diff --git a/.zed/rules b/.zed/rules index 06b3eb329..7e88196fe 100644 --- a/.zed/rules +++ b/.zed/rules @@ -175,7 +175,8 @@ src/orcapod/ join.py — Join (N-ary inner join, commutative) merge_join.py — MergeJoin (binary, colliding cols → sorted list[T]) semijoin.py — SemiJoin (binary, non-commutative) - batch.py — Batch (group rows, types become list[T]) + batch.py — Batch (group rows by count, types become list[T]) + group_by.py — GroupBy (many→one reduction keyed on tag values) column_selection.py — Select/Drop Tag/Data columns mappers.py — MapTags, MapData (rename columns) filters.py — PolarsFilter @@ -275,7 +276,16 @@ Prefixes are computed from SystemConstant in system_constants.py. 2. Name-extending — multi-input ops. System tag column name gets ::{pipeline_hash}:{canonical_position} appended. Commutative operators sort by pipeline_hash and sort system tag values per row. -3. Type-evolving — aggregation ops. Column type changes from str to list[str]. +3. Reducing — many→one ops (Batch, GroupBy). User tag and data columns become list[T], and + source-info columns become list[str] with one element per member. System tag columns must + stay scalar, because _build_record_id_preimage (core/nodes/function_node.py) hashes them + directly to derive record identity. They fold to a deterministic digest via + arrow_utils.fold_system_tag_values and their column name gains ::{pipeline_hash}. + + The fold is SHA-based (uuid5 for record_id, combine_hashes for source_id) and so is stable + across processes. Never use hash() or a set-based construction there: the digest becomes a + cache key, so a per-process value would look correct in a single-process test and miss the + cache on every new driver run. ### Key patterns @@ -296,4 +306,23 @@ Prefixes are computed from SystemConstant in system_constants.py. - DerivedSource before run() → raises ValueError (no computed records). - Join requires non-overlapping data columns; raises InputValidationError on collision. - MergeJoin requires colliding columns to have identical types; merges into sorted list[T]. -- Operators predict output schema (including system tag names) without computation. +- Operators predict output schema (including system tag names) without computation. Verify a + prediction against unary_static_process(stream).output_schema(...), NOT against process(...) + (circular — it delegates back to the pod) and NOT against as_table() (legitimately differs: + output_schema ignores columns.source, so _source_*, _content_hash and _context_key appear in + the table but never in the schema). +- GroupBy requires every column in `by` to be a scalar tag column; raises InputValidationError + for unknown, data, or list-valued columns (list-valued usually means the stream came from + Batch — group before batching, not after). Members sort by non-key tag values with system + record_id as tiebreaker; groups emit in key order, so reordered input yields a + byte-identical table. +- Data source-info values may be str, None, or list[...]; Arrow/Python types are derived from + the value, with None mapping to large_string. A many→one operator must emit a list for EVERY + row of a list-valued column — mixing in a bare None makes per-row as_table() schemas diverge + and pa.concat_tables fail on the barrier path. +- Aggregating operators (Batch, GroupBy) must build list-valued columns via + arrow_utils.build_aggregated_table, never pa.list_(field.type). Arrow cannot embed an + extension type inside a list value field, so the naive call raises + ArrowNotImplementedError on any logical-typed column (e.g. a pod annotated -> Path). + The helper builds the list over the element's storage type and wraps it in the outer + list[] extension type from ListLogicalType. diff --git a/CLAUDE.md b/CLAUDE.md index 624086dea..9e958448e 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -185,7 +185,8 @@ src/orcapod/ │ ├── join.py # Join (N-ary inner join, commutative) │ ├── merge_join.py # MergeJoin (binary, colliding cols → sorted list[T]) │ ├── semijoin.py # SemiJoin (binary, non-commutative) -│ ├── batch.py # Batch (group rows, types become list[T]) +│ ├── batch.py # Batch (group rows by count, types become list[T]) +│ ├── group_by.py # GroupBy (many→one reduction keyed on tag values) │ ├── column_selection.py # Select/Drop Tag/Data columns │ ├── mappers.py # MapTags, MapData (rename columns) │ └── filters.py # PolarsFilter @@ -266,6 +267,10 @@ FunctionNode. | Synthesizes new values | No | Yes | | Stream arity | Configurable | Single in, single out | +`GroupBy` is the only operator that reduces row count many→one: it collapses N rows sharing a +tag tuple into one row with list-valued members. It still synthesizes no new data values — +every emitted element came from an input row — and it keys only on tags, never on data. + ### Two identity chains Every pipeline element has two parallel hashes: @@ -301,7 +306,17 @@ Prefixes are computed from `SystemConstant` in `system_constants.py`. The `const 2. **Name-extending** — multi-input ops (join, merge join). Each input's system tag column name gets `::{pipeline_hash}:{canonical_position}` appended. Commutative operators canonically order inputs by `pipeline_hash` and sort system tag values per row. -3. **Type-evolving** — aggregation ops (batch). Column type changes from `str` to `list[str]`. +3. **Reducing** — many→one ops (`Batch`, `GroupBy`). User tag and data columns become + `list[T]`, and source-info columns become `list[str]` with one element per member. System + tag columns must stay **scalar**, because `_build_record_id_preimage` + (`core/nodes/function_node.py`) hashes them directly to derive record identity. They fold + to a deterministic digest via `arrow_utils.fold_system_tag_values` and their column name + gains `::{pipeline_hash}` — a blend of rules 2 and 3. + + The fold is SHA-based (`uuid5` for `record_id`, `combine_hashes` for `source_id`) and so is + stable across processes. Never use `hash()` or a set-based construction there: the digest + becomes a cache key, so a per-process value would look correct in a single-process test and + miss the cache on every new driver run. ### Schema types and ColumnConfig @@ -337,4 +352,25 @@ and `as_table()` methods. `all_info=True` sets everything to True. - MergeJoin requires colliding data columns to have identical types; merges into sorted `list[T]` with source columns reordered to match. - Operators predict their output schema (including system tag column names) without - performing the actual computation. + performing the actual computation. Verify a prediction against + `unary_static_process(stream).output_schema(...)`, **not** against `process(...)` (which + delegates straight back to the pod, making the check circular) and **not** against + `as_table()` (which legitimately differs: `ArrowTableStream.output_schema` ignores + `columns.source`, so `_source_*`, `_content_hash`, and `_context_key` appear in the table + but never in the schema). +- `GroupBy` requires every column in `by` to be a scalar tag column; it raises + `InputValidationError` for unknown, data, or list-valued columns (the last usually means the + stream came from `Batch` — group before batching, not after). Members are sorted by non-key + tag values with the system `record_id` as tiebreaker, so the hashed lists are stable across + runs. Groups themselves are emitted in key order, so a reordered input yields a + byte-identical table. +- `Data` source-info values may be `str`, `None`, or `list[...]`; the Arrow and Python types + are derived from the value, with `None` mapping to `large_string`. A many→one operator must + emit a list for **every** row of a list-valued column — mixing a bare `None` into some rows + makes per-row `as_table()` schemas diverge and `pa.concat_tables` fail on the barrier path. +- Aggregating operators (`Batch`, `GroupBy`) must build list-valued columns via + `arrow_utils.build_aggregated_table`, never `pa.list_(field.type)`. Arrow cannot embed an + extension type inside a list value field, so the naive call raises + `ArrowNotImplementedError` on any logical-typed column (e.g. a pod annotated `-> Path`). + The helper builds the list over the element's storage type and wraps it in the outer + `list[]` extension type from `ListLogicalType`. diff --git a/DESIGN_ISSUES.md b/DESIGN_ISSUES.md index c2d5c8121..9938514d1 100644 --- a/DESIGN_ISSUES.md +++ b/DESIGN_ISSUES.md @@ -698,6 +698,148 @@ Three categories of improvement are planned: overrides in the future but require careful handling of Polars expression evaluation and system-tag evolution respectively. +`GroupBy` (NPIPE-204) is barrier-only by construction and is not a candidate for either category: +no group can be emitted before the input channel closes, because any row not yet seen could belong +to a group already started. Emitting early would require a guarantee that input arrives clustered +by group key, which orcapod streams do not carry. + +--- + +### O2 — Operators silently discard a non-default `_context_key` +**Status:** open +**Severity:** medium + +Every operator reads its input with +`stream.as_table(columns={"source": True, "system_tags": True})` — `batch.py:48`, +`merge_join.py:168`, `semijoin.py:60`, `column_selection.py:524`. That column set excludes +`_context_key`, so no operator ever sees the input's data context. `ArrowTableStream.__init__` +then finds no context column on the result and substitutes +`contexts.get_default_context_key()`. + +Consequence: a stream carrying a non-default context key reverts to the default at the first +operator it crosses, with no warning. Function pods are unaffected — only the operator layer. + +This is pre-existing and layer-wide rather than specific to any one operator. It may also be +intended (operators produce output in the ambient context rather than inheriting an input's), in +which case the fix is to document the rule rather than change behavior. Logged while working +NPIPE-204; deliberately not addressed there. + +Fix: decide whether operator output should inherit the input context. If yes, request +`context: True` and propagate it, erroring when inputs disagree. If no, document the reset +explicitly on `StaticOutputPod`. + +--- + +### O3 — `_materialize_to_stream` broadcasts row 0's provenance to every row +**Status:** open +**Severity:** high + +`StaticOutputOperatorPod._materialize_to_stream` (`core/operators/static_output_pod.py:223`) +reads the source info of the *first* row and applies it to the whole concatenated table: + +```python +# Preserve actual source_info provenance from the first row +# (all rows share the same data columns and source tokens). +source_info = rows[0][1].source_info() +``` + +The comment is half right. All rows do share the same data *column names*, but not the same +*tokens* — a provenance token embeds a per-row identifier (`...::row_0::value` versus +`...::row_1::value`, built by `_make_provenance_token` in `core/sources/stream_builder.py`). +So every row from index 1 onward is attributed to row 0. `sync_orchestrator.py:207,222` +repeats the pattern. + +Observed directly in a `MergeJoin` round-trip: row 1 of the output carried `row_0` tokens. + +This predates NPIPE-204 and is equally wrong for scalar tokens, so nothing in that change +caused it. NPIPE-204 did make it **more visible**: previously the list-valued case crashed at +`as_table` (see U1), so a broken operator never got far enough to mis-attribute provenance. +That crash is now fixed, which converts a loud failure into a quiet wrong answer. + +Fix: build the source-info mapping per row rather than once from row 0, or push the +provenance into the concatenated table before constructing the stream. + +--- + +### O4 — `GroupBy` member order is undefined for duplicate tag tuples +**Status:** open +**Severity:** medium + +`GroupBy` sorts a group's members by their non-group-key tag values so the emitted lists are +stable across runs — orcapod hashes those lists to build the cache key, so an unstable order +causes spurious recomputes. + +When `by` covers *every* tag column there is no non-key tag left to sort on, and the operator +falls back to the system `record_id`. That is a UUID v5 over the source id and a per-row +provenance token, so it encodes source *row position*. Permuting the source table therefore +changes member order, changes the list hash, and triggers a recompute. Measured: 79 distinct +member orders across all 120 permutations of a 5-row duplicate-tag input. + +The branch is reachable with more than one member only when the input contains **duplicate tag +tuples** — `sort_columns` is empty exactly when the group key is the full tag tuple, so a +multi-member group requires duplicates. + +This is not fixable inside `GroupBy`. The only remaining ordering signal is data content, and +the operator / function pod boundary forbids an operator from inspecting data. + +The real root cause is upstream: `DuplicateTagError` (`errors.py:23`) is defined but **never +raised anywhere** in `src/` or `tests/`, so duplicate tag tuples are not prevented in the first +place. Tag uniqueness is assumed throughout the operator layer and enforced nowhere. + +Fix: enforce tag uniqueness at stream construction (raising `DuplicateTagError`), which makes +this branch unreachable with more than one member. + +--- +### O5 — A reduction persists a result computed from an incomplete member set +**Status:** open +**Severity:** high + +When an upstream pod's output changes, the first `job.run()` afterwards emits only the rows +that were recomputed in that pass. For a per-row pod that emission is *complete* — each row is +independent, so processing just the changed one is correct. For a many→one operator it is +*partial*: the group is reduced over only the members present in that pass, the reducing pod +executes, and its result is persisted with nothing marking it incomplete. The next run emits +the full set and the group is recomputed correctly. + +Measured with two upstream pod stages (`source → sync_like → stringify → [GroupBy] → pod`), +one member of one group changed, same Delta store, fresh objects per run: + +``` +CONTROL (per-row, no GroupBy) + changed run1 sync_like(99), stringify(sync_99), pod(sync_99.parquet) <- converged + changed run2 [] + +GROUPED + changed run1 sync_like(99), stringify(sync_99), GROUP ['sync_99'] <- ONE member + changed run2 GROUP ['sync_99', 'sync_1'] <- correct + changed run3 [] +``` + +The control converges in a single run; only the grouped path needs a second one, and it +executes on an incomplete set first. + +Consequence: a reducing pod with a side effect writes a complete-looking artifact from partial +input. For the motivating consumer (`common_clock_op`) an intermediate run can produce an +`alignment.json` built from one of a session's probes. It is replaced on the next run, so a +driver that runs to convergence never observes it — but a consumer reading between runs gets a +wrong answer with no signal. + +`Batch` behaves identically, so this is pre-existing pipeline semantics surfaced by reduction +rather than something `GroupBy` introduced. It is more consequential for `GroupBy`, because +`GroupBy` exists specifically so a pod can reason over a *complete* group, whereas nobody +derives a result from `Batch` membership. + +Not fixable inside the operator: an operator sees whatever its upstream emits and has no way to +know whether a group is complete. A fix needs a completeness signal at the node level — either +the upstream emitting its full cached set on every pass, or a reducing node deferring execution +until its inputs are known settled. + +History: first reported as a cache-corruption bug (wrong group recomputing, tag/data +misalignment), then retracted as a propagation lag "identical with and without `GroupBy`". Both +framings are wrong. There is no tag/data misalignment and no cache defect — but the control +above shows the behaviour is *not* identical, because a partial emission is harmless per-row +and harmful for a reduction. + --- ## `src/orcapod/core/` — AddResult pod and Pod Groups @@ -1055,15 +1197,45 @@ The `normalize_extension_columns` utility landed in ITL-432. ## `src/orcapod/utils/` ### U1 — Source-info column type hard-coded to `large_string` -**Status:** open +**Status:** resolved (`tag_data.py` half), open (`arrow_utils.py` half) **Severity:** critical -In `add_source_info_to_table()` (`arrow_utils.py:604`), when source info is a collection it is -unconditionally cast to `pa.list_(pa.large_string())`: +Two sibling call sites assume source-info values are always scalar strings. + +**`Data._ensure_source_info_table()` (`core/datagrams/tag_data.py:330`)** builds the Arrow schema +as `pa.field(k, pa.large_string())` for every key, and `Data.schema()` (line ~384) reports `str` +for every `_source_*` column. Any operator that produces list-valued source info therefore fails +inside `job.run()`: + +``` +pyarrow.lib.ArrowTypeError: Expected bytes, got a 'list' object + core/datagrams/tag_data.py:342 _ensure_source_info_table +``` + +Two operators hit this: `MergeJoin`, which carries source columns along as parallel lists when +merging colliding data columns (`merge_join.py:262`), and `Batch`, which list-wraps every column. +Both were reproduced against `966d759a`. + +**Fix:** landed in NPIPE-204 (`1b570e68`, `1ed22ee4`). `_source_info_arrow_type` and +`_source_info_python_type` in `core/datagrams/tag_data.py` derive the Arrow and Python types +from the stored value — `str`/`None` → `large_string`, list → `large_list()` +recursively. The `SourceInfoValue` alias lives in `types.py` next to `TagValue`/`DataValue`, +which lets `DataProtocol` and `ArrowTableStream` reference it without inverting the +`protocols/` → `core/` layering. No pipeline-DB schema bump: a node's source-column type is +fixed by its own output schema, so existing nodes keep `large_string`. Job-level regression +coverage is in `tests/test_pipeline/test_aggregation_job.py`. See +`superpowers/specs/2026-08-07-npipe-204-batch-group-by-design.md`. + +A third site, `polars_data_utils.add_source_info` (line 119), forces `dtype=pl.String()`. It is +dead code — nothing in `src/` calls it, and the tests importing `add_source_info` import it from +`arrow_utils` — and it carries a latent shadowing bug where `source_column` is rebound to a +`pl.Series` inside the per-column loop. NPIPE-204 deletes it rather than fixing it. + +**Still open:** in `add_source_info_to_table()` (`arrow_utils.py:604`), when source info is a +collection it is unconditionally cast to `pa.list_(pa.large_string())`: ```python # TODO: this won't work other data types!!! ``` - Any non-string collection values will fail or silently corrupt data. The logic also has an unclear nested isinstance check (line ~602: `# TODO: clean up the logic here`). @@ -1198,7 +1370,7 @@ Open questions: --- -## `src/orcapod/extension_types/` +## `src/orcapod/logical_types/` ### ET1 — `make_polars_extension_type` cannot accept a storage type containing nested extension types **Status:** open @@ -1275,6 +1447,14 @@ resolves to a logical type, pointing to this entry and PLT-1732. Use a direct `T type carries the annotation into the schema, and `reconstruct_from_arrow` re-registers `T` transitively on read. +**Operator layer (NPIPE-204):** `Batch` and `GroupBy` previously called +`pa.list_(field.type)` directly, so they raised `ArrowNotImplementedError: extension` on any +extension-typed column even after ITL-173 landed `ListLogicalType`. Both now build their +output through `arrow_utils.build_aggregated_table`, which constructs the list over the +element's storage type and wraps it in the outer `list[]` extension type. A pod +annotated `-> Path` groups into `list[orcapod.path]` and the downstream pod receives real +`Path` objects. + **Planned fix (PLT-1732, target v0.2):** Introduce `ListLogicalType` / `ListLogicalTypeFactory` and `StructLogicalType` / `StructLogicalTypeFactory`. A `list[UUID]` top-level column would be wrapped as a new extension type diff --git a/docs/api/operators.md b/docs/api/operators.md index ef1d69dd2..ca7fc539a 100644 --- a/docs/api/operators.md +++ b/docs/api/operators.md @@ -10,6 +10,8 @@ Operators perform structural transformations on streams without inspecting or sy ::: orcapod.core.operators.Batch +::: orcapod.core.operators.GroupBy + ::: orcapod.core.operators.SelectTagColumns ::: orcapod.core.operators.SelectDataColumns diff --git a/docs/concepts/operators.md b/docs/concepts/operators.md index 2c45a5fb1..ed2347778 100644 --- a/docs/concepts/operators.md +++ b/docs/concepts/operators.md @@ -1,10 +1,10 @@ # Operators Operators are structural transforms that reshape [streams](streams.md) without inspecting or -synthesizing data values. They join, filter, batch, rename, and select columns -- operations -that affect the *structure* of the data (which rows exist, which columns are present, how -columns are named) but never compute new values from data content. This is the key -distinction from [function pods](function-pods.md), which do the opposite: they transform +synthesizing data values. They join, filter, batch, group, rename, and select columns -- +operations that affect the *structure* of the data (which rows exist, which columns are +present, how columns are named) but never compute new values from data content. This is the +key distinction from [function pods](function-pods.md), which do the opposite: they transform data values but never touch tags or stream structure. ## The operator / function pod boundary @@ -23,12 +23,17 @@ This boundary ensures that structural operations (joins, filters) and value comp (transformations, model inference) are cleanly separated, making pipelines easier to reason about and optimize. +`GroupBy` is the one operator that reduces row count many-to-one: it collapses N rows sharing +a tag tuple into a single row with list-valued members. It stays on the operator side of the +boundary because it synthesizes no new data values -- every element of every emitted list came +from an input row -- and it keys only on tags, never on data content. + ## Operator categories ### `UnaryOperator` -- single input -Takes one stream, produces one stream. Used for filtering, column selection, renaming, and -batching. +Takes one stream, produces one stream. Used for filtering, column selection, renaming, +batching, and grouping. ### `BinaryOperator` -- two inputs @@ -119,6 +124,56 @@ Pass `batch_size=N` to create fixed-size batches instead of grouping everything: batch = Batch(batch_size=10, drop_partial_batch=False) ``` +### GroupBy + +Collapses every set of rows that share a tag tuple into a single row -- the one many-to-one +operator. Where `Batch` partitions by row count, `GroupBy` partitions by tag *value*, which is +what lets a downstream function pod receive a complete logical unit (all of a recording +session's probes, say) in one call. + +```python +from orcapod.sources import DictSource +from orcapod.operators import GroupBy + +source = DictSource( + data=[ + {"subject_id": "mouse_01", "probe": "imec0", "spike_count": 1200}, + {"subject_id": "mouse_01", "probe": "imec1", "spike_count": 980}, + {"subject_id": "mouse_02", "probe": "imec0", "spike_count": 1450}, + ], + tag_columns=["subject_id", "probe"], +) + +grouped = GroupBy(by=["subject_id"]).process(source) +for tag, data in grouped.iter_data(): + print("Tags:", tag.as_dict(), "Data:", data.as_dict()) + # Tags: {'subject_id': 'mouse_01'} Data: {'probe': ['imec0', 'imec1'], 'spike_count': [1200, 980]} + # Tags: {'subject_id': 'mouse_02'} Data: {'probe': ['imec0'], 'spike_count': [1450]} +``` + +The output schema follows three rules: + +- **Group keys stay scalar** and remain the output's tag columns (`subject_id: str`). +- **Non-key tag columns are promoted to list-valued data columns** (`probe: list[str]`), so a + consumer can tell which member each list element came from. +- **Data columns become `list[T]`** (`spike_count: list[int]`). + +Ordering is deterministic by construction, because orcapod hashes the emitted lists to build +cache keys: members are sorted by their non-key tag values (with the internal record id as a +final tiebreaker), and groups are emitted in group-key order. A reordered input therefore +produces a byte-identical output table. + +Every name in `by` must be a **scalar tag column** of the input. `GroupBy` raises +`InputValidationError` for an unknown column, a data column, or a list-valued tag column -- +the last usually means the stream came from `Batch`, so group *before* batching rather than +after. + +Streams also expose this as a fluent method: + +```python +grouped = source.group_by(["subject_id"]) +``` + ### Column selection Four operators for including or excluding columns: diff --git a/docs/reference/hashing.md b/docs/reference/hashing.md index e10d3340e..886b1580a 100644 --- a/docs/reference/hashing.md +++ b/docs/reference/hashing.md @@ -8,8 +8,9 @@ orcapod-python framework. For conceptual background on the two identity chains ## Hash Site Index -The table below summarises all 14 hash computation sites across the six usage groups -documented in the sections that follow. +The table below summarises all 15 hash computation sites across the six usage groups +documented in the sections that follow. (Site 6b is numbered alongside site 6 rather than +appended, because it is a variant of the same system tag mechanism.) | # | Site | Algorithm | Output format | One-line guarantee | |---|------|-----------|---------------|--------------------| @@ -19,6 +20,7 @@ documented in the sections that follow. | 4 | Default `source_id` | `StarfixArrowHasher.hash_table(table).to_hex(n)` | Truncated hex `str` | Unique per raw table content; used as source identifier when none is provided | | 5 | Per-row `record_id` (system tag value) | `uuid.uuid5(NAMESPACE, f"{source_id}::{provenance_token}")` | `bytes` (16, UUID v5) | Deterministic per `(source_id, row_identity)`; stable across re-runs of the same source | | 6 | Join system tag suffix | `stream.pipeline_hash().to_hex(n)` + `:{idx}` appended to column name | Column name suffix | Unique per `(input topology, canonical join position)`; encodes full join lineage in column name | +| 6b | Aggregating operator system tag fold | `uuid.uuid5(NAMESPACE, joined_member_hex)` for `record_id`; `combine_hashes(*member_values)` otherwise | `bytes` (16) / hex `str`, plus a column name suffix | Deterministic per ordered member set; keeps system tags scalar through a many-to-one reduction | | 7 | `compute_base_entry_id()` | `StarfixArrowHasher.hash_table(system_tags + INPUT_DATA_HASH_COL)` | `bytes` (`b"method:digest"`) | Unique per `(node, tag lineage, input_data content)` across all recomputation attempts | | 8 | `compute_pipeline_entry_id()` | `StarfixArrowHasher.hash_table(system_tags + INPUT_DATA_HASH_COL + recomputation_index)` | `bytes` (`b"method:digest"`) | Unique per `(node, tag lineage, input_data content, recomputation attempt)` | | 9 | Side-effect `record_id` | `StarfixArrowHasher.hash_table(system_tags + INPUT_DATA_HASH_COL + recomputation_index=0)` | `bytes` | Unique per `(tag lineage, input_data content)`; pod-version scoped via table path (`uri` + `pipeline_hash()`) | @@ -32,7 +34,8 @@ documented in the sections that follow. ## 7. Worked Example -The pipeline below exercises all 14 hash sites. `source_id` values are set explicitly so +The pipeline below exercises hash sites 1--14; it contains no aggregating operator, so site 6b +is not covered. `source_id` values are set explicitly so every content hash is reproducible across runs. Run the script to regenerate the values in this section: @@ -298,7 +301,27 @@ topology suffix, encoding the full join lineage into each column name. | **Algorithm** | Streams are first sorted by `stream.pipeline_hash().to_string()` for determinism. For each input at canonical position `idx`, every existing system tag column name has `{BLOCK_SEPARATOR}{stream.pipeline_hash().to_hex(system_tag_n_char)}:{idx}` appended via `arrow_utils.append_to_system_tags()`. (`system_tag_n_char` from `OrcapodConfig.hashing`; default `None` = full digest.) | | **Output format** | Column name suffix; no separate value is stored | | **Uniqueness guarantee** | Each post-join system tag column name uniquely identifies `(original schema, input topology, canonical join position)`; no collision even when joining streams with identical schemas | -| **Known exclusions** | `SemiJoin` passes system tags through unchanged. `Batch` changes the column type from `str` to `list[str]` but preserves the column name. | +| **Known exclusions** | `SemiJoin` passes system tags through unchanged. Aggregating operators (`Batch`, `GroupBy`) use the fold described in site 6b instead. | + +--- + +### Site 6b — Aggregating operator system tag fold + +| Field | Value | +|---|---| +| **Inputs** | The ordered list of a group's member values for one system tag column, plus that column's name (which selects the fold) | +| **Algorithm** | `arrow_utils.fold_system_tag_values()`. For a `record_id` column: `uuid.uuid5(_AGGREGATED_RECORD_ID_NAMESPACE, BLOCK_SEPARATOR.join(v.hex() for v in values))`, where `_AGGREGATED_RECORD_ID_NAMESPACE = uuid.uuid5(NAMESPACE_URL, "https://orcapod.org/namespaces/aggregated-record-id")` is a fixed constant. For every other system tag column: `combine_hashes(*[str(v) for v in values], order=False)` — SHA-256 over the members concatenated in order. `None` members contribute the empty string. The folded column name then gains `{BLOCK_SEPARATOR}{stream.pipeline_hash().to_hex(system_tag_n_char)}` via `arrow_utils.append_to_system_tags()`. | +| **Output format** | `bytes` (16, UUID v5 bit pattern) for a `record_id` column; 64-character hex `str` otherwise — in both cases **scalar**, matching the input column's type | +| **Uniqueness guarantee** | Deterministic per ordered member set, and stable across processes (both primitives are SHA-based). `Batch` and `GroupBy` both order their members deterministically before folding, so an unchanged member set always folds to the same digest | +| **Known exclusions** | Member order is significant; a permuted member set folds differently. This is why `GroupBy` sorts members by their non-key tag values with `record_id` as a tiebreaker rather than relying on upstream emission order | + +> **Why fold rather than list-wrap.** A many-to-one operator turns user tag and data columns +> into `list[T]`, but system tag columns must stay scalar: `_build_record_id_preimage` +> (`core/nodes/function_node.py`) hashes those columns directly to derive record identity. +> Folding rather than dropping them preserves the Merkle link, so an upstream recompute that +> yields identical data still invalidates downstream records. Because the digest becomes a +> cache key, the fold must never use `hash()` or a set-based construction — a per-process value +> would look correct in a single-process test and miss the cache on every new driver run. > **Key insight for entry IDs:** Because `compute_base_entry_id()` (§3) calls > `tag.as_table(columns={"system_tags": True})`, its preimage captures the full set of diff --git a/src/orcapod/core/datagrams/tag_data.py b/src/orcapod/core/datagrams/tag_data.py index 8acafa1f4..e9cc3e993 100644 --- a/src/orcapod/core/datagrams/tag_data.py +++ b/src/orcapod/core/datagrams/tag_data.py @@ -8,10 +8,10 @@ requested via ``ColumnConfig(system_tags=True)``. ``Data`` - Extends ``Datagram`` with *source information*: provenance tokens (strings or None) - keyed by data-column name. Source-info keys are stored without the - ``constants.SOURCE_PREFIX`` internally and added back when serialising via - ``as_dict()`` / ``as_table()``. + Extends ``Datagram`` with *source information*: provenance tokens (strings, + None, or lists of tokens) keyed by data-column name. Source-info keys are + stored without the ``constants.SOURCE_PREFIX`` internally and added back when + serialising via ``as_dict()`` / ``as_table()``. """ from __future__ import annotations @@ -25,7 +25,15 @@ from orcapod.core.datagrams.datagram import Datagram from orcapod.semantic_types import infer_python_schema_from_pylist_data from orcapod.system_constants import constants -from orcapod.types import ColumnConfig, ContentHash, DataValue, Schema, SchemaLike +from orcapod.types import ( + ColumnConfig, + ContentHash, + DataType, + DataValue, + Schema, + SchemaLike, + SourceInfoValue, +) from orcapod.utils import arrow_utils from orcapod.utils.lazy_module import LazyModule @@ -237,11 +245,60 @@ def copy(self, include_cache: bool = True, preserve_id: bool = False) -> Self: # --------------------------------------------------------------------------- +# TODO(NPIPE-204): these two helpers only understand scalars and ``list[T]``. +# Structured source-info values (a struct/dataclass token, or a mapping) will need +# their own branches, and the recursion should dispatch on a logical type rather +# than on ``isinstance(value, list)``. Deriving the type from the first element is +# also a shortcut: it is correct for the homogeneous lists many->one operators +# produce, but not in general. Raised in review on PR #250. +def _source_info_arrow_type(value: "SourceInfoValue") -> "pa.DataType": + """Derive the Arrow type for a single source-info value. + + Scalars and unknowns map to ``large_string``; lists map to ``large_list`` + of their element type, recursively. An empty list defaults to + ``large_list(large_string)``. + + Args: + value: The stored provenance token. + + Returns: + The Arrow type to declare for this value. + """ + import pyarrow as _pa + + if isinstance(value, list): + if not value: + return _pa.large_list(_pa.large_string()) + return _pa.large_list(_source_info_arrow_type(value[0])) + return _pa.large_string() + + +def _source_info_python_type(value: "SourceInfoValue") -> DataType: + """Derive the Python type for a single source-info value. + + Mirrors ``_source_info_arrow_type`` for the ``Schema`` representation, and + shares its list-only limitation -- see the TODO above that function. + + Args: + value: The stored provenance token. + + Returns: + ``str`` for scalars and unknowns, ``list[...]`` for lists. + """ + if isinstance(value, list): + if not value: + return list[str] + return list[_source_info_python_type(value[0])] # type: ignore[misc] + return str + + class Data(Datagram): """ Datagram with source-information tracking. - Source info maps each data-column name to a provenance token (``str | None``). + Source info maps each data-column name to a provenance token + (``SourceInfoValue``: a string, ``None``, or -- for many->one operators -- a + list of tokens, one per aggregated member). Keys in ``_source_info`` are stored **without** the ``SOURCE_PREFIX``; the prefix is added transparently when serialising to dict or Arrow table. @@ -254,7 +311,7 @@ def __init__( self, data: "Mapping[str, DataValue] | pa.Table | pa.RecordBatch", meta_info: "Mapping[str, DataValue] | None" = None, - source_info: "Mapping[str, str | None] | None" = None, + source_info: "Mapping[str, SourceInfoValue] | None" = None, python_schema: "SchemaLike | None" = None, data_context: "str | contexts.DataContext | None" = None, record_uuid: "uuid.UUID | None" = None, @@ -293,7 +350,7 @@ def __init__( ) si_table = prefixed_tables[constants.SOURCE_PREFIX] if si_table.num_columns > 0 and si_table.num_rows > 0: - self._source_info: dict[str, str | None] = { + self._source_info: dict[str, SourceInfoValue] = { k.removeprefix(constants.SOURCE_PREFIX): v for k, v in si_table.to_pylist()[0].items() } @@ -306,7 +363,7 @@ def __init__( for k, v in data.items() if not k.startswith(constants.SOURCE_PREFIX) } - contained_source_info: dict[str, str | None] = { + contained_source_info: dict[str, SourceInfoValue] = { k.removeprefix(constants.SOURCE_PREFIX): v # type: ignore[misc] for k, v in data.items() if k.startswith(constants.SOURCE_PREFIX) @@ -337,7 +394,10 @@ def _ensure_source_info_table(self) -> "pa.Table": for k, v in self._source_info.items() } schema = _pa.schema( - [_pa.field(k, _pa.large_string()) for k in prefixed] + [ + _pa.field(k, _source_info_arrow_type(v)) + for k, v in prefixed.items() + ] ) self._source_info_table = _pa.Table.from_pylist( [prefixed], schema=schema @@ -350,11 +410,11 @@ def _ensure_source_info_table(self) -> "pa.Table": # Source-info API # ------------------------------------------------------------------ - def source_info(self) -> "dict[str, str | None]": + def source_info(self) -> "dict[str, SourceInfoValue]": """Return source info for all data-column keys (None for unknown).""" return {k: self._source_info.get(k) for k in self.keys()} - def with_source_info(self, **source_info: "str | None") -> Self: + def with_source_info(self, **source_info: "SourceInfoValue") -> Self: """Create a copy with updated source-information entries.""" current = dict(self._source_info) for key, value in source_info.items(): @@ -391,7 +451,9 @@ def schema( column_config = ColumnConfig.handle_config(columns, all_info=all_info) if column_config.source: for key in super().keys(): - schema[f"{constants.SOURCE_PREFIX}{key}"] = str + schema[f"{constants.SOURCE_PREFIX}{key}"] = _source_info_python_type( + self._source_info.get(key) + ) return Schema(schema) def arrow_schema( diff --git a/src/orcapod/core/operators/__init__.py b/src/orcapod/core/operators/__init__.py index 7728d37aa..b8e846d5f 100644 --- a/src/orcapod/core/operators/__init__.py +++ b/src/orcapod/core/operators/__init__.py @@ -6,6 +6,7 @@ SelectTagColumns, ) from .filters import PolarsFilter +from .group_by import GroupBy from .index import Index from .join import Join from .mappers import MapData, MapTags @@ -20,6 +21,7 @@ "MapTags", "MapData", "Batch", + "GroupBy", "SelectTagColumns", "SelectDataColumns", "DropTagColumns", diff --git a/src/orcapod/core/operators/batch.py b/src/orcapod/core/operators/batch.py index 0795ede74..d3419599d 100644 --- a/src/orcapod/core/operators/batch.py +++ b/src/orcapod/core/operators/batch.py @@ -7,7 +7,9 @@ from orcapod.core.operators.base import UnaryOperator from orcapod.core.streams import ArrowTableStream from orcapod.protocols.core_protocols import DataProtocol, StreamProtocol, TagProtocol +from orcapod.system_constants import constants from orcapod.types import ColumnConfig +from orcapod.utils import arrow_utils from orcapod.utils.lazy_module import LazyModule if TYPE_CHECKING: @@ -41,44 +43,71 @@ def validate_unary_input(self, stream: StreamProtocol) -> None: return None def unary_static_process(self, stream: StreamProtocol) -> StreamProtocol: - """ - This method should be implemented by subclasses to define the specific behavior of the binary operator. - It takes two streams as input and returns a new stream as output. + """Group rows into fixed-size batches, list-wrapping their values. + + Tag and data columns become list-valued. Source-info columns become + list-valued too, one element per batch member. System-tag columns are + folded to a scalar instead -- record identity hashes them directly, so + they must not become lists. + + Args: + stream: The upstream stream. + + Returns: + A stream with one row per batch. """ table = stream.as_table(columns={"source": True, "system_tags": True}) - tag_columns, data_columns = stream.keys() + tag_columns, _ = stream.keys() - data_list = table.to_pylist() + system_tag_columns = tuple( + c for c in table.column_names if c.startswith(constants.SYSTEM_TAG_PREFIX) + ) + member_columns = tuple( + c for c in table.column_names if c not in system_tag_columns + ) - batched_data = [] + data_list = table.to_pylist() - next_batch = {} + batches: list[list[dict[str, Any]]] = [] + next_batch: list[dict[str, Any]] = [] - i = 0 for entry in data_list: - i += 1 - for c in entry: - next_batch.setdefault(c, []).append(entry[c]) - - if self.batch_size > 0 and i >= self.batch_size: - batched_data.append(next_batch) - next_batch = {} - i = 0 - - if i > 0 and not self.drop_partial_batch: - batched_data.append(next_batch) - - # Build the target schema upfront: each field becomes list, nullable=False. - # Passing schema= directly avoids a second table construction. - batched_schema = pa.schema([ - pa.field(f.name, pa.list_(f.type), nullable=False) - for f in table.schema - ]) - batched_table = pa.Table.from_pylist(batched_data, schema=batched_schema) + next_batch.append(entry) + if self.batch_size > 0 and len(next_batch) >= self.batch_size: + batches.append(next_batch) + next_batch = [] + + if next_batch and not self.drop_partial_batch: + batches.append(next_batch) + + batched_data = [ + { + **{c: [m[c] for m in members] for c in member_columns}, + **{ + c: arrow_utils.fold_system_tag_values(c, [m[c] for m in members]) + for c in system_tag_columns + }, + } + for members in batches + ] + + batched_table = arrow_utils.build_aggregated_table( + batched_data, + table.schema, + member_columns, + stream.data_context.type_converter, + ) + + n_char = self.orcapod_config.hashing.system_tag_n_char + batched_table = arrow_utils.append_to_system_tags( + batched_table, stream.pipeline_hash().to_hex(n_char) + ) + return ArrowTableStream( batched_table, tag_columns=tag_columns, + data_context=stream.data_context, ) def unary_output_schema( @@ -88,17 +117,35 @@ def unary_output_schema( columns: ColumnConfig | dict[str, Any] | None = None, all_info: bool = False, ) -> tuple[Schema, Schema]: + """Predict the batched output schemas without batching. + + Every user tag, data, and source column becomes ``list[T]``. System + tag columns keep their scalar type and gain a ``::{pipeline_hash}`` + name suffix. + + Args: + stream: The upstream stream. + columns: Column inclusion config. + all_info: Include all info columns. + + Returns: + A ``(tag_schema, data_schema)`` tuple. """ - This method should be implemented by subclasses to return the schemas of the input and output streams. - It takes two streams as input and returns a tuple of schemas. - """ - tag_types, data_types = stream.output_schema( - columns=columns, all_info=all_info - ) - batched_tag_types = {k: list[v] for k, v in tag_types.items()} + tag_types, data_types = stream.output_schema(columns=columns, all_info=all_info) + n_char = self.orcapod_config.hashing.system_tag_n_char + suffix = stream.pipeline_hash().to_hex(n_char) + + batched_tag_types: dict[str, Any] = {} + for name, col_type in tag_types.items(): + if name.startswith(constants.SYSTEM_TAG_PREFIX): + batched_tag_types[f"{name}{constants.BLOCK_SEPARATOR}{suffix}"] = ( + col_type + ) + else: + batched_tag_types[name] = list[col_type] + batched_data_types = {k: list[v] for k, v in data_types.items()} - # TODO: check if this is really necessary return Schema(batched_tag_types), Schema(batched_data_types) async def async_execute( diff --git a/src/orcapod/core/operators/group_by.py b/src/orcapod/core/operators/group_by.py new file mode 100644 index 000000000..3b8331926 --- /dev/null +++ b/src/orcapod/core/operators/group_by.py @@ -0,0 +1,272 @@ +"""GroupBy operator — many->one reduction keyed on tag values.""" + +from __future__ import annotations + +from collections.abc import Collection +from typing import TYPE_CHECKING, Any, get_origin + +from orcapod.core.operators.base import UnaryOperator +from orcapod.core.streams import ArrowTableStream +from orcapod.errors import InputValidationError +from orcapod.protocols.core_protocols import StreamProtocol +from orcapod.system_constants import constants +from orcapod.types import ColumnConfig, Schema +from orcapod.utils import arrow_utils +from orcapod.utils.lazy_module import LazyModule + +if TYPE_CHECKING: + import pyarrow as pa +else: + pa = LazyModule("pyarrow") + + +class GroupBy(UnaryOperator): + """Reduce rows sharing a tag tuple into one packet with list-valued members. + + This is the only many->one operator. Every other operator preserves one + row per tag; ``GroupBy`` collapses N rows into one, which is what lets a + downstream pod receive a whole group at once (for example, all of a + recording session's per-probe result parquets). + + Given tags ``(subject, date, probe)`` and data ``(path)``, grouping by + ``["subject", "date"]`` emits one row per distinct ``(subject, date)``: + + * ``subject`` and ``date`` stay scalar and remain the output's tag columns + * ``probe`` becomes a list-valued **data** column, so a consumer can tell + which member each list element came from + * ``path`` becomes list-valued + * the ``_source_*`` column of each original data column becomes + list-valued, one element per member. A promoted tag column such as + ``probe`` had no provenance token to begin with, so its ``_source_probe`` + is the scalar null the stream fills in, not a list. + * system-tag columns fold to a scalar digest and gain a + ``::{pipeline_hash}`` name suffix + + Members are sorted by their non-group-key tag values, with ``record_id`` + appended as a final tiebreaker, so the emitted lists are stable across + runs. This matters because orcapod hashes those lists to build the cache + key -- an unsorted list would make an identical member set hash differently + and trigger a spurious recompute. Groups themselves are emitted in + group-key order for the same reason: a reordered input must produce an + identical output. + + Tag tuples are expected to be unique within a stream, but nothing enforces + that. If the input does contain duplicate tag tuples, the members sharing + a tuple are separated only by ``record_id``, which is fixed when the source + is materialized -- so their order is stable for a given source table, but + undefined if the source table itself is permuted. + + Contrast with ``Batch``, which partitions by row count for throughput and + keeps its tag columns as list-valued tags. + + Args: + by: Tag column names to group on. Must be non-empty, free of + duplicates, and must all be scalar tag columns of the input stream. + """ + + def __init__(self, by: Collection[str], **kwargs: Any) -> None: + by_tuple = tuple(by) + if not by_tuple: + raise ValueError("GroupBy requires at least one column in `by`.") + duplicates = sorted({c for c in by_tuple if by_tuple.count(c) > 1}) + if duplicates: + raise ValueError( + f"GroupBy `by` contains duplicate column names: {duplicates}. " + f"Got {list(by_tuple)}." + ) + self.by = by_tuple + super().__init__(**kwargs) + + def identity_structure(self) -> Any: + return (self.__class__.__name__, self.by) + + def to_config(self) -> dict[str, Any]: + """Serialize this GroupBy operator to a config dict. + + ``by`` is emitted as a list rather than a tuple so the config stays + JSON-serializable; ``__init__`` normalizes it back to a tuple. + + Returns: + A dict with ``class_name``, ``module_path``, and ``config`` keys, + where ``config`` contains ``by``. + """ + config = super().to_config() + config["config"] = {"by": list(self.by)} + return config + + # ------------------------------------------------------------------ + # Validation + # ------------------------------------------------------------------ + + def validate_unary_input(self, stream: StreamProtocol) -> None: + """Verify every grouping column is a scalar tag column of the input. + + Args: + stream: The upstream stream to validate. + + Raises: + InputValidationError: If any name in ``by`` is not a tag column, or + names a list-valued tag column. + """ + tag_columns, data_columns = stream.keys() + missing = [c for c in self.by if c not in tag_columns] + if missing: + raise InputValidationError( + f"GroupBy: {missing} are not tag columns of the input stream. " + f"Available tag columns: {list(tag_columns)}. " + f"(Data columns cannot be grouping keys: {list(data_columns)})" + ) + + # A list-valued tag -- what `Batch` produces -- is unhashable, so it + # cannot key a group. Check the schema rather than hashing a value so + # the error names the operator and column instead of surfacing a bare + # `TypeError: unhashable type: 'list'` from the grouping loop. + tag_types, _ = stream.output_schema() + non_scalar = { + c: tag_types.get(c) for c in self.by if get_origin(tag_types.get(c)) is list + } + if non_scalar: + raise InputValidationError( + f"GroupBy: {list(non_scalar)} are list-valued tag columns and " + f"cannot be grouping keys. Column types: {non_scalar}. " + "A list-valued tag usually means the stream came from `Batch`; " + "group before batching rather than after." + ) + + # ------------------------------------------------------------------ + # Processing + # ------------------------------------------------------------------ + + def unary_static_process(self, stream: StreamProtocol) -> StreamProtocol: + """Partition rows by group key and emit one row per group. + + Args: + stream: The upstream stream. + + Returns: + A stream with one row per distinct group-key tuple. + """ + table = stream.as_table(columns={"source": True, "system_tags": True}) + tag_columns, _ = stream.keys() + + system_tag_columns = tuple( + c for c in table.column_names if c.startswith(constants.SYSTEM_TAG_PREFIX) + ) + member_columns = tuple( + c + for c in table.column_names + if c not in self.by and c not in system_tag_columns + ) + # Non-key user tags order the members of a group. The `record_id` + # columns are appended as a final tiebreaker: tag tuples are supposed + # to be unique within a stream, but nothing enforces that, and without + # the tiebreaker duplicate tuples would fall back to emission order -- + # which Ray scheduling and DB fetch order make nondeterministic. + # `record_id` is fixed when the source is materialized, so it is immune + # to that shuffling. + # + # A joined stream carries one `record_id` per canonical input position, + # and the broadcast side of a fan-out repeats its id on every row it + # was joined into. So *every* record_id column has to take part -- + # consulting only the first would tie on exactly the fan-out rows the + # tiebreaker exists for. They are sorted by name so the key order is + # itself independent of column order in the input table. + sort_columns = tuple(c for c in tag_columns if c not in self.by) + sort_columns += tuple( + sorted( + c + for c in system_tag_columns + if c.startswith(constants.SYSTEM_TAG_RECORD_ID_PREFIX) + ) + ) + + groups: dict[tuple[Any, ...], list[dict[str, Any]]] = {} + for row in table.to_pylist(): + groups.setdefault(tuple(row[c] for c in self.by), []).append(row) + + grouped_rows: list[dict[str, Any]] = [] + # Emit groups in key order rather than first-appearance order, so a + # reordered input produces a byte-identical output table. + for key, members in sorted( + groups.items(), key=lambda kv: tuple((v is None, v) for v in kv[0]) + ): + if sort_columns: + # The leading bool keeps a null comparable against a real + # value of any type, and sorts nulls last. + members.sort( + key=lambda r: tuple((r[c] is None, r[c]) for c in sort_columns) + ) + + grouped_rows.append({ + **dict(zip(self.by, key)), + **{c: [m[c] for m in members] for c in member_columns}, + **{ + c: arrow_utils.fold_system_tag_values(c, [m[c] for m in members]) + for c in system_tag_columns + }, + }) + + grouped_table = arrow_utils.build_aggregated_table( + grouped_rows, + table.schema, + member_columns, + stream.data_context.type_converter, + ) + + n_char = self.orcapod_config.hashing.system_tag_n_char + grouped_table = arrow_utils.append_to_system_tags( + grouped_table, stream.pipeline_hash().to_hex(n_char) + ) + + return ArrowTableStream( + grouped_table, + tag_columns=self.by, + data_context=stream.data_context, + ) + + # ------------------------------------------------------------------ + # Schema prediction + # ------------------------------------------------------------------ + + def unary_output_schema( + self, + stream: StreamProtocol, + *, + columns: ColumnConfig | dict[str, Any] | None = None, + all_info: bool = False, + ) -> tuple[Schema, Schema]: + """Predict the grouped output schemas without grouping. + + Args: + stream: The upstream stream. + columns: Column inclusion config. + all_info: Include all info columns. + + Returns: + A ``(tag_schema, data_schema)`` tuple. Group keys stay scalar in + the tag schema; promoted non-key tags and list-wrapped data columns + land in the data schema. + """ + tag_types, data_types = stream.output_schema(columns=columns, all_info=all_info) + n_char = self.orcapod_config.hashing.system_tag_n_char + suffix = stream.pipeline_hash().to_hex(n_char) + + out_tag_types: dict[str, Any] = {} + out_data_types: dict[str, Any] = {} + + for name, col_type in tag_types.items(): + if name.startswith(constants.SYSTEM_TAG_PREFIX): + out_tag_types[f"{name}{constants.BLOCK_SEPARATOR}{suffix}"] = col_type + elif name in self.by: + out_tag_types[name] = col_type + else: + # Promoted to a list-valued data column. No `_source_*` entry + # is emitted: `ArrowTableStream.output_schema` returns the data + # schema unconditionally, so `columns.source` is a no-op at the + # schema level even though `as_table` does add those columns. + out_data_types[name] = list[col_type] + + for name, col_type in data_types.items(): + out_data_types[name] = list[col_type] + + return Schema(out_tag_types), Schema(out_data_types) diff --git a/src/orcapod/core/streams/arrow_table_stream.py b/src/orcapod/core/streams/arrow_table_stream.py index 09bfc5995..ea87df4a2 100644 --- a/src/orcapod/core/streams/arrow_table_stream.py +++ b/src/orcapod/core/streams/arrow_table_stream.py @@ -11,7 +11,7 @@ from orcapod.protocols.core_protocols import PodProtocol, StreamProtocol, TagProtocol from orcapod.protocols.hashing_protocols import PipelineElementProtocol from orcapod.system_constants import constants -from orcapod.types import ColumnConfig, Schema +from orcapod.types import ColumnConfig, Schema, SourceInfoValue from orcapod.utils import arrow_utils from orcapod.utils.lazy_module import LazyModule @@ -41,7 +41,7 @@ def __init__( table: "pa.Table", tag_columns: Collection[str] = (), system_tag_columns: Collection[str] = (), - source_info: dict[str, str | None] | None = None, + source_info: dict[str, SourceInfoValue] | None = None, producer: PodProtocol | None = None, upstreams: tuple[StreamProtocol, ...] = (), **kwargs, diff --git a/src/orcapod/core/streams/base.py b/src/orcapod/core/streams/base.py index 2a18355ab..b210ac756 100644 --- a/src/orcapod/core/streams/base.py +++ b/src/orcapod/core/streams/base.py @@ -148,6 +148,27 @@ def batch( self, label=label ) + def group_by( + self, + by: Collection[str], + label: str | None = None, + ) -> StreamBase: + """Reduce rows sharing a tag tuple into one packet per group. + + Group-key columns stay scalar tags; every other column becomes + list-valued. See ``orcapod.core.operators.GroupBy``. + + Args: + by: Tag column names to group on. + label: Optional node label for the pipeline graph. + + Returns: + A stream with one row per distinct group-key tuple. + """ + from orcapod.core.operators import GroupBy + + return GroupBy(by=by)(self, label=label) + def polars_filter( self, *predicates: Any, diff --git a/src/orcapod/pipeline/serialization.py b/src/orcapod/pipeline/serialization.py index f255634ac..40e9de5f0 100644 --- a/src/orcapod/pipeline/serialization.py +++ b/src/orcapod/pipeline/serialization.py @@ -167,6 +167,7 @@ def _build_operator_registry() -> dict[str, type]: Batch, DropDataColumns, DropTagColumns, + GroupBy, Join, MapData, MapTags, @@ -182,6 +183,7 @@ def _build_operator_registry() -> dict[str, type]: "MergeJoin": MergeJoin, "SemiJoin": SemiJoin, "Batch": Batch, + "GroupBy": GroupBy, "SelectTagColumns": SelectTagColumns, "DropTagColumns": DropTagColumns, "SelectDataColumns": SelectDataColumns, diff --git a/src/orcapod/protocols/core_protocols/datagrams.py b/src/orcapod/protocols/core_protocols/datagrams.py index 831574b7a..3fd9b4f27 100644 --- a/src/orcapod/protocols/core_protocols/datagrams.py +++ b/src/orcapod/protocols/core_protocols/datagrams.py @@ -14,7 +14,7 @@ ContentIdentifiableProtocol, DataContextAwareProtocol, ) -from orcapod.types import ColumnConfig, DataValue, Schema +from orcapod.types import ColumnConfig, DataValue, Schema, SourceInfoValue if TYPE_CHECKING: import pyarrow as pa @@ -652,7 +652,7 @@ class DataProtocol(DatagramProtocol, Protocol): data flow: Tags provide context, Datas provide content. """ - def source_info(self) -> dict[str, str | None]: + def source_info(self) -> dict[str, SourceInfoValue]: """ Return metadata about the data's source/origin. @@ -664,13 +664,15 @@ def source_info(self) -> dict[str, str | None]: - Processing pipeline information Returns: - dict[str, str | None]: Source information for each data column as key-value pairs. + dict[str, SourceInfoValue]: Source information for each data column as + key-value pairs. A value is a provenance token string, ``None`` when + unknown, or a list of tokens for many-to-one operators. """ ... def with_source_info( self, - **source_info: str | None, + **source_info: SourceInfoValue, ) -> Self: """ Create a new data with updated source information. diff --git a/src/orcapod/types.py b/src/orcapod/types.py index 9c1ff170c..6cba0269c 100644 --- a/src/orcapod/types.py +++ b/src/orcapod/types.py @@ -47,6 +47,11 @@ arbitrarily nested collection thereof. Tags are used to label and organise data and datagrams.""" +SourceInfoValue: TypeAlias = str | None | list["SourceInfoValue"] +"""A per-column provenance token: a string, ``None`` when the provenance is +unknown, or -- for many-to-one operators such as ``GroupBy`` and +``MergeJoin`` -- a list of tokens, one per aggregated member.""" + PathSet: TypeAlias = PathLike | Collection[PathLike | None] """A single path or an arbitrarily nested collection of paths (with optional ``None`` entries). Used when operations need to address multiple files at diff --git a/src/orcapod/utils/arrow_utils.py b/src/orcapod/utils/arrow_utils.py index f3a1ca057..45c233bcc 100644 --- a/src/orcapod/utils/arrow_utils.py +++ b/src/orcapod/utils/arrow_utils.py @@ -1,9 +1,11 @@ from __future__ import annotations +import uuid from collections import defaultdict -from collections.abc import Mapping, Collection +from collections.abc import Mapping, Collection, Sequence from typing import Any, TYPE_CHECKING +from orcapod.hashing.hash_utils import combine_hashes from orcapod.system_constants import constants from orcapod.types import ColumnConfig from orcapod.utils.lazy_module import LazyModule @@ -1177,6 +1179,121 @@ def append_to_system_tags(table: "pa.Table", value: str) -> "pa.Table": return table.rename_columns(column_name_map) +# Fixed namespace for aggregated record IDs produced by many->one operators. +# Mirrors _SOURCE_RECORD_ID_NAMESPACE in core/sources/stream_builder.py. +# Computed value: uuid.UUID('96411bfc-d3ba-5395-ba6f-5bb5726f18ad') +_AGGREGATED_RECORD_ID_NAMESPACE = uuid.uuid5( + uuid.NAMESPACE_URL, + "https://orcapod.org/namespaces/aggregated-record-id", +) + + +def fold_system_tag_values(column_name: str, values: Sequence[Any]) -> str | bytes: + """Fold a group's system-tag values into one scalar of the same type. + + Many->one operators must emit scalar system tags, because + ``_build_record_id_preimage`` (``core/nodes/function_node.py``) hashes + those columns directly to derive a record's identity. Each column folds + independently over its own ordered member values. + + Both digests are SHA-based and therefore stable across processes. Never + substitute ``hash()`` or a set-based construction: orcapod uses the result + as a cache key, so a per-process digest would miss the cache on every new + driver run while looking correct in a single-process test. + + Member order is significant -- it matches the order of the list-valued + data columns the folded tag accompanies. + + Args: + column_name: The system-tag column name, used to select the fold. + Names starting with ``constants.SYSTEM_TAG_RECORD_ID_PREFIX`` fold + to ``binary(16)``; everything else folds to a hex string. + values: The group's member values, in emission order. + + Returns: + 16 raw bytes for a record_id column, a 64-character hex string + otherwise. + """ + if column_name.startswith(constants.SYSTEM_TAG_RECORD_ID_PREFIX): + name = constants.BLOCK_SEPARATOR.join( + "" if v is None else v.hex() for v in values + ) + return uuid.uuid5(_AGGREGATED_RECORD_ID_NAMESPACE, name).bytes + return combine_hashes( + *["" if v is None else str(v) for v in values], order=False + ) + + +def build_aggregated_table( + rows: "Sequence[Mapping[str, Any]]", + input_schema: "pa.Schema", + member_columns: "Collection[str]", + type_converter: Any, +) -> "pa.Table": + """Build a many->one operator's output table, list-wrapping member columns. + + Columns named in ``member_columns`` become list-valued, one element per + group member; every other column keeps its input field unchanged (that is + how scalar group keys and folded system tags pass through). + + Logical element types are preserved. Arrow cannot embed an extension type + inside a list value field, so ``pa.list_(extension_type)`` raises + ``ArrowNotImplementedError`` (see ``DESIGN_ISSUES`` ET1/ET2). For such a + column the list is built over the element's *storage* type and then wrapped + in the outer ``list[]`` extension type supplied by + ``ListLogicalType``. A pod annotated ``-> Path`` therefore groups into + ``list[orcapod.path]`` rather than losing the type or failing. + + Args: + rows: One mapping per output row, values already aggregated into lists + for every member column. + input_schema: Schema of the operator's input table, used to derive + element types and to pass non-member fields through unchanged. + member_columns: Names of the columns to list-wrap. + type_converter: The stream's type converter, used to resolve an + extension type's Python element type and the matching outer list + extension type. + + Returns: + The aggregated ``pa.Table``. + """ + member_columns = set(member_columns) + fields: list[pa.Field] = [] + # Extension columns are built as plain storage lists, then re-wrapped. + extension_overrides: dict[str, Any] = {} + + for field in input_schema: + if field.name not in member_columns: + fields.append(field) + continue + if isinstance(field.type, pa.ExtensionType): + element_python_type = type_converter.arrow_type_to_python_type(field.type) + list_type = type_converter.python_type_to_arrow_type( + list[element_python_type] + ) + extension_overrides[field.name] = list_type + fields.append( + pa.field(field.name, list_type.storage_type, nullable=False) + ) + else: + fields.append(pa.field(field.name, pa.list_(field.type), nullable=False)) + + table = pa.Table.from_pylist(list(rows), schema=pa.schema(fields)) + + for name, list_type in extension_overrides.items(): + index = table.schema.get_field_index(name) + storage = table.column(name) + wrapped = pa.chunked_array( + [pa.ExtensionArray.from_storage(list_type, chunk) for chunk in storage.chunks], + type=list_type, + ) + table = table.set_column( + index, pa.field(name, list_type, nullable=False), wrapped + ) + + return table + + def _parse_system_tag_column( col_name: str, ) -> tuple[str, str, str] | None: diff --git a/src/orcapod/utils/polars_data_utils.py b/src/orcapod/utils/polars_data_utils.py index a6ca778c5..443614ed7 100644 --- a/src/orcapod/utils/polars_data_utils.py +++ b/src/orcapod/utils/polars_data_utils.py @@ -88,38 +88,3 @@ def append_to_system_tags(df: "pl.DataFrame", value: str) -> "pl.DataFrame": if c.startswith(constants.SYSTEM_TAG_PREFIX) } return df.rename(column_name_map) - - -def add_source_info( - df: "pl.DataFrame", - source_info: str | Collection[str] | None, - exclude_prefixes: Collection[str] = ( - constants.META_PREFIX, - constants.DATAGRAM_PREFIX, - ), - exclude_columns: Collection[str] = (), -) -> "pl.DataFrame": - """Add source information to an Arrow table.""" - # Create a new column with the source information - if source_info is None or isinstance(source_info, str): - source_column = [source_info] * df.height - elif isinstance(source_info, Collection): - if len(source_info) != df.height: - raise ValueError( - "Length of source_info collection must match number of rows in the table." - ) - source_column = source_info - - # identify columns for which source columns should be created - - for col in df.columns: - if col.startswith(tuple(exclude_prefixes)) or col in exclude_columns: - continue - source_column = pl.Series( - f"{constants.SOURCE_PREFIX}{col}", - [f"{source_val}::{col}" for source_val in source_column], - dtype=pl.String(), - ) - df = df.with_columns(source_column) - - return df diff --git a/superpowers/plans/2026-08-07-npipe-204-groupby-operator.md b/superpowers/plans/2026-08-07-npipe-204-groupby-operator.md new file mode 100644 index 000000000..8a9eecedc --- /dev/null +++ b/superpowers/plans/2026-08-07-npipe-204-groupby-operator.md @@ -0,0 +1,1822 @@ +# NPIPE-204 — `GroupBy` Operator + Type-Aware Source Info Implementation Plan + +> **For agentic workers:** REQUIRED SUB-SKILL: Use superpowers:subagent-driven-development (recommended) or superpowers:executing-plans to implement this plan task-by-task. Steps use checkbox (`- [ ]`) syntax for tracking. + +**Goal:** Add a `GroupBy` operator that reduces N rows sharing a tag tuple into one packet with list-valued members, and fix the `Data` datagram so list-valued provenance no longer crashes `job.run()`. + +**Architecture:** Three independent layers. (1) `Data` derives `_source_*` Arrow/Python types from the stored value instead of hard-coding `large_string`/`str` — this alone repairs the already-shipped `MergeJoin`. (2) A shared `arrow_utils` helper folds N members' system-tag values into one scalar using SHA-based primitives already in the codebase. (3) `GroupBy` partitions by tag values, sorts members deterministically, list-wraps everything except the group keys, and folds the system tags. `Batch` is semantically untouched — it only picks up the provenance fix. + +**Tech Stack:** Python 3.11+, PyArrow, Polars, pytest, `uv run` for all commands, Delta Lake (`DeltaTableDatabase`) for job-level tests. + +**Spec:** `superpowers/specs/2026-08-07-npipe-204-batch-group-by-design.md` + +**Linear:** NPIPE-204 (status already *In Progress*). Branch `arnoldb/npipe-204-orcapod-add-tag-based-grouping-to-batch-manyone-reduction` is already checked out. + +--- + +## Background an implementer needs + +**The bug, minimally.** A `Data` datagram whose source info holds a list crashes on `as_table`: + +```python +data.source_info() # {'probe': None, 'path': ['s0', 's1']} +data.schema(columns={"source": True}) # '_source_path': str <- WRONG +data.as_table(columns={"source": True}) # ArrowTypeError: Expected bytes, got a 'list' object +``` + +Two operators already produce this state: `MergeJoin` (carries source columns as parallel lists when merging colliding data columns, `merge_join.py:262`) and `Batch` (list-wraps every column). Neither is caught by existing tests because all operator tests stop at `op.process(stream)` + `as_table()`; the crash happens in `StaticOutputOperatorPod._materialize_to_stream`, only reached via `job.run()`. + +**Column vocabulary.** For a stream with tag `subject` and data `path`, `stream.as_table(columns={"source": True, "system_tags": True})` yields exactly: + +``` +subject # user tag +path # data +_tag_source_id:: # system tag, large_string +_tag_record_id:: # system tag, binary(16) +_source_path # provenance, large_string +``` + +Note `_context_key` is **absent** — no operator sees it. Do not add it. + +**`ArrowTableStream` fills in missing source columns.** If the output table has a data column with no matching `_source_`, the stream creates one with value `None` and type `large_string`. Verified. So `GroupBy` does not need to synthesize source columns for the tag columns it promotes to data — but `unary_output_schema` must still *predict* them. + +--- + +## File Structure + +| File | Responsibility | Action | +|---|---|---| +| `src/orcapod/core/datagrams/tag_data.py` | `Data` source-info type derivation | Modify | +| `src/orcapod/utils/polars_data_utils.py` | Delete dead `add_source_info` | Modify | +| `src/orcapod/utils/arrow_utils.py` | Shared system-tag fold helper | Modify | +| `src/orcapod/core/operators/batch.py` | Provenance fix only | Modify | +| `src/orcapod/core/operators/group_by.py` | The new operator | **Create** | +| `src/orcapod/core/operators/__init__.py` | Export | Modify | +| `src/orcapod/pipeline/serialization.py` | Operator registry | Modify | +| `src/orcapod/core/streams/base.py` | `.group_by()` fluent method | Modify | +| `tests/test_core/datagrams/test_data_source_info_types.py` | Part 1 unit tests | **Create** | +| `tests/test_utils/test_arrow_utils.py` | Fold helper + cross-process stability | Modify | +| `tests/test_core/operators/test_group_by.py` | Operator-level tests | **Create** | +| `tests/test_pipeline/test_aggregation_job.py` | Job-level tests | **Create** | +| `CLAUDE.md` + `.zed/rules` | Doc updates | Modify | +| `DESIGN_ISSUES.md` | U1 resolution note | Modify | + +--- + +## Task 1: Type-aware source info in `Data` + +**Files:** +- Modify: `src/orcapod/core/datagrams/tag_data.py` +- Test: `tests/test_core/datagrams/test_data_source_info_types.py` (create) + +- [ ] **Step 1: Write the failing test** + +Create `tests/test_core/datagrams/test_data_source_info_types.py`: + +```python +"""Source-info values may be lists, not just scalar strings. + +Many->one operators (GroupBy, MergeJoin) produce one provenance token per +member. `Data` must represent those without collapsing or crashing. +""" + +from __future__ import annotations + +import pyarrow as pa + +from orcapod.core.datagrams import Data + + +def _data_with_mixed_source_info() -> Data: + """Data with one list-valued and one scalar-null source token.""" + return Data( + {"probe": [0, 1], "path": ["a", "b"]}, + source_info={"probe": None, "path": ["s0", "s1"]}, + ) + + +class TestListValuedSourceInfo: + def test_schema_reports_list_type_for_list_valued_token(self): + data = _data_with_mixed_source_info() + schema = data.schema(columns={"source": True}) + assert schema["_source_path"] == list[str] + + def test_schema_reports_str_for_none_token(self): + data = _data_with_mixed_source_info() + schema = data.schema(columns={"source": True}) + assert schema["_source_probe"] is str + + def test_as_table_round_trips_list_valued_token(self): + data = _data_with_mixed_source_info() + table = data.as_table(columns={"source": True}) + assert table.schema.field("_source_path").type == pa.large_list( + pa.large_string() + ) + assert table.column("_source_path").to_pylist() == [["s0", "s1"]] + + def test_as_table_keeps_none_token_as_large_string(self): + data = _data_with_mixed_source_info() + table = data.as_table(columns={"source": True}) + assert table.schema.field("_source_probe").type == pa.large_string() + assert table.column("_source_probe").to_pylist() == [None] + + def test_empty_list_token_defaults_to_list_of_string(self): + data = Data({"path": ["a"]}, source_info={"path": []}) + table = data.as_table(columns={"source": True}) + assert table.schema.field("_source_path").type == pa.large_list( + pa.large_string() + ) + + def test_scalar_token_unchanged(self): + """Existing scalar behavior must not regress.""" + data = Data({"path": "a"}, source_info={"path": "src::row_0::path"}) + table = data.as_table(columns={"source": True}) + assert table.schema.field("_source_path").type == pa.large_string() + assert data.schema(columns={"source": True})["_source_path"] is str + + def test_arrow_table_construction_recovers_list_token(self): + """Data built from an Arrow table keeps list-valued source info.""" + table = pa.table({ + "path": pa.array([["a", "b"]], pa.list_(pa.large_string())), + "_source_path": pa.array([["s0", "s1"]], pa.list_(pa.large_string())), + }) + data = Data(table) + assert data.source_info()["path"] == ["s0", "s1"] + assert data.as_table(columns={"source": True}).column( + "_source_path" + ).to_pylist() == [["s0", "s1"]] +``` + +- [ ] **Step 2: Run test to verify it fails** + +```bash +uv run pytest tests/test_core/datagrams/test_data_source_info_types.py -v +``` + +Expected: `test_schema_reports_list_type_for_list_valued_token` FAILS (`assert str == list[str]`), and the four `as_table` tests FAIL with `ArrowTypeError: Expected bytes, got a 'list' object`. `test_schema_reports_str_for_none_token` and `test_scalar_token_unchanged` should already PASS. + +- [ ] **Step 3: Add the type-derivation helpers** + +In `src/orcapod/core/datagrams/tag_data.py`, immediately above `class Data(Datagram):` (currently line 240), add: + +```python +# A provenance token is a string, an unknown (None), or — for many->one +# operators such as GroupBy and MergeJoin — a list of tokens, one per member. +SourceInfoValue = str | None | list["SourceInfoValue"] + + +def _source_info_arrow_type(value: "SourceInfoValue") -> "pa.DataType": + """Derive the Arrow type for a single source-info value. + + Scalars and unknowns map to ``large_string``; lists map to ``large_list`` + of their element type, recursively. An empty list defaults to + ``large_list(large_string)``. + + Args: + value: The stored provenance token. + + Returns: + The Arrow type to declare for this value. + """ + import pyarrow as _pa + + if isinstance(value, (list, tuple)): + if not value: + return _pa.large_list(_pa.large_string()) + return _pa.large_list(_source_info_arrow_type(value[0])) + return _pa.large_string() + + +def _source_info_python_type(value: "SourceInfoValue") -> type: + """Derive the Python type for a single source-info value. + + Mirrors ``_source_info_arrow_type`` for the ``Schema`` representation. + + Args: + value: The stored provenance token. + + Returns: + ``str`` for scalars and unknowns, ``list[...]`` for lists. + """ + if isinstance(value, (list, tuple)): + if not value: + return list[str] + return list[_source_info_python_type(value[0])] # type: ignore[misc] + return str +``` + +- [ ] **Step 4: Use the helpers in `_ensure_source_info_table`** + +Replace the body of `_ensure_source_info_table` (currently `tag_data.py:330-347`): + +```python + def _ensure_source_info_table(self) -> "pa.Table": + if self._source_info_table is None: + import pyarrow as _pa + + if self._source_info: + prefixed = { + f"{constants.SOURCE_PREFIX}{k}": v + for k, v in self._source_info.items() + } + schema = _pa.schema( + [ + _pa.field(k, _source_info_arrow_type(v)) + for k, v in prefixed.items() + ] + ) + self._source_info_table = _pa.Table.from_pylist( + [prefixed], schema=schema + ) + else: + self._source_info_table = _pa.table({}) + return self._source_info_table +``` + +- [ ] **Step 5: Use the helper in `Data.schema`** + +In `Data.schema` (currently `tag_data.py:384-395`), replace the `if column_config.source:` block: + +```python + if column_config.source: + for key in super().keys(): + schema[f"{constants.SOURCE_PREFIX}{key}"] = _source_info_python_type( + self._source_info.get(key) + ) +``` + +- [ ] **Step 6: Widen the type annotations** + +Four annotation sites in `tag_data.py`, all mechanical — no logic change: + +1. `Data.__init__` signature (line ~257): + `source_info: "Mapping[str, str | None] | None" = None` + → `source_info: "Mapping[str, SourceInfoValue] | None" = None` +2. Arrow-path assignment (line ~296): + `self._source_info: dict[str, str | None] = {` + → `self._source_info: dict[str, SourceInfoValue] = {` +3. Dict-path local (line ~309): + `contained_source_info: dict[str, str | None] = {` + → `contained_source_info: dict[str, SourceInfoValue] = {` +4. `source_info()` return (line ~353) and `with_source_info()` kwargs (line ~357): + `def source_info(self) -> "dict[str, str | None]":` + → `def source_info(self) -> "dict[str, SourceInfoValue]":` + `def with_source_info(self, **source_info: "str | None") -> Self:` + → `def with_source_info(self, **source_info: "SourceInfoValue") -> Self:` + +- [ ] **Step 7: Run the test to verify it passes** + +```bash +uv run pytest tests/test_core/datagrams/test_data_source_info_types.py -v +``` + +Expected: all 7 PASS. + +- [ ] **Step 8: Verify the MergeJoin crash is gone** + +This is the payoff — an already-shipped operator that was broken. + +```bash +uv run pytest tests/test_core/ tests/test_utils/ -x -q +``` + +Expected: no failures, no new errors. + +- [ ] **Step 9: Commit** + +```bash +git add src/orcapod/core/datagrams/tag_data.py tests/test_core/datagrams/test_data_source_info_types.py +git commit -m "fix(datagrams): derive source-info column types from value (NPIPE-204) + +_ensure_source_info_table and Data.schema hard-coded large_string / str for +every _source_* field, so any operator producing a list-valued provenance +token crashed with ArrowTypeError inside job.run(). MergeJoin already does +this (merge_join.py:262) and was silently broken. + +Types are now derived from the stored value. None still maps to +large_string, so unknown-provenance columns are unchanged. No pipeline-DB +schema bump: a node's source-column type is fixed by its own output schema. + +Refs DESIGN_ISSUES U1. +NPIPE-204" +``` + +--- + +## Task 2: Delete the dead `polars_data_utils.add_source_info` + +**Files:** +- Modify: `src/orcapod/utils/polars_data_utils.py` + +Third site with the same hard-coded assumption (`dtype=pl.String()`, line 119). It has **zero callers** in `src/` and no test coverage — the tests that import `add_source_info` import it from `arrow_utils`. It also has a latent shadowing bug: `source_column` is rebound to a `pl.Series` inside the per-column loop, so from the second column onward it formats `f"{}::{col}"`. + +- [ ] **Step 1: Confirm it is unreferenced** + +```bash +grep -rn "polars_data_utils" --include=*.py src/ tests/ +grep -rn "add_source_info" --include=*.py src/ tests/ +``` + +Expected: the only `polars_data_utils` references are `drop_system_columns` (from `data_frame_source.py:57` and `polling_source.py:454`). Every `add_source_info` reference resolves to `arrow_utils`. If either expectation fails, **stop** and report — do not delete. + +- [ ] **Step 2: Delete the function** + +Remove the entire `def add_source_info(...)` definition from `src/orcapod/utils/polars_data_utils.py` (starts line 93, ends at the `return df` around line 125). Leave every other function in the module untouched. + +- [ ] **Step 3: Verify nothing broke** + +```bash +uv run pytest tests/ -q -x +``` + +Expected: full suite passes. + +- [ ] **Step 4: Commit** + +```bash +git add src/orcapod/utils/polars_data_utils.py +git commit -m "chore(utils): delete dead polars_data_utils.add_source_info (NPIPE-204) + +Zero callers in src/, no test coverage, and a latent shadowing bug where +source_column is rebound to a pl.Series inside the per-column loop. It was +also a third site hard-coding a scalar string type for source info. + +Refs DESIGN_ISSUES U1. +NPIPE-204" +``` + +--- + +## Task 3: Shared system-tag fold helper + +**Files:** +- Modify: `src/orcapod/utils/arrow_utils.py` +- Test: `tests/test_utils/test_arrow_utils.py` + +`_build_record_id_preimage` (`core/nodes/function_node.py:82`) hashes system-tag columns directly, so they must be scalar. A many→one operator must fold N members' values into one, preserving each column's Arrow type. + +**Both primitives are reused from the codebase and are SHA-based, therefore stable across processes.** Never use `hash()` or set iteration order here — a fold using either would pass every same-process test and then miss the cache on every new driver run. + +- [ ] **Step 1: Write the failing test** + +Append to `tests/test_utils/test_arrow_utils.py`: + +```python +# --------------------------------------------------------------------------- +# fold_system_tag_values +# --------------------------------------------------------------------------- + + +class TestFoldSystemTagValues: + """Folding N members' system-tag values into one scalar (NPIPE-204). + + The expected digests below are hard-coded on purpose. A fold that used + hash() or set-iteration order would still be self-consistent within one + process; pinning the values is what catches it. + """ + + SOURCE_COL = "_tag_source_id::abc123" + RECORD_COL = "_tag_record_id::abc123" + + RIDS = [ + bytes.fromhex("0102030405060708090a0b0c0d0e0f10"), + bytes.fromhex("1112131415161718191a1b1c1d1e1f20"), + ] + EXPECTED_RID = bytes.fromhex("853be16a3f38565f8ced039f84fdbea6") + EXPECTED_SID = ( + "7916442d59841140bedf6c1f5dcc1304ae9fce0ba885765c06e511086b85da2e" + ) + + def test_record_id_folds_to_16_bytes(self): + from orcapod.utils.arrow_utils import fold_system_tag_values + + result = fold_system_tag_values(self.RECORD_COL, self.RIDS) + assert isinstance(result, bytes) + assert len(result) == 16 + + def test_record_id_digest_is_pinned(self): + from orcapod.utils.arrow_utils import fold_system_tag_values + + assert fold_system_tag_values(self.RECORD_COL, self.RIDS) == self.EXPECTED_RID + + def test_source_id_digest_is_pinned(self): + from orcapod.utils.arrow_utils import fold_system_tag_values + + result = fold_system_tag_values(self.SOURCE_COL, ["src_a", "src_b"]) + assert result == self.EXPECTED_SID + + def test_order_matters(self): + """Member order is part of the identity, matching the data lists.""" + from orcapod.utils.arrow_utils import fold_system_tag_values + + forward = fold_system_tag_values(self.RECORD_COL, self.RIDS) + reverse = fold_system_tag_values(self.RECORD_COL, list(reversed(self.RIDS))) + assert forward != reverse + + def test_single_member_is_still_folded(self): + """A one-member group folds rather than passing the value through.""" + from orcapod.utils.arrow_utils import fold_system_tag_values + + result = fold_system_tag_values(self.RECORD_COL, self.RIDS[:1]) + assert isinstance(result, bytes) and len(result) == 16 + assert result != self.RIDS[0] + + def test_none_members_are_tolerated(self): + from orcapod.utils.arrow_utils import fold_system_tag_values + + assert isinstance( + fold_system_tag_values(self.RECORD_COL, [None, self.RIDS[0]]), bytes + ) + assert isinstance( + fold_system_tag_values(self.SOURCE_COL, [None, "src_a"]), str + ) + + def test_digest_is_stable_across_processes(self): + """Fresh interpreter, therefore fresh PYTHONHASHSEED. + + This is the test that catches a fold built on hash() or set order: + such a fold is self-consistent within one process and only diverges + on a new driver run. + """ + import subprocess + import sys + + script = ( + "from orcapod.utils.arrow_utils import fold_system_tag_values\n" + "rids = [bytes.fromhex('0102030405060708090a0b0c0d0e0f10'),\n" + " bytes.fromhex('1112131415161718191a1b1c1d1e1f20')]\n" + "print(fold_system_tag_values('_tag_record_id::abc123', rids).hex())\n" + "print(fold_system_tag_values('_tag_source_id::abc123', ['src_a','src_b']))\n" + ) + out = subprocess.run( + [sys.executable, "-c", script], + capture_output=True, + text=True, + check=True, + ).stdout.split() + assert out[0] == self.EXPECTED_RID.hex() + assert out[1] == self.EXPECTED_SID +``` + +- [ ] **Step 2: Run test to verify it fails** + +```bash +uv run pytest tests/test_utils/test_arrow_utils.py::TestFoldSystemTagValues -v +``` + +Expected: all FAIL with `ImportError: cannot import name 'fold_system_tag_values'`. + +- [ ] **Step 3: Implement the helper** + +In `src/orcapod/utils/arrow_utils.py`, add immediately after `append_to_system_tags` (currently ends line 1157): + +```python +# Fixed namespace for aggregated record IDs produced by many->one operators. +# Mirrors _SOURCE_RECORD_ID_NAMESPACE in core/sources/stream_builder.py. +# Computed value: uuid.UUID('96411bfc-d3ba-5395-ba6f-5bb5726f18ad') +_AGGREGATED_RECORD_ID_NAMESPACE = uuid.uuid5( + uuid.NAMESPACE_URL, + "https://orcapod.org/namespaces/aggregated-record-id", +) + + +def fold_system_tag_values( + column_name: str, values: "Sequence[Any]" +) -> "str | bytes": + """Fold a group's system-tag values into one scalar of the same type. + + Many->one operators must emit scalar system tags, because + ``_build_record_id_preimage`` (``core/nodes/function_node.py``) hashes + those columns directly to derive a record's identity. Each column folds + independently over its own ordered member values. + + Both digests are SHA-based and therefore stable across processes. Never + substitute ``hash()`` or a set-based construction: orcapod uses the result + as a cache key, so a per-process digest would miss the cache on every new + driver run while looking correct in a single-process test. + + Member order is significant — it matches the order of the list-valued data + columns the folded tag accompanies. + + Args: + column_name: The system-tag column name, used to select the fold. Names + starting with ``constants.SYSTEM_TAG_RECORD_ID_PREFIX`` fold to + ``binary(16)``; everything else folds to a hex string. + values: The group's member values, in emission order. + + Returns: + 16 raw bytes for a record_id column, a 64-character hex string + otherwise. + """ + if column_name.startswith(constants.SYSTEM_TAG_RECORD_ID_PREFIX): + name = constants.BLOCK_SEPARATOR.join( + "" if v is None else v.hex() for v in values + ) + return uuid.uuid5(_AGGREGATED_RECORD_ID_NAMESPACE, name).bytes + return combine_hashes( + *["" if v is None else str(v) for v in values], order=False + ) +``` + +Add the imports at the top of `arrow_utils.py` if not already present: + +```python +import uuid + +from orcapod.hashing.hash_utils import combine_hashes +``` + +If importing `combine_hashes` at module scope creates a circular import, move it inside the function body instead and note why. + +- [ ] **Step 4: Run test to verify it passes** + +```bash +uv run pytest tests/test_utils/test_arrow_utils.py::TestFoldSystemTagValues -v +``` + +Expected: all 7 PASS, including `test_digest_is_stable_across_processes`. + +- [ ] **Step 5: Commit** + +```bash +git add src/orcapod/utils/arrow_utils.py tests/test_utils/test_arrow_utils.py +git commit -m "feat(arrow_utils): add fold_system_tag_values for many-to-one operators (NPIPE-204) + +System tags must stay scalar because _build_record_id_preimage hashes them +directly. Each column folds independently over its ordered member values: +record_id via uuid5 (matching stream_builder._make_record_id), source_id via +combine_hashes. Both SHA-based, so the digest is stable across processes -- +pinned by a subprocess test, since a hash()-based fold would look correct +within one process and miss the cache on every new driver run. + +NPIPE-204" +``` + +--- + +## Task 4: `Batch` provenance fix + +**Files:** +- Modify: `src/orcapod/core/operators/batch.py:43-102` +- Test: `tests/test_core/operators/test_operators.py` + +`Batch`'s partitioning, its list-valued tag columns, and its streaming `async_execute` all stay exactly as they are. The only change: system-tag columns fold to scalars instead of being list-wrapped, and `_source_*` columns keep their list form (now representable thanks to Task 1). + +- [ ] **Step 1: Write the failing test** + +Add to `tests/test_core/operators/test_operators.py`, inside `class TestBatchBehavior`: + +```python + def test_batch_system_tags_are_scalar(self): + """System tags must stay scalar -- record identity hashes them directly.""" + from orcapod.core.sources import ArrowTableSource + from orcapod.system_constants import constants + + table = pa.table({ + "animal": ["cat", "dog"], + "weight": [4.0, 12.0], + }) + source = ArrowTableSource(table, tag_columns=["animal"], infer_nullable=True) + out = Batch(batch_size=0).process(source) + result = out.as_table(columns={"source": True, "system_tags": True}) + + sys_cols = [ + c for c in result.column_names + if c.startswith(constants.SYSTEM_TAG_PREFIX) + ] + assert sys_cols, "expected system tag columns on the batched output" + for col in sys_cols: + assert not pa.types.is_list(result.schema.field(col).type) + assert not pa.types.is_large_list(result.schema.field(col).type) + + def test_batch_source_columns_are_lists(self): + """Provenance stays per-member rather than collapsing.""" + from orcapod.core.sources import ArrowTableSource + from orcapod.system_constants import constants + + table = pa.table({ + "animal": ["cat", "dog"], + "weight": [4.0, 12.0], + }) + source = ArrowTableSource(table, tag_columns=["animal"], infer_nullable=True) + out = Batch(batch_size=0).process(source) + result = out.as_table(columns={"source": True, "system_tags": True}) + + src_col = f"{constants.SOURCE_PREFIX}weight" + assert src_col in result.column_names + assert len(result.column(src_col).to_pylist()[0]) == 2 +``` + +- [ ] **Step 2: Run test to verify it fails** + +```bash +uv run pytest tests/test_core/operators/test_operators.py::TestBatchBehavior -v +``` + +Expected: `test_batch_system_tags_are_scalar` FAILS (system tag columns are `list<...>`). `test_batch_source_columns_are_lists` should already PASS. + +- [ ] **Step 3: Rewrite `unary_static_process`** + +Replace `Batch.unary_static_process` (`batch.py:43-82`) with: + +```python + def unary_static_process(self, stream: StreamProtocol) -> StreamProtocol: + """Group rows into fixed-size batches, list-wrapping their values. + + Tag and data columns become list-valued. Source-info columns become + list-valued too, one element per batch member. System-tag columns are + folded to a scalar instead — record identity hashes them directly, so + they must not become lists. + + Args: + stream: The upstream stream. + + Returns: + A stream with one row per batch. + """ + table = stream.as_table(columns={"source": True, "system_tags": True}) + + tag_columns, _ = stream.keys() + + system_tag_columns = tuple( + c + for c in table.column_names + if c.startswith(constants.SYSTEM_TAG_PREFIX) + ) + member_columns = tuple( + c for c in table.column_names if c not in system_tag_columns + ) + + data_list = table.to_pylist() + + batches: list[list[dict[str, Any]]] = [] + next_batch: list[dict[str, Any]] = [] + + for entry in data_list: + next_batch.append(entry) + if self.batch_size > 0 and len(next_batch) >= self.batch_size: + batches.append(next_batch) + next_batch = [] + + if next_batch and not self.drop_partial_batch: + batches.append(next_batch) + + batched_data = [ + { + **{c: [m[c] for m in members] for c in member_columns}, + **{ + c: arrow_utils.fold_system_tag_values(c, [m[c] for m in members]) + for c in system_tag_columns + }, + } + for members in batches + ] + + input_fields = {f.name: f for f in table.schema} + batched_schema = pa.schema( + [ + pa.field(c, pa.list_(input_fields[c].type), nullable=False) + if c in member_columns + else input_fields[c] + for c in table.column_names + ] + ) + batched_table = pa.Table.from_pylist(batched_data, schema=batched_schema) + + n_char = self.orcapod_config.hashing.system_tag_n_char + batched_table = arrow_utils.append_to_system_tags( + batched_table, stream.pipeline_hash().to_hex(n_char) + ) + + return ArrowTableStream( + batched_table, + tag_columns=tag_columns, + data_context=stream.data_context, + ) +``` + +Add these imports to the top of `batch.py`: + +```python +from orcapod.system_constants import constants +from orcapod.utils import arrow_utils +``` + +- [ ] **Step 4: Update `unary_output_schema` to mirror it** + +Replace `Batch.unary_output_schema` (`batch.py:84-102`): + +```python + def unary_output_schema( + self, + stream: StreamProtocol, + *, + columns: ColumnConfig | dict[str, Any] | None = None, + all_info: bool = False, + ) -> tuple[Schema, Schema]: + """Predict the batched output schemas without batching. + + Every user tag, data, and source column becomes ``list[T]``. System + tag columns keep their scalar type and gain a ``::{pipeline_hash}`` + name suffix. + + Args: + stream: The upstream stream. + columns: Column inclusion config. + all_info: Include all info columns. + + Returns: + A ``(tag_schema, data_schema)`` tuple. + """ + tag_types, data_types = stream.output_schema(columns=columns, all_info=all_info) + n_char = self.orcapod_config.hashing.system_tag_n_char + suffix = stream.pipeline_hash().to_hex(n_char) + + batched_tag_types: dict[str, Any] = {} + for name, col_type in tag_types.items(): + if name.startswith(constants.SYSTEM_TAG_PREFIX): + batched_tag_types[ + f"{name}{constants.BLOCK_SEPARATOR}{suffix}" + ] = col_type + else: + batched_tag_types[name] = list[col_type] + + batched_data_types = {k: list[v] for k, v in data_types.items()} + + return Schema(batched_tag_types), Schema(batched_data_types) +``` + +- [ ] **Step 5: Run the tests to verify they pass** + +```bash +uv run pytest tests/test_core/operators/ -v +``` + +Expected: all PASS, including the pre-existing `TestBatchBehavior` cases (`test_batch_groups_rows`, `test_batch_drop_partial`, `test_batch_output_lineage`, `test_batch_size_zero_returns_single_batch`, `test_negative_batch_size_raises`). + +- [ ] **Step 6: Commit** + +```bash +git add src/orcapod/core/operators/batch.py tests/test_core/operators/test_operators.py +git commit -m "fix(batch): fold system tags to scalar instead of list-wrapping (NPIPE-204) + +Batch list-wrapped every column including the system tags, but record +identity hashes those columns directly, so they must stay scalar. They now +fold via fold_system_tag_values and gain a ::{pipeline_hash} name suffix, +mirroring the name-extending rule joins already use. + +Partitioning, list-valued tag columns, and the streaming async_execute path +are unchanged. + +NPIPE-204" +``` + +--- + +## Task 5: The `GroupBy` operator + +**Files:** +- Create: `src/orcapod/core/operators/group_by.py` +- Test: `tests/test_core/operators/test_group_by.py` (create) + +- [ ] **Step 1: Write the failing test** + +Create `tests/test_core/operators/test_group_by.py`: + +```python +"""Tests for the GroupBy operator — many->one reduction keyed on tag values.""" + +from __future__ import annotations + +import pyarrow as pa +import pytest + +from orcapod.core.operators import GroupBy +from orcapod.core.sources import ArrowTableSource +from orcapod.errors import InputValidationError +from orcapod.system_constants import constants + + +@pytest.fixture +def session_table() -> pa.Table: + """Two sessions x two probes: the common-clock shape from NPIPE-204.""" + return pa.table({ + "subject": ["G", "G", "G", "G"], + "date": ["d1", "d1", "d2", "d2"], + "probe": [1, 0, 1, 0], + "path": ["b", "a", "d", "c"], + }) + + +@pytest.fixture +def session_source(session_table) -> ArrowTableSource: + return ArrowTableSource( + session_table, + tag_columns=["subject", "date", "probe"], + infer_nullable=True, + ) + + +class TestGroupByShape: + def test_one_row_per_distinct_key(self, session_source): + out = GroupBy(by=["subject", "date"]).process(session_source) + assert len(out.as_table()) == 2 + + def test_group_keys_are_scalar_tags(self, session_source): + op = GroupBy(by=["subject", "date"]) + out = op.process(session_source) + tag_cols, _ = out.keys() + assert set(tag_cols) == {"subject", "date"} + assert out.as_table().column("subject").to_pylist() == ["G", "G"] + + def test_non_key_tags_promoted_to_list_data(self, session_source): + out = GroupBy(by=["subject", "date"]).process(session_source) + _, data_cols = out.keys() + assert "probe" in data_cols + assert out.as_table().column("probe").to_pylist() == [[0, 1], [0, 1]] + + def test_data_columns_are_lists(self, session_source): + out = GroupBy(by=["subject", "date"]).process(session_source) + assert out.as_table().column("path").to_pylist() == [["a", "b"], ["c", "d"]] + + def test_source_columns_are_lists(self, session_source): + out = GroupBy(by=["subject", "date"]).process(session_source) + table = out.as_table(columns={"source": True}) + assert len(table.column(f"{constants.SOURCE_PREFIX}path").to_pylist()[0]) == 2 + + def test_system_tags_are_scalar_and_renamed(self, session_source): + out = GroupBy(by=["subject", "date"]).process(session_source) + table = out.as_table(columns={"system_tags": True}) + sys_cols = [ + c for c in table.column_names + if c.startswith(constants.SYSTEM_TAG_PREFIX) + ] + assert sys_cols + for col in sys_cols: + field_type = table.schema.field(col).type + assert not pa.types.is_list(field_type) + assert not pa.types.is_large_list(field_type) + # name-extended: original "::" plus "::" + assert col.count(constants.BLOCK_SEPARATOR) >= 2 + + +class TestGroupByOrdering: + def test_members_sorted_by_non_key_tags(self, session_source): + """probe=[1,0] on input must emit as [0,1].""" + out = GroupBy(by=["subject", "date"]).process(session_source) + assert out.as_table().column("probe").to_pylist()[0] == [0, 1] + + def test_row_order_does_not_affect_output(self, session_table): + """Same rows, shuffled, must produce a byte-identical table.""" + shuffled = session_table.take([3, 1, 2, 0]) + + def run(tbl): + src = ArrowTableSource( + tbl, tag_columns=["subject", "date", "probe"], infer_nullable=True + ) + return GroupBy(by=["subject", "date"]).process(src).as_table() + + assert run(session_table).equals(run(shuffled)) + + def test_falls_back_to_record_id_when_key_covers_all_tags(self): + """by covering every tag leaves no non-key tag to sort on.""" + table = pa.table({"subject": ["G", "G"], "path": ["b", "a"]}) + src = ArrowTableSource(table, tag_columns=["subject"], infer_nullable=True) + out = GroupBy(by=["subject"]).process(src) + assert len(out.as_table()) == 1 + assert sorted(out.as_table().column("path").to_pylist()[0]) == ["a", "b"] + + +class TestGroupByValidation: + def test_empty_by_raises(self): + with pytest.raises(ValueError, match="at least one"): + GroupBy(by=[]) + + def test_unknown_column_raises(self, session_source): + op = GroupBy(by=["subject", "nonexistent"]) + with pytest.raises(InputValidationError, match="nonexistent"): + op.process(session_source) + + def test_data_column_as_key_raises(self, session_source): + """Grouping on a data column is not allowed -- keys must be tags.""" + op = GroupBy(by=["path"]) + with pytest.raises(InputValidationError, match="path"): + op.process(session_source) + + +class TestGroupByEmptyInput: + def test_empty_input_yields_zero_groups(self): + table = pa.table({ + "subject": pa.array([], pa.large_string()), + "path": pa.array([], pa.large_string()), + }) + from orcapod.core.streams import ArrowTableStream + + stream = ArrowTableStream(table, tag_columns=["subject"]) + out = GroupBy(by=["subject"]).process(stream) + assert len(out.as_table()) == 0 + + +class TestGroupByIdentity: + def test_identity_structure_includes_by(self): + assert GroupBy(by=["a"]).identity_structure() != GroupBy( + by=["b"] + ).identity_structure() + + def test_to_config_round_trip(self): + op = GroupBy(by=["subject", "date"]) + config = op.to_config() + # A list, not a tuple, so the config stays JSON-serializable. + assert config["config"]["by"] == ["subject", "date"] + rebuilt = GroupBy.from_config(config) + assert rebuilt.identity_structure() == op.identity_structure() + + def test_to_config_is_json_serializable(self): + import json + + json.dumps(GroupBy(by=["subject", "date"]).to_config()["config"]) +``` + +- [ ] **Step 2: Run test to verify it fails** + +```bash +uv run pytest tests/test_core/operators/test_group_by.py -v +``` + +Expected: all FAIL at collection with `ImportError: cannot import name 'GroupBy'`. + +- [ ] **Step 3: Create the operator** + +**Do not write an `async_execute` override.** `UnaryOperator.async_execute` +(`core/operators/base.py:71`) already collects the full input before calling +`static_process`, which is exactly the barrier `GroupBy` needs — no group can be emitted +before the input channel closes, because any row not yet seen could belong to a group +already started. Task 7 adds a test asserting the override is absent. + +Create `src/orcapod/core/operators/group_by.py`: + +```python +"""GroupBy operator — many->one reduction keyed on tag values.""" + +from __future__ import annotations + +import logging +from collections.abc import Collection +from typing import TYPE_CHECKING, Any + +from orcapod.core.operators.base import UnaryOperator +from orcapod.core.streams import ArrowTableStream +from orcapod.errors import InputValidationError +from orcapod.protocols.core_protocols import StreamProtocol +from orcapod.system_constants import constants +from orcapod.types import ColumnConfig, Schema +from orcapod.utils import arrow_utils +from orcapod.utils.lazy_module import LazyModule + +if TYPE_CHECKING: + import pyarrow as pa +else: + pa = LazyModule("pyarrow") + +logger = logging.getLogger(__name__) + + +class GroupBy(UnaryOperator): + """Reduce rows sharing a tag tuple into one packet with list-valued members. + + This is the only many->one operator. Every other operator preserves one + row per tag; ``GroupBy`` collapses N rows into one, which is what lets a + downstream pod receive a whole group at once (for example, all of a + recording session's per-probe result parquets). + + Given tags ``(subject, date, probe)`` and data ``(path)``, grouping by + ``["subject", "date"]`` emits one row per distinct ``(subject, date)``: + + * ``subject`` and ``date`` stay scalar and remain the output's tag columns + * ``probe`` becomes a list-valued **data** column, so a consumer can tell + which member each list element came from + * ``path`` becomes list-valued + * ``_source_*`` columns become list-valued, one element per member + * system-tag columns fold to a scalar digest and gain a + ``::{pipeline_hash}`` name suffix + + Members are sorted by their non-group-key tag values, so the emitted lists + are stable across runs. This matters because orcapod hashes those lists to + build the cache key — an unsorted list would make an identical member set + hash differently and trigger a spurious recompute. + + Contrast with ``Batch``, which partitions by row count for throughput and + keeps its tag columns as list-valued tags. + + Args: + by: Tag column names to group on. Must be non-empty and must all be + tag columns of the input stream. + """ + + def __init__(self, by: Collection[str], **kwargs: Any) -> None: + by_tuple = tuple(by) + if not by_tuple: + raise ValueError("GroupBy requires at least one column in `by`.") + self.by = by_tuple + super().__init__(**kwargs) + + def identity_structure(self) -> Any: + return (self.__class__.__name__, self.by) + + def to_config(self) -> dict[str, Any]: + """Serialize this GroupBy operator to a config dict. + + ``by`` is emitted as a list rather than a tuple so the config stays + JSON-serializable; ``__init__`` normalizes it back to a tuple. + + Returns: + A dict with ``class_name``, ``module_path``, and ``config`` keys, + where ``config`` contains ``by``. + """ + config = super().to_config() + config["config"] = {"by": list(self.by)} + return config + + # ------------------------------------------------------------------ + # Validation + # ------------------------------------------------------------------ + + def validate_unary_input(self, stream: StreamProtocol) -> None: + """Verify every grouping column is a tag column of the input. + + Args: + stream: The upstream stream to validate. + + Raises: + InputValidationError: If any name in ``by`` is not a tag column. + """ + tag_columns, data_columns = stream.keys() + missing = [c for c in self.by if c not in tag_columns] + if missing: + raise InputValidationError( + f"GroupBy: {missing} are not tag columns of the input stream. " + f"Available tag columns: {list(tag_columns)}. " + f"(Data columns cannot be grouping keys: {list(data_columns)})" + ) + + # ------------------------------------------------------------------ + # Processing + # ------------------------------------------------------------------ + + def unary_static_process(self, stream: StreamProtocol) -> StreamProtocol: + """Partition rows by group key and emit one row per group. + + Args: + stream: The upstream stream. + + Returns: + A stream with one row per distinct group-key tuple. + """ + table = stream.as_table(columns={"source": True, "system_tags": True}) + tag_columns, _ = stream.keys() + + system_tag_columns = tuple( + c + for c in table.column_names + if c.startswith(constants.SYSTEM_TAG_PREFIX) + ) + member_columns = tuple( + c + for c in table.column_names + if c not in self.by and c not in system_tag_columns + ) + # Non-key user tags give a total order within a group: tags are unique + # within a stream. When `by` covers every tag, fall back to record_id. + sort_columns = tuple(c for c in tag_columns if c not in self.by) + record_id_column = next( + ( + c + for c in system_tag_columns + if c.startswith(constants.SYSTEM_TAG_RECORD_ID_PREFIX) + ), + None, + ) + + groups: dict[tuple[Any, ...], list[dict[str, Any]]] = {} + for row in table.to_pylist(): + groups.setdefault(tuple(row[c] for c in self.by), []).append(row) + + grouped_rows: list[dict[str, Any]] = [] + for key, members in groups.items(): + if sort_columns: + members.sort( + key=lambda r: tuple( + (r[c] is None, r[c]) for c in sort_columns + ) + ) + elif record_id_column is not None: + members.sort(key=lambda r: r[record_id_column] or b"") + + grouped_rows.append({ + **dict(zip(self.by, key)), + **{c: [m[c] for m in members] for c in member_columns}, + **{ + c: arrow_utils.fold_system_tag_values( + c, [m[c] for m in members] + ) + for c in system_tag_columns + }, + }) + + input_fields = {f.name: f for f in table.schema} + grouped_schema = pa.schema([ + pa.field(c, pa.list_(input_fields[c].type), nullable=False) + if c in member_columns + else input_fields[c] + for c in table.column_names + ]) + grouped_table = pa.Table.from_pylist(grouped_rows, schema=grouped_schema) + + n_char = self.orcapod_config.hashing.system_tag_n_char + grouped_table = arrow_utils.append_to_system_tags( + grouped_table, stream.pipeline_hash().to_hex(n_char) + ) + + return ArrowTableStream( + grouped_table, + tag_columns=self.by, + data_context=stream.data_context, + ) + + # ------------------------------------------------------------------ + # Schema prediction + # ------------------------------------------------------------------ + + def unary_output_schema( + self, + stream: StreamProtocol, + *, + columns: ColumnConfig | dict[str, Any] | None = None, + all_info: bool = False, + ) -> tuple[Schema, Schema]: + """Predict the grouped output schemas without grouping. + + Args: + stream: The upstream stream. + columns: Column inclusion config. + all_info: Include all info columns. + + Returns: + A ``(tag_schema, data_schema)`` tuple. Group keys stay scalar in + the tag schema; promoted non-key tags and list-wrapped data columns + land in the data schema. + """ + column_config = ColumnConfig.handle_config(columns, all_info=all_info) + tag_types, data_types = stream.output_schema( + columns=columns, all_info=all_info + ) + n_char = self.orcapod_config.hashing.system_tag_n_char + suffix = stream.pipeline_hash().to_hex(n_char) + + out_tag_types: dict[str, Any] = {} + out_data_types: dict[str, Any] = {} + + for name, col_type in tag_types.items(): + if name.startswith(constants.SYSTEM_TAG_PREFIX): + out_tag_types[ + f"{name}{constants.BLOCK_SEPARATOR}{suffix}" + ] = col_type + elif name in self.by: + out_tag_types[name] = col_type + else: + # Promoted to a list-valued data column. + out_data_types[name] = list[col_type] + if column_config.source: + # Promoted columns carry no provenance token; the stream + # fills in a scalar null. + out_data_types[f"{constants.SOURCE_PREFIX}{name}"] = str + + for name, col_type in data_types.items(): + out_data_types[name] = list[col_type] + + return Schema(out_tag_types), Schema(out_data_types) +``` + +- [ ] **Step 4: Export it** + +In `src/orcapod/core/operators/__init__.py`, add the import alphabetically after `from .filters import PolarsFilter`: + +```python +from .group_by import GroupBy +``` + +and add `"GroupBy",` to `__all__` after `"Batch",`. + +- [ ] **Step 5: Run the tests to verify they pass** + +```bash +uv run pytest tests/test_core/operators/test_group_by.py -v +``` + +Expected: all PASS. If `test_row_order_does_not_affect_output` fails, the sort key is wrong — check that `sort_columns` excludes the `by` columns and that `members.sort` runs before the dict comprehension. + +- [ ] **Step 6: Commit** + +```bash +git add src/orcapod/core/operators/group_by.py src/orcapod/core/operators/__init__.py tests/test_core/operators/test_group_by.py +git commit -m "feat(operators): add GroupBy for many-to-one tag-keyed reduction (NPIPE-204) + +Every other operator preserves one row per tag; GroupBy collapses N rows +sharing a tag tuple into one packet with list-valued members, which is what +lets a downstream pod receive a whole recording session at once. + +Group keys stay scalar tags; non-key tags are promoted to list-valued data +columns rather than dropped, so consumers can tell which member each element +came from. Members sort by non-key tag values so the hashed lists are stable +across runs. + +NPIPE-204" +``` + +--- + +## Task 6: Schema-mirroring test + +**Files:** +- Test: `tests/test_core/operators/test_group_by.py` + +`unary_output_schema` predicting something the table does not match is the classic operator bug — it surfaces far downstream, as a DB write failure rather than a schema error. + +- [ ] **Step 1: Write the failing test** + +Append to `tests/test_core/operators/test_group_by.py`: + +```python +class TestGroupBySchemaMirror: + """unary_output_schema must match what the produced stream reports.""" + + @pytest.mark.parametrize( + "config", + [ + {}, + {"source": True}, + {"system_tags": True}, + {"source": True, "system_tags": True}, + ], + ids=["bare", "source", "system_tags", "source+system_tags"], + ) + def test_predicted_schema_matches_produced_stream(self, session_source, config): + op = GroupBy(by=["subject", "date"]) + + predicted_tags, predicted_data = op.unary_output_schema( + session_source, columns=config + ) + produced = op.unary_static_process(session_source) + actual_tags, actual_data = produced.output_schema(columns=config) + + assert dict(predicted_tags) == dict(actual_tags) + assert dict(predicted_data) == dict(actual_data) +``` + +**Two comparisons that look right and are not.** Both were tried during review and both silently pass: + +- **Against `op.process(...)`** — `DynamicPodStream.output_schema` (`static_output_pod.py:321-332`) delegates straight back to the pod, so this compares the prediction against itself and can never fail. +- **Against `as_table(columns=config).column_names`** — `output_schema` and `as_table` legitimately disagree. `ArrowTableStream.output_schema` (`arrow_table_stream.py:183-204`) returns `self._data_schema` unconditionally, so `columns.source` is a documented no-op there (`source_node.py:178-180` states this). `_source_*`, `_content_hash`, and `_context_key` appear in the table but never in the schema. The same gap exists on a plain un-batched `ArrowTableSource`, so it is a property of the schema layer, not of any operator. + +Compare the prediction against the materialized `ArrowTableStream` from `unary_static_process`. That is the contract that actually binds. + +- [ ] **Step 2: Run the test** + +```bash +uv run pytest tests/test_core/operators/test_group_by.py::TestGroupBySchemaMirror -v +``` + +Expected: PASS if Task 5 was implemented correctly. If a case fails, the assertion message names the exact mismatched columns — fix `unary_output_schema` to match the table, not the other way round. + +- [ ] **Step 3: Commit** + +```bash +git add tests/test_core/operators/test_group_by.py +git commit -m "test(group_by): assert output_schema mirrors as_table (NPIPE-204) + +NPIPE-204" +``` + +--- + +## Task 7: Registration — serialization and fluent API + +**Files:** +- Modify: `src/orcapod/pipeline/serialization.py:160-192` +- Modify: `src/orcapod/core/streams/base.py:135-149` +- Test: `tests/test_core/operators/test_group_by.py` + +Without the registry entry, a pipeline containing a `GroupBy` cannot be deserialized. + +- [ ] **Step 1: Write the failing test** + +Append to `tests/test_core/operators/test_group_by.py`: + +```python +class TestGroupByRegistration: + def test_in_operator_registry(self): + from orcapod.pipeline.serialization import _build_operator_registry + + assert _build_operator_registry()["GroupBy"] is GroupBy + + def test_stream_fluent_method(self, session_source): + out = session_source.group_by(["subject", "date"]) + assert len(out.as_table()) == 2 + + +class TestGroupByAsyncIsBarrier: + """GroupBy must NOT override async_execute. + + ``UnaryOperator.async_execute`` (``core/operators/base.py:71``) already + collects the full input before calling ``static_process``, which is exactly + the barrier GroupBy needs: no group can be emitted before the input channel + closes, because any row not yet seen could belong to a group already + started. Adding an override would duplicate that logic and risk drifting + from it. + """ + + def test_does_not_override_async_execute(self): + from orcapod.core.operators.base import UnaryOperator + + assert "async_execute" not in GroupBy.__dict__ + assert GroupBy.async_execute is UnaryOperator.async_execute +``` + +- [ ] **Step 2: Run test to verify it fails** + +```bash +uv run pytest tests/test_core/operators/test_group_by.py::TestGroupByRegistration -v +``` + +Expected: `test_in_operator_registry` FAILS with `KeyError: 'GroupBy'`; `test_stream_fluent_method` FAILS with `AttributeError`. + +- [ ] **Step 3: Add the registry entry** + +In `src/orcapod/pipeline/serialization.py`, inside `_build_operator_registry`, add `GroupBy` to the import list (alphabetically, after `DropTagColumns`): + +```python + from orcapod.core.operators import ( + Batch, + DropDataColumns, + DropTagColumns, + GroupBy, + Join, + MapData, + MapTags, + MergeJoin, + PolarsFilter, + SelectDataColumns, + SelectTagColumns, + SemiJoin, + ) +``` + +and add the entry to the returned dict, after `"Batch": Batch,`: + +```python + "GroupBy": GroupBy, +``` + +- [ ] **Step 4: Add the fluent method** + +In `src/orcapod/core/streams/base.py`, immediately after the existing `batch` method (ends line 149), add: + +```python + def group_by( + self, + by: Collection[str], + label: str | None = None, + ) -> StreamBase: + """Reduce rows sharing a tag tuple into one packet per group. + + Group-key columns stay scalar tags; every other column becomes + list-valued. See ``orcapod.core.operators.GroupBy``. + + Args: + by: Tag column names to group on. + label: Optional node label for the pipeline graph. + + Returns: + A stream with one row per distinct group-key tuple. + """ + from orcapod.core.operators import GroupBy + + return GroupBy(by=by)(self, label=label) +``` + +`Collection` is already imported at the top of `base.py`. Verify with `grep -n "from collections.abc" src/orcapod/core/streams/base.py` and add it to that import if absent. + +- [ ] **Step 5: Run the tests to verify they pass** + +```bash +uv run pytest tests/test_core/operators/test_group_by.py -v +uv run pytest tests/test_pipeline/test_serialization.py -q +``` + +Expected: all PASS. + +- [ ] **Step 6: Commit** + +```bash +git add src/orcapod/pipeline/serialization.py src/orcapod/core/streams/base.py tests/test_core/operators/test_group_by.py +git commit -m "feat(operators): register GroupBy in serialization and stream API (NPIPE-204) + +Without the registry entry a pipeline containing a GroupBy cannot be +deserialized. + +NPIPE-204" +``` + +--- + +## Task 8: Job-level tests + +**Files:** +- Create: `tests/test_pipeline/test_aggregation_job.py` + +This is the gap that let the provenance crash ship. Every existing `Batch` test stops at `op.process(stream)` + `as_table()`; the crash only happens inside `job.run()`, in `StaticOutputOperatorPod._materialize_to_stream`. + +- [ ] **Step 1: Write the failing test** + +Create `tests/test_pipeline/test_aggregation_job.py`: + +```python +"""Job-level tests for aggregating operators (NPIPE-204). + +Operator-level tests (`op.process(stream)` then `as_table()`) never reach +`StaticOutputOperatorPod._materialize_to_stream`, which is where list-valued +provenance used to crash. These run the full `job.run()` path against a real +Delta Lake store. +""" + +from __future__ import annotations + +from pathlib import Path + +import pyarrow as pa +import pytest + +from orcapod.core.data_function import PythonDataFunction +from orcapod.core.function_pod import FunctionPod +from orcapod.core.operators import Batch, GroupBy, MergeJoin +from orcapod.core.sources import ArrowTableSource +from orcapod.databases import DeltaTableDatabase +from orcapod.pipeline import PipelineJob + + +@pytest.fixture +def store(tmp_path: Path) -> DeltaTableDatabase: + return DeltaTableDatabase(str(tmp_path / "store")) + + +@pytest.fixture +def session_source_factory(): + """Build a 2-group source; `paths` lets a test mutate one group's data.""" + + def _make(paths: list[str] | None = None) -> ArrowTableSource: + table = pa.table({ + "subject": ["G", "G", "G", "G"], + "date": ["d1", "d1", "d2", "d2"], + "probe": [0, 1, 0, 1], + "path": paths or ["a", "b", "c", "d"], + }) + return ArrowTableSource( + table, + tag_columns=["subject", "date", "probe"], + infer_nullable=True, + ) + + return _make + + +class _CountingFunction: + """Records every invocation so memoization can be asserted on. + + `PythonDataFunction` binds data columns to parameter names via + `inspect.signature`, which resolves a callable instance to `__call__`. + The parameter must therefore be named after the data column — hence the + separate `_CountingV` below for MergeJoin's merged `v` column. + """ + + def __init__(self) -> None: + self.calls: list[list[str]] = [] + + def __call__(self, path: list[str]) -> int: + self.calls.append(list(path)) + return len(path) + + +class _CountingV: + """Same as `_CountingFunction`, bound to a column named `v`.""" + + def __init__(self) -> None: + self.calls: list[list[str]] = [] + + def __call__(self, v: list[str]) -> int: + self.calls.append(list(v)) + return len(v) + + +def _run(store, source, operator, fn, name): + pod = FunctionPod(PythonDataFunction(fn, output_keys="n", function_name="count")) + job = PipelineJob(name=name, store=store) + with job: + pod(operator(source, label="agg"), label="counter") + return job.run() + + +class TestGroupByInJob: + def test_group_by_completes(self, store, session_source_factory): + fn = _CountingFunction() + _run(store, session_source_factory(), GroupBy(by=["subject", "date"]), fn, "gb") + assert len(fn.calls) == 2 + assert sorted(fn.calls) == [["a", "b"], ["c", "d"]] + + def test_batch_completes(self, store, session_source_factory): + """The provenance fix is independent of grouping.""" + fn = _CountingFunction() + _run(store, session_source_factory(), Batch(batch_size=2), fn, "b") + assert len(fn.calls) == 2 + + +class TestGroupByMemoization: + def test_identical_runs_hit_cache(self, store, session_source_factory): + fn = _CountingFunction() + _run(store, session_source_factory(), GroupBy(by=["subject", "date"]), fn, "m") + assert len(fn.calls) == 2 + + fn2 = _CountingFunction() + _run(store, session_source_factory(), GroupBy(by=["subject", "date"]), fn2, "m") + assert fn2.calls == [], "second identical run must not recompute" + + def test_changed_member_invalidates_only_its_group( + self, store, session_source_factory + ): + """Two groups; change one member of the first only. + + With a single group this assertion would be vacuous -- it must show + that the untouched group stays cached. + """ + fn = _CountingFunction() + _run(store, session_source_factory(), GroupBy(by=["subject", "date"]), fn, "i") + assert len(fn.calls) == 2 + + fn2 = _CountingFunction() + changed = session_source_factory(["a", "B_CHANGED", "c", "d"]) + _run(store, changed, GroupBy(by=["subject", "date"]), fn2, "i") + + assert fn2.calls == [["a", "B_CHANGED"]], ( + "only the changed group should recompute; " + f"got {fn2.calls}" + ) + + +class TestMergeJoinRegression: + def test_merge_join_completes_in_job(self, store): + """MergeJoin carries source columns as parallel lists (merge_join.py:262). + + It crashed with the same ArrowTypeError before the Data fix. + """ + left = ArrowTableSource( + pa.table({"id": ["a", "b"], "v": ["l1", "l2"]}), + tag_columns=["id"], + infer_nullable=True, + ) + right = ArrowTableSource( + pa.table({"id": ["a", "b"], "v": ["r1", "r2"]}), + tag_columns=["id"], + infer_nullable=True, + ) + + fn = _CountingV() + pod = FunctionPod( + PythonDataFunction(fn, output_keys="n", function_name="count_v") + ) + job = PipelineJob(name="mj", store=store) + with job: + pod(MergeJoin()(left, right, label="mj"), label="counter") + job.run() + + assert len(fn.calls) == 2 + # MergeJoin merges colliding `v` columns into a sorted 2-element list. + assert sorted(fn.calls) == [["l1", "r1"], ["l2", "r2"]] +``` + +- [ ] **Step 2: Run test to verify it fails, and how** + +```bash +uv run pytest tests/test_pipeline/test_aggregation_job.py -v +``` + +Expected before Tasks 1–7: `ArrowTypeError: Expected bytes, got a 'list' object`. After Tasks 1–7: these should pass. If a memoization test fails, that is a real finding — report the actual `fn2.calls` value rather than relaxing the assertion. + +- [ ] **Step 3: Fix any failures found** + +If `test_changed_member_invalidates_only_its_group` shows *both* groups recomputing, the likely cause is that `fold_system_tag_values` is being fed members in a different order between runs — verify the sort in `unary_static_process` runs before the fold. If *neither* recomputes, the input data hash is not reaching the record preimage; inspect `_build_record_id_preimage` output for both runs. + +- [ ] **Step 4: Run the full suite** + +```bash +uv run pytest tests/ -q +``` + +Expected: no failures. + +- [ ] **Step 5: Commit** + +```bash +git add tests/test_pipeline/test_aggregation_job.py +git commit -m "test(pipeline): job-level coverage for GroupBy, Batch, MergeJoin (NPIPE-204) + +Every existing aggregating-operator test stopped at op.process() + as_table(), +which never reaches _materialize_to_stream -- the reason the list-valued +provenance crash shipped. These run the full job.run() path against a Delta +Lake store and assert memoization holds across identical runs and invalidates +only the group whose member changed. + +NPIPE-204" +``` + +--- + +## Task 9: Documentation + +**Files:** +- Modify: `CLAUDE.md` +- Modify: `.zed/rules` +- Modify: `DESIGN_ISSUES.md` + +Per `CLAUDE.md`, agent instructions must be updated in **both** `CLAUDE.md` and `.zed/rules`. + +- [ ] **Step 1: Update the project layout tree** + +In `CLAUDE.md`, in the `src/orcapod/core/operators/` block, add after the `batch.py` line: + +``` +│ ├── group_by.py # GroupBy (many→one reduction keyed on tag values) +``` + +- [ ] **Step 2: Update the operator / function pod boundary table** + +`GroupBy` synthesizes no new values but does change row count. Add a note under that table in `CLAUDE.md`: + +```markdown +`GroupBy` is the only operator that changes row count in a many→one direction: it +reduces N rows sharing a tag tuple to one row with list-valued members. It still +synthesizes no new data values — every emitted element came from an input row. +``` + +- [ ] **Step 3: Update the system tag evolution rules** + +In `CLAUDE.md`, the "System tag evolution rules" section currently has three rules, and rule 3 describes `Batch` as purely type-evolving (`str` → `list[str]`). That is now wrong. Replace rule 3 with: + +```markdown +3. **Reducing** — many→one ops (`Batch`, `GroupBy`). User tag and data columns become + `list[T]`; source-info columns become `list[str]`, one element per member. System tag + columns must stay **scalar** (record identity hashes them directly), so they fold to a + deterministic digest via `arrow_utils.fold_system_tag_values` and their column name gains + `::{pipeline_hash}`. The fold is SHA-based and stable across processes — never use + `hash()` there, since the digest becomes a cache key. +``` + +- [ ] **Step 4: Add the `Important implementation details` entries** + +In `CLAUDE.md`, append to that list: + +```markdown +- `GroupBy` requires every column in `by` to be a tag column; raises `InputValidationError` + otherwise. Members are sorted by non-group-key tag values (falling back to system + `record_id`) so the hashed lists are stable across runs. +- `Data` source-info values may be `str`, `None`, or `list[...]`. Types are derived from the + value; `None` maps to `large_string`. +``` + +- [ ] **Step 5: Mirror every change into `.zed/rules`** + +`.zed/rules` carries the same instructions for Zed AI and must stay in sync — this is +required by `CLAUDE.md` itself. Apply Steps 1–4 verbatim to `.zed/rules`, then confirm the +two files agree on this content: + +```bash +diff CLAUDE.md .zed/rules +``` + +Expected: no differences in any section touched by Steps 1–4. If the files already diverge +elsewhere for unrelated reasons, leave those differences alone — only verify that every +line you added appears in both. + +- [ ] **Step 6: Resolve the `tag_data.py` half of U1** + +In `DESIGN_ISSUES.md`, update the U1 status line: + +```markdown +**Status:** resolved (`tag_data.py` half), open (`arrow_utils.py` half) +``` + +and append to the `**Fix (NPIPE-204):**` paragraph: + +```markdown +**Fix:** landed in NPIPE-204. `_source_info_arrow_type` / `_source_info_python_type` in +`core/datagrams/tag_data.py` derive the Arrow and Python types from the stored value. +The dead `polars_data_utils.add_source_info` was deleted. `add_source_info_to_table()` +in `arrow_utils.py` is untouched and remains open. +``` + +- [ ] **Step 7: Verify docs are consistent** + +```bash +uv run pytest tests/ -q +grep -n "group_by.py" CLAUDE.md .zed/rules +grep -n "GroupBy" CLAUDE.md .zed/rules +``` + +Expected: suite passes; both files mention `group_by.py` and `GroupBy`. + +- [ ] **Step 8: Commit** + +```bash +git add CLAUDE.md .zed/rules DESIGN_ISSUES.md +git commit -m "docs: document GroupBy and the reducing system-tag rule (NPIPE-204) + +Rule 3 described Batch as purely type-evolving, which was only ever true of +user tag and data columns -- system tags fold to a scalar and extend their +name. Renames the rule to 'reducing' and covers both operators. + +Resolves the tag_data.py half of DESIGN_ISSUES U1. +NPIPE-204" +``` + +--- + +## Task 10: Final verification and PR + +**Files:** none + +- [ ] **Step 1: Full suite** + +```bash +uv run pytest tests/ -q +``` + +Expected: all pass. Record the exact pass/fail counts — do not claim success without reading the output. + +- [ ] **Step 2: Confirm both original crashes are gone** + +```bash +uv run pytest tests/test_pipeline/test_aggregation_job.py -v +``` + +Expected: `TestGroupByInJob`, `TestGroupByMemoization`, and `TestMergeJoinRegression` all pass. + +- [ ] **Step 3: Review the whole diff** + +```bash +git diff main...HEAD --stat +git diff main...HEAD +``` + +Check: no debug prints, no `sys.modules` manipulation, no backward-compat shims (all forbidden by `CLAUDE.md`), Google-style docstrings with no ReST roles. + +- [ ] **Step 4: Push and open the PR** + +```bash +git push -u origin arnoldb/npipe-204-orcapod-add-tag-based-grouping-to-batch-manyone-reduction +``` + +PR body must include `Fixes NPIPE-204` so Linear's GitHub integration links it, and must call out the API deviation: + +```markdown +## Summary + +Adds a `GroupBy` operator for many→one reduction keyed on tag values, and fixes the +`Data` datagram so list-valued provenance no longer crashes `job.run()`. + +**API note:** NPIPE-204 specifies `Batch(group_by=[...])`. This PR adds a separate +`GroupBy(by=[...])` instead, so each operator keeps one output contract and `Batch`'s +streaming `async_execute` path stays untouched. See +`superpowers/specs/2026-08-07-npipe-204-batch-group-by-design.md`. + +The provenance fix also repairs `MergeJoin`, which produced list-valued `_source_*` +columns (`merge_join.py:262`) and was silently broken inside `job.run()`. + +Fixes NPIPE-204 + +🤖 Generated with [Claude Code](https://claude.com/claude-code) +``` + +- [ ] **Step 5: Update the Linear issue** + +NPIPE-204's "What to build" section and the `common_clock_op` wiring example both still say `Batch(group_by=...)`. Update them to `GroupBy(by=...)` so the downstream orcapod-sync-and-qc change is written against the real API. + +**Confirm with the user before editing the Linear issue** — it is an outward-facing change to a shared record. + +--- + +## Downstream (not part of this plan) + +The rev bump must land in **both** orcapod-sync-and-qc and orcapod-spikesorting together — they are deliberately kept on the same orcapod rev because they target the same Ray cluster. `common_clock_op` wiring becomes: + +```python +grouped = op.operators.GroupBy(by=["subject", "date"]).process(sync_out) +common_clock_op.pod(grouped) +``` diff --git a/superpowers/specs/2026-08-07-npipe-204-batch-group-by-design.md b/superpowers/specs/2026-08-07-npipe-204-batch-group-by-design.md new file mode 100644 index 000000000..07cd91c99 --- /dev/null +++ b/superpowers/specs/2026-08-07-npipe-204-batch-group-by-design.md @@ -0,0 +1,363 @@ +# NPIPE-204 — Tag-based grouping via a `GroupBy` operator (many→one reduction) + +**Linear:** [NPIPE-204](https://linear.app/metamorphic/issue/NPIPE-204/orcapod-add-tag-based-grouping-to-batch-manyone-reduction-fix-batch) +**Branch:** `arnoldb/npipe-204-orcapod-add-tag-based-grouping-to-batch-manyone-reduction` +**Base rev:** `966d759a` + +> **API deviation from the issue.** NPIPE-204 specifies `Batch(group_by=[...])`. This spec +> introduces a separate `GroupBy(by=[...])` operator instead — see *Why a separate operator*. +> The Linear issue and the downstream `common_clock_op` wiring both need updating to match. + +## Overview + +Every existing orcapod pipeline fans *out*. Nothing reduces. The consumer that motivates this +change — `common_clock_op` in orcapod-sync-and-qc — produces one `AlignmentResult` per recording +session, computed from *all* of that session's spikeglx-sync result parquets at once. Expressing +that requires a many→one operator keyed on tag values, which orcapod does not have. + +Two independent problems, verified against `966d759a`: + +1. **No tag-based reduction exists.** `Batch` groups by row count only; there is no way to say + "one packet per `(subject, date)`". +2. **List-valued source info breaks `job.run()`.** The `Data` datagram cannot represent a + non-scalar `_source_*` value. + +Reproduction of (2): + +``` +pyarrow.lib.ArrowTypeError: Expected bytes, got a 'list' object + core/nodes/operator_node.py:946 execute + core/operators/static_output_pod.py:214 _materialize_to_stream + core/datagrams/tag_data.py:433 as_table + core/datagrams/tag_data.py:342 _ensure_source_info_table +``` + +### Problem 2 is not Batch-specific + +`MergeJoin` merges colliding data columns into `list[T]` and carries their `_source_*` columns +along as parallel lists (`merge_join.py:262`). It therefore fails with the identical +`ArrowTypeError` inside `job.run()`, reproduced independently of `Batch`. The root cause is in the +`Data` datagram, not in either operator. + +This is the already-logged `DESIGN_ISSUES.md` **U1 — Source-info column type hard-coded to +`large_string`** (severity: critical), whose recorded location was the sibling call site +`arrow_utils.add_source_info_to_table()`. Part 1 fixes the `tag_data.py` half; the `arrow_utils.py` +half stays open. + +### Why a separate operator + +`Batch` exists for throughput and pipelining — its `async_execute` docstring is explicit that +batching lets "downstream consumers start processing before all input is consumed" — and its git +history contains only refactors. Semantic reduction was never its intent. + +Folding grouping into `Batch` would give one class two meaningfully different output contracts, +selected by which kwarg the caller passed, guarded by a mutually-exclusive-args check. Splitting +gives each class one contract: + +| | `Batch` | `GroupBy` | +|---|---|---| +| Partitions by | row count | tag values | +| Purpose | throughput / pipelining | many→one reduction | +| Output tags | list-valued, unchanged | scalar group keys | +| Async | streams when `batch_size > 0` | always a barrier | + +`Batch` is then semantically untouched by this change — it receives only the provenance fix. All +new reduction semantics live in `GroupBy`, and its barrier behavior is structural rather than a +conditional inside a streaming operator. + +## Goals & Success Criteria + +* `GroupBy(by=["subject", "date"])` emits one packet per distinct tag tuple, with the group keys as + scalar tag columns and the members as list-valued data columns. +* `GroupBy`, `Batch`, and `MergeJoin` all survive `job.run()` end to end. +* Memoization over a grouped stream holds across two identical runs and invalidates when a single + member's data changes. +* Group member order and the folded provenance digests are deterministic **across processes**, so + an unchanged member set never produces a different list hash or a cache miss on a new driver run. +* No pipeline-DB schema version bump. + +## Scope & Boundaries + +In scope: +* `core/datagrams/tag_data.py` — type-aware source info on `Data`. +* `core/operators/group_by.py` — new operator. +* `core/operators/batch.py` — provenance handling only. +* `utils/arrow_utils.py` — shared system-tag fold helper. +* `utils/polars_data_utils.py` — delete dead `add_source_info`. +* Operator-level and job-level tests, including a `MergeJoin` job-level regression test. + +Out of scope: +* `arrow_utils.add_source_info_to_table()` (the other half of U1). +* Operators discarding a non-default `_context_key` (new DESIGN_ISSUES entry, not fixed here). +* Incremental emission for `GroupBy` — see *Async execution*. +* The downstream rev bump in orcapod-sync-and-qc and orcapod-spikesorting. + +--- + +## Part 1 — Type-aware source info in `Data` + +`Data` stores per-data-column provenance tokens in `self._source_info`, surfaced as `_source_*` +columns. Two places assume those values are always scalar strings: + +* `_ensure_source_info_table()` (`tag_data.py:330`) builds the Arrow schema as + `pa.field(k, pa.large_string())` for every key. +* `Data.schema()` (`tag_data.py:384`) sets `schema[f"{SOURCE_PREFIX}{key}"] = str` for every key. + +Both become derived from the stored value: + +| Stored value | Arrow type | Python type | +|---|---|---| +| `str` | `large_string` | `str` | +| `None` | `large_string` | `str` | +| `list[str]` | `large_list(large_string)` | `list[str]` | +| nested list | `large_list()`, recursive | `list[...]` | +| `[]` | `large_list(large_string)` | `list[str]` | + +`None` keeps mapping to `large_string`, so unknown-provenance columns behave exactly as today. + +`self._source_info`'s annotation widens from `dict[str, str | None]` to a recursive +`SourceInfoValue = str | None | list["SourceInfoValue"]`, with `source_info()`, +`with_source_info()`, `rename()`, and `with_columns()` following. The dict and table construction +paths in `__init__` already pass values through untouched — the table path recovers lists correctly +via `to_pylist()`. + +**No schema version bump.** A node's `_source_*` column type is fixed by that node's own output +schema. A `FunctionNode` downstream of a `GroupBy` writes list-typed source columns from its first +record; every pre-existing node keeps `large_string`. Nothing re-reads an old table under a new +type. + +This part alone fixes `MergeJoin` inside `job.run()`. + +### Third hard-coded site: delete it + +`polars_data_utils.add_source_info` (line 119) forces `dtype=pl.String()`, making it a third site +with the same assumption. It is dead code — nothing in `src/` calls it (only `drop_system_columns` +is imported from that module), and the tests that import `add_source_info` import it from +`arrow_utils`. It also carries a latent shadowing bug: `source_column` is rebound to a `pl.Series` +inside the per-column loop, so from the second column onward it formats `f"{}::{col}"`. + +Delete the function rather than fix it. Greenfield pre-v0.1.0 means no back-compat obligation, and +removing it eliminates the site outright instead of leaving a trap for whoever wires it up later. + +--- + +## Part 2a — `Batch` provenance fix + +`Batch`'s partitioning, its list-valued tag columns, and its streaming `async_execute` are all +unchanged. The only change: `_source_*` columns keep their list-valued form (now representable +thanks to Part 1), and the system-tag columns fold to scalars via the shared helper below instead of +being list-wrapped. + +That is the minimum needed to make plain `Batch(batch_size=N)` work inside `job.run()`. + +--- + +## Part 2b — The `GroupBy` operator + +New module `core/operators/group_by.py`. A `UnaryOperator`. + +### Constructor and validation + +```python +def __init__(self, by: Collection[str], **kwargs): +``` + +* Empty `by` → `ValueError`. +* `validate_unary_input` raises `InputValidationError` if any name in `by` is not a tag column of + the input stream. + +### Partitioning + +Rows are keyed on the tuple of their group-key tag values and accumulated into a plain dict, so +first-seen group order is preserved in the output. + +### Member ordering + +Within a group, members are sorted by the tuple of their **non-group-key tag values**, falling back +to the member's system `record_id` bytes when `by` covers every tag column. Tags are unique within a +stream, so this is a total order, and it does not depend on which data column happens to hold a +path. Sort keys wrap each value as `(v is None, v)` so nulls order consistently without comparing +against `None`. + +This matters because orcapod hashes the emitted list to build the cache key. Upstream emission order +is not stable across runs (Ray executor scheduling, DB fetch order), so an unsorted list would make +an identical member set hash differently and trigger a spurious recompute. + +### Column treatment + +| Column class | Result | +|---|---| +| group-key tags | **scalar**, remain tag columns | +| other user tags | list-valued **data** columns | +| data columns | list-valued | +| `_source_*` | list-valued, one element per member | +| `_tag_source_id` / `_tag_record_id` | **scalar digest**, name extended | + +Non-key tags become list-valued *data* rather than being dropped, so a consumer can tell which +member each list element came from. Those promoted columns have no provenance token, so +`Data.source_info()` reports `None` for them — its existing behavior for unknown keys. + +`_context_key` does not appear: `stream.as_table(columns={"source": True, "system_tags": True})` +returns only user tags, data, system tags, and `_source_*`. Every operator in the repo uses that +same column set, and `ArrowTableStream` re-defaults the context key on construction. `GroupBy` +therefore cannot silently pick one member's context. That layer-wide behavior is logged separately. + +Nullability: list-wrapped columns are `nullable=False` (a group always has at least one member, so +the list itself is never null). Scalar group-key columns inherit the input column's nullable flag. + +### System tag folding + +`_build_record_id_preimage` (`core/nodes/function_node.py:82`) computes a record's identity from the +system-tag columns plus a hash of the input data. System tags must therefore be scalar. A many→one +operator has to define how N members' system tags collapse into one record's provenance. + +**Rule:** fold each system-tag column independently over its own ordered member values, preserving +the column's declared Arrow type, then extend the column name with +`{BLOCK_SEPARATOR}{pipeline_hash}` via the existing `arrow_utils.append_to_system_tags` — mirroring +the name-extending rule already used by joins. + +Both primitives are reused from the codebase rather than invented, and both are **SHA-based and +stable across processes**. Neither `hash()` nor set iteration order may appear anywhere in the fold: + +| Column | Arrow type | Fold | +|---|---|---| +| `_tag_source_id::` | `large_string` | `hash_utils.combine_hashes(*member_source_ids, order=False)` — SHA-256 hex | +| `_tag_record_id::` | `binary(16)` | `uuid.uuid5(_GROUP_RECORD_ID_NAMESPACE, "::".join(rid.hex() for rid in member_record_ids)).bytes` | + +`order=False` preserves member order, which is already deterministic from the sort above, so the +digest reflects both the member set and its order — matching the data lists it accompanies. +`uuid5` is the same construction `stream_builder._make_record_id` (line 55) already uses to mint +record IDs, and `.bytes` is exactly the 16 bytes `pa.binary(16)` wants. Record IDs are globally +unique (uuid5 over `source_id::token`), so folding them alone captures full member identity. + +Why this rule: + +* **Invalidation is correct.** Any member whose record identity changes changes the digest, so the + downstream record_id changes and the cache misses. This is strictly stronger than relying on the + input-data hash alone, which would miss an upstream recompute that produced identical data. +* **Nothing else has to change.** `function_node.py` and the pdb schema are untouched. + `_tag_record_id` is consumed only for identity and sorting (`join.py:687`, + `arrow_utils.py:1101/1169/1235`), never as a key to look a record up, so synthesizing a value + cannot break a lookup. +* **Member identities remain recoverable** from the list-valued `_source_*` columns, which Part 1 + now preserves per member. + +The extended name `_tag_source_id::::` has no trailing `:position`, so +`_parse_system_tag_column` returns `None` for it and `sort_system_tag_values` skips it. That is +correct — `GroupBy` is unary, so there is no cross-input commutativity to normalize. + +Two alternatives were rejected. Having `GroupBy` mint a fresh source-like identity severs the Merkle +link to the member records. Keeping system tags list-valued and teaching `_build_record_id_preimage` +to hash lists is the most information-preserving option but changes the record-identity machinery +and the pdb column types, requiring a v1→v2 migration on top of an already cross-repo change. + +The fold lives in `arrow_utils` as a shared helper, next to `append_to_system_tags` and +`sort_system_tag_values`, since both `Batch` and `GroupBy` need it. + +### Output schema + +`unary_output_schema` must mirror the table exactly: scalar group keys in the tag schema, promoted +non-key tags and list-wrapped data columns in the data schema, list-typed `_source_*` entries when +`columns={"source": True}`, and renamed scalar system-tag entries when +`columns={"system_tags": True}`. The operator predicts this without performing the computation, +consistent with every other operator — including for empty input, which yields zero groups but the +same schema. + +### Registration + +Four sites, mirroring `Batch`: + +* `core/operators/group_by.py` — the module. +* `core/operators/__init__.py` — import and `__all__`. +* `pipeline/serialization.py:_build_operator_registry` — import and registry entry. +* `core/streams/base.py` — a `group_by(by, label=None)` fluent method alongside `batch()`. + +`to_config()` and `identity_structure()` both include `by`. `from_config` needs no change — it +already forwards `config["config"]` as kwargs. + +### Async execution + +`async_execute` is barrier-only: no group can be emitted before the input channel closes, because +any row not yet seen could belong to a group already started. + +Under `AsyncPipelineOrchestrator` this stalls one node, not the pipeline. Upstream nodes still run +concurrently and stream into the `GroupBy`; downstream resumes full concurrency once the barrier +releases, fanning the N groups out to N concurrent invocations. The unavoidable cost is that the +last straggler member in *any* group delays the first downstream invocation for *every* group. + +This is not a regression relative to what already exists: `Batch(batch_size=0)` takes the barrier +path today (`batch.py:118`), and `SyncPipelineOrchestrator` is node-at-a-time regardless. `Batch`'s +`batch_size=N>0` streaming path is untouched by this change. + +Emitting groups early would require a guarantee that input arrives clustered by group key. orcapod +streams carry no ordering guarantee, so that would need an unverifiable `assume_grouped=True` +opt-in. Deliberately not built. + +--- + +## Part 3 — Tests + +### Operator-level (`tests/test_core/operators/`) + +* `GroupBy` produces one row per distinct tag tuple; group keys scalar, members list-valued. +* Non-key tags are promoted to list-valued data columns, not dropped. +* Empty `by` raises `ValueError`; `by` naming a non-tag column raises `InputValidationError`. +* Members are sorted by non-key tags — the same rows fed in two different orders yield + byte-identical output tables. +* Fallback ordering by `record_id` when `by` covers every tag column. +* `unary_output_schema` matches `as_table().schema` for every `ColumnConfig` combination. +* System tag columns are scalar and renamed; `_source_*` columns are list-valued with one element + per member. +* Empty input → zero groups, schema still predicted correctly. +* Existing `TestBatchBehavior` cases continue to pass unchanged, plus new assertions that `Batch` + emits scalar system tags and list-valued `_source_*`. + +### Cross-process digest stability + +The failure this guards against: a fold that accidentally used `hash()` or set-iteration order +would pass every same-process test — including the reordering test, which only checks byte-identity +*within* one run — and then miss the cache on every new driver run. + +A test runs the fold in a `subprocess` (fresh interpreter, therefore fresh `PYTHONHASHSEED`) over a +fixed member list and asserts the digests equal hard-coded expected values. + +### Job-level (new) + +The gap that let problem 2 ship: `Batch`'s existing tests are all `op.process(stream)` followed by +`as_table()`. None run through `job.run()`. + +Backed by a `DeltaTableDatabase` store: + +* `GroupBy` completes end to end inside `job.run()` feeding a `@function_pod` that takes a + `list[str]` parameter. +* `Batch(batch_size=N)` likewise — the provenance fix is independent of grouping. +* **Memoization holds:** two identical `job.run()` calls; the pod body executes only on the first. +* **Memoization invalidates:** a fixture with **at least two groups**; change one file in one + group's list and assert only that group's pod invocation re-runs while the other stays cached. + With a single group the assertion would be vacuous. +* **`MergeJoin` regression:** a `MergeJoin` with colliding data columns completes inside + `job.run()`. + +--- + +## Dependencies & Risks + +* `orcapod-python` is upstream (`walkerlab/orcapod-python`); orcapod-sync-and-qc pins + `nauticalab/orcapod-python @ 966d759a`. This needs an upstream PR followed by a coordinated rev + bump in **both** orcapod-sync-and-qc and orcapod-spikesorting, which are deliberately kept on the + same rev because they target the same Ray cluster. +* The `GroupBy` split means NPIPE-204's stated API and the documented `common_clock_op` wiring both + need updating before the downstream change is written. +* The schedule risk is the upstream PR and two-repo rev bump, not the code. + +## Resources & References + +* `DESIGN_ISSUES.md` §U1 — source-info type hard-coding (this change fixes the `tag_data.py` half + and deletes the dead `polars_data_utils` site). +* `DESIGN_ISSUES.md` §O2 — operators discard a non-default `_context_key` (logged by this change, + not fixed). +* `CLAUDE.md` §"System tag evolution rules" — needs updating. Rule 3 currently describes `Batch` as + purely type-evolving (`str` → `list[str]`), which is true only of user tag and data columns; + system tags fold to a scalar and extend their name. `GroupBy` needs adding as a fourth, + reducing category, along with a row in the operator/function-pod boundary table and the project + layout tree. diff --git a/test-objective/integration/test_provenance.py b/test-objective/integration/test_provenance.py index a7b668262..9f4a41bd7 100644 --- a/test-objective/integration/test_provenance.py +++ b/test-objective/integration/test_provenance.py @@ -174,14 +174,34 @@ def test_join_sorts_system_tag_values_for_commutativity(self): # =================================================================== -# Type-evolving (aggregation ops) +# Reducing (many->one aggregation ops) # =================================================================== -class TestTypeEvolving: - """Per design: batch operation changes system tag type from str to list[str].""" +class TestReducing: + """Per design: a many->one op list-wraps user columns but NOT system tags. - def test_batch_evolves_system_tag_type(self): + System tag columns must stay scalar, because ``_build_record_id_preimage`` + hashes them directly to derive record identity. They fold to a + deterministic digest and their column name gains ``::{pipeline_hash}``. + """ + + def test_batch_list_wraps_user_columns(self): + source = _make_source( + {"group": pa.array(["a", "a", "b"], type=pa.large_string())}, + {"value": pa.array([1, 2, 3], type=pa.int64())}, + ["group"], + ) + result_table = Batch().process(source).as_table(all_info=True) + + for col_name in ("group", "value"): + col_type = result_table.schema.field(col_name).type + assert pa.types.is_list(col_type) or pa.types.is_large_list(col_type), ( + f"Expected list type for user column {col_name} after batch, " + f"got {col_type}" + ) + + def test_batch_keeps_system_tags_scalar(self): source = _make_source( {"group": pa.array(["a", "a", "b"], type=pa.large_string())}, {"value": pa.array([1, 2, 3], type=pa.int64())}, @@ -198,12 +218,18 @@ def test_batch_evolves_system_tag_type(self): # System tag columns should exist in output assert len(result_tag_cols) == len(source_tag_cols) - # The type should have evolved to list + # They must stay scalar -- record identity hashes them directly. for col_name in result_tag_cols: col_type = result_table.schema.field(col_name).type - assert pa.types.is_list(col_type) or pa.types.is_large_list( + assert not pa.types.is_list(col_type) and not pa.types.is_large_list( col_type - ), f"Expected list type for {col_name} after batch, got {col_type}" + ), f"Expected scalar type for {col_name} after batch, got {col_type}" + + # ...and their names gain a ::{pipeline_hash} block. + for src_col, out_col in zip(sorted(source_tag_cols), sorted(result_tag_cols)): + assert out_col.count("::") > src_col.count("::"), ( + f"Expected name-extension on {out_col}" + ) # =================================================================== @@ -234,7 +260,7 @@ def test_full_chain(self): filt = PolarsFilter(constraints={"group": "x"}) filtered = filt.process(joined) - # Step 3: Batch (type-evolving) + # Step 3: Batch (reducing) batch = Batch() batched = batch.process(filtered) @@ -244,7 +270,16 @@ def test_full_chain(self): # After all three stages, system tags should exist assert len(tag_cols) > 0 - # After batch, types should be lists + # After batch, system tags stay scalar (record identity hashes them), + # while the user columns are the ones that become lists. for col_name in tag_cols: col_type = table.schema.field(col_name).type - assert pa.types.is_list(col_type) or pa.types.is_large_list(col_type) + assert not pa.types.is_list(col_type) and not pa.types.is_large_list( + col_type + ), f"Expected scalar system tag {col_name}, got {col_type}" + + for col_name in ("group", "a", "b"): + col_type = table.schema.field(col_name).type + assert pa.types.is_list(col_type) or pa.types.is_large_list(col_type), ( + f"Expected list type for user column {col_name}, got {col_type}" + ) diff --git a/tests/test_core/datagrams/test_data_source_info_types.py b/tests/test_core/datagrams/test_data_source_info_types.py new file mode 100644 index 000000000..93d547b83 --- /dev/null +++ b/tests/test_core/datagrams/test_data_source_info_types.py @@ -0,0 +1,87 @@ +"""Source-info values may be lists, not just scalar strings. + +Many->one operators (GroupBy, MergeJoin) produce one provenance token per +member. `Data` must represent those without collapsing or crashing. +""" + +from __future__ import annotations + +import pyarrow as pa + +from orcapod.core.datagrams import Data + + +def _data_with_mixed_source_info() -> Data: + """Data with one list-valued and one scalar-null source token.""" + return Data( + {"probe": [0, 1], "path": ["a", "b"]}, + source_info={"probe": None, "path": ["s0", "s1"]}, + ) + + +class TestListValuedSourceInfo: + def test_schema_reports_list_type_for_list_valued_token(self): + data = _data_with_mixed_source_info() + schema = data.schema(columns={"source": True}) + assert schema["_source_path"] == list[str] + + def test_schema_reports_str_for_none_token(self): + data = _data_with_mixed_source_info() + schema = data.schema(columns={"source": True}) + assert schema["_source_probe"] is str + + def test_as_table_round_trips_list_valued_token(self): + data = _data_with_mixed_source_info() + table = data.as_table(columns={"source": True}) + assert table.schema.field("_source_path").type == pa.large_list( + pa.large_string() + ) + assert table.column("_source_path").to_pylist() == [["s0", "s1"]] + + def test_as_table_keeps_none_token_as_large_string(self): + data = _data_with_mixed_source_info() + table = data.as_table(columns={"source": True}) + assert table.schema.field("_source_probe").type == pa.large_string() + assert table.column("_source_probe").to_pylist() == [None] + + def test_empty_list_token_defaults_to_list_of_string(self): + data = Data({"path": ["a"]}, source_info={"path": []}) + table = data.as_table(columns={"source": True}) + assert table.schema.field("_source_path").type == pa.large_list( + pa.large_string() + ) + + def test_nested_list_token_recurses(self): + """Type derivation descends into nested lists rather than flattening. + + A single level of list-wrapping would be satisfied by returning + ``large_list(large_string)`` unconditionally; this pins the recursion. + """ + data = Data({"path": ["a", "b"]}, source_info={"path": [["s0"], ["s1"]]}) + assert data.schema(columns={"source": True})["_source_path"] == list[list[str]] + table = data.as_table(columns={"source": True}) + assert table.schema.field("_source_path").type == pa.large_list( + pa.large_list(pa.large_string()) + ) + assert table.column("_source_path").to_pylist() == [[["s0"], ["s1"]]] + + def test_scalar_token_unchanged(self): + """Existing scalar behavior must not regress.""" + data = Data({"path": "a"}, source_info={"path": "src::row_0::path"}) + table = data.as_table(columns={"source": True}) + assert table.schema.field("_source_path").type == pa.large_string() + assert data.schema(columns={"source": True})["_source_path"] is str + + def test_arrow_table_construction_recovers_list_token(self): + """Data built from an Arrow table keeps list-valued source info.""" + table = pa.table( + { + "path": pa.array([["a", "b"]], pa.list_(pa.large_string())), + "_source_path": pa.array([["s0", "s1"]], pa.list_(pa.large_string())), + } + ) + data = Data(table) + assert data.source_info()["path"] == ["s0", "s1"] + assert data.as_table(columns={"source": True}).column( + "_source_path" + ).to_pylist() == [["s0", "s1"]] diff --git a/tests/test_core/operators/test_group_by.py b/tests/test_core/operators/test_group_by.py new file mode 100644 index 000000000..7e961ced7 --- /dev/null +++ b/tests/test_core/operators/test_group_by.py @@ -0,0 +1,397 @@ +"""Tests for the GroupBy operator — many->one reduction keyed on tag values.""" + +from __future__ import annotations + +from pathlib import Path + +import pyarrow as pa +import pytest + +from orcapod.core.data_function import PythonDataFunction +from orcapod.core.function_pod import FunctionPod +from orcapod.core.operators import Batch, GroupBy, Join +from orcapod.core.sources import ArrowTableSource +from orcapod.core.streams import ArrowTableStream +from orcapod.errors import InputValidationError +from orcapod.protocols.core_protocols import StreamProtocol +from orcapod.system_constants import constants + + +@pytest.fixture +def session_table() -> pa.Table: + """Two sessions x two probes: the common-clock shape from NPIPE-204.""" + return pa.table({ + "subject": ["G", "G", "G", "G"], + "date": ["d1", "d1", "d2", "d2"], + "probe": [1, 0, 1, 0], + "path": ["b", "a", "d", "c"], + }) + + +@pytest.fixture +def session_source(session_table) -> ArrowTableSource: + return ArrowTableSource( + session_table, + tag_columns=["subject", "date", "probe"], + infer_nullable=True, + ) + + +class TestGroupByShape: + def test_one_row_per_distinct_key(self, session_source): + out = GroupBy(by=["subject", "date"]).process(session_source) + assert len(out.as_table()) == 2 + + def test_group_keys_are_scalar_tags(self, session_source): + op = GroupBy(by=["subject", "date"]) + out = op.process(session_source) + tag_cols, _ = out.keys() + assert set(tag_cols) == {"subject", "date"} + assert out.as_table().column("subject").to_pylist() == ["G", "G"] + + def test_non_key_tags_promoted_to_list_data(self, session_source): + out = GroupBy(by=["subject", "date"]).process(session_source) + _, data_cols = out.keys() + assert "probe" in data_cols + assert out.as_table().column("probe").to_pylist() == [[0, 1], [0, 1]] + + def test_data_columns_are_lists(self, session_source): + out = GroupBy(by=["subject", "date"]).process(session_source) + assert out.as_table().column("path").to_pylist() == [["a", "b"], ["c", "d"]] + + def test_source_columns_are_lists(self, session_source): + out = GroupBy(by=["subject", "date"]).process(session_source) + table = out.as_table(columns={"source": True}) + assert len(table.column(f"{constants.SOURCE_PREFIX}path").to_pylist()[0]) == 2 + + def test_system_tags_are_scalar_and_renamed(self, session_source): + out = GroupBy(by=["subject", "date"]).process(session_source) + table = out.as_table(columns={"system_tags": True}) + sys_cols = [ + c for c in table.column_names if c.startswith(constants.SYSTEM_TAG_PREFIX) + ] + assert sys_cols + for col in sys_cols: + field_type = table.schema.field(col).type + assert not pa.types.is_list(field_type) + assert not pa.types.is_large_list(field_type) + # name-extended: original "::" plus "::" + assert col.count(constants.BLOCK_SEPARATOR) >= 2 + + +class TestGroupByOrdering: + """Member and group ordering. + + Determinism here is load-bearing rather than cosmetic: orcapod hashes the + emitted lists to build the memoization key, so an unstable order makes an + unchanged member set hash differently and triggers a spurious recompute. + Upstream emission order is not stable across runs (Ray executor scheduling, + DB fetch order), so the operator cannot inherit it. + + *That* ordering must be deterministic is a requirement; *which* order to + use, and whether to promise one publicly, is an open design question -- as + is empty-group handling (see ``TestGroupByEmptyInput``). These tests pin + current behaviour so a change is deliberate, not accidental. + """ + + # TODO(ITL-630): decide whether GroupBy guarantees an ordering or only + # determinism, and which sort key. Deferred per review on PR #250. + + def test_members_sorted_by_non_key_tags(self, session_source): + """probe=[1,0] on input must emit as [0,1].""" + out = GroupBy(by=["subject", "date"]).process(session_source) + assert out.as_table().column("probe").to_pylist()[0] == [0, 1] + + def test_row_order_does_not_affect_output(self, session_table): + """Same rows, shuffled, must produce a byte-identical table.""" + shuffled = session_table.take([3, 1, 2, 0]) + + def run(tbl): + src = ArrowTableSource( + tbl, tag_columns=["subject", "date", "probe"], infer_nullable=True + ) + return GroupBy(by=["subject", "date"]).process(src).as_table() + + assert run(session_table).equals(run(shuffled)) + + def test_falls_back_to_record_id_when_key_covers_all_tags(self): + """by covering every tag leaves no non-key tag to sort on. + + record_id is a uuid5 of source content, so the order it imposes is not + human-predictable but IS stable across processes. Asserted exactly -- + sorting the result before comparing would make this blind to the very + ordering it exists to pin down. + """ + table = pa.table({"subject": ["G", "G"], "path": ["b", "a"]}) + src = ArrowTableSource(table, tag_columns=["subject"], infer_nullable=True) + out = GroupBy(by=["subject"]).process(src) + assert len(out.as_table()) == 1 + assert out.as_table().column("path").to_pylist()[0] == ["a", "b"] + + def test_record_id_breaks_ties_between_duplicate_tags(self): + """Duplicate tag tuples must not fall back to raw emission order.""" + table = pa.table({"s": ["G", "G"], "p": [1, 1], "path": ["z", "a"]}) + src = ArrowTableSource(table, tag_columns=["s", "p"], infer_nullable=True) + out = GroupBy(by=["s"]).process(src) + # Emission order would give ["z", "a"]; record_id imposes ["a", "z"]. + assert out.as_table().column("path").to_pylist()[0] == ["a", "z"] + + def test_join_fan_out_ties_broken_by_every_record_id_column(self): + """A joined stream carries one record_id per canonical input position. + + The broadcast side of a fan-out repeats its record_id on every row it + was joined into, so it cannot separate those rows; only the fanned-out + side's record_id can. Consulting just the first record_id column in + table order therefore leaves the sort key fully tied whenever the + broadcast side happens to sort first canonically, and the members fall + back to emission order -- which DB fetch order and Ray scheduling make + nondeterministic. Every record_id column has to take part. + + The literals are load-bearing: with this pair of schemas the broadcast + input sorts to canonical position 0, and its record_ids order the + members *against* emission order, so emission order and record_id + order are distinguishable. Sorting the result before comparing would + make the test blind to the ordering it exists to pin down. + """ + left = ArrowTableSource( + pa.table({"k": ["t", "t"], "a": [10, 20]}), + tag_columns=["k"], + infer_nullable=True, + ) + right = ArrowTableSource( + pa.table({"k": ["t"], "b": [1]}), + tag_columns=["k"], + infer_nullable=True, + ) + joined = Join().process(left, right) + out = GroupBy(by=["k"]).process(joined).as_table() + # Emission order would give [10, 20]; the fanned-out side's record_id + # imposes [20, 10]. + assert out.column("a").to_pylist() == [[20, 10]] + + +class TestGroupByValidation: + def test_empty_by_raises(self): + with pytest.raises(ValueError, match="at least one"): + GroupBy(by=[]) + + def test_unknown_column_raises(self, session_source): + op = GroupBy(by=["subject", "nonexistent"]) + with pytest.raises(InputValidationError, match="nonexistent"): + op.process(session_source) + + def test_data_column_as_key_raises(self, session_source): + """Grouping on a data column is not allowed -- keys must be tags.""" + op = GroupBy(by=["path"]) + with pytest.raises(InputValidationError, match="path"): + op.process(session_source) + + def test_duplicate_by_raises(self): + with pytest.raises(ValueError, match="duplicate"): + GroupBy(by=["subject", "subject"]) + + def test_list_valued_tag_as_key_raises(self): + """Batch list-wraps its tags; those are unhashable and cannot key a group. + + Without an explicit check this surfaced as a bare + ``TypeError: unhashable type: 'list'`` naming neither operator nor column. + """ + table = pa.table({"s": ["G", "G"], "v": [1, 2]}) + src = ArrowTableSource(table, tag_columns=["s"], infer_nullable=True) + batched = Batch(batch_size=2).process(src) + + with pytest.raises(InputValidationError, match="list-valued"): + GroupBy(by=["s"]).process(batched) + + +class TestGroupByEmptyInput: + """Empty-input behaviour. + + Current behaviour is "zero rows in, zero groups out" -- the least + surprising default, but never designed. Whether a reduction should instead + emit nothing, emit an empty group, or raise is open. This test pins the + status quo. + """ + + # TODO(ITL-630): decide empty-input behaviour for reductions. Deferred per + # review on PR #250, same follow-up as ordering. + + def test_empty_input_yields_zero_groups(self): + table = pa.table({ + "subject": pa.array([], pa.large_string()), + "path": pa.array([], pa.large_string()), + }) + stream = ArrowTableStream(table, tag_columns=["subject"]) + out = GroupBy(by=["subject"]).process(stream) + assert len(out.as_table()) == 0 + + +class TestGroupByOutputSchema: + """The predicted schema must match what grouping actually produces. + + Compared against the materialized ``ArrowTableStream`` from + ``unary_static_process``, never against ``process(...)``: a + ``DynamicPodStream.output_schema`` delegates straight back to the pod, so + that comparison is circular and passes no matter how wrong the prediction is. + """ + + @pytest.mark.parametrize( + "config", + [ + {}, + {"source": True}, + {"system_tags": True}, + {"source": True, "system_tags": True}, + ], + ids=["plain", "source", "system_tags", "source+system_tags"], + ) + def test_predicted_schema_matches_materialized(self, session_source, config): + op = GroupBy(by=["subject", "date"]) + + pred_tag, pred_data = op.unary_output_schema(session_source, columns=config) + actual = op.unary_static_process(session_source) + act_tag, act_data = actual.output_schema(columns=config) + + def diff(label, predicted, actual_schema): + predicted, actual_schema = dict(predicted), dict(actual_schema) + only_pred = set(predicted) - set(actual_schema) + only_act = set(actual_schema) - set(predicted) + assert predicted == actual_schema, ( + f"{label} schema mismatch for columns={config}: " + f"predicted-only={sorted(only_pred)}, " + f"actual-only={sorted(only_act)}, " + f"predicted={predicted}, actual={actual_schema}" + ) + + diff("tag", pred_tag, act_tag) + diff("data", pred_data, act_data) + + +class TestGroupByIdentity: + def test_identity_structure_includes_by(self): + assert ( + GroupBy(by=["a"]).identity_structure() + != GroupBy(by=["b"]).identity_structure() + ) + + def test_to_config_round_trip(self): + op = GroupBy(by=["subject", "date"]) + config = op.to_config() + # A list, not a tuple, so the config stays JSON-serializable. + assert config["config"]["by"] == ["subject", "date"] + rebuilt = GroupBy.from_config(config) + assert rebuilt.identity_structure() == op.identity_structure() + + def test_to_config_is_json_serializable(self): + import json + + json.dumps(GroupBy(by=["subject", "date"]).to_config()["config"]) + + +class TestGroupByRegistration: + def test_in_operator_registry(self): + from orcapod.pipeline.serialization import _build_operator_registry + + assert _build_operator_registry()["GroupBy"] is GroupBy + + def test_stream_fluent_method(self, session_source): + out = session_source.group_by(["subject", "date"]) + assert len(out.as_table()) == 2 + + +class TestGroupByAsyncIsBarrier: + """GroupBy must NOT override async_execute. + + ``UnaryOperator.async_execute`` (``core/operators/base.py:71``) already + collects the full input before calling ``static_process``, which is exactly + the barrier GroupBy needs: no group can be emitted before the input channel + closes, because any row not yet seen could belong to a group already + started. Adding an override would duplicate that logic and risk drifting + from it. + """ + + def test_does_not_override_async_execute(self): + from orcapod.core.operators.base import UnaryOperator + + assert "async_execute" not in GroupBy.__dict__ + assert GroupBy.async_execute is UnaryOperator.async_execute + + +# --------------------------------------------------------------------------- +# Logical (extension) element types — NPIPE-204 / ET2 +# --------------------------------------------------------------------------- + +# Defined at module level on purpose. This file uses `from __future__ import +# annotations`, so a pod function nested inside a test method would have its +# annotations stringified with no resolvable scope, and `Path` would fail to +# resolve. Module scope keeps `-> Path` resolvable, matching how real pipelines +# declare pods. +def _make_path(seed: str) -> Path: + return Path(f"/data/sync_{seed}.parquet") + + +def _path_stream() -> StreamProtocol: + """A stream whose data column is ``extension``.""" + source = ArrowTableSource( + pa.table({ + "date": ["d1", "d1", "d2"], + "probe": [0, 1, 0], + "seed": ["a", "b", "c"], + }), + tag_columns=["date", "probe"], + infer_nullable=True, + ) + return FunctionPod( + PythonDataFunction(_make_path, output_keys="result_path") + )(source) + + +class TestAggregationLogicalTypes: + """Aggregating operators must preserve logical (extension) element types. + + A pod annotated ``-> Path`` emits an ``extension`` column. + Naively wrapping that in ``pa.list_()`` raises ``ArrowNotImplementedError``: + Arrow cannot embed an extension type inside a list value field + (DESIGN_ISSUES ET1/ET2). The list must be built over the element's + *storage* type and wrapped in the ``list[orcapod.path]`` extension type at + the outermost level, which is what ``ListLogicalType`` provides. + + Every other fixture in this file uses plain ``large_string``/``int64``, + which is why this case was originally missed. + """ + + def test_upstream_column_really_is_extension_typed(self): + """Guard the fixture -- if this stops holding, the rest is vacuous.""" + field = _path_stream().as_table().schema.field("result_path") + assert isinstance(field.type, pa.ExtensionType) + assert field.type.storage_type == pa.large_string() + + def test_group_by_preserves_path_element_type(self): + out = GroupBy(by=["date"]).process(_path_stream()) + field = out.as_table().schema.field("result_path") + + assert isinstance(field.type, pa.ExtensionType), ( + f"expected a list extension type, got {field.type}" + ) + assert field.type.storage_type == pa.large_list(pa.large_string()) + + def test_group_by_output_schema_is_list_of_path(self): + out = GroupBy(by=["date"]).process(_path_stream()) + _, data_schema = out.output_schema() + assert data_schema["result_path"] == list[Path] + + def test_group_by_preserves_path_values(self): + out = GroupBy(by=["date"]).process(_path_stream()) + assert out.as_table().column("result_path").to_pylist() == [ + ["/data/sync_a.parquet", "/data/sync_b.parquet"], + ["/data/sync_c.parquet"], + ] + + def test_batch_preserves_path_element_type(self): + """Batch has the same defect, independent of GroupBy.""" + out = Batch(batch_size=0).process(_path_stream()) + field = out.as_table().schema.field("result_path") + assert isinstance(field.type, pa.ExtensionType), ( + f"expected a list extension type, got {field.type}" + ) + assert field.type.storage_type == pa.large_list(pa.large_string()) diff --git a/tests/test_core/operators/test_operators.py b/tests/test_core/operators/test_operators.py index d519d1c0b..cfd3f7e59 100644 --- a/tests/test_core/operators/test_operators.py +++ b/tests/test_core/operators/test_operators.py @@ -433,6 +433,48 @@ def test_negative_batch_size_raises(self): with pytest.raises(ValueError, match="non-negative"): Batch(batch_size=-1) + def test_batch_system_tags_are_scalar(self): + """System tags must stay scalar -- record identity hashes them directly.""" + from orcapod.core.sources import ArrowTableSource + from orcapod.system_constants import constants + + table = pa.table( + { + "animal": ["cat", "dog"], + "weight": [4.0, 12.0], + } + ) + source = ArrowTableSource(table, tag_columns=["animal"], infer_nullable=True) + out = Batch(batch_size=0).process(source) + result = out.as_table(columns={"source": True, "system_tags": True}) + + sys_cols = [ + c for c in result.column_names if c.startswith(constants.SYSTEM_TAG_PREFIX) + ] + assert sys_cols, "expected system tag columns on the batched output" + for col in sys_cols: + assert not pa.types.is_list(result.schema.field(col).type) + assert not pa.types.is_large_list(result.schema.field(col).type) + + def test_batch_source_columns_are_lists(self): + """Provenance stays per-member rather than collapsing.""" + from orcapod.core.sources import ArrowTableSource + from orcapod.system_constants import constants + + table = pa.table( + { + "animal": ["cat", "dog"], + "weight": [4.0, 12.0], + } + ) + source = ArrowTableSource(table, tag_columns=["animal"], infer_nullable=True) + out = Batch(batch_size=0).process(source) + result = out.as_table(columns={"source": True, "system_tags": True}) + + src_col = f"{constants.SOURCE_PREFIX}weight" + assert src_col in result.column_names + assert len(result.column(src_col).to_pylist()[0]) == 2 + class TestJoinBehavior: def test_join_combines_streams_on_shared_tags(self, simple_stream, disjoint_stream): diff --git a/tests/test_pipeline/test_aggregation_job.py b/tests/test_pipeline/test_aggregation_job.py new file mode 100644 index 000000000..e7403862b --- /dev/null +++ b/tests/test_pipeline/test_aggregation_job.py @@ -0,0 +1,297 @@ +"""Job-level tests for aggregating operators (NPIPE-204). + +Operator-level tests (`op.process(stream)` then `as_table()`) never reach +`StaticOutputOperatorPod._materialize_to_stream`, which is where list-valued +provenance used to crash. These run the full `job.run()` path against a real +Delta Lake store. +""" + +from __future__ import annotations + +from collections.abc import Iterator +from pathlib import Path + +import pyarrow as pa +import pytest + +from orcapod.core.data_function import PythonDataFunction +from orcapod.core.function_pod import FunctionPod +from orcapod.core.operators import Batch, GroupBy, MergeJoin +from orcapod import File +from orcapod.core.sources import ArrowTableSource, DictSource +from orcapod.databases import DeltaTableDatabase +from orcapod.pipeline import PipelineJob + +# --------------------------------------------------------------------------- +# Invocation recording +# +# `PythonDataFunction` hashes the wrapped callable via `inspect.getsource` and +# `func.__code__`, so the callable must be a real function -- a callable class +# instance is rejected. The invocation log therefore lives at module level +# rather than on a per-test callable object, which has the useful side effect +# of keeping the function identity byte-for-byte identical across the two runs +# a memoization test performs. +# --------------------------------------------------------------------------- + +_CALLS: dict[str, list] = {"path": [], "v": [], "cfg": []} + + +def count_paths(path: list[str]) -> int: + """Record the batch's `path` members and return the batch size. + + Used with `Batch`, which keeps `probe` as a list-valued *tag*, so `path` + is the only data column reaching the pod. + """ + _CALLS["path"].append(list(path)) + return len(path) + + +def count_group(probe: list[int], path: list[str]) -> int: + """Record the group's `path` members and return the group size. + + Used with `GroupBy`, which promotes the non-key tag `probe` to a + list-valued *data* column. `FunctionPod` requires the pod signature to + cover every incoming data column, so `probe` must be a parameter even + though only `path` is asserted on. + """ + _CALLS["path"].append(list(path)) + return len(path) + + +def count_v(v: list[str]) -> int: + """Record the merged `v` members and return the merged-list size.""" + _CALLS["v"].append(list(v)) + return len(v) + + +@pytest.fixture(autouse=True) +def _reset_calls() -> Iterator[None]: + """Clear the shared invocation log before every test.""" + _CALLS["path"].clear() + _CALLS["v"].clear() + yield + + +@pytest.fixture +def store(tmp_path: Path) -> DeltaTableDatabase: + return DeltaTableDatabase(base_path=tmp_path / "store") + + +@pytest.fixture +def session_source_factory(): + """Build a 2-group source; `paths` lets a test mutate one group's data.""" + + def _make(paths: list[str] | None = None) -> ArrowTableSource: + table = pa.table({ + "subject": ["G", "G", "G", "G"], + "date": ["d1", "d1", "d2", "d2"], + "probe": [0, 1, 0, 1], + "path": paths or ["a", "b", "c", "d"], + }) + return ArrowTableSource( + table, + tag_columns=["subject", "date", "probe"], + infer_nullable=True, + ) + + return _make + + +def _run(store, source, operator, name, function=count_group): + """Record and run `source -> operator -> function` as a PipelineJob.""" + pod = FunctionPod(PythonDataFunction(function, output_keys="n")) + job = PipelineJob(name=name, store=store) + with job: + pod(operator(source, label="agg"), label="counter") + return job.run() + + +class TestGroupByInJob: + def test_group_by_completes(self, store, session_source_factory): + _run(store, session_source_factory(), GroupBy(by=["subject", "date"]), "gb") + assert len(_CALLS["path"]) == 2 + assert sorted(_CALLS["path"]) == [["a", "b"], ["c", "d"]] + + def test_batch_completes(self, store, session_source_factory): + """The provenance fix is independent of grouping.""" + _run( + store, + session_source_factory(), + Batch(batch_size=2), + "b", + function=count_paths, + ) + assert len(_CALLS["path"]) == 2 + + +class TestGroupByMemoization: + def test_identical_runs_hit_cache(self, store, session_source_factory): + _run(store, session_source_factory(), GroupBy(by=["subject", "date"]), "m") + assert len(_CALLS["path"]) == 2 + + _CALLS["path"].clear() + result = _run( + store, session_source_factory(), GroupBy(by=["subject", "date"]), "m" + ) + assert _CALLS["path"] == [], "second identical run must not recompute" + # A cache hit must still surface the results, not an empty stream. + table = result.nodes["counter"].as_table() + assert table.num_rows == 2 + assert sorted(table.column("n").to_pylist()) == [2, 2] + + def test_fresh_store_recomputes(self, tmp_path, session_source_factory): + """Control for `test_identical_runs_hit_cache`. + + Same two runs, but the second points at a *different* store. If this + did not recompute, the cache-hit assertion above would be vacuous -- + it would be passing because the pod never runs, not because the record + was found. + """ + store_a = DeltaTableDatabase(base_path=tmp_path / "store_a") + store_b = DeltaTableDatabase(base_path=tmp_path / "store_b") + + _run(store_a, session_source_factory(), GroupBy(by=["subject", "date"]), "m") + assert len(_CALLS["path"]) == 2 + + _CALLS["path"].clear() + _run(store_b, session_source_factory(), GroupBy(by=["subject", "date"]), "m") + assert len(_CALLS["path"]) == 2, ( + "a fresh store has no cached records, so both groups must recompute" + ) + + def test_changed_member_invalidates_only_its_group( + self, store, session_source_factory + ): + """Two groups; change one member of the first only. + + With a single group this assertion would be vacuous -- it must show + that the untouched group stays cached. + """ + _run(store, session_source_factory(), GroupBy(by=["subject", "date"]), "i") + assert len(_CALLS["path"]) == 2 + + _CALLS["path"].clear() + changed = session_source_factory(["a", "B_CHANGED", "c", "d"]) + result = _run(store, changed, GroupBy(by=["subject", "date"]), "i") + + assert _CALLS["path"] == [["a", "B_CHANGED"]], ( + f"only the changed group should recompute; got {_CALLS['path']}" + ) + # Both groups must still be present in the output — the recomputed one + # and the one served from cache. + table = result.nodes["counter"].as_table() + assert table.num_rows == 2 + assert sorted(table.column("n").to_pylist()) == [2, 2] + + +class TestMergeJoinRegression: + def test_merge_join_completes_in_job(self, store): + """MergeJoin carries source columns as parallel lists. + + It crashed with the same ArrowTypeError before the Data fix. + """ + left = ArrowTableSource( + pa.table({"id": ["a", "b"], "v": ["l1", "l2"]}), + tag_columns=["id"], + infer_nullable=True, + ) + right = ArrowTableSource( + pa.table({"id": ["a", "b"], "v": ["r1", "r2"]}), + tag_columns=["id"], + infer_nullable=True, + ) + + pod = FunctionPod(PythonDataFunction(count_v, output_keys="n")) + job = PipelineJob(name="mj", store=store) + with job: + pod(MergeJoin()(left, right, label="mj"), label="counter") + job.run() + + assert len(_CALLS["v"]) == 2 + # MergeJoin merges colliding `v` columns into a sorted 2-element list. + assert sorted(_CALLS["v"]) == [["l1", "r1"], ["l2", "r2"]] + + +# --------------------------------------------------------------------------- +# list[File] content hashing through a reduction — ITL-627 Defect 2 +# --------------------------------------------------------------------------- + + +def reduce_configs(probe: list[int], cfg: list[File]) -> int: + """Record each member's file *contents* so a stale cache is visible.""" + _CALLS["cfg"].append([Path(c).read_text().strip() for c in cfg]) + return len(cfg) + + +def read_config(cfg: File) -> int: + """Per-row control: scalar File columns are known to hash by content.""" + _CALLS["cfg"].append(Path(cfg).read_text().strip()) + return 1 + + +class TestFileContentHashingThroughReduction: + """A `list[File]` column must invalidate on content change, like a scalar one. + + A `File` column exists so that orcapod hashes the file's *contents* — editing + a file re-runs its consumer even though the path is unchanged. ITL-627 + Defect 2 was that list-wrapping lost this: `build_aggregated_table` produced + a `list[orcapod.file]` hashed by its storage values (JSON path strings), so a + content edit silently failed to invalidate. + + The unit-level hashing contract is covered by + `tests/test_hashing/test_extension_type_hashing.py::TestListExtensionHashing`. + This test covers the job level, which is where the failure was observed and + where it fails *silently* rather than raising. + """ + + @staticmethod + def _source(cfg_paths): + return DictSource( + [ + {"date": "d1", "probe": i, "cfg": File(p)} + for i, p in enumerate(cfg_paths) + ], + tag_columns=["date", "probe"], + data_schema={"date": str, "probe": int, "cfg": File}, + source_id="cfgsrc", + ) + + def _run(self, store, cfg_paths, *, grouped): + _CALLS["cfg"].clear() + source = self._source(cfg_paths) + job = PipelineJob(name="filehash", store=store) + with job: + if grouped: + pod = FunctionPod( + PythonDataFunction(reduce_configs, output_keys="n") + ) + pod(GroupBy(by=["date"])(source, label="agg"), label="cc") + else: + pod = FunctionPod(PythonDataFunction(read_config, output_keys="n")) + pod(source, label="cc") + job.run() + return list(_CALLS["cfg"]) + + @pytest.mark.parametrize("grouped", [False, True], ids=["per_row", "grouped"]) + def test_content_edit_invalidates(self, tmp_path, grouped): + store = DeltaTableDatabase(base_path=tmp_path / "store") + a = tmp_path / "a.toml" + b = tmp_path / "b.toml" + a.write_text("entities = 'A'\n") + b.write_text("entities = 'B'\n") + + assert self._run(store, [a, b], grouped=grouped), "cold run must execute" + assert self._run(store, [a, b], grouped=grouped) == [], ( + "identical re-run must hit the cache" + ) + + # Same path, new contents. + a.write_text("entities = 'A_CHANGED'\n") + after = self._run(store, [a, b], grouped=grouped) + + assert after, ( + "a content edit must invalidate; an empty call log means the " + "File column was hashed by path rather than by contents (ITL-627)" + ) + flattened = after[0] if grouped else after + assert "entities = 'A_CHANGED'" in flattened diff --git a/tests/test_utils/test_arrow_utils.py b/tests/test_utils/test_arrow_utils.py index 385c399d6..a35154f4f 100644 --- a/tests/test_utils/test_arrow_utils.py +++ b/tests/test_utils/test_arrow_utils.py @@ -792,3 +792,107 @@ def test_round_trips_through_arrow_table_stream(self): _, data_schema = stream.output_schema() assert data_schema["score"] == (int | None) +# fold_system_tag_values +# --------------------------------------------------------------------------- + + +class TestFoldSystemTagValues: + """Folding N members' system-tag values into one scalar (NPIPE-204). + + The expected digests below are hard-coded on purpose. A fold that used + hash() or set-iteration order would still be self-consistent within one + process; pinning the values is what catches it. + """ + + SOURCE_COL = "_tag_source_id::abc123" + RECORD_COL = "_tag_record_id::abc123" + + RIDS = [ + bytes.fromhex("0102030405060708090a0b0c0d0e0f10"), + bytes.fromhex("1112131415161718191a1b1c1d1e1f20"), + ] + EXPECTED_RID = bytes.fromhex("853be16a3f38565f8ced039f84fdbea6") + EXPECTED_SID = ( + "7916442d59841140bedf6c1f5dcc1304ae9fce0ba885765c06e511086b85da2e" + ) + + def test_record_id_folds_to_16_bytes(self): + from orcapod.utils.arrow_utils import fold_system_tag_values + + result = fold_system_tag_values(self.RECORD_COL, self.RIDS) + assert isinstance(result, bytes) + assert len(result) == 16 + + def test_record_id_digest_is_pinned(self): + from orcapod.utils.arrow_utils import fold_system_tag_values + + assert fold_system_tag_values(self.RECORD_COL, self.RIDS) == self.EXPECTED_RID + + def test_source_id_digest_is_pinned(self): + from orcapod.utils.arrow_utils import fold_system_tag_values + + result = fold_system_tag_values(self.SOURCE_COL, ["src_a", "src_b"]) + assert result == self.EXPECTED_SID + + def test_order_matters(self): + """Member order is part of the identity, matching the data lists.""" + from orcapod.utils.arrow_utils import fold_system_tag_values + + forward = fold_system_tag_values(self.RECORD_COL, self.RIDS) + reverse = fold_system_tag_values(self.RECORD_COL, list(reversed(self.RIDS))) + assert forward != reverse + + def test_single_member_is_still_folded(self): + """A one-member group folds rather than passing the value through.""" + from orcapod.utils.arrow_utils import fold_system_tag_values + + result = fold_system_tag_values(self.RECORD_COL, self.RIDS[:1]) + assert isinstance(result, bytes) and len(result) == 16 + assert result != self.RIDS[0] + + def test_none_members_are_tolerated(self): + from orcapod.utils.arrow_utils import fold_system_tag_values + + assert isinstance( + fold_system_tag_values(self.RECORD_COL, [None, self.RIDS[0]]), bytes + ) + assert isinstance( + fold_system_tag_values(self.SOURCE_COL, [None, "src_a"]), str + ) + + def test_digest_is_stable_across_processes(self): + """Fresh interpreter, therefore fresh PYTHONHASHSEED. + + This is the test that catches a fold built on hash() or set order: + such a fold is self-consistent within one process and only diverges + on a new driver run. + """ + import os + import subprocess + import sys + from pathlib import Path + + script = ( + "from orcapod.utils.arrow_utils import fold_system_tag_values\n" + "rids = [bytes.fromhex('0102030405060708090a0b0c0d0e0f10'),\n" + " bytes.fromhex('1112131415161718191a1b1c1d1e1f20')]\n" + "print(fold_system_tag_values('_tag_record_id::abc123', rids).hex())\n" + "print(fold_system_tag_values('_tag_source_id::abc123', ['src_a','src_b']))\n" + ) + # The subprocess does not inherit pytest.ini's `pythonpath = src`, so + # pass it explicitly rather than relying on orcapod being installed in + # the active environment. + src_dir = Path(__file__).resolve().parents[2] / "src" + env = dict(os.environ) + env["PYTHONPATH"] = os.pathsep.join( + [str(src_dir), env["PYTHONPATH"]] if env.get("PYTHONPATH") else [str(src_dir)] + ) + out = subprocess.run( + [sys.executable, "-c", script], + capture_output=True, + text=True, + check=True, + env=env, + ).stdout.split() + assert out[0] == self.EXPECTED_RID.hex() + assert out[1] == self.EXPECTED_SID