feat(stats): stats_request runs units for a time budget, with per-client cursors and the final assignment (rows-first s5)

[paddymul]

Stacked on #1026 (feat/rowsfirst-s4-stat-units), which is stacked on #1024, #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, #1024 and #1026 until they merge. The commits of this phase are 2a73c8ed (failing tests) and ad26fdf4 (implementation); read only those two (git diff origin/feat/rowsfirst-s4-stat-units...HEAD).

Problem

stats_request (#1024) is a whole run: one synchronous call that computes every stat before it replies. On a long xorq batch it holds the loop for as long as that takes, and no client sees a stat before the last one is done. #1026 cut a run into units that can be run one at a time, and added the per-session StatRun and the per-connection StatCursor, but nothing in the server calls them. The request branch has to run units for a bounded time per request, answer each client with the fragments it has not seen, and do the final assignment when the last unit completes.

Phase and plan references

Rows-first s5, from buckaroo2-reports/plans/: plan 2 (02-rows-first-xorq-and-lazy-polars.md) section 7 "Phase 3" with sections 4.1 (Request, Final assignment) and 4.3 (reload, telemetry), and plan 1 (01-rows-first-stats-separate-plumbing.md) section 8 (telemetry).

Approach

• Two shapes of stats_request. {stats_gen, scope, columns?} is still the whole run and answers as in feat(server): stats_request, stats_update and df_meta.stats on deferred sessions (rows-first s3) #1024. {stats_gen, scope, columns?, incremental: true} is a time-boxed step. A value other than a boolean for incremental is stats_aborted with reason bad_request, and is not run as a whole run.
• An incremental step (stats_wire._serve_units). It takes the session's StatRun for the current generation (start_stat_run, planning only). If this client's cursor is behind the run, the reply carries the fragments it has not seen and no unit runs. Otherwise run_units runs units until STATS_BUDGET_S (75 ms) is spent, at least one, so a unit that is a query is the only one of its request and units that are snapshot-cache hits share a request. columns are the grid's column names (a, b, c) and put the units that cover them first. The reply is stats_update {stats_gen, scope, tier, final: false, remaining, payload, elapsed_ms}. payload is an inline wide DFEnvelope of the columns the unseen fragments cover, assembled the way merged_sd is (CustomizableDataflow._assemble_merged_sd(running_sd), the body of the merged_sd observer with the run's accumulating sd standing in for the scopes that share the filt chain), so init_sd, a processing step's sd and the cleaned_* and filtered_* layers apply. A client merges it key by key.
• Final assignment (complete_stats, shared by every path). When the last unit has run, in whichever request ran it: write the sd into summary_stats_cache under the full-tier key, then assign summary_sd; refresh the session snapshot through refresh_session_snapshot and set the status to complete in the same step; free the run. A state whose full sd is already in the cache runs no unit. The reply that follows is final: true, remaining: 0, carries the complete all_stats, and carries the rebuilt df_display_args when its digest differs from the one the client holds (a float column's minWidth is 114 with full stats and 78 without in the test fixture). The digest of the config a stats-free frame carried is kept on the handler (display_args_hash, set by build_state_message_for). The client's own search highlight is applied to the config it gets, since the config replaces the one that carries it.
• Whole run and legacy clients. With no StatRun for the generation, complete_stats is one _get_summary_sd call, as before: it runs the same units in the same order, and it is the hook the xorq dataflow customizes. With a run that a client has made progress on, it runs the units left and assigns the run's results, so no unit runs twice. A legacy client's complete frame does that too.
• Per-handler cursors. Each connection reads the shared fragment list at its own pace. Two clients that take turns each get the other's fragments as a catch-up that runs nothing, and each unit runs once.
• Failure. A unit that raises fails the generation (stats_status error, reason stats_failed, runs dropped). It is not retried by the next request, and the next generation starts clean. A run that cannot be planned (a post-processor that fails leaves an error frame, which the xorq stats class cannot plan) falls back to the whole-run path, with a warning.
• Telemetry. DataStreamHandler._handle_stats_request already binds session.tele_sink. New: a stats.unit span per unit (session, stats_gen, unit, phase, cost, columns); firstpull.stats_total once per completed run, now with units, run_secs and, on xorq, the cache_status, cache_hits, cache_misses, cache_secs, cache_snapshots, cache_bytes and cache_write_errors that firstpull.summary_stats carries for a whole run; the stats.request span gains final, remaining and units. StatRun sums elapsed_s over its units.
• Reload. /reload_expr bumps stats_gen and drops the run (begin_stats_generation, from feat(server): stats_request, stats_update and df_meta.stats on deferred sessions (rows-first s3) #1024 and feat(stats): resumable stat units and the StatRun store (rows-first s4) #1026); a request for the old generation gets stats_aborted with no query, and the new generation starts from its first unit.

What changes

• buckaroo/server/stats_wire.py: STATS_BUDGET_S, run_units, partial_payload, display_args_hash, highlighted_display_args (the overlay's loop, moved here so the final reply can use it), assign_full_stats(dataflow, computed=None), complete_stats finishing a run, _serve_units, handle_stats_request(session, msg, client=None), the bad_request abort, and the digest kept in build_state_message_for.
• buckaroo/server/stat_run.py: StatRun.elapsed_s and StatRun.errs().
• buckaroo/dataflow/dataflow.py: _assemble_merged_sd(running_sd=None); the _merged_sd observer calls it.
• buckaroo/server/websocket_handler.py: passes the handler to handle_stats_request, initializes display_args_hash, and _send_highlight_overlay uses highlighted_display_args.

Tests

Two commits: 2a73c8ed (failing tests) and ad26fdf4 (implementation).

• test_load_expr.py, TestStatsWire (xorq, over WebSocket): one unit per request and the complete state at the end; the last unit assigns the stats and frees the run; a long budget runs every unit in one request; the columns hint orders units; a stale generation mid-run; client B opening after client A finished; two interleaved clients end with full stats and run no unit twice; the final reply carries the rebuilt df_display_args when the digest differs, carries none when it does not (an int and string table), and keeps the client's highlight; a filtered state ends as the inline session's does; a legacy client connecting mid-run finishes only the units left; a request without the flag finishes a run in progress; a concurrent infinite_request is served between units, with every request under plan 1's 250 ms ceiling (a fake unit that sleeps 40 ms); the first initial_state and the first rows precede every data stat query and rows are served before stats complete; /reload_expr mid-run; the stats.unit, firstpull.stats_total and stats.request spans on the session sink; the completion span's cache outcome (a miss, then a hit after a reload on a cache_storage_path); a unit that raises; a non-boolean flag.
• test_data_loading_polars.py (pandas and polars): TestRunUnits (units that cost little run together, a request stops once the budget is spent, a unit over the budget is alone in its request, a cold unit after warm ones ends its request, one unit runs whatever the budget, no budget runs every unit, the hint in rewritten names, a raising unit, the spans, elapsed_s, errs()), run on real units with a fake clock and a fake cost; TestPartialPayload and TestAssembledMergedSd (the payload of every fragment equals the dataflow's all_stats, only the covered columns, filtered_* keys for a search, init_sd wins); TestHighlightedDisplayArgs and TestDisplayArgsHash.
• Three guard tests were added with the implementation because they pass on the earlier code: an incremental request for a cached state runs no unit, an incremental request that cannot plan units takes the whole-run path (and logs why), and a columns hint that is not a list of names is only a hint.
• Two things in the failing-tests commit were corrected in the implementation commit, both mistakes in the tests: next_unit takes plan order among the preferred columns, not the order of prefer (test_the_columns_hint_is_read_in_the_client_s_rewritten_names assumed the latter), and a decoded all_stats row has a level_0 key that the column sets must skip, as StatsChannel's merge does. One test also dropped a legacy-client check that could not hold (every send to a legacy client completes the session, so with one connected the session is never pending after a push).

On 2a73c8ed (tests only) all nine Python / Test jobs failed (3.11 to 3.14, Max Versions 3.11 to 3.14, Windows), all 28 checks completed. Locally 63 of the new tests fail on the earlier code: 43 in test_data_loading_polars.py, each on a missing attribute, and 20 in TestStatsWire, each on an assertion about the reply (the earlier code ignores incremental and answers with the whole run) or a missing key. CI logs were not read (the REST log API is rate-limited for this account), so the reasons rest on the local run of the same commit. The 174 tests that were in the two files before pass unchanged.

CI on ad26fdf4: 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.

Locally, on ad26fdf4, the full unit suite (pytest ./tests/unit -m "not slow") gives 1459 passed and 5 skipped, the 1393 of #1026 plus 66 new tests, and the same 1459 in an environment resolved with the newest versions and pandas 3 (pandas 3.0.6, polars 1.44.2, xorq 0.4.5, numpy 2.5.3). basedpyright on the scoped files reports 0 errors and 0 warnings. Nineteen deliberate regressions of the implementation each make at least one test fail: the budget ignored, the budget checked before the first unit, no catch-up guard, complete_stats ignoring the run, the display config never attached, always attached, the run not freed, the partial payload not restricted to the covered columns, no stats.unit span, a failure that keeps the runs, no flag validation, no cache check before planning, the running sd not used for the raw scope, the final assignment recomputing, the digest not recorded, the columns hint ignored, elapsed_s not summed, a missing units attribute on stats.request, and the wrong tier on a partial reply.

Measurements

Each incremental request timed through handle_stats_request (no WebSocket), from a schema-tier session to the final reply, on the 27-column synthetic parquet fixture of #1026 (4 ints, 6 floats with nulls, 2 low-cardinality ints, 8 strings of 5 to 5,000 distinct values, 3 booleans, 4 timestamps), the default 75 ms budget, an Apple M4 Pro with other agents running. The final request includes its own units, the final assignment and the snapshot refresh.

Fixture	Batch	Requests	Median non-final request	Longest non-final request	Final request	Sum
52,814 rows	one batch	4	0.087 s	0.087 s	0.084 s	0.341 s
52,814 rows	stat_chunk_cells=1M	5	0.079 s	0.095 s	0.038 s	0.369 s
2,000,000 rows	one batch	10	0.090 s	0.860 s (the batch)	0.055 s	1.636 s
2,000,000 rows	stat_chunk_cells=12M (6-column chunks)	13	0.103-0.105 s	0.252-0.253 s (two runs)	0.029-0.031 s	1.636 s

A request is as long as its longest unit plus the partial payload, so the bound a waiting infinite_request sees is the unit. With the batch whole at 2M rows that is 0.86 s; with 6-column chunks it is 0.25 s, at plan 1's proposed 250 ms ceiling and a few milliseconds over it. The final request's own work (the assignment and the cascade it fires) is small here, 30 to 80 ms with its units; a state with a cleaning op or a search computes more in it (see "Not in this PR").

Why default behaviour is unchanged

• Everything new is opt-in. Only a client that sends ?caps=stats_update to a session loaded with stats_delivery="deferred" can send a stats_request that matters, and only incremental: true runs units for a budget. An inline session is never pending, so none of this runs for it. Splitting the merged_sd observer into _assemble_merged_sd changes no value: the observer calls it with no running sd, and every existing test of merged_sd passes unchanged.
• A request without the flag answers as in feat(server): stats_request, stats_update and df_meta.stats on deferred sessions (rows-first s3) #1024: stats_update {type, stats_gen, scope, tier, final: true, payload, elapsed_ms}, with two additive fields, remaining: 0 and, for a client whose pending frame carried a different display config, df_display_args. A client that does not read them ignores them.
• _send_highlight_overlay builds the same message from the same loop, which now lives in highlighted_display_args.
• No existing test changed. The edits to the test files are added tests and the corrections named above to tests added in this PR.

Deviations from the plan

• The plan has every stats_request run units for 50 to 100 ms. Here only incremental: true does, because the whole-run request exists in feat(server): stats_request, stats_update and df_meta.stats on deferred sessions (rows-first s3) #1024 with its own tests and clients, and the maintainer's rule is that default behaviour does not change. The unflagged request is the whole run, as the phase asks ("the whole-run request becomes run all units"): it runs every unit still to run.
• With no StatRun the whole run is one _get_summary_sd call (the same units in the same order, shown by feat(stats): resumable stat units and the StatRun store (rows-first s4) #1026) instead of a loop over a run's units. _get_summary_sd is the hook the xorq dataflow overrides (the pandas error-frame result, the debug-mode raise), and the feat(server): stats_request, stats_update and df_meta.stats on deferred sessions (rows-first s3) #1024 tests patch it. With a run that has progress it runs the units left.
• The final reply carries the complete all_stats, not only what lies beyond the client's cursor. The final assignment also fills layers that no fragment carries (the other scopes' cleaned_* and filtered_* keys, an init_sd or processing sd override), and a client that missed a frame then still ends complete.
• The budget is a server constant (STATS_BUDGET_S = 75 ms). The plan's 250 ms ceiling is a bound on the longest request, and a request cannot be cut short inside a unit, so the ceiling is what a test asserts, not something the server enforces.
• columns are the grid's rewritten names, and the first unit that covers any of them runs first (plan order among them, as StatRun.next_unit does), not in the order the client lists them.
• stats_aborted gains the reason bad_request.
• firstpull.stats_total is not emitted as firstpull.summary_stats. That span's duration is the time spent computing stats, which a run spread over several requests has no single value for. The attributes are the same and sit on firstpull.stats_total, whose duration is the request that completed the run, with run_secs for the sum of the units.
• The /load_expr snapshot copy (the third site the plan names) is not replaced by refresh_session_snapshot. The helper applies the session's stored component_config, and that copy applies only the one in the POST, so replacing it would change what a re-POST that omits the field sends. feat(server): stats_request, stats_update and df_meta.stats on deferred sessions (rows-first s3) #1024 left it for the same reason.
• Counters for updates sent, requests dropped as stale and errors are the outcome attribute of stats.request, as in feat(server): stats_request, stats_update and df_meta.stats on deferred sessions (rows-first s3) #1024.

Not in this PR

• Making deferred the default, the filtered-count memoization and the sort guards (plan 3 phase 7), the scalar and policy tiers, client code, a child process.
• A state with a cleaning op or a search filter still computes its raw and clean scopes whole inside the final assignment, because assigning summary_sd fills the other scopes' cache entries through the cascade. For xorq that is a full stats pass per scope in the request that finishes the run. Units for those scopes belong with plan 3's scoped runs.
• A measurement of the final assignment's own cost (the cascade it triggers) on a wide table; it is one request, but the figure is not taken here.

Stack

#1021, #1022, #1024 and #1026 are underneath; their commits are in this diff until they merge. Commits of this phase: 2a73c8ed (failing tests) and ad26fdf4 (implementation).

🤖 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.dev37188764309

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.dev37188764309

MCP server for Claude Code

claude mcp add buckaroo-table -- uvx --from "buckaroo[mcp]==0.15.9.dev37188764309" --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]

Closing. This is the client-pull loop (incremental stats_request, a 75 ms budget per request, per-connection cursors) that the revised D2 of ADR-002 (#1043, 7943412) replaces with a server push after the first row reply.

One piece carries over: attaching the rebuilt df_display_args to the final reply when its hash changed (display_args_hash, highlighted_display_args), which the revised D5 still calls for. The branch is kept for that.