feat(stats): resumable stat units and the StatRun store (rows-first s4)

[paddymul]

Stacked on #1024 (feat/rowsfirst-s3-stats-wire-server), which is stacked on #1022 and #1021. This PR is based on main so the repo's Checks workflow runs on it (its pull_request trigger uses branches: "*", which does not match a base branch containing a slash). The diff therefore includes the commits of #1021, #1022 and #1024 until they merge. The commits of this phase are dde83c16 (a characterization test that passes on the base), 667ae34d (failing tests) and 799583f7 (implementation), then 746ae8c5 (failing tests) and fa05adba (fix) from a review round on column names; read only those five.

Problem

The stats classes compute everything in one call: DfStatsV2 and PlDfStatsV2 run StatPipeline.process_df in their constructor, and XorqDfStatsV2 runs XorqStatPipeline.process_table. #1024's stats_request is therefore a whole run, one synchronous call that blocks the loop for as long as the stats take. On xorq that is one batch aggregate over every column plus one histogram query per column, and the batch alone took 162.5 s on the 78M-row telemetry entry. To bound the length of a request, a run has to be cut into units that can be run one at a time. A session also has to keep the pieces that have finished, so that two clients, or a client that reconnects, do not recompute them or lose them.

Phase and plan references

Rows-first s4, from buckaroo2-reports/plans/: plan 1 (01-rows-first-stats-separate-plumbing.md) section 9 "Phase 3" and plan 2 (02-rows-first-xorq-and-lazy-polars.md) section 7 "Phase 2", with plan 2 sections 4.1 (units, run object) and 4.2 (the xorq batch split), and gate 2 of measurements-phase0.md.

Approach

• Units. A new module, pluggable_analysis_framework/stat_units.py, defines StatUnit (id, columns, phase, prerequisite unit ids, cost class), StatState (the frame or expression, the columns that get no unit, an optional column group, columns to plan first, an optional row count, and the namespace the group and priority names are written in), the fragment type {orig_col: {stat: value}}, StatAccumulator, and three helpers: merge_fragments, resolve_names, which turns requested column names into original names (see "Column names in a request"), and rewrite_sd, which re-keys an original-name dict to the rewritten a, b, c names the way XorqDataflow._get_summary_sd did inline. StatPipeline and XorqStatPipeline each gain plan(state), new_accumulator(state) and run(unit, acc), and so do the three stats classes (DfStatsV2, PlDfStatsV2, XorqDfStatsV2), which delegate to their pipeline. plan and new_accumulator compute nothing. process_df and process_table are these three called in order, so the inline path and a caller running units one at a time cannot differ.
• pandas and polars. One unit per column that is not skipped, in column order, run through the existing process_column.
• xorq. The scalar batch is the first unit (length, null_count, min, max, mean, std, approx_median, distinct_count, plus the stats that are pure functions of them, so histogram_bins comes with it and color_map needs no histogram query). Then one histogram GROUP BY per column, with the columns the caller names first. A stat that needs an XorqExpr or XorqExecute, or reads one that does, belongs to its column's histogram unit, so a user-defined stat that queries the table lands there without a special case. The two phases of the old _process_table_impl are the two kinds of unit. The histogram queries are built and executed as before, so their snapshot-cache keys are unchanged.
• Column-chunk split of the batch. XorqStatPipeline(..., chunk_cells=N) cuts the batch into one aggregate per chunk of columns, sized by cells (rows x columns, with at least one column per chunk). It is off unless a host sets stat_chunk_cells on XorqServerDataflow, and it needs the row count (StatState.rows; the dataflow passes its cached _expr_count, and without a count the batch stays whole). Phase 0 gate 2 found the split by column cheap on a parquet scan (the chunks sum to 0.90-1.16x the single batch at 10-12M rows, peak RSS halves, the longest unit falls from 1.1-1.4 s to 0.20-0.26 s for 6-column chunks), not worth it by stat class (1.3-1.4x, every piece rescans) and costly on a CSV scan (2.1-2.2x, every chunk reparses the file). So only the column split exists.
• What decides that a source may be split. The maintainer removed an op-tree scan classifier on purpose: 17f031e9 (fix(stats): materialize the stat source only for chained sources (#915) #916) gated source materialization on the op tree, and fde5213b (refactor(stats): remove buckaroo-managed source materialization #929) deleted the materialization with it, because the worthiness heuristic and the machinery around it "made the system hard to reason about". That classifier chose behaviour for every source. This one only vetoes a host's request. The host opts in with stat_chunk_cells, which is its statement that it wants chunks for this source. XorqStatPipeline.chunk_refusal(table) then refuses the split unless the expression is a Read of a parquet file, or a projection of plain columns of one. Anything else (a join, an aggregate, a diff, a filter, a derived column, a cache node, a CSV read, a memtable, and a table registered on a backend, whose format is not recorded) keeps the single batch and logs why at info level. The guard cannot turn chunking on, it has one allow-list entry, and an expression it does not recognise is refused. If the host's declaration and the guard disagree, the guard wins. Nothing classifies an expression to decide what a default does.
• StatRun. buckaroo/server/stat_run.py holds the run a session keeps for one (stats_gen, scope): the planned units, an append-only list of the fragments in the order they finished, the accumulator the units read, and a status (pending, complete, error). run_next(prefer=(), namespace="any") runs one unit, the first whose prerequisites have run, or the first of those that covers a column named in prefer, and returns its fragment. It has no thread, timer or callback; nothing runs unless a caller asks. A unit that raises fails the run (status == "error", the exception propagates once, and the run is not retried); a stat that fails inside a unit is not that, it is an entry in acc.errors and the unit finishes. StatCursor is one client's position in the list: take(run) returns the fragments it has not seen and advances, and a cursor handed a different run starts again at the beginning. DataStreamHandler creates one per connection (stats_cursor), so the work is done once and each client reads the whole list at its own pace.
• On the session. SessionState.stat_runs maps (stats_gen, scope) to its run. begin_stats_generation clears it, so a run lives until the generation changes. start_stat_run(session, scope="raw") in stats_wire.py returns the run for the current generation, creating it from the dataflow if absent (planning only, no query). It analyzes the frame assign_full_stats does. Nothing in the server calls it yet.
• Dataflow. CustomizableDataflow.build_stats(processed_df, run=True) builds the stats class the way _get_summary_sd always did, and _get_summary_sd now calls it; run=False builds it without running, for plan/run. XorqDataflow overrides it to pass cache_storage, chunk_cells and rows. The three stats classes take run=True, and build_stats passes run only when it is False, so a DFStatsClass written before the keyword existed still builds.

Column names in a request

A column group, StatState.priority and StatRun's prefer hint are written in one of two kinds of name: the original column names, or the rewritten a, b, c names a client holds. The first version of this PR read each name as both. That breaks when an original name equals another column's rewritten name. A frame with columns c, b, a rewrites them to a, b, c, so c is the original name of the first column and the rewritten name of the last. A group of ('c',) planned the units of both, and prefer=('a',) ran the unit of the first column. A review found this; no server path calls these yet, but the columns hint of the next phase's stats_request will.

resolve_names(pairs, names, namespace) now reads each name once, in the namespace the caller states:

• rewritten: every name is a rewritten name. This is what the request branch passes for a client's hint.
• original: every name is an original name.
• any (the default): a name that is an original column name picks that column, and only a name that is not one is read as a rewritten name. Names that do not collide pick what they picked before, so the existing tests are unchanged. A name that does collide is read as the original, which is wrong for a client; that is why a client says rewritten.

StatState(namespace=...) sets it for the group and the priority names, and StatRun.next_unit and run_next take it for prefer. An unknown namespace raises ValueError. Priority and prefer names are resolved against every column of the frame, not only the columns in the group, so a name picks the same column whatever group is asked for. prioritized takes the state instead of the priority tuple because of this. skip_columns matching is not part of this change; it keeps the matching each backend had on main.

What changes

• buckaroo/pluggable_analysis_framework/stat_units.py (new), stat_pipeline.py (plan, new_accumulator, run; process_df over them), xorq_stat_pipeline.py (the same, XorqAccumulator, chunk_cells, chunk_refusal; process_table over them), df_stats_v2.py (run=False, state, plan/run).
• buckaroo/dataflow/dataflow.py (build_stats, the DfStats protocol), buckaroo/xorq_buckaroo.py (build_stats, stat_chunk_cells, rewrite_sd), buckaroo/server/xorq_loading.py (the stat_chunk_cells keyword).
• buckaroo/server/stat_run.py (new), session.py (stat_runs, cleared in begin_stats_generation), stats_wire.py (start_stat_run), websocket_handler.py (stats_cursor).
• Review round on column names: stat_units.py (resolve_names, StatState.namespace, prioritized(state, pairs)), stat_pipeline.py and xorq_stat_pipeline.py (the prioritized call), stat_run.py (namespace on next_unit and run_next).

Tests

Five commits: dde83c16 (characterization tests that pass on the base), 667ae34d (failing tests) and 799583f7 (implementation), then 746ae8c5 (failing tests for the column-name namespaces) and fa05adba (the fix).

• Characterization, in test_paf_v2.py and test_xorq_stats_v2.py: process_df equals the process_column loop it replaced (pandas and polars, with and without a skipped column, plus the per-column error list), and each xorq histogram query keeps its snapshot-cache key, checked against keys of queries built in the test in the shape histogram documents (the batch plus one per column, nothing else).

• Units, in the files that already test each backend: test_paf_v2.py (pandas and polars pipelines: one unit per column, a fragment keyed by original name, the union of fragments equals process_df, a column group by original or rewritten name, columns planned first, an empty frame, an error recorded on the accumulator), test_xorq_stats_v2.py (the batch first then one histogram unit per column, visible columns first, plan sends no query, the batch is one query and each histogram unit one more, the batch fragment carries histogram_bins and no histogram, the union equals process_table, a failing stat is reported once and the run goes on, a column group), skip_columns_test.py (a skipped column gets no unit on the three backends and still has its entry).

• The column-chunk split, in test_xorq_stats_v2.py (TestColumnChunkSplit): off by default, chunks sized by cells with at least one column, an unknown row count keeps one batch, batches precede histograms, visible columns fill the first chunk, the split equals the single batch (and queries are one per chunk), the histogram cache keys keep their keys and order, and nine kinds of source are refused (a join, an aggregate, a diff, a CSV read, a filter, a derived column, a cache node, a memtable and a backend table) while a plain parquet scan and a column projection of it are allowed. A refusal is logged and gives the single batch and the same stats.

• Through a dataflow, in test_data_loading_polars.py (pandas and polars) and test_load_expr.py (xorq): build_stats(run=False) computes nothing; the default run goes through the units (every stat query is issued inside a unit); the fragments of each scope's units, assembled by assemble_merged_sd with init_sd, a cleaning op, a search filter, column_config_overrides and a post-processing sd override, equal the dataflow's merged_sd; a column group returns only those columns; skip_stat_columns columns get no unit; stat_chunk_cells splits a parquet-backed dataflow's batch and a joined one stays whole. The helpers shared by those two files sit beside _scope_inputs in scoped_summary_stats_test.py.

• StatRun and the session: test_data_loading_polars.py (key, unit plan, fragments appended in completion order and never changed, prefer, completion, the accumulator equals the merged fragments, raw_sd equals the full-stats summary_sd, cursors read one list each at their own pace, a cursor starts again on another run, no thread, timer or IOLoop callback is created, a raising unit fails the run) and TestStatsWire in test_load_expr.py (a run started from a deferred session is keyed by the generation and sends no query, running it ends at the summary_sd an inline session computes, a dataflow-field change and /reload_expr drop it while a typed search term does not, each connection has its own cursor).

• Column-name namespaces (746ae8c5), in the files that already test each layer. test_paf_v2.py uses a frame with columns c, b, a: a group or a priority name that is an original name and also another column's rewritten name picks the original column only, rewritten and original read every name in their namespace, a rewritten name that is no original name still picks its column, and an unknown namespace raises. test_xorq_stats_v2.py checks the same on the xorq plan (the batch columns and one histogram unit per column picked). test_data_loading_polars.py checks next_unit and run_next on pandas and polars. Of the 32 new tests, 27 fail on 746ae8c5 alone (the wrong columns picked, or the missing namespace argument) and 5 pin names that do not collide and pass on both commits. All nine Python / Test jobs failed on 746ae8c5.

On 667ae34d (tests only) all nine Python / Test jobs failed, including 3.14 and Windows, where the xorq test file is skipped but the pandas and polars tests are not; 86 of the new tests fail locally, each on a missing module or attribute (stat_units, stat_run, plan, run, build_stats, chunk_refusal, chunk_cells, start_stat_run, stat_chunk_cells). CI logs were not read (the REST log API is rate-limited for this account), so the reason rests on the local run of the same commit. test_the_batch_is_one_query_by_default passes on that commit by design: it pins the default the flag must not change.

CI on 799583f7: all 28 checks completed, 27 succeeded (one of them the Read the Docs status) and deploy was skipped. That includes all nine Python / Test jobs (3.11 to 3.14, Max Versions 3.11 to 3.14, Windows), Python / Lint, Python / Typecheck, the JS job, the wheel build and the Playwright jobs.

CI on fa05adba: all 28 checks completed, 27 succeeded (one of them the Read the Docs status) and deploy was skipped, including all nine Python / Test jobs.

Locally the full unit suite (pytest ./tests/unit -m "not slow") gives 1361 passed and 5 skipped, which is the 1268 of #1024 plus 6 characterization tests and 87 new ones. In a Max Versions environment (pandas 3.0.6, polars 1.44.2, xorq 0.4.5, numpy 2.5.3, resolved as CI does) it gives the same 1361 passed. basedpyright on the scoped files reports 0 errors and 0 warnings. After fa05adba the unit suite gives 1393 passed and 5 skipped locally, the 1361 plus the 32 namespace tests.

Eight deliberate regressions of the implementation each make at least one test fail: the batch no longer first, a chunk guard that always allows, begin_stats_generation not dropping runs, a cursor that does not advance, a skipped column's length not set, fragments that alias the accumulator's dicts, priority ignored, and process_df bypassing the units. A ninth, StatRun.next_unit ignoring a unit's prerequisites, survives: no test asks for a unit whose prerequisite has not run (the prefer tests run pandas units, which have none). I did not add one after the implementation because it would not have been seen failing on CI first; the request branch of the next phase, which drives a xorq run with a column hint, is where it belongs.

Measurements

The longest single unit, each unit timed alone through plan/run as a request would run it, three repeats, Apple M4 Pro with other agents running (load average about 4). The fixture is synthetic: 27 columns (4 ints, 6 floats with nulls, 2 low-cardinality ints, 8 strings of 5 to 5,000 distinct values, 3 booleans, 4 timestamps). xorq reads a parquet file. pandas and polars analyze the 50,000-row sample their stats classes take above 1M cells, so their units do not grow with the file.

Fixture	Backend	Units	Longest unit (3 runs)	Whole run, median
small, 52,814 rows (1.4M cells)	pandas	27	0.010-0.011 s (column:float3)	0.145 s
	polars	27	0.004-0.005 s	0.074 s
	xorq, one batch	28	0.061-0.069 s (the batch)	0.325 s
	xorq, chunk_cells=1M (18 columns per chunk)	29	0.048-0.050 s	0.323 s
medium, 2,000,000 rows (54M cells)	pandas	27	0.011 s	0.175 s
	polars	27	0.005 s	0.078 s
	xorq, one batch	28	0.858-0.899 s (the batch)	1.624 s
	xorq, chunk_cells=60M	28	0.848-0.866 s (still one batch: 27 columns x 2M rows is under 60M cells)	1.618 s
	xorq, chunk_cells=12M (6 columns per chunk)	32	0.243-0.245 s (batch:0)	1.610 s

On xorq the longest histogram unit is 0.038-0.039 s on the medium fixture (one outlier run at 0.093 s), so with the batch cut into 6-column chunks the longest unit is a batch chunk at about 0.24 s, and the chunks add up to what the single batch costs (1.61 s against 1.62 s for the whole run). That matches gate 2: 0.20-0.26 s for 6-column chunks at 10-12M rows. Against the 250 ms budget the plans propose, a 6-column chunk fits at 2M rows, just under it; the plans' 60M cells per unit is a size where this fixture needs no split. These are the figures for parquet; a CSV-backed source is refused by the guard and keeps its single batch, whose cost is the 54-215 s the telemetry shows, which is plan 3's tier policy and not this phase's.

Why default behaviour is unchanged

• Everything new is opt-in or not called. stat_chunk_cells and chunk_cells default to off; run=True is the default of the stats classes; start_stat_run, StatRun and stats_cursor are not used by any request path yet, so stats_request and complete_stats behave as in feat(server): stats_request, stats_update and df_meta.stats on deferred sessions (rows-first s3) #1024, through the same units.
• The default path computes what it did. I compared the summary dicts of DfStatsV2, PlDfStatsV2 and XorqDfStatsV2 (on a parquet scan and on a memtable, each with and without a skipped column, on a mixed-dtype fixture with nulls, strings, booleans, timestamps and a categorical) and the xorq snapshot-cache keys of a run, between 9bfebfb9 and this branch, with a script outside the repo. All were identical, modulo the order of tied histogram entries and polars' mode over tied counts, which differ between two runs of the same code (two runs of the base disagree there, so the script sorts histograms and drops mode). The characterization tests above pin the same on the base.
• process_df keeps its early return for an empty frame and its perf summary; process_table keeps the stat.xorq.total and stat.xorq.batch_aggregate spans, the cache counters and the firstpull.summary_stats span on the dataflow, which the existing telemetry tests cover.
• No existing test changed. One import line in test_data_loading_polars.py grew (_scope_inputs plus the new shared helpers); everything else in the test files is added.

Deviations from the plan

• Fragments are keyed by original column name on all three backends, as the plan writes them, but process_df still returns its summary keyed by rewritten name and process_table by original name, as they always did. rewrite_sd converts a run's results to the rewritten keys the dataflow carries, and XorqDataflow._get_summary_sd now uses it instead of its inline loop.
• plan(state) takes a StatState, which the plan does not define. It carries what the plan lists as the unit's inputs, plus the skipped columns, an optional column group, columns to plan first and the row count. A group, the priority columns and prefer are written in a namespace (any, original or rewritten, default any; see "Column names in a request"), because a client only knows the rewritten names; skip_columns keeps the matching each backend always had (original or rewritten on pandas and polars, original on xorq).
• A skipped column gets no unit, but it keeps its entry in acc.sd(), with the structural keys the pipeline always gave it (orig_col_name and rewritten_col_name, and on xorq dtype, the batch's length and the None placeholders), because process_df and process_table always returned one. It is in no fragment.
• Visible-columns-first is StatState.priority for the plan and StatRun.next_unit(prefer) for a run, since the plan leaves the source of "visible" to the request.
• The xorq batch also computes the stats that are pure functions of its results (histogram_bins, non_null_count, nan_per, distinct_per), so the batch fragment is complete for everything that needs no histogram query, as plan 2 section 4.2 describes. With chunk_cells every chunk precedes every histogram unit, as the plan's unit order says, and each histogram unit names its chunk as its prerequisite.
• The flag is stat_chunk_cells on XorqServerDataflow (and chunk_cells on the pipeline and the stats class), a cell count and not a boolean, because chunks are sized by cells. It is not a /load_expr body field yet.
• The guard is an allow-list of one source kind, where the plan's wording ("scan-backed sources") invites a classifier. See "What decides that a source may be split".
• StatRun has the plan's unit list, fragment list, accumulator, creation time and status, but not the cache and timing counters of plan 2 section 4.5, which belong with the request branch's spans.
• build_stats and the DfStats protocol gain run and plan/new_accumulator/run (the plan says "each stats class gains two methods"; new_accumulator is the third, because run(unit, acc) needs somewhere to start).

Not in this PR

• The request branch that runs units under a time budget, the remaining field of stats_update and final per client: stats_request still runs the whole run, through the same units.
• Final assignment and snapshot refresh generalized to a run that finished in several requests, and the config upgrade.
• Per-unit telemetry spans (stats.unit) and the run's cache and timing counters.
• The client, the count memoization and the scalar-tier unit sets of plan 3 phase 4.
• A /load_expr body field for stat_chunk_cells, and any default for it.
• A test for StatRun.next_unit's prerequisite filter (see Tests).

🤖 Generated with Claude Code

[github-actions]

📦 TestPyPI package published

pip install --index-strategy unsafe-best-match --index-url https://test.pypi.org/simple/ --extra-index-url https://pypi.org/simple/ buckaroo==0.15.9.dev37186186027

or with uv:

uv pip install --index-strategy unsafe-best-match --index-url https://test.pypi.org/simple/ --extra-index-url https://pypi.org/simple/ buckaroo==0.15.9.dev37186186027

MCP server for Claude Code

claude mcp add buckaroo-table -- uvx --from "buckaroo[mcp]==0.15.9.dev37186186027" --index-strategy unsafe-best-match --index-url https://test.pypi.org/simple/ --extra-index-url https://pypi.org/simple/ buckaroo-table

📖 Docs preview

🎨 Storybook preview

[paddymul]

Review round on this PR: one finding, fixed in fa05adba with its tests in 746ae8c5.

Finding (medium): a column name was read as an original name and as a rewritten name at once. A column group, StatState.priority and StatRun's prefer hint matched each name against both the original and the rewritten (a, b, c) column names. When an original name equals another column's rewritten name, one name picked two columns. On a frame with columns c, b, a (rewritten a, b, c), StatState(df, columns=('c',)) planned the units of c and a, priority=('a',) planned c before a, and StatRun.next_unit(prefer=('a',)) returned the unit of c. A frame b, a, c with columns=('a',) returned both b and a. The values in the fragments were right; the selection and the order were wrong. I reproduced it with the reviewer's prefer script and its edge script before changing anything. The PR body described accepting both kinds of name but not this collision, and it now has a section on it. No request path calls these yet, but the columns hint of the next phase's stats_request will.

How it was addressed.

• 746ae8c5 adds the tests alone: 32 new tests in test_paf_v2.py, test_xorq_stats_v2.py and test_data_loading_polars.py. Locally 27 of them failed on the tests-only commit (the wrong columns picked, or the new namespace argument missing) and 5 pass because they pin names that do not collide. All nine Python / Test jobs failed on that commit before the fix was pushed.
• fa05adba adds resolve_names, which reads each name once, in a namespace the caller states: rewritten (every name is a rewritten name, which is what a client's hint is), original, or any. In any, the default, a name that is an original column name picks that column and only a name that is not one is read as a rewritten name, so names that never collided pick what they picked before and no existing test changed. StatState(namespace=...) sets it for the group and priority names, and StatRun.next_unit and run_next take namespace for prefer. An unknown namespace raises ValueError. Priority and prefer names are resolved against every column of the frame rather than the group's columns, so a name picks the same column whatever group is asked for. prioritized now takes the state because of that.

I chose an explicit namespace over a fixed rule because neither order of precedence is right for both callers: a client's c means the rewritten name, a Python caller's c means the column called c. The request branch of the next phase should pass namespace="rewritten" for the client's hint.

Not changed. skip_columns still matches original or rewritten names on pandas and polars and original names on xorq. That is the matching those backends have on main, so it is outside this PR.

Verification. After the fix, the reviewer's two scripts pick the single column asked for (prefer picks a for rewritten a, a group of ('c',) plans only c, a group of ('a',) on b, a, c plans only a). The full unit suite gives 1393 passed and 5 skipped locally, the previous 1361 plus the 32 new tests, with ruff and paddy_format clean. On fa05adba all 28 checks completed: 27 succeeded (one is the Read the Docs status) and deploy was skipped, including all nine Python / Test jobs. Nothing was declined.