From a9419af0c1d11f8137b61d0f8246cdc484954268 Mon Sep 17 00:00:00 2001 From: Paddy Mullen Date: Sat, 3 Oct 2026 22:51:32 -0400 Subject: [PATCH 01/18] test(dataflow): characterize widget trait order and merged_sd before the _handle_widget_change split (rows-first s1) A default-tier XorqBuckarooWidget and BuckarooWidget publish df_data_dict, then df_display_args, then the rest of the widget_args_tuple observers on a search change, and merged_sd carries the full stat set. These pass on main and pin the behaviour the split must keep. Co-Authored-By: Claude Sonnet 5.5 --- tests/unit/basic_widget_test.py | 35 +++++++++++++ tests/unit/test_xorq_buckaroo_widget.py | 69 +++++++++++++++++++++++++ 2 files changed, 104 insertions(+) diff --git a/tests/unit/basic_widget_test.py b/tests/unit/basic_widget_test.py index a028751ed..9e63765e1 100644 --- a/tests/unit/basic_widget_test.py +++ b/tests/unit/basic_widget_test.py @@ -288,6 +288,41 @@ def counting_handle(*args, **kwargs): ) +def test_default_tier_widget_change_sequence_and_merged_sd(): + """Characterization for the ``_handle_widget_change`` split (rows-first + s1): a default-tier pandas widget keeps the order in which a search change + reaches the dataflow's and widget's traits, and its merged_sd.""" + bw = BuckarooWidget(pd.DataFrame({ + 'price': [12.5, 18.9, 7.4, 22.1, 14.0], 'qty': [1, 2, 1, 3, 2], + 'category': ['a', 'b', 'a', 'c', 'b']})) + msd = bw.dataflow.merged_sd + assert {col: (sd['length'], sd['_type'], sd['dtype']) for col, sd in msd.items()} == { + 'a': (5, 'float', 'float64'), 'b': (5, 'integer', 'int64'), + 'c': (5, 'string', 'object')} + assert msd['a']['mean'] == pytest.approx(14.98) + + seen = [] + bw.dataflow.observe(lambda c: seen.append(('dataflow', c['name'])), names=[ + 'summary_sd', 'merged_sd', 'widget_args_tuple', 'df_data_dict', + 'df_display_args', 'df_meta', 'operations', 'raw_sd_key', 'clean_sd_key', + 'filt_sd_key', 'summary_stats_cache', 'processed_result', 'cleaned']) + bw.observe(lambda c: seen.append(('widget', c['name'])), + names=['df_data_dict', 'df_display_args', 'df_meta', 'operations']) + bw.dataflow.quick_command_args = {'search': ['a']} + assert seen == [ + ('widget', 'df_data_dict'), ('dataflow', 'df_data_dict'), + ('widget', 'df_display_args'), ('dataflow', 'df_display_args'), + ('dataflow', 'widget_args_tuple'), ('dataflow', 'merged_sd'), + ('dataflow', 'summary_sd'), ('dataflow', 'processed_result'), + ('widget', 'df_meta'), ('dataflow', 'df_meta'), ('dataflow', 'cleaned'), + ('dataflow', 'summary_stats_cache'), ('dataflow', 'widget_args_tuple'), + ('dataflow', 'merged_sd'), ('dataflow', 'filt_sd_key'), + ('widget', 'operations'), ('dataflow', 'operations')] + assert {col: (sd['length'], sd['filtered_length']) + for col, sd in bw.dataflow.merged_sd.items()} == { + 'a': (5, 2), 'b': (5, 2), 'c': (5, 2)} + + def test_multi_index_cols() -> None: df = get_multiindex_cols_df() bw = BuckarooWidget(df) diff --git a/tests/unit/test_xorq_buckaroo_widget.py b/tests/unit/test_xorq_buckaroo_widget.py index 196175748..9e8f4486f 100644 --- a/tests/unit/test_xorq_buckaroo_widget.py +++ b/tests/unit/test_xorq_buckaroo_widget.py @@ -60,6 +60,75 @@ def counting(self, q): ) +_CHANGE_ORDER_TRAITS = [ + "summary_sd", "merged_sd", "widget_args_tuple", "df_data_dict", + "df_display_args", "df_meta", "operations", "raw_sd_key", "clean_sd_key", + "filt_sd_key", "summary_stats_cache", "processed_result", "cleaned"] + +# Order in which the dataflow's observers see a search change on a default +# (full-tier) XorqBuckarooWidget, recorded before ``_handle_widget_change`` is +# split into separately callable halves. The split must leave it unchanged: +# df_data_dict, then df_display_args, then widget_args_tuple's other observers. +_SEARCH_CHANGE_ORDER = [ + ("widget", "df_data_dict"), ("dataflow", "df_data_dict"), + ("widget", "df_display_args"), ("dataflow", "df_display_args"), + ("dataflow", "widget_args_tuple"), ("dataflow", "merged_sd"), + ("dataflow", "summary_sd"), ("dataflow", "processed_result"), + ("widget", "df_meta"), ("dataflow", "df_meta"), ("dataflow", "cleaned"), + ("dataflow", "summary_stats_cache"), ("dataflow", "widget_args_tuple"), + ("dataflow", "merged_sd"), ("dataflow", "filt_sd_key"), + ("widget", "operations"), ("dataflow", "operations")] + +_FULL_TIER_STAT_KEYS = { + "a": ["_type", "distinct_count", "distinct_per", "dtype", "histogram", + "histogram_bins", "is_bool", "is_datetime", "is_float", "is_integer", + "is_numeric", "is_string", "length", "max", "mean", "median", "min", + "nan_per", "non_null_count", "null_count", "orig_col_name", + "rewritten_col_name", "std"]} + + +class TestWidgetChangeSequence: + """Characterization of what a default-tier widget publishes and in what + order, so the ``_handle_widget_change`` split (rows-first s1) can be shown + to leave a full-stats widget unchanged.""" + + def _widget(self): + return XorqBuckarooWidget(xo.memtable( + {"price": [12.5, 18.9, 7.4, 22.1, 14.0], "qty": [1, 2, 1, 3, 2], + "category": ["a", "b", "a", "c", "b"]})) + + def test_search_change_keeps_trait_order(self): + w = self._widget() + seen = [] + w.dataflow.observe( + lambda c: seen.append(("dataflow", c["name"])), names=_CHANGE_ORDER_TRAITS) + w.observe( + lambda c: seen.append(("widget", c["name"])), + names=["df_data_dict", "df_display_args", "df_meta", "operations"]) + w.dataflow.quick_command_args = {"search": ["a"]} + assert seen == _SEARCH_CHANGE_ORDER + + def test_merged_sd_and_display_state_unchanged(self): + w = self._widget() + msd = w.dataflow.merged_sd + assert sorted(msd["a"]) == _FULL_TIER_STAT_KEYS["a"] + assert {col: (sd["length"], sd["_type"], sd["dtype"]) for col, sd in msd.items()} == { + "a": (5, "float", "float64"), "b": (5, "integer", "int64"), + "c": (5, "string", "string")} + assert msd["a"]["mean"] == pytest.approx(14.98) + assert list(w.df_data_dict) == ["main", "all_stats", "empty"] + main = w.df_display_args["main"] + assert [pr["primary_key_val"] for pr in main["df_viewer_config"]["pinned_rows"]] == [ + "dtype", "histogram"] + assert [cc["col_name"] for cc in main["df_viewer_config"]["column_config"]] == [ + "a", "b", "c"] + + w.dataflow.quick_command_args = {"search": ["a"]} + msd = w.dataflow.merged_sd + assert {col: (sd["length"], sd["filtered_length"]) for col, sd in msd.items()} == { + "a": (5, 2), "b": (5, 2), "c": (5, 2)} + + class TestExprCountMemoization: """Regression for #795: ``_expr_count`` issued a fresh ``expr.count().execute()`` against the backend on every call, even From 9b96220122c42bdac52b0f081a856885c3bb030e Mon Sep 17 00:00:00 2001 From: Paddy Mullen Date: Sat, 3 Oct 2026 22:56:48 -0400 Subject: [PATCH 02/18] test(xorq): failing tests for the xorq schema stats tier and its load-handler fields (rows-first s1) A schema-tier XorqServerDataflow should match full stats on pinned_rows, data_key, summary_stats_key and (except the stats-derived minWidth) column_config, issue no data query besides the cached count, and keep init_sd hints and sorted windows working. A pending state must write nothing under a full-tier cache key, and a later full assignment must reach merged_sd for the raw, clean and filt scopes. assemble_merged_sd must equal merged_sd, and _handle_widget_change must be built from separately callable all_stats and display-args builders. /load_expr and /reload_expr accept stats_tier and stats_delivery, replay them on reload, and keep them out of the warm short-circuit's has_config tuple. Co-Authored-By: Claude Sonnet 5.5 --- .../dataflow/customizable_dataflow_test.py | 19 + .../dataflow/scoped_summary_stats_test.py | 67 +++ tests/unit/server/test_load_expr.py | 432 ++++++++++++++++++ 3 files changed, 518 insertions(+) diff --git a/tests/unit/dataflow/customizable_dataflow_test.py b/tests/unit/dataflow/customizable_dataflow_test.py index 09b36eb70..909302c7a 100644 --- a/tests/unit/dataflow/customizable_dataflow_test.py +++ b/tests/unit/dataflow/customizable_dataflow_test.py @@ -75,6 +75,25 @@ def test_widget_instatiation(): assert dfc.df_data_dict['main'] == BASIC_DF_JSON_DATA assert dfc.df_display_args['main']['df_viewer_config'] == DFVIEWER_CONFIG_DEFAULT +def test_all_stats_and_display_args_are_built_separately(monkeypatch): + """``_handle_widget_change`` is the composition of separately callable + builders (rows-first s1): the stats payload can be rebuilt without + rebuilding column_config, and column_config without serializing stats.""" + dfc = ACDFC(BASIC_DF) + _unused, processed_df, merged_sd = dfc.widget_args_tuple + serialized = [] + original = dfc._sd_to_jsondf + monkeypatch.setattr(dfc, '_sd_to_jsondf', lambda sd: serialized.append(sd) or original(sd)) + + display_args = dfc._build_df_display_args(processed_df, merged_sd) + assert serialized == [], "building display args must not serialize all_stats" + assert display_args == dfc.df_display_args + + data_dict = dfc._build_df_data_dict(processed_df, merged_sd) + assert serialized == [merged_sd] + assert data_dict == dfc.df_data_dict + + def test_widget_operations_instatiation(): dfc = ACDFC(BASIC_DF) # dfc starts with operations of [{'meta':'no-op'], but the first diff --git a/tests/unit/dataflow/scoped_summary_stats_test.py b/tests/unit/dataflow/scoped_summary_stats_test.py index b5648ab56..a8b0a5339 100644 --- a/tests/unit/dataflow/scoped_summary_stats_test.py +++ b/tests/unit/dataflow/scoped_summary_stats_test.py @@ -15,6 +15,8 @@ observable in ``merged_sd``. """ +import copy + import pandas as pd from typing import Any, TypedDict @@ -23,6 +25,7 @@ from buckaroo.customizations.pd_stats_v2 import PD_ANALYSIS_V2, PD_AUTOCLEAN_DEFAULT_V2, cleaning_gen_ops from buckaroo.dataflow.autocleaning import (AutocleaningConfig, PandasAutocleaning) from buckaroo.dataflow.dataflow import CustomizableDataflow, StylingAnalysis +from buckaroo.dataflow.sd_cache import split_chain_by_scope from buckaroo.pluggable_analysis_framework.stat_func import stat @@ -249,3 +252,67 @@ def test_analysis_klasses_change_invalidates_scoped_sd(): f"got the same key {key1} for both — likely the cache key still " f"omits analysis_klasses identity (Codex P2)" ) + + +def _scope_inputs(dfc): + """The pieces ``_merged_sd`` assembles, read back from the dataflow.""" + cache = dfc.summary_stats_cache + return dict( + init_sd=dfc.init_sd, cleaned_sd=dfc.cleaned_sd, raw_sd=cache[dfc.raw_sd_key], + processed_sd=dfc.processed_sd, processed_df=dfc.processed_df, + chains=split_chain_by_scope(dfc.operations), + clean_sd=cache[dfc.clean_sd_key], filt_sd=cache[dfc.filt_sd_key]) + + +def _assembly_dataflow(): + df = pd.DataFrame({'amount': ['10', '20', '30', '40', '50', '60', '70'], + 'word': ['foo', 'bar', 'foo', 'baz', 'foo', 'bar', 'foo']}) + # init_sd is keyed by the original column name; 'amount' is rewritten to 'a'. + dfc = ScopedDataflow(df, init_sd={ + 'amount': {'displayer_args': {'displayer': 'string', 'max_length': 99}, 'init_only': 1}}) + dfc.cleaning_method = 'default' + dfc.quick_command_args = {'search': ['foo']} + return dfc + + +def test_assemble_merged_sd_equals_merged_sd_with_init_sd_and_cleaning(): + """``assemble_merged_sd`` is the body of the ``merged_sd`` observer as a + pure function, so a server can assemble wire stats from an accumulating + raw sd without going through the traitlets cascade.""" + from buckaroo.dataflow.dataflow import assemble_merged_sd + dfc = _assembly_dataflow() + sd = dfc.merged_sd['a'] + assert sd['init_only'] == 1, "init_sd override must reach merged_sd under the rewritten name" + assert 'cleaned_mean' in sd and 'filtered_mean' in sd, "cleaning and filter layers must be active" + + assembled = assemble_merged_sd(**_scope_inputs(dfc)) + + assert assembled == dfc.merged_sd + + +def test_assemble_merged_sd_without_a_processed_df_merges_the_unrewritten_sds(): + from buckaroo.dataflow.dataflow import assemble_merged_sd + assembled = assemble_merged_sd( + init_sd={'x': {'init': 1}}, cleaned_sd={'x': {'cleaned': 2}}, raw_sd={'x': {'raw': 3}}, + processed_sd={'x': {'processed': 4}}, processed_df=None, + chains={'raw': [], 'clean': [], 'filt': []}) + assert assembled == {'x': {'init': 1, 'cleaned': 2, 'raw': 3, 'processed': 4}} + + +def test_assemble_merged_sd_does_not_mutate_its_inputs(): + from buckaroo.dataflow.dataflow import assemble_merged_sd + dfc = _assembly_dataflow() + inputs = _scope_inputs(dfc) + sd_names = ('init_sd', 'cleaned_sd', 'raw_sd', 'processed_sd', 'clean_sd', 'filt_sd') + + def key_sets(): + return {name: {col: sorted(stats) for col, stats in inputs[name].items()} for name in sd_names} + + before, chains_before = key_sets(), copy.deepcopy(inputs['chains']) + + assemble_merged_sd(**inputs) + + # Layering cleaned_* / filtered_* into a result column that shared its dict + # with an input would add keys to that input. + assert key_sets() == before + assert inputs['chains'] == chains_before diff --git a/tests/unit/server/test_load_expr.py b/tests/unit/server/test_load_expr.py index bdbb47db6..03feb571e 100644 --- a/tests/unit/server/test_load_expr.py +++ b/tests/unit/server/test_load_expr.py @@ -25,6 +25,9 @@ from xorq.common.utils.provenance_utils import read_parquet_provenance # noqa: E402 from xorq.expr.relations import CachedNode # noqa: E402 +from buckaroo.dataflow.sd_cache import split_chain_by_scope # noqa: E402 +from buckaroo.jlisp.lisp_utils import s as lisp_sym # noqa: E402 +from buckaroo.pluggable_analysis_framework.col_analysis import ColAnalysis # noqa: E402 from buckaroo.server import telemetry, xorq_loading # noqa: E402 from buckaroo.server.app import make_app as _make_app # noqa: E402 @@ -910,6 +913,435 @@ def tracking_default_backend(): shutil.rmtree(builds_root, ignore_errors=True) +def _stats_tier_expr(): + """A 5-row table whose float column has a 1e9 maximum, so its estimated + column width depends on the min and max that only full stats supply.""" + return xo.memtable({ + "price": [12.5, 18.9, 7.4, 22.1, 1e9], + "qty": [1, 2, 1, 3, 2], + "category": ["a", "b", "a", "c", "b"]}) + + +def _build_dataflow(expr=None, **kwargs): + return xorq_loading.XorqServerDataflow( + _stats_tier_expr() if expr is None else expr, skip_main_serial=True, **kwargs) + + +def _three_scope_dataflow(expr=None, **kwargs): + """A dataflow with a user op (clean scope) and a search (filt scope) active, + so raw, clean and filt each hold their own summary-stats cache entry.""" + dataflow = _build_dataflow(expr, **kwargs) + dataflow.operations = [[lisp_sym("fillna"), {"symbol": "df"}, "qty", 0]] + dataflow.quick_command_args = {"search": ["a"]} + return dataflow + + +def _spy_data_queries(monkeypatch): + """Record the outermost op of every materialisation (``execute`` and + ``to_pyarrow``), and every stat query XorqStatPipeline sends.""" + from xorq.vendor.ibis.expr.types.core import Expr + + from buckaroo.pluggable_analysis_framework.xorq_stat_pipeline import XorqStatPipeline + ops, stat_queries = [], [] + original_execute, original_to_pyarrow = Expr.execute, Expr.to_pyarrow + original_stat_execute = XorqStatPipeline._execute + + def spy_execute(self, *args, **kwargs): + ops.append(type(self.op()).__name__) + return original_execute(self, *args, **kwargs) + + def spy_to_pyarrow(self, *args, **kwargs): + ops.append(type(self.op()).__name__) + return original_to_pyarrow(self, *args, **kwargs) + + def spy_stat_execute(self, query): + stat_queries.append(type(query.op()).__name__) + return original_stat_execute(self, query) + + monkeypatch.setattr(Expr, "execute", spy_execute) + monkeypatch.setattr(Expr, "to_pyarrow", spy_to_pyarrow) + monkeypatch.setattr(XorqStatPipeline, "_execute", spy_stat_execute) + return ops, stat_queries + + +def _without_min_width(column_config): + return [{**cc, "ag_grid_specs": {k: v for k, v in cc["ag_grid_specs"].items() if k != "minWidth"}} + for cc in column_config] + + +def _as_json(sd): + """Comparable form of an sd: json equates NaN with NaN where == would not.""" + return json.dumps(sd, sort_keys=True, default=str) + + +class _NoopPostProcessing(ColAnalysis): + provides_defaults = {} + post_processing_method = "noop_post" + + @classmethod + def post_process_df(cls, expr): + return [expr, {}] + + +class TestStatsTierSchema: + """``XorqServerDataflow(..., stats_tier="schema")`` (rows-first s1): the + dataflow publishes identity and typing for every column, with no data query + beyond the cached row count.""" + + def test_matches_full_stats_display_state(self): + expr = _stats_tier_expr() + full = _build_dataflow(expr) + schema = _build_dataflow(expr, stats_tier="schema") + assert schema.df_display_args.keys() == full.df_display_args.keys() + for name, full_arg in full.df_display_args.items(): + schema_arg = schema.df_display_args[name] + assert schema_arg["data_key"] == full_arg["data_key"] + assert schema_arg["summary_stats_key"] == full_arg["summary_stats_key"] + full_cfg, schema_cfg = full_arg["df_viewer_config"], schema_arg["df_viewer_config"] + assert schema_cfg["pinned_rows"] == full_cfg["pinned_rows"] + assert (_without_min_width(schema_cfg["column_config"]) + == _without_min_width(full_cfg["column_config"])) + + def test_min_width_is_the_stats_derived_difference(self): + expr = _stats_tier_expr() + widths = {} + for tier in ("full", "schema"): + cfg = _build_dataflow(expr, stats_tier=tier).df_display_args["main"]["df_viewer_config"] + widths[tier] = {cc["header_name"]: cc["ag_grid_specs"]["minWidth"] + for cc in cfg["column_config"]} + # price's 1e9 maximum widens it under full stats; without min and max + # the estimate falls back to a one-digit value. + assert widths["schema"]["price"] < widths["full"]["price"] + + def test_schema_sd_is_the_schema_keys_of_the_full_sd(self): + expr = _stats_tier_expr() + full_sd = _build_dataflow(expr).merged_sd + schema_sd = _build_dataflow(expr, stats_tier="schema").merged_sd + assert schema_sd.keys() == full_sd.keys() + for col, stats in schema_sd.items(): + assert {"orig_col_name", "rewritten_col_name", "dtype", "_type", "is_numeric", + "is_integer", "is_float", "is_bool", "is_datetime", "is_string", + "length"} <= stats.keys() + assert "mean" not in stats and "histogram" not in stats + assert stats == {k: full_sd[col][k] for k in stats} + + def test_issues_no_data_query_besides_the_cached_count(self, monkeypatch): + ops, stat_queries = _spy_data_queries(monkeypatch) + dataflow = _build_dataflow(stats_tier="schema") + dataflow.quick_command_args = {"search": ["a"]} + dataflow.add_analysis(_NoopPostProcessing) + assert stat_queries == [] + assert ops and set(ops) == {"CountStar"} + + def test_the_spy_sees_stat_queries_at_the_full_tier(self, monkeypatch): + ops, stat_queries = _spy_data_queries(monkeypatch) + _build_dataflow() + assert stat_queries, "the spy would not notice a stat query" + assert set(ops) - {"CountStar"} + + def test_init_sd_hints_and_overrides_still_apply(self): + dataflow = _build_dataflow( + stats_tier="schema", + init_sd={"qty": {"displayer_args": {"displayer": "string", "max_length": 200}}}, + column_config_overrides={ + "category": {"displayer_args": {"displayer": "string", "max_length": 5000}}}) + cfg = dataflow.df_display_args["main"]["df_viewer_config"]["column_config"] + by_header = {cc["header_name"]: cc for cc in cfg} + assert by_header["qty"]["displayer_args"]["max_length"] == 200 + assert by_header["category"]["displayer_args"]["max_length"] == 5000 + + def test_sorted_infinite_request_works(self): + dataflow = _build_dataflow(stats_tier="schema") + qty = next(k for k, v in dataflow.merged_sd.items() if v["orig_col_name"] == "qty") + resp, parquet = xorq_loading.handle_infinite_request_xorq( + dataflow, {"start": 0, "end": 5, "sourceName": "default", + "sort": qty, "sort_direction": "desc"}) + assert "error_info" not in resp + assert resp["length"] == 5 + assert pq.read_table(io.BytesIO(parquet)).column(qty).to_pylist() == [3, 2, 2, 1, 1] + + def test_pending_state_writes_no_full_tier_cache_key(self): + dataflow = _three_scope_dataflow(stats_tier="schema") + chains = split_chain_by_scope(dataflow.operations) + full_keys = {dataflow._scope_cache_key(chain, tier="full") for chain in chains.values()} + assert len(full_keys) == 3, "raw, clean and filt must each have their own key" + assert dataflow.summary_stats_cache + assert not full_keys & dataflow.summary_stats_cache.keys() + + def test_later_full_assignment_reaches_merged_sd_for_all_scopes(self): + expr = _stats_tier_expr() + full = _three_scope_dataflow(expr) + dataflow = _three_scope_dataflow(expr, stats_tier="schema") + assert "mean" not in dataflow.merged_sd["a"] + assert {"mean", "cleaned_mean", "filtered_mean"} <= full.merged_sd["a"].keys() + + dataflow.stats_tier = "full" + dataflow.summary_sd = full.summary_sd + + assert _as_json(dataflow.merged_sd) == _as_json(full.merged_sd) + + def test_a_summary_sd_from_another_tier_is_not_cached_under_the_new_tier(self): + """An sd computed at the schema tier must not become the full tier's + entry because the tier flipped before the next cascade: a present key + is a hit, so it would never be repaired.""" + dataflow = _three_scope_dataflow(stats_tier="schema") + stale = {col: {**stats, "stale": True} for col, stats in dataflow.summary_sd.items()} + + dataflow.stats_tier = "full" + dataflow.summary_sd = stale + + filt_sd = dataflow.summary_stats_cache[dataflow.filt_sd_key] + assert "mean" in filt_sd["a"] and "stale" not in filt_sd["a"] + assert "filtered_mean" in dataflow.merged_sd["a"] + + def test_full_scope_sds_cached_first_are_used_without_a_stat_query(self, monkeypatch): + expr = _stats_tier_expr() + full = _three_scope_dataflow(expr) + dataflow = _three_scope_dataflow(expr, stats_tier="schema") + cache = dict(dataflow.summary_stats_cache) + for scope, chain in split_chain_by_scope(dataflow.operations).items(): + cache[dataflow._scope_cache_key(chain, tier="full")] = full.summary_stats_cache[ + getattr(full, f"{scope}_sd_key")] + dataflow.summary_stats_cache = cache + _ops, stat_queries = _spy_data_queries(monkeypatch) + + dataflow.stats_tier = "full" + dataflow.summary_sd = full.summary_sd + + assert stat_queries == [] + assert _as_json(dataflow.merged_sd) == _as_json(full.merged_sd) + + +class TestLoadExprStatsPolicy(tornado.testing.AsyncHTTPTestCase): + """``stats_tier`` and ``stats_delivery`` on POST /load_expr and + /reload_expr (rows-first s1). The pair is stored beside dataflow_kwargs, + replayed by /reload_expr, and kept out of the has_config tuple.""" + + def get_app(self): + return make_app() + + def _session(self, sid): + return self._app.settings["sessions"].get(sid) + + @tornado.testing.gen_test + async def test_defaults_are_full_and_inline(self): + builds_root = tempfile.mkdtemp() + try: + build_path = _build_expr_dir(builds_root) + resp = await _post(self.get_http_port(), "/load_expr", + {"session": "sp-default", "build_dir": build_path}) + self.assertEqual(resp.code, 200) + session = self._session("sp-default") + self.assertEqual((session.stats_tier, session.stats_delivery), ("full", "inline")) + self.assertEqual(session.xorq_dataflow.stats_tier, "full") + self.assertIn("mean", session.xorq_dataflow.merged_sd["a"]) + finally: + shutil.rmtree(builds_root, ignore_errors=True) + + @tornado.testing.gen_test + async def test_schema_tier_builds_a_schema_dataflow(self): + builds_root = tempfile.mkdtemp() + try: + build_path = _build_expr_dir(builds_root) + resp = await _post(self.get_http_port(), "/load_expr", + {"session": "sp-schema", "build_dir": build_path, "stats_tier": "schema"}) + self.assertEqual(resp.code, 200) + self.assertEqual(json.loads(resp.body)["rows"], 10) + session = self._session("sp-schema") + self.assertEqual((session.stats_tier, session.stats_delivery), ("schema", "inline")) + self.assertEqual(session.xorq_dataflow.stats_tier, "schema") + self.assertNotIn("mean", session.xorq_dataflow.merged_sd["a"]) + finally: + shutil.rmtree(builds_root, ignore_errors=True) + + @tornado.testing.gen_test + async def test_deferred_delivery_publishes_a_schema_dataflow_and_serves_rows(self): + builds_root = tempfile.mkdtemp() + try: + build_path = _build_expr_dir(builds_root) + resp = await _post(self.get_http_port(), "/load_expr", + {"session": "sp-deferred", "build_dir": build_path, "stats_delivery": "deferred"}) + self.assertEqual(resp.code, 200) + session = self._session("sp-deferred") + # The session still targets full stats; the dataflow is built at + # the schema tier until a later phase delivers them. + self.assertEqual((session.stats_tier, session.stats_delivery), ("full", "deferred")) + self.assertEqual(session.xorq_dataflow.stats_tier, "schema") + + ws = await tornado.websocket.websocket_connect( + f"ws://localhost:{self.get_http_port()}/ws/sp-deferred") + initial = json.loads(await ws.read_message()) + self.assertEqual(initial["type"], "initial_state") + self.assertEqual(initial["df_meta"]["total_rows"], 10) + ws.write_message(json.dumps({"type": "infinite_request", + "payload_args": {"start": 0, "end": 10, "sourceName": "default", "origEnd": 10}})) + rows = json.loads(await ws.read_message()) + self.assertEqual(rows["length"], 10) + self.assertNotIn("error_info", rows) + table = pq.read_table(io.BytesIO(await ws.read_message())) + self.assertEqual(table.num_rows, 10) + ws.close() + finally: + shutil.rmtree(builds_root, ignore_errors=True) + + @tornado.testing.gen_test + async def test_deferred_delivery_runs_no_stat_query(self): + from buckaroo.pluggable_analysis_framework.xorq_stat_pipeline import XorqStatPipeline + builds_root = tempfile.mkdtemp() + stat_queries = [] + original = XorqStatPipeline._execute + + def spy(pipeline, query): + stat_queries.append(query) + return original(pipeline, query) + + try: + build_path = _build_expr_dir(builds_root) + with patch.object(XorqStatPipeline, "_execute", spy): + resp = await _post(self.get_http_port(), "/load_expr", + {"session": "sp-deferred-q", "build_dir": build_path, + "stats_delivery": "deferred"}) + self.assertEqual(resp.code, 200) + self.assertEqual(stat_queries, []) + finally: + shutil.rmtree(builds_root, ignore_errors=True) + + @tornado.testing.gen_test + async def test_invalid_values_are_a_400(self): + builds_root = tempfile.mkdtemp() + try: + build_path = _build_expr_dir(builds_root) + for field, code in (("stats_tier", "invalid_stats_tier"), + ("stats_delivery", "invalid_stats_delivery")): + resp = await _post(self.get_http_port(), "/load_expr", + {"session": "sp-bad", "build_dir": build_path, field: "sometimes"}) + self.assertEqual(resp.code, 400) + self.assertEqual(json.loads(resp.body)["error_code"], code) + self.assertIsNone(self._session("sp-bad")) + finally: + shutil.rmtree(builds_root, ignore_errors=True) + + @tornado.testing.gen_test + async def test_warm_repost_with_an_unchanged_pair_short_circuits(self): + builds_root = tempfile.mkdtemp() + try: + build_path = _build_expr_dir(builds_root) + body = {"session": "sp-warm", "build_dir": build_path, + "stats_tier": "schema", "stats_delivery": "deferred"} + self.assertEqual((await _post(self.get_http_port(), "/load_expr", body)).code, 200) + with patch.object(xorq_loading, "load_expr_build_dir", + side_effect=AssertionError("an unchanged pair must take the warm exit")): + same = await _post(self.get_http_port(), "/load_expr", body) + omitted = await _post(self.get_http_port(), "/load_expr", + {"session": "sp-warm", "build_dir": build_path}) + self.assertEqual(same.code, 200) + self.assertEqual(omitted.code, 200, "omitting the pair keeps the session's") + session = self._session("sp-warm") + self.assertEqual((session.stats_tier, session.stats_delivery), ("schema", "deferred")) + finally: + shutil.rmtree(builds_root, ignore_errors=True) + + @tornado.testing.gen_test + async def test_the_default_pair_sent_on_every_post_does_not_defeat_the_warm_exit(self): + """has_config tests truthiness, so a host that always sends the pair + would rebuild on every POST if the fields were in it (#944).""" + builds_root = tempfile.mkdtemp() + try: + build_path = _build_expr_dir(builds_root) + body = {"session": "sp-warm-default", "build_dir": build_path, + "stats_tier": "full", "stats_delivery": "inline"} + self.assertEqual((await _post(self.get_http_port(), "/load_expr", body)).code, 200) + with patch.object(xorq_loading, "load_expr_build_dir", + side_effect=AssertionError("the warm exit must not see the pair as config")): + resp = await _post(self.get_http_port(), "/load_expr", body) + self.assertEqual(resp.code, 200) + finally: + shutil.rmtree(builds_root, ignore_errors=True) + + @tornado.testing.gen_test + async def test_warm_repost_with_a_changed_pair_rebuilds(self): + builds_root = tempfile.mkdtemp() + try: + build_path = _build_expr_dir(builds_root) + sid = "sp-changed" + original = xorq_loading.load_expr_build_dir + calls = [] + + def counting_loader(bd, **kwargs): + calls.append(bd) + return original(bd, **kwargs) + + with patch.object(xorq_loading, "load_expr_build_dir", side_effect=counting_loader): + await _post(self.get_http_port(), "/load_expr", + {"session": sid, "build_dir": build_path}) + self.assertEqual(len(calls), 1) + resp = await _post(self.get_http_port(), "/load_expr", + {"session": sid, "build_dir": build_path, "stats_delivery": "deferred"}) + self.assertEqual(resp.code, 200) + self.assertEqual(len(calls), 2, "a changed delivery must rebuild") + self.assertEqual(self._session(sid).xorq_dataflow.stats_tier, "schema") + resp = await _post(self.get_http_port(), "/load_expr", + {"session": sid, "build_dir": build_path, + "stats_tier": "schema", "stats_delivery": "deferred"}) + self.assertEqual(len(calls), 3, "a changed tier must rebuild") + resp = await _post(self.get_http_port(), "/load_expr", + {"session": sid, "build_dir": build_path, + "stats_tier": "full", "stats_delivery": "inline"}) + self.assertEqual(len(calls), 4, "returning to the defaults must rebuild") + session = self._session(sid) + self.assertEqual((session.stats_tier, session.stats_delivery), ("full", "inline")) + self.assertEqual(session.xorq_dataflow.stats_tier, "full") + self.assertIn("mean", session.xorq_dataflow.merged_sd["a"]) + finally: + shutil.rmtree(builds_root, ignore_errors=True) + + @tornado.testing.gen_test + async def test_reload_expr_replays_the_stored_pair(self): + builds_root = tempfile.mkdtemp() + project_root = tempfile.mkdtemp() + try: + build_path = _build_expr_dir(builds_root) + sid = "sp-reload" + await _post(self.get_http_port(), "/load_expr", + {"session": sid, "build_dir": build_path, "project_root": project_root, + "stats_delivery": "deferred"}) + before = self._session(sid).xorq_dataflow + resp = await _post(self.get_http_port(), f"/reload_expr/{sid}", {}) + self.assertEqual(resp.code, 200) + session = self._session(sid) + self.assertIsNot(session.xorq_dataflow, before) + self.assertEqual(session.xorq_dataflow.stats_tier, "schema") + self.assertEqual((session.stats_tier, session.stats_delivery), ("full", "deferred")) + finally: + shutil.rmtree(builds_root, ignore_errors=True) + shutil.rmtree(project_root, ignore_errors=True) + + @tornado.testing.gen_test + async def test_reload_expr_accepts_a_new_pair_and_stores_it(self): + builds_root = tempfile.mkdtemp() + project_root = tempfile.mkdtemp() + try: + build_path = _build_expr_dir(builds_root) + sid = "sp-reload-new" + await _post(self.get_http_port(), "/load_expr", + {"session": sid, "build_dir": build_path, "project_root": project_root, + "stats_delivery": "deferred"}) + resp = await _post(self.get_http_port(), f"/reload_expr/{sid}", + {"stats_delivery": "inline"}) + self.assertEqual(resp.code, 200) + session = self._session(sid) + self.assertEqual((session.stats_tier, session.stats_delivery), ("full", "inline")) + self.assertEqual(session.xorq_dataflow.stats_tier, "full") + bad = await _post(self.get_http_port(), f"/reload_expr/{sid}", + {"stats_tier": "sometimes"}) + self.assertEqual(bad.code, 400) + self.assertEqual(json.loads(bad.body)["error_code"], "invalid_stats_tier") + self.assertEqual((session.stats_tier, session.stats_delivery), ("full", "inline")) + finally: + shutil.rmtree(builds_root, ignore_errors=True) + shutil.rmtree(project_root, ignore_errors=True) + + class TestReloadExpr(tornado.testing.AsyncHTTPTestCase): def get_app(self): return make_app() From 7f56c2019f9fe79b0bc0285b49373faa3a58b9cb Mon Sep 17 00:00:00 2001 From: Paddy Mullen Date: Sat, 3 Oct 2026 23:14:52 -0400 Subject: [PATCH 03/18] feat(xorq): stats tier machinery and the xorq schema tier (rows-first s1) Add a dataflow-level stats_tier ("full" default, "schema"). The xorq schema tier builds identity and typing for every column from the expression's schema, with no data query beyond the cached row count, so a dataflow constructs in milliseconds rather than the stats' hundreds. The tier is part of _scope_cache_key, so a schema entry is never read as a full one, and _populate_sd_cache stores summary_sd under the filt key only if it was computed for the current frame, klass list and tier. add_analysis no longer builds DFStatsClass outside the hook when the tier is not full. The merged_sd observer body is extracted as the pure assemble_merged_sd, and _handle_widget_change is split into _build_df_data_dict and _build_df_display_args. /load_expr and /reload_expr accept stats_tier and stats_delivery, stored on the session beside dataflow_kwargs and replayed on reload. They stay out of the has_config tuple; the warm short-circuit compares the stored pair. stats_delivery="deferred" builds the schema-tier dataflow and publishes it. Defaults (full, inline) leave behaviour unchanged. Co-Authored-By: Claude Sonnet 5.5 --- buckaroo/customizations/xorq_stats_v2.py | 10 + buckaroo/dataflow/dataflow.py | 235 +++++++++++++++-------- buckaroo/server/handlers.py | 73 ++++++- buckaroo/server/session.py | 18 ++ buckaroo/xorq_buckaroo.py | 19 +- 5 files changed, 273 insertions(+), 82 deletions(-) diff --git a/buckaroo/customizations/xorq_stats_v2.py b/buckaroo/customizations/xorq_stats_v2.py index fe613f302..7f5fcaa63 100644 --- a/buckaroo/customizations/xorq_stats_v2.py +++ b/buckaroo/customizations/xorq_stats_v2.py @@ -138,6 +138,16 @@ def _type(is_bool: bool, is_integer: bool, is_float: bool, is_datetime: bool, is return "obj" +def schema_stats(dtype: str) -> dict: + """The stats a column's dtype alone determines: ``dtype``, the ``is_*`` + flags and ``_type``. This is what the xorq dataflow's schema tier + publishes, derived through the same stats the pipeline runs.""" + flags = typing_stats(dtype) + return {"dtype": dtype, **flags, "_type": _type( + flags["is_bool"], flags["is_integer"], flags["is_float"], flags["is_datetime"], + flags["is_string"])} + + # ============================================================ # Batched aggregates — one ibis.Expr each, folded into the batch query # ============================================================ diff --git a/buckaroo/dataflow/dataflow.py b/buckaroo/dataflow/dataflow.py index 271e86ca4..68b7bd416 100644 --- a/buckaroo/dataflow/dataflow.py +++ b/buckaroo/dataflow/dataflow.py @@ -27,6 +27,63 @@ from .df_types import DataFrameT from .sd_cache import hash_chain, split_chain_by_scope +# Content levels a dataflow can be built at. ``full`` runs the summary-stats +# pipeline. ``schema`` publishes identity and typing keys derived from the +# schema alone and runs no data query (backends opt in, see ``_get_schema_sd``). +STATS_TIERS = ("full", "schema") + + +def assemble_merged_sd(init_sd: InitSD, cleaned_sd: SDType, raw_sd: SDType, processed_sd: SDType, + processed_df: TAny, chains: TDict[str, List[TAny]], + clean_sd: Optional[SDType] = None, filt_sd: Optional[SDType] = None) -> SDType: + """Build the wire ``merged_sd`` from the pieces ``_merged_sd`` reads. + + Bare keys come from the raw scope's sd (``raw_sd``, computed on + ``sampled_df``) merged under ``init_sd`` (keyed by original column name, + so rewritten here), ``cleaned_sd`` and ``processed_sd``. ``cleaned_*`` keys + are layered on top from ``clean_sd`` when cleaning is active, and + ``filtered_*`` keys from ``filt_sd`` when a search filter is active. + + Pure: the inputs are not mutated, so a caller holding an accumulating + ``raw_sd`` can assemble it repeatedly without the traitlets cascade. + ``chains`` is ``split_chain_by_scope(operations)``. + """ + if processed_df is None: + # on initial startup + return merge_sds(init_sd, cleaned_sd, raw_sd, processed_sd) + + # Gate each prefixed layer on a chain-shape diff between scopes. + # ``cleaned_*`` fires when the clean chain has ops the raw chain + # doesn't (cleaning is on); ``filtered_*`` fires when the filt + # chain has ops the clean chain doesn't (at least one + # quick-command op is present, i.e. a search filter is on). + # Keying off ``filt_sd_key != raw_sd_key`` would also fire + # ``filtered_*`` for cleaning-only states, mislabelling cleaned + # stats as filtered. + cleaning_active = chains['clean'] != chains['raw'] + filter_active = chains['filt'] != chains['clean'] + + #we do this to get rewrtten keys for init_sd + rewritten_init_sd = merge_sd_overrides({}, processed_df, init_sd) + intermediate_sd = merge_sds(rewritten_init_sd, cleaned_sd, raw_sd) + base = merge_sd_overrides(intermediate_sd, processed_df, processed_sd) + + # Layer ``cleaned_*`` keys on top when cleaning is active. + if cleaning_active and clean_sd: + for col, stats in clean_sd.items(): + col_dict = base.setdefault(col, {}) + for stat_name, val in stats.items(): + col_dict[f'cleaned_{stat_name}'] = val + + # Layer ``filtered_*`` keys on top when a filter is active. + if filter_active and filt_sd: + for col, stats in filt_sd.items(): + col_dict = base.setdefault(col, {}) + for stat_name, val in stats.items(): + col_dict[f'filtered_{stat_name}'] = val + + return base + class DfTrait(Any): """Any-trait for values that may contain a pandas/polars DataFrame. @@ -161,6 +218,10 @@ def __init__(self, raw_df): # are id()-based cache keys and the dead "foo"/"bar" test branches below, # all of which treat None and [] alike). analysis_klasses: List[Type[ColAnalysis]] = [] + # The content level the summary-stats cascade runs at (see STATS_TIERS). + # Instance-level on CustomizableDataflow; ``full`` here so the bare + # pipeline's dedupe key and cache keys can read it. + stats_tier: str = "full" summary_sd = Any() df_meta = Any() @@ -312,7 +373,7 @@ def _get_summary_sd(self, processed_df: DataFrameT) -> Tuple[SDType, TAny]: ret_summary[col] = {} return ret_summary, {} - _summary_sd_cache_key = (None, None) + _summary_sd_cache_key = (None, None, None) @observe('processed_result', 'analysis_klasses') @exception_protect('summary_sd-protector') @@ -322,7 +383,8 @@ def _summary_sd(self, change): # tuple wrapper, which makes this observer fire twice per widget # construction even when processed_df identity is unchanged. # Skip when neither the dataframe nor analysis_klasses has actually - # changed since the last run. See issue #709. + # changed since the last run. See issue #709. The tier is part of the + # key so a tier change reruns the summary even on the same frame. df = self.processed_df if df is None: # Nothing to summarize before the pipeline has produced a frame. @@ -330,9 +392,9 @@ def _summary_sd(self, change): # narrows Optional[DataFrameT] -> DataFrameT for _get_summary_sd. return klasses = self.analysis_klasses - if (id(df), id(klasses)) == self._summary_sd_cache_key: + if (id(df), id(klasses), self.stats_tier) == self._summary_sd_cache_key: return - self._summary_sd_cache_key = (id(df), id(klasses)) + self._summary_sd_cache_key = (id(df), id(klasses), self.stats_tier) result_summary_sd, errs = self._get_summary_sd(df) self.summary_sd = result_summary_sd self.errs = errs @@ -390,7 +452,11 @@ def __init__(self, orig_df, debug=False, pinned_rows:Union[Literal[None], PinnedRowConfig]=None, extra_grid_config=None, component_config:Union[Literal[None], ComponentConfig]=None, init_sd:Union[Literal[None], InitSD]=None, skip_main_serial=False, - skip_stat_columns=None): + skip_stat_columns=None, stats_tier:str="full"): + if stats_tier not in STATS_TIERS: + raise ValueError(f"stats_tier must be one of {STATS_TIERS}, got {stats_tier!r}") + # Set before super().__init__ — assigning raw_df runs the whole cascade. + self.stats_tier = stats_tier self.init_sd: InitSD if init_sd is None: self.init_sd = {} @@ -520,43 +586,9 @@ def _merged_sd(self, change): if filt_sd is None: filt_sd = self.summary_sd or {} - # Gate each prefixed layer on a chain-shape diff between scopes. - # ``cleaned_*`` fires when the clean chain has ops the raw chain - # doesn't (cleaning is on); ``filtered_*`` fires when the filt - # chain has ops the clean chain doesn't (at least one - # quick-command op is present, i.e. a search filter is on). - # Keying off ``filt_sd_key != raw_sd_key`` would also fire - # ``filtered_*`` for cleaning-only states, mislabelling cleaned - # stats as filtered. - chains = split_chain_by_scope(self.operations) - cleaning_active = chains['clean'] != chains['raw'] - filter_active = chains['filt'] != chains['clean'] - - if self.processed_df is None: - #on initial startup - self.merged_sd = merge_sds(self.init_sd, self.cleaned_sd, raw_sd, self.processed_sd) - return - - #we do this to get rewrtten keys for init_sd - rewritten_init_sd = merge_sd_overrides({}, self.processed_df, self.init_sd) - intermediate_sd = merge_sds(rewritten_init_sd, self.cleaned_sd, raw_sd) - base = merge_sd_overrides(intermediate_sd, self.processed_df, self.processed_sd) - - # Layer ``cleaned_*`` keys on top when cleaning is active. - if cleaning_active and clean_sd: - for col, stats in clean_sd.items(): - col_dict = base.setdefault(col, {}) - for stat_name, val in stats.items(): - col_dict[f'cleaned_{stat_name}'] = val - - # Layer ``filtered_*`` keys on top when a filter is active. - if filter_active and filt_sd: - for col, stats in filt_sd.items(): - col_dict = base.setdefault(col, {}) - for stat_name, val in stats.items(): - col_dict[f'filtered_{stat_name}'] = val - - self.merged_sd = base + self.merged_sd = assemble_merged_sd( + self.init_sd, self.cleaned_sd, raw_sd, self.processed_sd, self.processed_df, + split_chain_by_scope(self.operations), clean_sd=clean_sd, filt_sd=filt_sd) def _compute_scope_df(self, scope: str): """Return the df that scope's SD should be computed against. @@ -598,14 +630,14 @@ def _compute_scope_df(self, scope: str): return base return pp_result[0] if pp_result else base - def _scope_cache_key(self, chain): + def _scope_cache_key(self, chain, tier=None): """Hash that identifies a scope's SD-input identity. Includes the op chain *and* an identifier for the source dataframe (``id(sampled_df)``) *and* the post-processing method - *and* the analysis-klasses identity — all four are inputs to - the scope's SD, and a cache hit must mean "same SD-producing - inputs" not just "same chain". + *and* the analysis-klasses identity *and* the stats tier — all + five are inputs to the scope's SD, and a cache hit must mean + "same SD-producing inputs" not just "same chain". - sampled_df identity addresses codex P1 on #783: a ``raw_df`` swap with an unchanged chain must invalidate. @@ -617,11 +649,18 @@ def _scope_cache_key(self, chain): - analysis_klasses identity addresses codex P2 on #783: a klass-list swap with an unchanged chain must invalidate so new stat klasses surface in ``merged_sd``. + - the tier keeps a schema-tier sd from being read as a full one: + without it, assigning full stats after a schema-tier + construction hits the schema entries and never reaches + ``merged_sd``. ``tier`` defaults to the dataflow's own, and is + explicit so a caller can address the full-tier entry while the + dataflow is still at the schema tier. """ sampled_id = id(self.sampled_df) if self.sampled_df is not None else 0 pp = self.post_processing_method or '' klasses_id = id(self.analysis_klasses) - return hash_chain(chain, extra=f"{sampled_id}|{pp}|{klasses_id}") + tier = tier or self.stats_tier + return hash_chain(chain, extra=f"{sampled_id}|{pp}|{klasses_id}|{tier}") @observe('summary_sd', 'operations', 'analysis_klasses') @exception_protect('sd-cache-protector') @@ -653,9 +692,18 @@ def _populate_sd_cache(self, _change): new_cache = dict(self.summary_stats_cache) cache_grew = False - # filt scope reuses the SD that _summary_sd just produced. + # filt scope reuses the SD that _summary_sd just produced, but only if + # it was produced for this frame, klass list and tier. An sd left over + # from another (a tier changed since it ran, or an assignment ahead of + # the cascade) would be stored under this key and never repaired, since + # a present key is a hit. Compute the scope's SD instead. if keys['filt'] not in new_cache: - new_cache[keys['filt']] = dict(self.summary_sd or {}) + current = (id(self.processed_df), id(self.analysis_klasses), self.stats_tier) + if self._summary_sd_cache_key == current: + new_cache[keys['filt']] = dict(self.summary_sd or {}) + else: + filt_sd, _errs = self._get_summary_sd(self.processed_df) + new_cache[keys['filt']] = filt_sd cache_grew = True # raw + clean: fresh compute, but only on cache miss. @@ -711,6 +759,8 @@ def _build_error_dataframe(self, e) -> pd.DataFrame: #TAny closer to some error type @override def _get_summary_sd(self, processed_df: DataFrameT) -> Tuple[SDType, ErrDict]: + if self.stats_tier == "schema": + return self._get_schema_sd(processed_df), {} stats = self.DFStatsClass( processed_df, self.analysis_klasses, @@ -728,6 +778,16 @@ def _get_summary_sd(self, processed_df: DataFrameT) -> Tuple[SDType, ErrDict]: return sdf, {} + def _get_schema_sd(self, processed_df: DataFrameT) -> SDType: + """Identity and typing keys for every column, from the schema alone. + + ``orig_col_name``, ``rewritten_col_name``, ``dtype``, the ``is_*`` + flags, ``_type`` and ``length``: the keys column styling needs to pick + displayers, with no data query beyond the row count. A backend opts in + to ``stats_tier="schema"`` by overriding this. + """ + raise NotImplementedError(f"{type(self).__name__} has no schema stats tier") + # ### end summary stats block def _sd_to_jsondf(self, sd:SDType): @@ -745,35 +805,43 @@ def _sd_to_jsondf(self, sd:SDType): def _df_to_obj(self, df: DataFrameT) -> List[TDict[str, TAny]]: return pd_to_obj(self.sampling_klass.serialize_sample(df)) + def _analysis_klasses_with(self, analysis_klass:Type[ColAnalysis]) -> List[Type[ColAnalysis]]: + """The klass list the stats pipelines' ``add_stat`` builds: a class of the + same name is replaced, anything else is appended.""" + klasses = list(self.analysis_klasses) + if isinstance(analysis_klass, type): + klasses = [k for k in klasses + if not (isinstance(k, type) and k.__name__ == analysis_klass.__name__)] + klasses.append(analysis_klass) + return klasses + def add_analysis(self, analysis_klass:Type[ColAnalysis]) -> None: """ same as get_summary_sd, call whatever to set summary_sd and trigger further comps """ - stats = self.DFStatsClass( - self.processed_df, - self.analysis_klasses, - self.df_name, debug=self.debug) - stats.add_analysis(analysis_klass) - - self.analysis_klasses = stats.ap.ordered_a_objs + if self.stats_tier == "full": + stats = self.DFStatsClass( + self.processed_df, + self.analysis_klasses, + self.df_name, debug=self.debug) + stats.add_analysis(analysis_klass) + self.analysis_klasses = stats.ap.ordered_a_objs + else: + # Building DFStatsClass here would run the stats this tier exists + # to skip (twice, since the constructor and add_analysis each + # process the frame), and it does not go through _get_summary_sd. + self.analysis_klasses = self._analysis_klasses_with(analysis_klass) self.DFStatsClass.verify_analysis_objects(self.analysis_klasses) self.setup_options_from_analysis() #force recomputation self._handle_widget_change({}) - #final processing block - @observe('widget_args_tuple') - def _handle_widget_change(self, change): - """ - put together df_dict for consumption by the frontend - """ - # Tuple[TAny, pd.DataFrame, SDType] - _unused, processed_df, merged_sd = self.widget_args_tuple - if processed_df is None: - return - + def _build_df_data_dict(self, processed_df, merged_sd): + """The ``df_data_dict`` trait value: the main rows, the ``all_stats`` + payload and the empty placeholder. ``all_stats`` is the 5-11 ms part, + which is why it is built apart from ``_build_df_display_args``.""" # df_data_dict is still hardcoded for now # eventually processed_df will be able to add or alter values of df_data_dict # correlation would be added, filtered would probably be altered @@ -781,15 +849,15 @@ def _handle_widget_change(self, change): # to expedite processing maybe future provided dfs from # postprcoessing could default to empty until that is # selected, optionally - if self.skip_main_serial: - self.df_data_dict = {'main': [], - 'all_stats': self._sd_to_jsondf(merged_sd), - 'empty': []} - else: - self.df_data_dict = {'main': self._df_to_obj(processed_df), - 'all_stats': self._sd_to_jsondf(merged_sd), - 'empty': []} - + main = [] if self.skip_main_serial else self._df_to_obj(processed_df) + return {'main': main, + 'all_stats': self._sd_to_jsondf(merged_sd), + 'empty': []} + + def _build_df_display_args(self, processed_df, merged_sd): + """The ``df_display_args`` trait value: each display's column_config, + pinned_rows and keys, with the overrides applied. Reads ``merged_sd`` + but serializes none of it.""" temp_display_args = {} for display_name, A_Klass in self.df_display_klasses.items(): df_viewer_config = A_Klass.get_dfviewer_config(merged_sd, processed_df) @@ -807,8 +875,21 @@ def _handle_widget_change(self, change): temp_display_args['main']['df_viewer_config']['extra_grid_config'] = self.extra_grid_config if self.component_config: temp_display_args['main']['df_viewer_config']['component_config'] = self.component_config + return temp_display_args + + #final processing block + @observe('widget_args_tuple') + def _handle_widget_change(self, change): + """ + put together df_dict for consumption by the frontend + """ + # Tuple[TAny, pd.DataFrame, SDType] + _unused, processed_df, merged_sd = self.widget_args_tuple + if processed_df is None: + return - self.df_display_args = temp_display_args + self.df_data_dict = self._build_df_data_dict(processed_df, merged_sd) + self.df_display_args = self._build_df_display_args(processed_df, merged_sd) """ diff --git a/buckaroo/server/handlers.py b/buckaroo/server/handlers.py index c60d99062..b844bd3e0 100644 --- a/buckaroo/server/handlers.py +++ b/buckaroo/server/handlers.py @@ -14,7 +14,9 @@ from buckaroo.compare import col_join_dfs from buckaroo.df_util import old_col_new_col from buckaroo.server.focus import find_or_create_session_window -from buckaroo.server.session import build_state_message +from buckaroo.dataflow.dataflow import STATS_TIERS +from buckaroo.server.session import ( + STATS_DELIVERIES, build_state_message, dataflow_stats_tier) from buckaroo.server import telemetry from buckaroo.pluggable_analysis_framework import perf_log @@ -380,6 +382,28 @@ async def post(self): self.write({"session": session_id, "server_pid": os.getpid(), "browser_action": browser_action, **metadata}) +def _stats_policy_from_body(body: dict, current_tier: str, current_delivery: str): + """Read ``stats_tier`` and ``stats_delivery`` from a request body. + + A field the body omits (or sends as null) keeps the current value (the + session's, or the default for a new session), as ``cache_dir`` does. Returns + ``(stats_tier, stats_delivery, None)``, or ``(None, None, error)`` with the + 400 response body for an unknown value.""" + stats_tier = body.get("stats_tier") + if stats_tier is None: + stats_tier = current_tier + if stats_tier not in STATS_TIERS: + return None, None, {"error_code": "invalid_stats_tier", + "message": f"stats_tier must be one of {list(STATS_TIERS)}, got {stats_tier!r}"} + stats_delivery = body.get("stats_delivery") + if stats_delivery is None: + stats_delivery = current_delivery + if stats_delivery not in STATS_DELIVERIES: + return None, None, {"error_code": "invalid_stats_delivery", + "message": f"stats_delivery must be one of {list(STATS_DELIVERIES)}, got {stats_delivery!r}"} + return stats_tier, stats_delivery, None + + class LoadExprHandler(tornado.web.RequestHandler): """POST /load_expr — load a xorq/ibis expression from a build dir and serve it via the xorq-backed buckaroo dataflow. @@ -438,7 +462,10 @@ async def post(self): # Config-bearing fields that change how the result is computed or # rendered. If the caller passes any of these on a warm POST we must # re-run the pipeline — returning cached metadata would silently - # ignore the new config. + # ignore the new config. stats_tier and stats_delivery are not in this + # tuple: it tests truthiness, and a host that sends the pair on every + # POST would never get the warm exit (#944). They are compared with the + # session's stored pair below instead. has_config = any(body.get(k) for k in ( "component_config", "column_config_overrides", "extra_grid_config", "init_sd", "skip_stat_columns")) @@ -468,11 +495,22 @@ async def post(self): # ~/.cache/xorq and recomputing the embedder's snapshots there. if cache_dir is None and existing is not None: cache_dir = existing.cache_dir + # The stats policy carries over the same way: a re-POST that omits it + # keeps the session's, and one that changes it rebuilds. + stats_tier, stats_delivery, policy_error = _stats_policy_from_body( + body, + existing.stats_tier if existing is not None else "full", + existing.stats_delivery if existing is not None else "inline") + if policy_error is not None: + self.set_status(400) + self.write(policy_error) + return # /load swaps a session to pandas without clearing build_dir, so the # backend is checked too — else its pandas metadata comes back here. if (not force_reload and not has_config and existing and existing.backend == "xorq" and existing.build_dir == build_dir - and existing.cache_dir == cache_dir and existing.metadata): + and existing.cache_dir == cache_dir and existing.metadata + and (existing.stats_tier, existing.stats_delivery) == (stats_tier, stats_delivery)): # The pipeline is skipped, but the refreshed page still opens a new # WS and pulls a fresh time-to-first-rows. Re-arm first-pull telemetry # on the existing session — rebind this request's sink and reset the @@ -553,6 +591,7 @@ async def post(self): with perf_log.perf_span("firstpull.dataflow_construct", session=session_id): xorq_dataflow = xorq_loading.XorqServerDataflow( expr, skip_main_serial=True, extra_klasses=extra_klasses, + stats_tier=dataflow_stats_tier(stats_tier, stats_delivery), **dataflow_kwargs) # Spanning metadata too leaves only the small klass-load step # unmeasured inside the outer firstpull.load_expr total. @@ -578,6 +617,8 @@ async def post(self): session.cache_dir = cache_dir session.project_root = project_root session.dataflow_kwargs = dataflow_kwargs + session.stats_tier = stats_tier + session.stats_delivery = stats_delivery session.tele_sink = tele_sink session.xorq_dataflow = xorq_dataflow # Clear pandas-side state left by a prior /load on the same @@ -807,8 +848,22 @@ class ReloadExprHandler(tornado.web.RequestHandler): ``cache_storage_path`` store are cache hits, and column overrides, extra grid config, init_sd and skip_stat_columns survive the reload. + The session's stored ``stats_tier`` / ``stats_delivery`` are replayed too. + The body is optional; a pair in it replaces the stored one (and is stored) + when the reload succeeds. + Returns 404 when the session does not exist, 400 when it is not a xorq - session or has no project_root recorded, 501 when xorq is not installed.""" + session, has no project_root recorded or carries an invalid stats policy, + 501 when xorq is not installed.""" + + def _optional_body(self) -> dict: + """The JSON object body, or ``{}`` when there is none or it is not an + object: the endpoint took no body before the stats policy fields.""" + try: + body = json.loads(self.request.body) + except (json.JSONDecodeError, TypeError): + return {} + return body if isinstance(body, dict) else {} async def post(self, session_id): sessions = self.application.settings["sessions"] @@ -831,6 +886,13 @@ async def post(self, session_id): "message": "Session has no project_root — pass project_root to /load_expr first"}) return + stats_tier, stats_delivery, policy_error = _stats_policy_from_body( + self._optional_body(), session.stats_tier, session.stats_delivery) + if policy_error is not None: + self.set_status(400) + self.write(policy_error) + return + try: from buckaroo.server import xorq_loading except ImportError: @@ -847,6 +909,7 @@ async def post(self, session_id): + xorq_loading.load_project_display_klasses(session.project_root)) xorq_dataflow = xorq_loading.XorqServerDataflow( session.expr, skip_main_serial=True, extra_klasses=extra_klasses, + stats_tier=dataflow_stats_tier(stats_tier, stats_delivery), **session.dataflow_kwargs) except Exception: tb = traceback.format_exc() @@ -874,6 +937,8 @@ async def post(self, session_id): refreshed = get_buckaroo_display_state(xorq_dataflow) session.xorq_dataflow = xorq_dataflow + session.stats_tier = stats_tier + session.stats_delivery = stats_delivery session.df_display_args = refreshed["df_display_args"] session.df_data_dict = refreshed["df_data_dict"] session.df_meta = refreshed["df_meta"] diff --git a/buckaroo/server/session.py b/buckaroo/server/session.py index 563ce28fc..bafa5a870 100644 --- a/buckaroo/server/session.py +++ b/buckaroo/server/session.py @@ -11,6 +11,17 @@ _DEFAULT_SESSION_TTL_S = 3600.0 # 1 hour idle before eviction _DEFAULT_EVICTION_INTERVAL_S = 300.0 # check every 5 minutes +# How a session reaches its stats tier: ``inline`` computes them in the +# dataflow constructor (today's behaviour); ``deferred`` builds the schema tier +# first and leaves the rest to later requests. +STATS_DELIVERIES = ("inline", "deferred") + + +def dataflow_stats_tier(stats_tier: str, stats_delivery: str) -> str: + """The tier a session's dataflow is constructed at. A deferred session + starts at the schema tier whatever tier it is headed for.""" + return "schema" if stats_delivery == "deferred" else stats_tier + @dataclass class SessionState: @@ -43,6 +54,13 @@ class SessionState: # extra_grid_config, init_sd, skip_stat_columns), replayed by /reload_expr so # a reload keeps stat caching and column config (#957). dataflow_kwargs: dict = field(default_factory=dict) + # The /load_expr stats policy. ``stats_tier`` is the tier the session is + # headed for (see dataflow.STATS_TIERS) and ``stats_delivery`` how it gets + # there (see STATS_DELIVERIES). Kept apart from dataflow_kwargs, which is + # splatted into the dataflow constructor, and replayed by /reload_expr. The + # tier the dataflow is built at follows from the pair (dataflow_stats_tier). + stats_tier: str = "full" + stats_delivery: str = "inline" # Companion telemetry sink (#943): a fire-and-forget POST callable, built # once from the /load_expr payload's telemetry_url on the IOLoop (where # make_http_sink captures AsyncHTTPClient/IOLoop.current()). Stored here so diff --git a/buckaroo/xorq_buckaroo.py b/buckaroo/xorq_buckaroo.py index 6b9149c40..b10cb2e6d 100644 --- a/buckaroo/xorq_buckaroo.py +++ b/buckaroo/xorq_buckaroo.py @@ -27,7 +27,7 @@ from .buckaroo_widget import BuckarooInfiniteWidget, BuckarooWidget from .customizations.styling import DefaultMainStyling, DefaultSummaryStatsStyling from .customizations.xorq_autoclean_conf import NoCleaningConfXorq -from .customizations.xorq_stats_v2 import XORQ_STATS_V2 +from .customizations.xorq_stats_v2 import XORQ_STATS_V2, schema_stats from .dataflow.autocleaning import PandasAutocleaning from .dataflow.dataflow import CustomizableDataflow from .dataflow.dataflow_extras import Sampling @@ -173,6 +173,8 @@ def _get_summary_sd(self, processed_df: "XorqExpr | pd.DataFrame"): 'orig_col_name': orig_col, 'rewritten_col_name': rewritten_col} return empty, {} + if self.stats_tier == "schema": + return self._get_schema_sd(processed_df), {} cache_storage = getattr(self, 'cache_storage', None) # The owning widget injects XorqDfStatsV2 as DFStatsClass (via its # InnerDataFlow subclass); the cast exposes cache_run_stats() below. @@ -213,6 +215,21 @@ def _get_summary_sd(self, processed_df: "XorqExpr | pd.DataFrame"): rewritten[rewritten_col] = col_meta return rewritten, errs + def _get_schema_sd(self, processed_df: "XorqExpr | pd.DataFrame") -> dict: + """Identity, dtype, typing flags and row count for every column of an + expression, with no query beyond ``_expr_count`` (cached per + expression). Carries the same values full stats give those keys, so + styling that reads only them renders the same. ``_get_summary_sd`` + handles a pandas frame before it gets here.""" + expr = cast("XorqExpr", processed_df) + schema = expr.schema() + length = _expr_count(expr) + return { + rewritten_col: { + 'orig_col_name': orig_col, 'rewritten_col_name': rewritten_col, + **schema_stats(str(schema[str(orig_col)])), 'length': length} + for orig_col, rewritten_col in old_col_new_col(expr)} + _XORQ_ANALYSIS_KLASSES = list(XORQ_STATS_V2) + [DefaultSummaryStatsStyling, DefaultMainStyling] From 21cd73e16843527d71cbb03f60654535f7b431b8 Mon Sep 17 00:00:00 2001 From: Paddy Mullen Date: Sat, 3 Oct 2026 23:24:11 -0400 Subject: [PATCH 04/18] test(dataflow): accept pandas 3's str dtype in the pandas characterization test (rows-first s1) The Max Versions jobs resolve pandas 3, which reports a string column's dtype as 'str' where pandas 2 says 'object'. The characterization test asserted 'object'. Verified in a Max Versions environment (pandas 3.0.6, polars 1.44.2, xorq 0.4.5): the unit suite passes. Co-Authored-By: Claude Sonnet 5.5 --- tests/unit/basic_widget_test.py | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) diff --git a/tests/unit/basic_widget_test.py b/tests/unit/basic_widget_test.py index 9e63765e1..8852b8cdd 100644 --- a/tests/unit/basic_widget_test.py +++ b/tests/unit/basic_widget_test.py @@ -296,9 +296,11 @@ def test_default_tier_widget_change_sequence_and_merged_sd(): 'price': [12.5, 18.9, 7.4, 22.1, 14.0], 'qty': [1, 2, 1, 3, 2], 'category': ['a', 'b', 'a', 'c', 'b']})) msd = bw.dataflow.merged_sd - assert {col: (sd['length'], sd['_type'], sd['dtype']) for col, sd in msd.items()} == { - 'a': (5, 'float', 'float64'), 'b': (5, 'integer', 'int64'), - 'c': (5, 'string', 'object')} + assert {col: (sd['length'], sd['_type']) for col, sd in msd.items()} == { + 'a': (5, 'float'), 'b': (5, 'integer'), 'c': (5, 'string')} + # pandas 3 names the string dtype 'str' where pandas 2 says 'object'. + assert [msd[col]['dtype'] for col in 'ab'] == ['float64', 'int64'] + assert msd['c']['dtype'] in ('object', 'str') assert msd['a']['mean'] == pytest.approx(14.98) seen = [] From 1c89cfd6a82f5b3ee2c2e366784e43a0a3ab028d Mon Sep 17 00:00:00 2001 From: Paddy Mullen Date: Sat, 3 Oct 2026 23:49:42 -0400 Subject: [PATCH 05/18] test(xorq): failing test for the schema tier overriding init_sd typing on a skipped column (rows-first s1) A column in skip_stat_columns gets only name, dtype and length from the full-tier pipeline, so its _type comes from init_sd. The schema tier layers the schema-derived _type and is_* keys over it, so an int64 column that init_sd types as float merges as integer and renders with zero fraction digits instead of the float displayer init_sd asked for. Co-Authored-By: Claude Sonnet 5.5 --- tests/unit/server/test_load_expr.py | 23 +++++++++++++++++++++++ 1 file changed, 23 insertions(+) diff --git a/tests/unit/server/test_load_expr.py b/tests/unit/server/test_load_expr.py index 03feb571e..36557afb7 100644 --- a/tests/unit/server/test_load_expr.py +++ b/tests/unit/server/test_load_expr.py @@ -1050,6 +1050,29 @@ def test_init_sd_hints_and_overrides_still_apply(self): assert by_header["qty"]["displayer_args"]["max_length"] == 200 assert by_header["category"]["displayer_args"]["max_length"] == 5000 + def test_skipped_column_keeps_init_sd_typing(self): + """A column in ``skip_stat_columns`` gets only name, dtype and length + from the pipeline, so its ``_type`` comes from ``init_sd``. The schema + tier must not layer the schema-derived typing keys over it.""" + expr = _stats_tier_expr() + kwargs = { + "init_sd": {"qty": {"_type": "float", "mean": 1.8, "min": 1, "max": 3}}, + "skip_stat_columns": ["qty"]} + full = _build_dataflow(expr, **kwargs) + schema = _build_dataflow(expr, stats_tier="schema", **kwargs) + + def merged(dataflow, orig_col): + return next(v for v in dataflow.merged_sd.values() if v["orig_col_name"] == orig_col) + + assert merged(full, "qty")["_type"] == "float" + assert merged(schema, "qty")["_type"] == "float" + assert not {"is_numeric", "is_integer", "is_float"} & merged(schema, "qty").keys() + assert merged(schema, "price")["_type"] == "float" + assert merged(schema, "price")["is_float"] is True + full_cfg = full.df_display_args["main"]["df_viewer_config"]["column_config"] + schema_cfg = schema.df_display_args["main"]["df_viewer_config"]["column_config"] + assert _without_min_width(schema_cfg) == _without_min_width(full_cfg) + def test_sorted_infinite_request_works(self): dataflow = _build_dataflow(stats_tier="schema") qty = next(k for k, v in dataflow.merged_sd.items() if v["orig_col_name"] == "qty") From 10d9d015e08048ed5c911d83a96750f8cb026642 Mon Sep 17 00:00:00 2001 From: Paddy Mullen Date: Sat, 3 Oct 2026 23:56:33 -0400 Subject: [PATCH 06/18] fix(xorq): schema tier leaves a skipped column's typing to init_sd (rows-first s1) _get_schema_sd never read skip_stat_columns, so a skipped column's schema-derived _type and is_* keys overrode init_sd's _type once merged. The full tier gives a skipped column only name, dtype and length. The schema tier now does the same, so init_sd's _type decides the displayer at both tiers. Co-Authored-By: Claude Sonnet 5.5 --- buckaroo/xorq_buckaroo.py | 18 ++++++++++++------ 1 file changed, 12 insertions(+), 6 deletions(-) diff --git a/buckaroo/xorq_buckaroo.py b/buckaroo/xorq_buckaroo.py index b10cb2e6d..e1aad2315 100644 --- a/buckaroo/xorq_buckaroo.py +++ b/buckaroo/xorq_buckaroo.py @@ -219,16 +219,22 @@ def _get_schema_sd(self, processed_df: "XorqExpr | pd.DataFrame") -> dict: """Identity, dtype, typing flags and row count for every column of an expression, with no query beyond ``_expr_count`` (cached per expression). Carries the same values full stats give those keys, so - styling that reads only them renders the same. ``_get_summary_sd`` - handles a pandas frame before it gets here.""" + styling that reads only them renders the same. A column in + ``skip_stat_columns`` gets only name, dtype and length, as at the full + tier, so its typing comes from ``init_sd``. ``_get_summary_sd`` handles + a pandas frame before it gets here.""" expr = cast("XorqExpr", processed_df) schema = expr.schema() length = _expr_count(expr) - return { - rewritten_col: { + skip = getattr(self, 'skip_stat_columns', None) or () + sd: dict = {} + for orig_col, rewritten_col in old_col_new_col(expr): + dtype = str(schema[str(orig_col)]) + typing = {'dtype': dtype} if orig_col in skip else schema_stats(dtype) + sd[rewritten_col] = { 'orig_col_name': orig_col, 'rewritten_col_name': rewritten_col, - **schema_stats(str(schema[str(orig_col)])), 'length': length} - for orig_col, rewritten_col in old_col_new_col(expr)} + **typing, 'length': length} + return sd _XORQ_ANALYSIS_KLASSES = list(XORQ_STATS_V2) + [DefaultSummaryStatsStyling, DefaultMainStyling] From 9cac279a2a839ea24621b93e12539f66e2ac101b Mon Sep 17 00:00:00 2001 From: Paddy Mullen Date: Sun, 4 Oct 2026 00:18:19 -0400 Subject: [PATCH 07/18] test(server): failing tests for the pandas and polars schema stats tier (rows-first s2) ServerDataflow and PolarsServerDataflow with stats_tier="schema" should publish the display state full stats give (column_config without stats-derived keys, pinned_rows including a host-supplied one, data_key, summary_stats_key) with no stat computed on the data, still apply init_sd, serve sorted windows, take a later full assignment into merged_sd for every scope, and assemble to the same sd as merged_sd. Both backends run through the same parametrized class, plus a pandas test that pins how an object column is typed from its dtype. Co-Authored-By: Claude Sonnet 5.5 --- tests/unit/server/test_data_loading_polars.py | 337 +++++++++++++++++- 1 file changed, 336 insertions(+), 1 deletion(-) diff --git a/tests/unit/server/test_data_loading_polars.py b/tests/unit/server/test_data_loading_polars.py index 58f06f5bb..6ed441642 100644 --- a/tests/unit/server/test_data_loading_polars.py +++ b/tests/unit/server/test_data_loading_polars.py @@ -9,15 +9,34 @@ * ``.json`` files are read as standard JSON arrays (matches ``pd.read_json`` default ``lines=False`` so the same file loads under either backend). + +It also pins the schema stats tier (rows-first s2) for the pandas and +polars server dataflows side by side: ``stats_tier="schema"`` publishes +the display state full stats give, without reading a value. """ +import base64 +import datetime +import decimal +import io import json import os import tempfile +import pandas as pd import polars as pl +import pyarrow.parquet as pq import pytest -from buckaroo.server.data_loading_polars import (create_polars_dataflow, handle_infinite_request_buckaroo_polars, load_file_polars) +from buckaroo.dataflow.dataflow import assemble_merged_sd +from buckaroo.dataflow.sd_cache import split_chain_by_scope +from buckaroo.jlisp.lisp_utils import s as lisp_sym +from buckaroo.pluggable_analysis_framework.col_analysis import ColAnalysis +from buckaroo.pluggable_analysis_framework.stat_pipeline import StatPipeline +from buckaroo.pluggable_analysis_framework.utils import PERVERSE_DF +from buckaroo.server.data_loading import ServerDataflow, handle_infinite_request_buckaroo +from buckaroo.server.data_loading_polars import (PolarsServerDataflow, create_polars_dataflow, handle_infinite_request_buckaroo_polars, load_file_polars) +from buckaroo.styling_helpers import float_, obj_ +from tests.unit.dataflow.scoped_summary_stats_test import _scope_inputs def _payload(start=0, end=100): @@ -89,3 +108,319 @@ def test_load_file_polars_ndjson_still_works(): def test_load_file_polars_unsupported_extension(): with pytest.raises(ValueError, match="Unsupported file format"): load_file_polars("/tmp/foo.xyz") + + +# --------------------------------------------------------------------------- +# stats_tier="schema" on the pandas and polars server dataflows (rows-first s2) +# --------------------------------------------------------------------------- + +_BACKENDS = { + "pandas": (ServerDataflow, pd.DataFrame, handle_infinite_request_buckaroo), + "polars": (PolarsServerDataflow, pl.DataFrame, handle_infinite_request_buckaroo_polars)} + +# price has a 1e9 maximum, so its estimated column width depends on the min and +# max that only full stats supply. +_TIER_DATA = { + "price": [12.5, 18.9, 7.4, 22.1, 1e9], + "qty": [1, 2, 1, 3, 2], + "category": ["a", "b", "a", "c", "b"]} + +# The same columns with a different number of rows per value in every column. +# polars does not keep the order of equal counts, so on _TIER_DATA the mode of +# two full-stats runs over the same frame can differ. +_UNTIED_DATA = { + "price": [12.5, 12.5, 12.5, 18.9, 18.9, 1e9], + "qty": [1, 1, 1, 2, 2, 3], + "category": ["a", "a", "a", "b", "b", "c"]} + +# Stats that read values. The schema tier must not publish any of them. +_DATA_STAT_KEYS = { + "mean", "std", "median", "min", "max", "null_count", "non_null_count", + "distinct_count", "unique_count", "value_counts", "mode", "histogram", + "histogram_bins", "memory_usage"} + + +@pytest.fixture(params=list(_BACKENDS)) +def backend(request): + return request.param + + +def _frame(backend, data=None): + return _BACKENDS[backend][1](_TIER_DATA if data is None else data) + + +def _typed_frame(backend): + """One column per kind of dtype the typing stats tell apart.""" + days = [datetime.date(2020, 1, 1), datetime.date(2020, 1, 2), datetime.date(2020, 1, 3)] + if backend == "pandas": + return pd.DataFrame({ + "float": [1.5, 2.5, 3.5], "int": [1, 2, 3], + "nullable_int": pd.array([1, None, 3], dtype="Int64"), + "bool": [True, False, True], "str": ["a", "b", "c"], + "datetime": pd.to_datetime(days), "datetime_utc": pd.to_datetime(days, utc=True), + "timedelta": pd.to_timedelta([1, 2, 3], unit="D"), + "category": pd.Categorical(["a", "b", "a"])}) + return pl.DataFrame({ + "float": [1.5, 2.5, 3.5], "int32": pl.Series([1, 2, 3], dtype=pl.Int32), + "bool": [True, False, True], "str": ["a", "b", "c"], "date": days, + "datetime": [datetime.datetime(2020, 1, d) for d in (1, 2, 3)], + "duration": [datetime.timedelta(days=d) for d in (1, 2, 3)], + "time": [datetime.time(h) for h in (1, 2, 3)], + "decimal": pl.Series([decimal.Decimal("1.5")] * 3, dtype=pl.Decimal(10, 2)), + "binary": [b"a", b"b", b"c"], + "category": pl.Series(["a", "b", "a"], dtype=pl.Categorical)}) + + +def _build_dataflow(backend, frame=None, **kwargs): + frame = _frame(backend) if frame is None else frame + return _BACKENDS[backend][0](frame, skip_main_serial=True, **kwargs) + + +def _three_scope_dataflow(backend, data=None, **kwargs): + """A dataflow with a search (filt scope) active and, on pandas, a user op + (clean scope), so each scope holds its own summary-stats cache entry. The + polars conf ships no command to run outside a search, so there raw and + clean are one scope.""" + dataflow = _build_dataflow(backend, _frame(backend, data), **kwargs) + if backend == "pandas": + dataflow.operations = [[lisp_sym("fillna"), {"symbol": "df"}, "qty", 0]] + dataflow.quick_command_args = {"search": ["a"]} + return dataflow + + +def _scope_keys(dataflow, tier): + chains = split_chain_by_scope(dataflow.operations) + return {scope: dataflow._scope_cache_key(chain, tier=tier) for scope, chain in chains.items()} + + +def _spy_stat_pipeline(monkeypatch): + """Record the columns of every frame ``StatPipeline.process_df`` runs on, + apart from the ``PERVERSE_DF`` its DAG self-check uses, which has no part + of the loaded data in it.""" + frames = [] + original = StatPipeline.process_df + + def spy(self, df, *args, **kwargs): + if list(df.columns) != list(PERVERSE_DF.columns): + frames.append(list(df.columns)) + return original(self, df, *args, **kwargs) + + monkeypatch.setattr(StatPipeline, "process_df", spy) + return frames + + +def _without_min_width(column_config): + return [{**cc, "ag_grid_specs": {k: v for k, v in cc["ag_grid_specs"].items() if k != "minWidth"}} + for cc in column_config] + + +def _as_json(sd): + """Comparable form of an sd: json equates NaN with NaN where == would not.""" + return json.dumps(sd, sort_keys=True, default=str) + + +def _wire_stat_names(dataflow): + """The stat names in the ``all_stats`` payload sent to the client.""" + envelope = dataflow.df_data_dict["all_stats"] + table = pq.read_table(io.BytesIO(base64.b64decode(envelope["data"]))) + return {name.split("__", 1)[1] for name in table.column_names} + + +def _viewer_config(dataflow, display="main"): + return dataflow.df_display_args[display]["df_viewer_config"] + + +class _NoopPostProcessing(ColAnalysis): + provides_defaults = {} + post_processing_method = "noop_post" + + @classmethod + def post_process_df(cls, df): + return [df, {}] + + +class TestStatsTierSchema: + """``stats_tier="schema"`` on ``ServerDataflow`` and + ``PolarsServerDataflow``: every column is typed from its dtype alone, so + the dataflow publishes the display state full stats give with no stat + computed on the data.""" + + def test_matches_full_stats_display_state(self, backend): + full = _build_dataflow(backend) + schema = _build_dataflow(backend, stats_tier="schema") + assert schema.df_display_args.keys() == full.df_display_args.keys() + for name, full_arg in full.df_display_args.items(): + schema_arg = schema.df_display_args[name] + assert schema_arg["data_key"] == full_arg["data_key"] + assert schema_arg["summary_stats_key"] == full_arg["summary_stats_key"] + full_cfg, schema_cfg = full_arg["df_viewer_config"], schema_arg["df_viewer_config"] + assert schema_cfg["pinned_rows"] == full_cfg["pinned_rows"] + assert (_without_min_width(schema_cfg["column_config"]) + == _without_min_width(full_cfg["column_config"])) + + def test_host_supplied_pinned_rows_match_full_stats(self, backend): + """pinned_rows is configuration, so it is the same in every tier even + when it names a stat the schema tier does not compute.""" + pinned = [obj_("dtype"), float_("mean")] + full = _build_dataflow(backend, pinned_rows=pinned) + schema = _build_dataflow(backend, stats_tier="schema", pinned_rows=pinned) + assert _viewer_config(schema)["pinned_rows"] == _viewer_config(full)["pinned_rows"] == pinned + + def test_min_width_is_the_stats_derived_difference(self, backend): + widths = {} + for tier in ("full", "schema"): + cfg = _viewer_config(_build_dataflow(backend, stats_tier=tier)) + widths[tier] = {cc["header_name"]: cc["ag_grid_specs"]["minWidth"] + for cc in cfg["column_config"]} + # price's 1e9 maximum widens it under full stats; without min and max + # the estimate falls back to a one-digit value. + assert widths["schema"]["price"] < widths["full"]["price"] + + def test_typing_matches_full_stats_for_every_dtype(self, backend): + full = _build_dataflow(backend, _typed_frame(backend)) + schema = _build_dataflow(backend, _typed_frame(backend), stats_tier="schema") + assert schema.merged_sd.keys() == full.merged_sd.keys() + for col, stats in schema.merged_sd.items(): + assert {"orig_col_name", "rewritten_col_name", "dtype", "_type", "length"} <= stats.keys() + assert stats == {k: full.merged_sd[col][k] for k in stats} + assert {"duration", "categorical", "datetime"} <= { + stats["_type"] for stats in schema.merged_sd.values()} + + def test_publishes_no_data_derived_stat(self, backend): + full = _build_dataflow(backend) + schema = _build_dataflow(backend, stats_tier="schema") + for stats in schema.merged_sd.values(): + assert not _DATA_STAT_KEYS & stats.keys() + assert _wire_stat_names(schema) == {"dtype"} + # The same reads on the full tier see the stats, so they would notice. + assert _DATA_STAT_KEYS & full.merged_sd["a"].keys() + assert {"dtype", "histogram_bins"} <= _wire_stat_names(full) + + def test_computes_no_stat_on_the_data(self, backend, monkeypatch): + frames = _spy_stat_pipeline(monkeypatch) + dataflow = _build_dataflow(backend, stats_tier="schema") + dataflow.quick_command_args = {"search": ["a"]} + dataflow.add_analysis(_NoopPostProcessing) + assert frames == [] + assert "noop_post" in dataflow.buckaroo_options["post_processing"] + + def test_the_spy_sees_the_pipeline_at_the_full_tier(self, backend, monkeypatch): + frames = _spy_stat_pipeline(monkeypatch) + _build_dataflow(backend) + assert frames, "the spy would not notice a stat computed on the data" + + def test_init_sd_hints_and_overrides_still_apply(self, backend): + dataflow = _build_dataflow( + backend, stats_tier="schema", + init_sd={"qty": {"displayer_args": {"displayer": "string", "max_length": 200}}}, + column_config_overrides={ + "category": {"displayer_args": {"displayer": "string", "max_length": 5000}}}) + by_header = {cc["header_name"]: cc for cc in _viewer_config(dataflow)["column_config"]} + assert by_header["qty"]["displayer_args"]["max_length"] == 200 + assert by_header["category"]["displayer_args"]["max_length"] == 5000 + + def test_skipped_column_keeps_init_sd_typing(self, backend): + """A column in ``skip_stat_columns`` gets only its names from the + pipeline, so its ``_type`` comes from ``init_sd``. The schema tier + must not layer the dtype-derived typing keys over it.""" + kwargs = { + "init_sd": {"qty": {"_type": "float", "mean": 1.8, "min": 1, "max": 3}}, + "skip_stat_columns": ["qty"]} + full = _build_dataflow(backend, **kwargs) + schema = _build_dataflow(backend, stats_tier="schema", **kwargs) + + def merged(dataflow, orig_col): + return next(v for v in dataflow.merged_sd.values() if v["orig_col_name"] == orig_col) + + assert merged(full, "qty")["_type"] == "float" + assert merged(schema, "qty")["_type"] == "float" + assert not {"is_numeric", "is_integer", "is_float", "dtype"} & merged(schema, "qty").keys() + assert merged(schema, "price")["_type"] == "float" + assert merged(schema, "price")["is_float"] is True + assert (_without_min_width(_viewer_config(schema)["column_config"]) + == _without_min_width(_viewer_config(full)["column_config"])) + + def test_empty_frame_matches_full_stats(self, backend): + """Full stats publish nothing for a frame with no rows (the pipeline + returns an empty summary), so the schema tier does the same.""" + empty = {"x": [], "y": []} + full = _build_dataflow(backend, _frame(backend, empty)) + schema = _build_dataflow(backend, _frame(backend, empty), stats_tier="schema") + assert full.merged_sd == {} + assert schema.merged_sd == {} + assert schema.df_display_args == full.df_display_args + + def test_sorted_infinite_request_works(self, backend): + dataflow = _build_dataflow(backend, stats_tier="schema") + qty = next(k for k, v in dataflow.merged_sd.items() if v["orig_col_name"] == "qty") + handler = _BACKENDS[backend][2] + resp, parquet = handler(dataflow, {"start": 0, "end": 5, "sourceName": "default", + "sort": qty, "sort_direction": "desc"}) + assert "error_info" not in resp + assert resp["length"] == 5 + assert pq.read_table(io.BytesIO(parquet)).column(qty).to_pylist() == [3, 2, 2, 1, 1] + + def test_pending_state_writes_no_full_tier_cache_key(self, backend): + dataflow = _three_scope_dataflow(backend, stats_tier="schema") + full_keys = set(_scope_keys(dataflow, "full").values()) + assert len(full_keys) == (3 if backend == "pandas" else 2) + assert dataflow.summary_stats_cache + assert not full_keys & dataflow.summary_stats_cache.keys() + + def test_later_full_assignment_reaches_merged_sd_for_all_scopes(self, backend): + full = _three_scope_dataflow(backend, _UNTIED_DATA) + dataflow = _three_scope_dataflow(backend, _UNTIED_DATA, stats_tier="schema") + assert "mean" not in dataflow.merged_sd["a"] + assert {"mean", "filtered_mean"} <= full.merged_sd["a"].keys() + + dataflow.stats_tier = "full" + dataflow.summary_sd = full.summary_sd + + assert _as_json(dataflow.merged_sd) == _as_json(full.merged_sd) + + def test_full_scope_sds_cached_first_are_used_without_the_pipeline(self, backend, monkeypatch): + full = _three_scope_dataflow(backend) + dataflow = _three_scope_dataflow(backend, stats_tier="schema") + cache = dict(dataflow.summary_stats_cache) + for scope, key in _scope_keys(dataflow, "full").items(): + cache[key] = full.summary_stats_cache[getattr(full, f"{scope}_sd_key")] + dataflow.summary_stats_cache = cache + frames = _spy_stat_pipeline(monkeypatch) + + dataflow.stats_tier = "full" + dataflow.summary_sd = full.summary_sd + + assert frames == [] + assert _as_json(dataflow.merged_sd) == _as_json(full.merged_sd) + + def test_assembled_sd_equals_merged_sd_with_init_sd_and_a_user_op(self, backend): + # init_sd is keyed by the original column name; price is rewritten to "a". + dataflow = _three_scope_dataflow(backend, stats_tier="schema", init_sd={ + "price": {"displayer_args": {"displayer": "string", "max_length": 99}, "init_only": 1}}) + sd = dataflow.merged_sd["a"] + assert sd["init_only"] == 1, "init_sd must reach merged_sd under the rewritten name" + assert "filtered_length" in sd, "the filter layer must be active" + if backend == "pandas": + assert "cleaned_length" in sd, "the cleaning layer must be active" + + assembled = assemble_merged_sd(**_scope_inputs(dataflow)) + + assert _as_json(assembled) == _as_json(dataflow.merged_sd) + + +def test_object_column_is_typed_as_a_string_at_the_schema_tier(): + """pandas gives an object column no dtype to tell strings from other + values, and calls an empty one a string, so the schema tier types every + object column as one. Full stats read the values and call the mixed column + ``obj``; the gap closes when full stats arrive.""" + def frame(): + return pd.DataFrame({ + "words": pd.Series(["a", "b", "c"], dtype=object), + "mixed": pd.Series([1, "b", 2.5], dtype=object)}) + + types = {} + for tier in ("full", "schema"): + sd = ServerDataflow(frame(), skip_main_serial=True, stats_tier=tier).merged_sd + types[tier] = {v["orig_col_name"]: v["_type"] for v in sd.values()} + assert types["full"] == {"words": "string", "mixed": "obj"} + assert types["schema"] == {"words": "string", "mixed": "string"} From 334b73675fdf4f581ad4018056fccab4e05edeb7 Mon Sep 17 00:00:00 2001 From: Paddy Mullen Date: Sun, 4 Oct 2026 00:47:10 -0400 Subject: [PATCH 08/18] feat(server): pandas and polars schema stats tier (rows-first s2) ServerDataflow and PolarsServerDataflow now implement the _get_schema_sd hook, so stats_tier="schema" builds a dataflow that types every column from its dtype and runs no stat on the data. schema_sd (stat_pipeline.py) builds the sd the way process_df shapes it: an empty frame gives {}, a skipped column keeps only its names. pandas applies the existing typing_stats to a zero-row slice and derives _type through the _type stat; polars factors pl_dtype_typing out of pl_typing_stats and feeds it the dtype. Co-Authored-By: Claude Sonnet 5.5 --- buckaroo/customizations/pd_stats_v2.py | 19 ++++++++++++ buckaroo/customizations/pl_stats_v2.py | 23 ++++++++++---- .../stat_pipeline.py | 31 ++++++++++++++++++- buckaroo/server/data_loading.py | 6 +++- buckaroo/server/data_loading_polars.py | 5 +++ 5 files changed, 76 insertions(+), 8 deletions(-) diff --git a/buckaroo/customizations/pd_stats_v2.py b/buckaroo/customizations/pd_stats_v2.py index 68d281fb7..bb59792d2 100644 --- a/buckaroo/customizations/pd_stats_v2.py +++ b/buckaroo/customizations/pd_stats_v2.py @@ -101,6 +101,25 @@ def _type(is_bool: bool, is_numeric: bool, is_float: bool, return "obj" +def with_type(flags: dict) -> dict: + """``flags``, the ``is_*`` keys a typing stat returns, plus the ``_type`` + derived from them through the ``_type`` stat.""" + return {**flags, '_type': _type(**{key.name: flags[key.name] for key in _type._stat_func.requires})} + + +def schema_stats(ser: pd.Series) -> dict: + """The stats a column's dtype alone determines: ``dtype``, the ``is_*`` + flags and ``_type``. ``typing_stats`` runs on a zero-row slice, so no value + is read, and ``memory_usage`` is left out because it measures the data. + + An ``object`` column gives the typing nothing to go on, and pandas calls an + empty one a string, so every object column types as ``string`` here. Full + stats look at the values and may say ``obj``.""" + flags = typing_stats(ser.iloc[:0]) + del flags['memory_usage'] + return with_type(flags) + + # ============================================================ # Base Summary Stats (replaces DefaultSummaryStats ColAnalysis) # ============================================================ diff --git a/buckaroo/customizations/pl_stats_v2.py b/buckaroo/customizations/pl_stats_v2.py index 687080eae..5e5392a06 100644 --- a/buckaroo/customizations/pl_stats_v2.py +++ b/buckaroo/customizations/pl_stats_v2.py @@ -21,7 +21,7 @@ from buckaroo.pluggable_analysis_framework.column_filters import is_numeric_not_bool # Reused unchanged from pd_stats_v2 — operate on stat dict, not raw series -from buckaroo.customizations.pd_stats_v2 import (_type, computed_default_summary_stats, histogram, cleaning_gen_ops, BaseSummaryResult, NumericStatsResult, HistogramSeriesResult) +from buckaroo.customizations.pd_stats_v2 import (_type, computed_default_summary_stats, histogram, cleaning_gen_ops, with_type, BaseSummaryResult, NumericStatsResult, HistogramSeriesResult) # ============================================================ @@ -44,16 +44,27 @@ def pl_orig_col_name(ser: RawSeries) -> Any: 'is_time': bool, 'is_decimal': bool, 'is_binary': bool, 'memory_usage': int}) -@stat() -def pl_typing_stats(ser: RawSeries) -> PlTypingResult: - """Compute dtype and type flags for a polars column.""" - dt = ser.dtype +def pl_dtype_typing(dt: pl.DataType) -> dict: + """dtype string and type flags for a polars dtype. Needs a schema only, + which is what ``pl_schema_stats`` has to go on.""" return {'dtype': str(dt), 'is_numeric': dt.is_numeric() and dt.base_type() is not pl.Decimal, 'is_integer': dt.is_integer(), 'is_float': dt.is_float(), 'is_bool': dt == pl.Boolean, 'is_datetime': dt.is_temporal() and dt not in (pl.Duration, pl.Time), 'is_timedelta': dt == pl.Duration, 'is_string': dt in (pl.Utf8, pl.String), 'is_categorical': dt == pl.Categorical or isinstance(dt, pl.Enum), 'is_period': False, 'is_interval': False, 'is_time': dt == pl.Time, 'is_decimal': dt.base_type() is pl.Decimal, - 'is_binary': dt == pl.Binary, 'memory_usage': ser.estimated_size()} + 'is_binary': dt == pl.Binary} + + +@stat() +def pl_typing_stats(ser: RawSeries) -> PlTypingResult: + """Compute dtype and type flags for a polars column.""" + return {**pl_dtype_typing(ser.dtype), 'memory_usage': ser.estimated_size()} + + +def pl_schema_stats(ser: pl.Series) -> dict: + """The stats a polars column's dtype alone determines: ``dtype``, the + ``is_*`` flags and ``_type``. Reads ``ser.dtype`` and no value.""" + return with_type(pl_dtype_typing(ser.dtype)) # ============================================================ diff --git a/buckaroo/pluggable_analysis_framework/stat_pipeline.py b/buckaroo/pluggable_analysis_framework/stat_pipeline.py index 67e8c0f77..f121624e0 100644 --- a/buckaroo/pluggable_analysis_framework/stat_pipeline.py +++ b/buckaroo/pluggable_analysis_framework/stat_pipeline.py @@ -9,10 +9,11 @@ from __future__ import annotations import time -from typing import Any, Dict, List, Optional, Tuple +from typing import Any, Callable, Dict, List, Optional, Tuple import pandas as pd +from buckaroo.dataflow.df_types import DataFrameLike from buckaroo.df_util import old_col_new_col from . import perf_log @@ -432,3 +433,31 @@ def errors_to_errdict(errors: List[StatError]) -> ErrDict: err_key = (se.column, se.stat_func.name if se.stat_func else "unknown") errs[err_key] = (se.error, None) return errs + + +def schema_sd(df: DataFrameLike, column_typing: Callable[[Any], Dict[str, Any]], + skip_columns=None) -> SDType: + """The summary dict for ``df`` from its schema alone, with no stat run. + + Shaped like ``StatPipeline.process_df``'s output: one entry per column, + keyed by the rewritten name, holding ``orig_col_name``, + ``rewritten_col_name``, ``length`` and whatever ``column_typing`` returns + for the column. ``column_typing`` is handed the column's series and may read + its dtype only; it is what keeps this from touching a value. As in + ``process_df``, a frame with no rows gives an empty dict and a column in + ``skip_columns`` gets only its two names, so its typing comes from + ``init_sd``. + """ + length = len(df) + if length == 0: + return {} + skip = set(skip_columns or ()) + summary: SDType = {} + for orig_col_name, rewritten_col_name in old_col_new_col(df): + col_stats: Dict[str, Any] = { + 'orig_col_name': orig_col_name, 'rewritten_col_name': rewritten_col_name} + if orig_col_name not in skip and rewritten_col_name not in skip: + col_stats.update(column_typing(df[orig_col_name])) + col_stats['length'] = length + summary[rewritten_col_name] = col_stats + return summary diff --git a/buckaroo/server/data_loading.py b/buckaroo/server/data_loading.py index 1db3385cc..e162e6901 100644 --- a/buckaroo/server/data_loading.py +++ b/buckaroo/server/data_loading.py @@ -12,10 +12,11 @@ from buckaroo.dataflow.dataflow_extras import Sampling from buckaroo.dataflow.autocleaning import PandasAutocleaning from buckaroo.dataflow.styling_core import InitSD, StylingAnalysis -from buckaroo.customizations.pd_stats_v2 import PD_ANALYSIS_V2 +from buckaroo.customizations.pd_stats_v2 import PD_ANALYSIS_V2, schema_stats from buckaroo.customizations.styling import DefaultSummaryStatsStyling, DefaultMainStyling from buckaroo.customizations.pd_autoclean_conf import CleaningConf, NoCleaningConf from buckaroo.pluggable_analysis_framework.df_stats_v2 import DfStatsV2 +from buckaroo.pluggable_analysis_framework.stat_pipeline import schema_sd class ServerSampling(Sampling): @@ -49,6 +50,9 @@ def _df_to_obj(self, df): # No sampling — matches BuckarooInfiniteWidget._df_to_obj return pd_to_obj(df) + def _get_schema_sd(self, processed_df): + return schema_sd(processed_df, schema_stats, skip_columns=self.skip_stat_columns) + def create_dataflow(df: pd.DataFrame, column_config_overrides=None, extra_grid_config=None, init_sd: InitSD | None = None) -> ServerDataflow: diff --git a/buckaroo/server/data_loading_polars.py b/buckaroo/server/data_loading_polars.py index 64e211782..adadee42b 100644 --- a/buckaroo/server/data_loading_polars.py +++ b/buckaroo/server/data_loading_polars.py @@ -27,7 +27,9 @@ from buckaroo.dataflow.styling_core import InitSD from buckaroo.dataflow.autocleaning import PandasAutocleaning from buckaroo.customizations.pl_autocleaning_conf import NoCleaningConfPl +from buckaroo.customizations.pl_stats_v2 import pl_schema_stats from buckaroo.pluggable_analysis_framework.df_stats_v2 import PlDfStatsV2 +from buckaroo.pluggable_analysis_framework.stat_pipeline import schema_sd from buckaroo.polars_buckaroo import ( PLSampling, local_analysis_klasses, prepare_df_for_serialization) from buckaroo.serialization_utils import pd_to_obj, make_infinite_resp @@ -59,6 +61,9 @@ def _df_to_obj(self, df): return pd_to_obj(self.sampling_klass.serialize_sample(df)) return pd_to_obj(self.sampling_klass.serialize_sample(df.to_pandas())) + def _get_schema_sd(self, processed_df): + return schema_sd(processed_df, pl_schema_stats, skip_columns=self.skip_stat_columns) + def load_file_polars(path: str) -> pl.DataFrame: """Eager polars read. Extension dispatch mirrors :func:`load_file`. From d7ba7bd963efae445d3f8a235ed8178f191d8a04 Mon Sep 17 00:00:00 2001 From: Paddy Mullen Date: Sun, 4 Oct 2026 01:26:27 -0400 Subject: [PATCH 09/18] test(server): failing tests for the stats wire format and whole-run stats_request (rows-first s3) On a deferred /load_expr session a stats_request {stats_gen, scope} should return a stats_update with the matching stats_gen whose inline wide payload equals the all_stats an inline session sends, a stale stats_gen should get stats_aborted and run no query, and /load_expr and /reload_expr should bump the generation. df_meta.stats should be injected on every frame and survive a dataflow-field change, which returns the session to the schema tier. With a caps client and a legacy client on one session, the legacy client should keep getting complete messages through the websocket broadcast, the /load_expr, /reload_expr, /load and /load_compare pushes and the highlight overlay, while the caps client gets a stats-free frame and then pulls a stats_update. A spy telemetry sink should see a stats.request span. Co-Authored-By: Claude Sonnet 5.5 --- tests/unit/server/test_load_expr.py | 500 ++++++++++++++++++++++++++++ 1 file changed, 500 insertions(+) diff --git a/tests/unit/server/test_load_expr.py b/tests/unit/server/test_load_expr.py index 36557afb7..faf694d91 100644 --- a/tests/unit/server/test_load_expr.py +++ b/tests/unit/server/test_load_expr.py @@ -1,17 +1,21 @@ """End-to-end tests for POST /load_expr — server load path for XorqBuckarooInfiniteWidget over a xorq/ibis expression.""" +import datetime import io import json import os import shutil import sys import tempfile +from contextlib import contextmanager from pathlib import Path +from types import SimpleNamespace from unittest.mock import MagicMock, patch import pandas as pd import pyarrow.parquet as pq import pytest +import tornado.gen import tornado.httpclient import tornado.testing import tornado.websocket @@ -28,8 +32,11 @@ from buckaroo.dataflow.sd_cache import split_chain_by_scope # noqa: E402 from buckaroo.jlisp.lisp_utils import s as lisp_sym # noqa: E402 from buckaroo.pluggable_analysis_framework.col_analysis import ColAnalysis # noqa: E402 +from buckaroo.pluggable_analysis_framework.xorq_stat_pipeline import XorqStatPipeline # noqa: E402 +from buckaroo.serialization_utils import resolve_summary_stats_payload # noqa: E402 from buckaroo.server import telemetry, xorq_loading # noqa: E402 from buckaroo.server.app import make_app as _make_app # noqa: E402 +from buckaroo.server.websocket_handler import DataStreamHandler # noqa: E402 pytestmark = pytest.mark.skipif( sys.platform == "win32", @@ -1365,6 +1372,499 @@ async def test_reload_expr_accepts_a_new_pair_and_stores_it(self): shutil.rmtree(project_root, ignore_errors=True) +def _build_stats_wire_dir(builds_root): + """Build the ``_stats_tier_expr`` table to ``builds_root``. Its float column + has a 1e9 maximum, so full stats change that column's ``minWidth`` and a + message built from the schema tier differs from a complete one there.""" + expr = xo.memtable({ + "price": [12.5, 18.9, 7.4, 22.1, 1e9], + "qty": [1, 2, 1, 3, 2], + "category": ["a", "b", "a", "c", "b"]}, name="t") + return str(xo.build_expr(expr, builds_dir=builds_root)) + + +def _state_change(**changes): + new_state = {"post_processing": "", "cleaning_method": "", "quick_command_args": {}, + "df_display": "main", "show_commands": False, "sampled": False, "search_string": ""} + return json.dumps({"type": "buckaroo_state_change", "new_state": {**new_state, **changes}}) + + +def _stats_request(stats_gen, **fields): + return json.dumps({"type": "stats_request", "stats_gen": stats_gen, "scope": "raw", **fields}) + + +async def _read_json(ws, timeout=3.0): + """The next frame on ``ws``, decoded. Raises AssertionError, not a hang, + when none arrives.""" + try: + frame = await tornado.gen.with_timeout( + datetime.timedelta(seconds=timeout), ws.read_message()) + except tornado.gen.TimeoutError: + raise AssertionError(f"no frame within {timeout}s") from None + assert frame is not None, "the connection closed" + return json.loads(frame) + + +def _rows_by_stat(payload): + """The decoded ``all_stats`` rows of a payload, keyed by stat name.""" + return {row["index"]: row for row in resolve_summary_stats_payload(payload)} + + +def _comparable(frame): + """An ``initial_state`` frame as a JSON string, for comparing a legacy + client's frame with what an inline session sends: ``all_stats`` decoded + (its parquet bytes follow column order) and ``df_meta.stats`` dropped.""" + frame = json.loads(json.dumps(frame)) + frame["df_data_dict"]["all_stats"] = _rows_by_stat(frame["df_data_dict"]["all_stats"]) + frame["df_meta"].pop("stats", None) + return _as_json(frame) + + +@contextmanager +def _count_stat_queries(): + """Record every query ``XorqStatPipeline`` sends while the block runs.""" + queries = [] + original = XorqStatPipeline._execute + + def spy(pipeline, query): + queries.append(query) + return original(pipeline, query) + + with patch.object(XorqStatPipeline, "_execute", spy): + yield queries + + +class TestStatsWire(tornado.testing.AsyncHTTPTestCase): + """``stats_request``, ``stats_update`` and ``stats_aborted`` on a deferred + ``/load_expr`` session, the ``stats_gen`` counter and ``df_meta.stats``, and + the per-connection capability (``?caps=stats_update``) that keeps a client + without it on complete messages (rows-first s3).""" + + def get_app(self): + return make_app() + + def setUp(self): + super().setUp() + self.builds_root = tempfile.mkdtemp() + self.project_root = tempfile.mkdtemp() + self.addCleanup(shutil.rmtree, self.builds_root, ignore_errors=True) + self.addCleanup(shutil.rmtree, self.project_root, ignore_errors=True) + pp_dir = os.path.join(self.project_root, "post_processing") + os.makedirs(pp_dir) + with open(os.path.join(pp_dir, "first_three.py"), "w") as f: + f.write("def process(expr):\n return expr.limit(3)\n") + self.build_path = _build_stats_wire_dir(self.builds_root) + self.clients = [] + + def tearDown(self): + for ws in self.clients: + ws.close() + super().tearDown() + + def _session(self, sid): + return self._app.settings["sessions"].get(sid) + + async def _load(self, sid, **body): + resp = await _post(self.get_http_port(), "/load_expr", + {"session": sid, "build_dir": self.build_path, "project_root": self.project_root, **body}) + self.assertEqual(resp.code, 200, resp.body) + + async def _connect(self, sid, caps=None): + """Open a WebSocket, with ``caps`` as ``?caps=``. Returns it with its + first ``initial_state``.""" + suffix = f"?caps={caps}" if caps else "" + ws = await tornado.websocket.websocket_connect( + f"ws://localhost:{self.get_http_port()}/ws/{sid}{suffix}") + self.clients.append(ws) + return ws, await _read_json(ws) + + def _stats(self, frame): + stats = frame["df_meta"].get("stats") + self.assertIsNotNone(stats, "initial_state carries no df_meta.stats") + return stats + + async def _inline_frame(self, sid, **changes): + """What an inline session of the same build sends a client: the complete + message every complete message here must match. ``changes`` are sent as + a state change first, and the broadcast that follows is returned.""" + await self._load(sid) + ws, frame = await self._connect(sid) + if changes: + ws.write_message(_state_change(**changes)) + frame = await _read_json(ws) + return frame + + def _assert_pending(self, frame, gen): + stats = self._stats(frame) + self.assertEqual((stats["status"], stats["tier"], stats["gen"]), ("pending", "schema", gen)) + self.assertEqual(list(_rows_by_stat(frame["df_data_dict"]["all_stats"])), ["dtype"], + "a pending frame must carry the schema tier only") + + def _assert_complete(self, frame, gen, inline_frame): + stats = self._stats(frame) + self.assertEqual((stats["status"], stats["tier"], stats["gen"]), ("complete", "full", gen)) + self.assertIn("histogram_bins", _rows_by_stat(frame["df_data_dict"]["all_stats"])) + self.assertEqual(_comparable(frame), _comparable(inline_frame)) + + async def _pair(self, sid): + """A caps client, then a legacy client, on a deferred session. The + legacy client's connect completes the stats, so a later push is the + first thing to put the session back to pending.""" + await self._load(sid, stats_delivery="deferred") + a, a_open = await self._connect(sid, caps="stats_update") + b, b_open = await self._connect(sid) + return a, a_open, b, b_open + + @tornado.testing.gen_test + async def test_stats_request_returns_a_stats_update_that_completes_the_stats(self): + await self._load("sw-complete", stats_delivery="deferred") + inline = await self._inline_frame("sw-complete-inline") + ws, first = await self._connect("sw-complete", caps="stats_update") + gen = self._stats(first)["gen"] + self._assert_pending(first, gen) + self.assertEqual(first["protocol_version"], 1) + + ws.write_message(_stats_request(gen)) + update = await _read_json(ws) + + self.assertEqual(update["type"], "stats_update") + self.assertEqual(update["stats_gen"], gen) + self.assertEqual((update["scope"], update["tier"], update["final"]), ("raw", "full", True)) + self.assertEqual((update["payload"]["format"], update["payload"]["layout"]), + ("parquet_b64", "wide"), "the payload must be inline, so binary pairing stays single-slot") + self.assertEqual(_rows_by_stat(update["payload"]), + _rows_by_stat(inline["df_data_dict"]["all_stats"])) + self.assertIsInstance(update["elapsed_ms"], (int, float)) + + @tornado.testing.gen_test + async def test_rows_are_served_before_and_after_the_stats_request(self): + await self._load("sw-rows", stats_delivery="deferred") + ws, first = await self._connect("sw-rows", caps="stats_update") + window = {"start": 0, "end": 5, "sourceName": "default", "origEnd": 5} + ws.write_message(json.dumps({"type": "infinite_request", "payload_args": window})) + self.assertEqual((await _read_json(ws))["length"], 5) + pq.read_table(io.BytesIO(await ws.read_message())) + ws.write_message(_stats_request(self._stats(first)["gen"])) + self.assertEqual((await _read_json(ws))["type"], "stats_update") + ws.write_message(json.dumps({"type": "infinite_request", "payload_args": window})) + self.assertEqual((await _read_json(ws))["type"], "infinite_resp", + "a stats_update must leave no stray binary frame in the stream") + + @tornado.testing.gen_test + async def test_a_stale_stats_gen_gets_stats_aborted_and_runs_nothing(self): + await self._load("sw-stale", stats_delivery="deferred") + ws, first = await self._connect("sw-stale", caps="stats_update") + gen = self._stats(first)["gen"] + + with _count_stat_queries() as queries: + ws.write_message(_stats_request(gen - 1)) + aborted = await _read_json(ws) + self.assertEqual(aborted["type"], "stats_aborted") + self.assertEqual((aborted["stats_gen"], aborted["current_gen"], aborted["reason"]), + (gen - 1, gen, "stale")) + self.assertEqual(queries, [], "a stale request must run no stat query") + + ws.write_message(_stats_request(gen)) + self.assertEqual((await _read_json(ws))["type"], "stats_update", + "the stale request must not have consumed or failed the session's stats") + + @tornado.testing.gen_test + async def test_an_unsupported_scope_gets_stats_aborted(self): + await self._load("sw-scope", stats_delivery="deferred") + ws, first = await self._connect("sw-scope", caps="stats_update") + ws.write_message(_stats_request(self._stats(first)["gen"], scope="filt")) + aborted = await _read_json(ws) + self.assertEqual((aborted["type"], aborted["reason"], aborted["scope"]), + ("stats_aborted", "unsupported_scope", "filt")) + + @tornado.testing.gen_test + async def test_load_expr_and_reload_expr_bump_stats_gen(self): + await self._load("sw-gen", stats_delivery="deferred") + ws, first = await self._connect("sw-gen", caps="stats_update") + gen = self._stats(first)["gen"] + + await self._load("sw-gen", stats_delivery="deferred", force_reload=True) + pushed = await _read_json(ws) + self.assertGreater(self._stats(pushed)["gen"], gen, "/load_expr must bump stats_gen") + gen = self._stats(pushed)["gen"] + + resp = await _post(self.get_http_port(), "/reload_expr/sw-gen", {}) + self.assertEqual(resp.code, 200, resp.body) + pushed = await _read_json(ws) + self.assertGreater(self._stats(pushed)["gen"], gen, "/reload_expr must bump stats_gen") + + @tornado.testing.gen_test + async def test_df_meta_stats_survives_a_dataflow_field_change(self): + await self._load("sw-change", stats_delivery="deferred") + ws, first = await self._connect("sw-change", caps="stats_update") + gen = self._stats(first)["gen"] + ws.write_message(_stats_request(gen)) + self.assertEqual((await _read_json(ws))["type"], "stats_update") + + with _count_stat_queries() as queries: + ws.write_message(_state_change(quick_command_args={"search": ["a"]})) + changed = await _read_json(ws) + self._assert_pending(changed, gen + 1) + self.assertEqual(queries, [], "a dataflow-field change on a deferred session must run no stats") + # The dataflow rebuilt df_meta for the filtered state; its keys are still there. + self.assertEqual((changed["df_meta"]["total_rows"], changed["df_meta"]["filtered_rows"]), (5, 2)) + + ws.write_message(_stats_request(gen)) + aborted = await _read_json(ws) + self.assertEqual((aborted["type"], aborted["current_gen"], aborted["reason"]), + ("stats_aborted", gen + 1, "stale")) + ws.write_message(_stats_request(gen + 1)) + update = await _read_json(ws) + self.assertEqual((update["type"], update["stats_gen"], update["final"]), + ("stats_update", gen + 1, True)) + + @tornado.testing.gen_test + async def test_a_legacy_client_stays_complete_while_a_caps_client_gets_a_stats_free_frame(self): + sid = "sw-ab" + inline = await self._inline_frame("sw-ab-inline", post_processing="first_three") + await self._load(sid, stats_delivery="deferred") + a1, a1_open = await self._connect(sid, caps="other,stats_update") + a2, _ = await self._connect(sid, caps="stats_update") + # Connecting runs the missing stats for a legacy client, so the session + # is complete when B1 changes it. + b1, b1_open = await self._connect(sid, caps="unknown") + b2, _ = await self._connect(sid) + gen = self._stats(a1_open)["gen"] + self.assertEqual(self._stats(b1_open)["status"], "complete") + + with _count_stat_queries() as queries: + b1.write_message(_state_change(post_processing="first_three")) + frames = {name: await _read_json(ws) + for name, ws in (("a1", a1), ("a2", a2), ("b1", b1), ("b2", b2))} + for name in ("a1", "a2"): + self._assert_pending(frames[name], gen + 1) + for name in ("b1", "b2"): + self._assert_complete(frames[name], gen + 1, inline) + ran = len(queries) + self.assertGreater(ran, 0, "B's frame must have run the stats it carries") + + a1.write_message(_stats_request(gen + 1)) + update = await _read_json(a1) + self.assertEqual((update["type"], update["stats_gen"]), ("stats_update", gen + 1)) + self.assertEqual(_rows_by_stat(update["payload"]), + _rows_by_stat(frames["b1"]["df_data_dict"]["all_stats"])) + self.assertEqual(len(queries), ran, + "A's request must be answered from the stats B's frame computed") + + @tornado.testing.gen_test + async def test_load_expr_push_keeps_a_legacy_client_complete(self): + sid = "sw-push-load-expr" + inline = await self._inline_frame("sw-push-load-expr-inline") + a, a_open, b, _ = await self._pair(sid) + gen = self._stats(a_open)["gen"] + await self._load(sid, stats_delivery="deferred", force_reload=True) + self._assert_pending(await _read_json(a), gen + 1) + self._assert_complete(await _read_json(b), gen + 1, inline) + + @tornado.testing.gen_test + async def test_reload_expr_push_keeps_a_legacy_client_complete(self): + sid = "sw-push-reload" + inline = await self._inline_frame("sw-push-reload-inline") + a, a_open, b, _ = await self._pair(sid) + gen = self._stats(a_open)["gen"] + resp = await _post(self.get_http_port(), f"/reload_expr/{sid}", {}) + self.assertEqual(resp.code, 200, resp.body) + self._assert_pending(await _read_json(a), gen + 1) + self._assert_complete(await _read_json(b), gen + 1, inline) + + @tornado.testing.gen_test + async def test_load_push_leaves_both_clients_complete(self): + """/load swaps the session to pandas, which has no deferred stats: the + stored policy must not leave it looking pending.""" + sid = "sw-push-load" + a, _, b, _ = await self._pair(sid) + csv_fd, csv_path = tempfile.mkstemp(suffix=".csv") + os.close(csv_fd) + try: + pd.DataFrame({"x": [1, 2, 3], "y": ["p", "q", "r"]}).to_csv(csv_path, index=False) + resp = await _post(self.get_http_port(), "/load", + {"session": sid, "path": csv_path, "mode": "buckaroo"}) + self.assertEqual(resp.code, 200, resp.body) + finally: + os.unlink(csv_path) + frames = [await _read_json(a), await _read_json(b)] + for frame in frames: + self.assertNotIn("stats", frame["df_meta"]) + self.assertIn("histogram_bins", _rows_by_stat(frame["df_data_dict"]["all_stats"])) + self.assertEqual(_comparable(frames[0]), _comparable(frames[1])) + + @tornado.testing.gen_test + async def test_load_compare_push_leaves_both_clients_complete(self): + sid = "sw-push-compare" + a, _, b, _ = await self._pair(sid) + paths = [] + try: + for frame in (pd.DataFrame({"id": [1, 2], "v": [10, 20]}), + pd.DataFrame({"id": [1, 3], "v": [10, 30]})): + fd, path = tempfile.mkstemp(suffix=".csv") + os.close(fd) + frame.to_csv(path, index=False) + paths.append(path) + resp = await _post(self.get_http_port(), "/load_compare", + {"session": sid, "path1": paths[0], "path2": paths[1], "join_columns": ["id"]}) + self.assertEqual(resp.code, 200, resp.body) + finally: + for path in paths: + os.unlink(path) + frames = [await _read_json(a), await _read_json(b)] + for frame in frames: + self.assertNotIn("stats", frame["df_meta"]) + self.assertEqual(_comparable(frames[0]), _comparable(frames[1])) + + @tornado.testing.gen_test + async def test_highlight_overlay_is_complete_for_a_legacy_client_and_stats_free_for_a_caps_client(self): + sid = "sw-overlay" + await self._load(sid, stats_delivery="deferred") + a, a_open = await self._connect(sid, caps="stats_update") + gen = self._stats(a_open)["gen"] + + a.write_message(_state_change(search_string="ca")) + a_overlay = await _read_json(a) + self._assert_pending(a_overlay, gen) + + b, _ = await self._connect(sid) + b.write_message(_state_change(search_string="ca")) + b_overlay = await _read_json(b) + self.assertEqual(self._stats(b_overlay)["status"], "complete") + self.assertIn("histogram_bins", _rows_by_stat(b_overlay["df_data_dict"]["all_stats"])) + # The overlay's display config is the complete one with the highlight on top. + session_args = self._session(sid).df_display_args["main"]["df_viewer_config"]["column_config"] + by_col = {cc["col_name"]: cc for cc in b_overlay["df_display_args"]["main"]["df_viewer_config"]["column_config"]} + for expected in session_args: + self.assertEqual(by_col[expected["col_name"]]["ag_grid_specs"], expected["ag_grid_specs"]) + highlighted = [cc for cc in by_col.values() + if cc.get("displayer_args", {}).get("highlight_phrase") == ["ca"]] + self.assertTrue(highlighted, "the overlay must still carry the highlight") + + @tornado.testing.gen_test + async def test_overlay_for_a_legacy_client_is_built_after_its_stats_are_completed(self): + """Completing the stats replaces the session's display config, so an + overlay that copied it first would send a legacy client the schema + tier's. Every send completes a session that has a legacy client + connected, so only a direct call reaches a pending session here.""" + await self._load("sw-overlay-order", stats_delivery="deferred") + session = self._session("sw-overlay-order") + sent: list = [] + legacy = SimpleNamespace(search_string="ca", caps=frozenset(), session_id="sw-overlay-order", + write_message=sent.append) + self.assertEqual(getattr(session, "stats_status", None), "pending") + + DataStreamHandler._send_highlight_overlay(legacy, session) + + self.assertEqual(session.stats_status, "complete") + overlay = json.loads(sent[0])["df_display_args"]["main"]["df_viewer_config"]["column_config"] + complete = session.df_display_args["main"]["df_viewer_config"]["column_config"] + self.assertEqual({cc["col_name"]: cc["ag_grid_specs"] for cc in overlay}, + {cc["col_name"]: cc["ag_grid_specs"] for cc in complete}) + + @tornado.testing.gen_test + async def test_a_client_connecting_after_completion_gets_the_complete_state(self): + inline = await self._inline_frame("sw-late-inline") + await self._load("sw-late", stats_delivery="deferred") + a, a_open = await self._connect("sw-late", caps="stats_update") + gen = self._stats(a_open)["gen"] + a.write_message(_stats_request(gen)) + self.assertEqual((await _read_json(a))["type"], "stats_update") + + with _count_stat_queries() as queries: + _, late = await self._connect("sw-late", caps="stats_update") + _, legacy = await self._connect("sw-late") + self._assert_complete(late, gen, inline) + self._assert_complete(legacy, gen, inline) + self.assertEqual(queries, [], "stats computed once must not be computed again for a later client") + + @tornado.testing.gen_test + async def test_a_legacy_client_connecting_to_a_pending_session_completes_it(self): + inline = await self._inline_frame("sw-open-inline") + await self._load("sw-open", stats_delivery="deferred") + _, legacy = await self._connect("sw-open") + self._assert_complete(legacy, self._stats(legacy)["gen"], inline) + + @tornado.testing.gen_test + async def test_a_session_targeting_the_schema_tier_is_not_computed(self): + await self._load("sw-schema", stats_tier="schema", stats_delivery="deferred") + ws, first = await self._connect("sw-schema", caps="stats_update") + stats = self._stats(first) + self.assertEqual((stats["status"], stats["tier"], stats["reason"]), ("not_computed", "schema", "host")) + ws.write_message(_stats_request(stats["gen"])) + aborted = await _read_json(ws) + self.assertEqual((aborted["type"], aborted["reason"]), ("stats_aborted", "not_requestable")) + # Nothing is missing relative to the schema tier, so a legacy client + # gets the schema tier, as it does from an inline schema session. + _, legacy = await self._connect("sw-schema") + self.assertEqual(list(_rows_by_stat(legacy["df_data_dict"]["all_stats"])), ["dtype"]) + await self._load("sw-schema-inline", stats_tier="schema") + _, inline_schema = await self._connect("sw-schema-inline") + self.assertEqual(self._stats(inline_schema)["status"], "not_computed") + + @tornado.testing.gen_test + async def test_a_failed_stats_run_reports_error_until_the_next_generation(self): + await self._load("sw-error", stats_delivery="deferred") + ws, first = await self._connect("sw-error", caps="stats_update") + gen = self._stats(first)["gen"] + + with patch.object(xorq_loading.XorqServerDataflow, "_get_summary_sd", + side_effect=RuntimeError("stats query failed")): + ws.write_message(_stats_request(gen)) + aborted = await _read_json(ws) + self.assertEqual((aborted["type"], aborted["stats_gen"], aborted["reason"]), ("stats_aborted", gen, "error")) + + # The failure is the session's state for this generation: a client that + # connects now is told so, and a legacy client still gets its frame. + _, caps_frame = await self._connect("sw-error", caps="stats_update") + stats = self._stats(caps_frame) + self.assertEqual((stats["status"], stats["tier"], stats["gen"]), ("error", "schema", gen)) + self.assertIn("reason", stats) + _, legacy = await self._connect("sw-error") + self.assertEqual(legacy["type"], "initial_state") + self.assertEqual(list(_rows_by_stat(legacy["df_data_dict"]["all_stats"])), ["dtype"]) + + with _count_stat_queries() as queries: + ws.write_message(_stats_request(gen)) + self.assertEqual((await _read_json(ws))["reason"], "error") + self.assertEqual(queries, [], "a failed generation must not be retried by every request") + + ws.write_message(_state_change(quick_command_args={"search": ["a"]})) + self._assert_pending(await _read_json(ws), gen + 1) + ws.write_message(_stats_request(gen + 1)) + self.assertEqual((await _read_json(ws))["type"], "stats_update") + + @tornado.testing.gen_test + async def test_stats_request_emits_a_stats_request_span(self): + captured: list = [] + with patch.object(telemetry, "make_http_sink", lambda url, **kw: captured.append): + await self._load("sw-span", stats_delivery="deferred", + telemetry_url="http://companion.invalid/internal/telemetry") + ws, first = await self._connect("sw-span", caps="stats_update") + gen = self._stats(first)["gen"] + ws.write_message(_stats_request(gen - 1)) + await _read_json(ws) + ws.write_message(_stats_request(gen, columns=["price"])) + await _read_json(ws) + + stale, served = [r for r in captured if r["name"] == "stats.request"] + self.assertEqual(served["trace"], "sw-span") + self.assertEqual((served["attrs"]["stats_gen"], served["attrs"]["scope"], served["attrs"]["tier"]), + (gen, "raw", "full")) + self.assertEqual((served["attrs"]["outcome"], served["attrs"]["columns"]), ("update", 1)) + self.assertEqual((stale["attrs"]["outcome"], stale["attrs"]["stats_gen"]), ("stale", gen - 1)) + self.assertIn("firstpull.stats_total", [r["name"] for r in captured]) + + @tornado.testing.gen_test + async def test_inline_sessions_send_no_df_meta_stats(self): + """Default behaviour: a session with the default policy sends the + message it always has, and a client reads the absence as complete.""" + await self._load("sw-inline") + _, frame = await self._connect("sw-inline", caps="stats_update") + self.assertNotIn("stats", frame["df_meta"]) + self.assertIn("histogram_bins", _rows_by_stat(frame["df_data_dict"]["all_stats"])) + + class TestReloadExpr(tornado.testing.AsyncHTTPTestCase): def get_app(self): return make_app() From 6bd63ed7e14c6e31b33bbdf6c6910735e5e07ffe Mon Sep 17 00:00:00 2001 From: Paddy Mullen Date: Sun, 4 Oct 2026 01:40:09 -0400 Subject: [PATCH 10/18] test(server): failing tests for stats reuse, component_config and generation edge cases (rows-first s3) Four more cases for the stats wire format, kept in their own commit so each is seen failing on CI before the implementation lands. Returning to a state whose stats were completed once is answered from summary_stats_cache with no query. Completing the stats keeps the session's component_config on the refreshed display config. A warm /load_expr, which rebuilds nothing, leaves stats_gen alone. A stats_request on a session with no data is answered with stats_aborted. Co-Authored-By: Claude Sonnet 5.5 --- tests/unit/server/test_load_expr.py | 48 +++++++++++++++++++++++++++++ 1 file changed, 48 insertions(+) diff --git a/tests/unit/server/test_load_expr.py b/tests/unit/server/test_load_expr.py index faf694d91..5ab61e063 100644 --- a/tests/unit/server/test_load_expr.py +++ b/tests/unit/server/test_load_expr.py @@ -1855,6 +1855,54 @@ async def test_stats_request_emits_a_stats_request_span(self): self.assertEqual((stale["attrs"]["outcome"], stale["attrs"]["stats_gen"]), ("stale", gen - 1)) self.assertIn("firstpull.stats_total", [r["name"] for r in captured]) + @tornado.testing.gen_test + async def test_returning_to_a_completed_state_is_answered_from_the_cache(self): + await self._load("sw-revisit", stats_delivery="deferred") + ws, first = await self._connect("sw-revisit", caps="stats_update") + gen = self._stats(first)["gen"] + ws.write_message(_stats_request(gen)) + original = await _read_json(ws) + ws.write_message(_state_change(quick_command_args={"search": ["a"]})) + self._assert_pending(await _read_json(ws), gen + 1) + ws.write_message(_stats_request(gen + 1)) + self.assertEqual((await _read_json(ws))["type"], "stats_update") + + ws.write_message(_state_change(quick_command_args={})) + self._assert_pending(await _read_json(ws), gen + 2) + with _count_stat_queries() as queries: + ws.write_message(_stats_request(gen + 2)) + again = await _read_json(ws) + self.assertEqual(queries, [], "the first state's stats are in summary_stats_cache") + self.assertEqual(_rows_by_stat(again["payload"]), _rows_by_stat(original["payload"])) + + @tornado.testing.gen_test + async def test_completing_the_stats_keeps_component_config(self): + await self._load("sw-theme", stats_delivery="deferred", component_config={"className": "sw-themed"}) + ws, first = await self._connect("sw-theme", caps="stats_update") + ws.write_message(_stats_request(self._stats(first)["gen"])) + self.assertEqual((await _read_json(ws))["type"], "stats_update") + _, late = await self._connect("sw-theme", caps="stats_update") + dvc = late["df_display_args"]["main"]["df_viewer_config"] + self.assertEqual(dvc["component_config"]["className"], "sw-themed") + + @tornado.testing.gen_test + async def test_a_warm_load_expr_keeps_the_generation(self): + await self._load("sw-warm", stats_delivery="deferred") + _, first = await self._connect("sw-warm", caps="stats_update") + gen = self._stats(first)["gen"] + await self._load("sw-warm", stats_delivery="deferred") + _, again = await self._connect("sw-warm", caps="stats_update") + self.assertEqual(self._stats(again)["gen"], gen, "a warm exit rebuilds nothing, so the generation stands") + + @tornado.testing.gen_test + async def test_a_stats_request_with_no_data_loaded_is_aborted(self): + ws = await tornado.websocket.websocket_connect( + f"ws://localhost:{self.get_http_port()}/ws/sw-no-data?caps=stats_update") + self.clients.append(ws) + ws.write_message(_stats_request(1)) + aborted = await _read_json(ws) + self.assertEqual((aborted["type"], aborted["reason"]), ("stats_aborted", "no_data")) + @tornado.testing.gen_test async def test_inline_sessions_send_no_df_meta_stats(self): """Default behaviour: a session with the default policy sends the From 9bfebfb9a48b7d5e51015ace24212723fd8efb86 Mon Sep 17 00:00:00 2001 From: Paddy Mullen Date: Sun, 4 Oct 2026 01:52:34 -0400 Subject: [PATCH 11/18] feat(server): stats_request, stats_update and df_meta.stats on deferred sessions (rows-first s3) A client that advertises ?caps=stats_update gets a stats-free initial_state on a deferred /load_expr session (df_meta.stats.status "pending") and pulls the stats with stats_request {stats_gen, scope}. The reply is a stats_update carrying the dataflow's all_stats as an inline wide envelope, or stats_aborted when the generation is stale. The request is the whole run: one synchronous call that computes the full stats, writes the full-tier summary_stats_cache entry, assigns summary_sd and refreshes the session snapshot through one helper. stats_gen is a server-owned counter bumped by every load handler and by a state change that touches a dataflow field, which also returns a deferred session to the schema tier. df_meta.stats is injected by build_state_message from the session, since the dataflow rebuilds df_meta wholesale. Every send site goes through build_state_message_for, so a client without the capability still gets complete messages (its missing stats run synchronously first); broadcast_state replaces the five copies of the send loop and sends to capable clients first. The stats_request branch binds the session's telemetry sink and emits a stats.request span. Co-Authored-By: Claude Sonnet 5.5 --- buckaroo/server/handlers.py | 81 +++------- buckaroo/server/session.py | 64 +++++++- buckaroo/server/stats_wire.py | 231 +++++++++++++++++++++++++++ buckaroo/server/websocket_handler.py | 76 +++++---- 4 files changed, 364 insertions(+), 88 deletions(-) create mode 100644 buckaroo/server/stats_wire.py diff --git a/buckaroo/server/handlers.py b/buckaroo/server/handlers.py index b844bd3e0..48100161a 100644 --- a/buckaroo/server/handlers.py +++ b/buckaroo/server/handlers.py @@ -16,7 +16,8 @@ from buckaroo.server.focus import find_or_create_session_window from buckaroo.dataflow.dataflow import STATS_TIERS from buckaroo.server.session import ( - STATS_DELIVERIES, build_state_message, dataflow_stats_tier) + STATS_DELIVERIES, begin_stats_generation, dataflow_stats_tier) +from buckaroo.server.stats_wire import broadcast_state, refresh_session_snapshot from buckaroo.server import telemetry from buckaroo.pluggable_analysis_framework import perf_log @@ -194,17 +195,10 @@ def _push_state_to_clients(self, session, metadata: dict): if not session.ws_clients: return - for client in list(session.ws_clients): - try: - # Reset per-client live search first (dataset changed), - # then build the msg so the injected ``buckaroo_state`` - # mirrors the post-reset value. - client.search_string = "" - msg = build_state_message(session, metadata=metadata, - search_string=client.search_string) - client.write_message(json.dumps(msg)) - except Exception: - session.ws_clients.discard(client) + # Each client's search is reset first (dataset changed), then its + # message is built so the injected ``buckaroo_state`` mirrors the + # post-reset value. + broadcast_state(session, metadata=metadata, reset_search=True) def _handle_browser_window(self, session_id: str) -> str: """Handle browser window management.""" @@ -373,6 +367,11 @@ async def post(self): **component_config, } + # /load builds no deferred stats, whatever policy a prior /load_expr on + # this session left behind, and the generation moves on with the data. + session.stats_tier, session.stats_delivery = "full", "inline" + begin_stats_generation(session) + # Notify connected clients and open browser self._push_state_to_clients(session, metadata) browser_action = "skipped" if no_browser else self._handle_browser_window(session_id) @@ -660,17 +659,13 @@ async def post(self): **dvc.get("component_config", {}), **component_config} - if session.ws_clients: - for client in list(session.ws_clients): - try: - # Reset per-client live search (#851): a term from - # the prior expression would silently filter the new one. - client.search_string = "" - msg = build_state_message(session, metadata=metadata, - search_string=client.search_string) - client.write_message(json.dumps(msg)) - except Exception: - session.ws_clients.discard(client) + # A new expression is a new stats generation; a deferred session starts + # it pending. + begin_stats_generation(session) + + # Reset per-client live search (#851): a term from the prior + # expression would silently filter the new one. + broadcast_state(session, metadata=metadata, reset_search=True) if no_browser or not self.application.settings.get("open_browser", False): browser_action = "skipped" @@ -807,18 +802,12 @@ async def post(self): session.df_data_dict = display_state["df_data_dict"] session.df_meta = display_state["df_meta"] session.mode = "viewer" + # A viewer session has no dataflow and so no deferred stats. + session.stats_tier, session.stats_delivery = "full", "inline" + begin_stats_generation(session) - # Push to WebSocket clients - if session.ws_clients: - for client in list(session.ws_clients): - try: - # Reset per-client live search (#851). - client.search_string = "" - msg = build_state_message(session, - search_string=client.search_string) - client.write_message(json.dumps(msg)) - except Exception: - session.ws_clients.discard(client) + # Push to WebSocket clients. Reset per-client live search (#851). + broadcast_state(session, reset_search=True) # Browser window if no_browser or not self.application.settings.get("open_browser", False): @@ -935,31 +924,13 @@ async def post(self, session_id): if bs.get("quick_command_args"): xorq_dataflow.quick_command_args = bs["quick_command_args"] - refreshed = get_buckaroo_display_state(xorq_dataflow) session.xorq_dataflow = xorq_dataflow session.stats_tier = stats_tier session.stats_delivery = stats_delivery - session.df_display_args = refreshed["df_display_args"] - session.df_data_dict = refreshed["df_data_dict"] - session.df_meta = refreshed["df_meta"] - session.buckaroo_options = refreshed["buckaroo_options"] - session.command_config = refreshed["command_config"] + refresh_session_snapshot(session, xorq_dataflow) + begin_stats_generation(session) - if session.component_config and session.df_display_args: - for key in session.df_display_args: - dvc = session.df_display_args[key].get("df_viewer_config") - if dvc is not None: - dvc["component_config"] = { - **dvc.get("component_config", {}), - **session.component_config} - - for client in list(session.ws_clients): - try: - msg = build_state_message(session, - search_string=getattr(client, "search_string", "")) - client.write_message(json.dumps(msg)) - except Exception: - session.ws_clients.discard(client) + broadcast_state(session) klass_count = len(extra_klasses) log.info("reload_expr session=%s project_root=%s klasses=%d", diff --git a/buckaroo/server/session.py b/buckaroo/server/session.py index bafa5a870..dc09c377b 100644 --- a/buckaroo/server/session.py +++ b/buckaroo/server/session.py @@ -23,6 +23,25 @@ def dataflow_stats_tier(stats_tier: str, stats_delivery: str) -> str: return "schema" if stats_delivery == "deferred" else stats_tier +# What ``df_meta.stats.status`` says about a session's stats for its current +# ``stats_gen``: ``complete`` (the stats of the tier the session is headed for +# are in the snapshot), ``pending`` (a deferred session that has not produced +# them yet), ``not_computed`` (the session is headed for the schema tier, so +# none will arrive) and ``error`` (the run failed; sticky until the next +# generation). +STATS_STATUSES = ("complete", "pending", "not_computed", "error") + + +def initial_stats_status(stats_tier: str, stats_delivery: str) -> tuple[str, Optional[str]]: + """The status, and its reason, of a session whose stats generation has just + started, from the policy pair alone.""" + if stats_tier != "full": + return "not_computed", "host" + if stats_delivery == "deferred": + return "pending", None + return "complete", None + + @dataclass class SessionState: session_id: str @@ -61,6 +80,16 @@ class SessionState: # tier the dataflow is built at follows from the pair (dataflow_stats_tier). stats_tier: str = "full" stats_delivery: str = "inline" + # The stats generation: a counter the server owns, bumped whenever the + # dataflow state the stats describe changes (/load, /load_expr, /load_compare, + # /reload_expr, and a buckaroo_state_change that touches a dataflow field). + # It rides on df_meta.stats and on every stats message, so a client can drop + # a reply for a state it has left. Independent of any client-owned request + # sequence (state_seq). ``stats_status`` and ``stats_reason`` live here, not + # in ``df_meta``, because the dataflow rebuilds df_meta wholesale. + stats_gen: int = 0 + stats_status: str = "complete" + stats_reason: Optional[str] = None # Companion telemetry sink (#943): a fire-and-forget POST callable, built # once from the /load_expr payload's telemetry_url on the IOLoop (where # make_http_sink captures AsyncHTTPClient/IOLoop.current()). Stored here so @@ -98,6 +127,33 @@ def touch(self) -> None: this field is the runtime escape hatch.""" +def begin_stats_generation(session: "SessionState") -> None: + """Start a new stats generation: the session's dataflow state has changed, + so stats for the previous one no longer describe it. The status restarts + from the session's policy pair.""" + session.stats_gen += 1 + session.stats_status, session.stats_reason = initial_stats_status( + session.stats_tier, session.stats_delivery) + + +def stats_meta(session: "SessionState") -> Optional[dict]: + """The ``df_meta.stats`` value for a session: ``{status, tier, gen}`` plus + ``reason`` when there is one. ``tier`` is the tier reached so far, which is + the one the session is headed for only once it is complete. + + ``None`` for a session on the default policy (inline delivery, full tier, + complete), which sends the message it always has; a client reads a missing + ``stats`` as complete.""" + if session.stats_status == "complete" and session.stats_delivery != "deferred": + return None + stats: dict = {"status": session.stats_status, + "tier": session.stats_tier if session.stats_status == "complete" else "schema", + "gen": session.stats_gen} + if session.stats_reason: + stats["reason"] = session.stats_reason + return stats + + def build_state_message(session: "SessionState", metadata: dict | None = None, search_string: str = "") -> dict: """Build the full ``initial_state`` WebSocket payload from a session. @@ -116,10 +172,16 @@ def build_state_message(session: "SessionState", metadata: dict | None = None, Returns: A dict ready to be JSON-serialised and sent to WebSocket clients. """ + # The dataflow rebuilds df_meta wholesale, so the stats status is injected + # here, into a copy, rather than stored in it. + df_meta = session.df_meta + stats = stats_meta(session) + if stats is not None: + df_meta = {**df_meta, "stats": stats} msg: dict = {"type": "initial_state", "protocol_version": PROTOCOL_VERSION, "metadata": metadata if metadata is not None else session.metadata, "prompt": session.prompt, "df_display_args": session.df_display_args, "df_data_dict": session.df_data_dict, - "df_meta": session.df_meta, "mode": session.mode} + "df_meta": df_meta, "mode": session.mode} if session.mode == "buckaroo": # Per-client search_string overlay (#851 Codex P1): the snapshot # on the session is search-agnostic; we re-inject the recipient's diff --git a/buckaroo/server/stats_wire.py b/buckaroo/server/stats_wire.py new file mode 100644 index 000000000..7aeecf2d9 --- /dev/null +++ b/buckaroo/server/stats_wire.py @@ -0,0 +1,231 @@ +"""Server side of the stats wire protocol (rows-first s3). + +A deferred session (``stats_delivery="deferred"``) publishes its dataflow at the +schema tier and leaves the rest of the stats to later. Two kinds of client +reach them: + +* A client that advertised ``?caps=stats_update`` on its WebSocket URL gets a + stats-free ``initial_state`` (``df_meta.stats.status == "pending"``), pulls + the stats with ``stats_request {stats_gen, scope}`` and merges the + ``stats_update`` that answers it. A request for a generation the session has + left gets ``stats_aborted``. +* Any other client gets complete messages: ``build_state_message_for`` runs the + missing stats synchronously before it builds a message for one, which is + today's cost, paid on the loop. + +Every send site goes through ``build_state_message_for`` (or ``broadcast_state``, +which calls it per client), because the session holds one shared snapshot and +the client is known only to the handler that owns the connection. + +The request here is the whole run: one synchronous call that computes the full +stats and applies the final assignment (``assign_full_stats``, +``refresh_session_snapshot``). Resumable units generalize it later. +""" +import json +import logging +import time +import traceback +from typing import Any, Optional + +from buckaroo.dataflow.sd_cache import split_chain_by_scope +from buckaroo.pluggable_analysis_framework import perf_log +from buckaroo.server.data_loading import get_buckaroo_display_state +from buckaroo.server.session import SessionState, build_state_message + +log = logging.getLogger("buckaroo.server.stats_wire") + +# The capability a client advertises, as one of the comma-separated values of +# ``?caps=`` on the WebSocket URL: it merges ``stats_update`` messages. The +# connection is the only place it can be recorded, because ``open()`` sends the +# first message before the client has said anything. +STATS_UPDATE_CAP = "stats_update" + +# The scopes a ``stats_request`` can name. ``raw`` is the unfiltered table of +# the session's current generation; filtered stats exist only through a +# dataflow-field change (``quick_command_args.search``), which bumps the +# generation. +STATS_SCOPES = ("raw",) + + +def parse_caps(raw: str) -> frozenset: + """The capabilities in a ``?caps=a,b`` value.""" + return frozenset(cap for cap in (part.strip() for part in raw.split(",")) if cap) + + +def client_has_cap(client: Any, cap: str) -> bool: + """Whether a WebSocket handler recorded ``cap`` at open. Anything else with + no ``caps`` (a test double, a future non-WebSocket client) has none.""" + return cap in getattr(client, "caps", ()) + + +def session_dataflow(session: Optional[SessionState]) -> Any: + """The dataflow behind a buckaroo-mode session, or ``None`` (viewer and lazy + sessions have none).""" + if session is None or session.mode != "buckaroo": + return None + return session.xorq_dataflow if session.backend == "xorq" else session.dataflow + + +def refresh_session_snapshot(session: SessionState, dataflow: Any) -> None: + """Copy the dataflow's display state onto the session snapshot that new + clients and every push read, and re-apply ``component_config`` so theme + settings survive. The dataflow reaches no client until this runs.""" + refreshed = get_buckaroo_display_state(dataflow) + session.df_display_args = refreshed["df_display_args"] + session.df_data_dict = refreshed["df_data_dict"] + session.df_meta = refreshed["df_meta"] + session.buckaroo_options = refreshed["buckaroo_options"] + session.command_config = refreshed["command_config"] + if session.component_config and session.df_display_args: + for key in session.df_display_args: + dvc = session.df_display_args[key].get("df_viewer_config") + if dvc is not None: + dvc["component_config"] = {**dvc.get("component_config", {}), **session.component_config} + + +def assign_full_stats(dataflow: Any) -> None: + """Take a schema-tier dataflow to the full tier: compute the current + state's full stats (or find them in ``summary_stats_cache``, where an + earlier visit to the same state left them), write them under the full-tier + key, then assign ``summary_sd``. Assigning runs the cascade, which fills the + other scopes' entries it still lacks and rebuilds ``merged_sd``, + ``df_data_dict`` and ``df_display_args``. + + The cache entry goes in first because ``_populate_sd_cache`` skips a key it + finds, so the filt scope is not computed a second time. The tier is a + plain attribute, flipped before the compute so ``_get_summary_sd`` runs the + full pipeline, and put back if anything raises.""" + tier = dataflow.stats_tier + filt_chain = split_chain_by_scope(dataflow.operations)["filt"] + key = dataflow._scope_cache_key(filt_chain, tier="full") + dataflow.stats_tier = "full" + try: + sd = dataflow.summary_stats_cache.get(key) + errs = {} + if sd is None: + sd, errs = dataflow._get_summary_sd(dataflow.processed_df) + dataflow.summary_stats_cache = {**dataflow.summary_stats_cache, key: sd} + dataflow.summary_sd = sd + dataflow.errs = errs + except Exception: + dataflow.stats_tier = tier + raise + + +def complete_stats(session: SessionState) -> bool: + """Run the stats a pending session is missing, in one synchronous call, and + publish them: the final assignment (``assign_full_stats``), then the session + snapshot refreshed and the status set to ``complete`` in the same step, so + ``all_stats`` and ``df_meta.stats`` cannot disagree. Returns whether the + session is complete. + + A failure is the session's state for this generation (``error``, reason + ``stats_failed``) and is not retried by the next request; the next + generation starts clean.""" + if session.stats_status == "complete": + return True + dataflow = session_dataflow(session) + if session.stats_status != "pending" or dataflow is None: + return False + with ( + perf_log.telemetry_context(session.session_id, session.tele_sink), + perf_log.perf_span("firstpull.stats_total", session=session.session_id, stats_gen=session.stats_gen), + ): + try: + assign_full_stats(dataflow) + refresh_session_snapshot(session, dataflow) + except Exception: + log.error("stats run failed session=%s stats_gen=%s: %s", session.session_id, session.stats_gen, + traceback.format_exc()) + session.stats_status, session.stats_reason = "error", "stats_failed" + return False + session.stats_status, session.stats_reason = "complete", None + return True + + +def build_state_message_for(session: SessionState, client: Any, metadata: Optional[dict] = None) -> dict: + """The ``initial_state`` message for one client. + + A client without the ``stats_update`` capability on a pending deferred + session gets its missing stats run first, so its message is complete; a + capable client gets the session snapshot as it is (stats-free while the + session is pending) and pulls the rest. The search term is the recipient's + own (#851).""" + if (session.stats_delivery == "deferred" and session.stats_status == "pending" + and not client_has_cap(client, STATS_UPDATE_CAP)): + complete_stats(session) + return build_state_message(session, metadata=metadata, search_string=getattr(client, "search_string", "")) + + +def broadcast_state(session: SessionState, metadata: Optional[dict] = None, reset_search: bool = False) -> None: + """Send every connected client its own ``initial_state``. A client whose + write fails is dropped from the session. + + ``reset_search`` clears each client's live search term first (#851), for a + push that replaces the dataset. Clients that merge ``stats_update`` go + first: a legacy client's message completes the session's stats, and a + message built after that would carry them, so the capable client would + never see the pending state its own frame is meant to describe.""" + for client in sorted(session.ws_clients, key=lambda c: not client_has_cap(c, STATS_UPDATE_CAP)): + try: + if reset_search: + client.search_string = "" + client.write_message(json.dumps(build_state_message_for(session, client, metadata=metadata))) + except Exception: + session.ws_clients.discard(client) + + +def _aborted(stats_gen: Any, scope: Any, reason: str, session: Optional[SessionState] = None) -> dict: + """``stats_aborted``: the request was not run. ``stats_gen`` echoes the + request so the client can pair them; ``current_gen`` is the session's, so a + stale client can ask again without waiting for the next ``initial_state``.""" + msg: dict = {"type": "stats_aborted", "stats_gen": stats_gen, "scope": scope, "reason": reason} + if session is not None: + msg["current_gen"] = session.stats_gen + return msg + + +def _answer_stats_request(session: Optional[SessionState], stats_gen: Any, scope: Any, started: float) -> dict: + dataflow = session_dataflow(session) + if session is None or dataflow is None: + return _aborted(stats_gen, scope, "no_data") + if stats_gen != session.stats_gen: + return _aborted(stats_gen, scope, "stale", session) + if scope not in STATS_SCOPES: + return _aborted(stats_gen, scope, "unsupported_scope", session) + if session.stats_status == "not_computed": + return _aborted(stats_gen, scope, "not_requestable", session) + if session.stats_status == "error" or (session.stats_status == "pending" and not complete_stats(session)): + return _aborted(stats_gen, scope, "error", session) + # Complete: from here the answer is the dataflow's own all_stats, with no + # query, whether this request ran the stats or an earlier one did. + return {"type": "stats_update", "stats_gen": stats_gen, "scope": scope, "tier": dataflow.stats_tier, + "final": True, "payload": session.df_data_dict["all_stats"], + "elapsed_ms": round((time.perf_counter() - started) * 1000, 1)} + + +def handle_stats_request(session: Optional[SessionState], msg: dict) -> dict: + """Answer a ``stats_request {stats_gen, scope, columns?}``: a + ``stats_update`` carrying the complete ``all_stats`` as an inline wide + ``DFEnvelope`` (self-contained, so binary pairing stays single-slot), or a + ``stats_aborted``. ``columns`` is a hint a whole-run request ignores. + + The ``stats.request`` span records the request and its ``outcome`` + (``update`` or the abort reason), which is where updates sent, requests + dropped as stale and errors are counted. The caller binds the session's + telemetry sink around this call.""" + stats_gen, scope, columns = msg.get("stats_gen"), msg.get("scope", "raw"), msg.get("columns") + started = time.perf_counter() + with perf_log.perf_span("stats.request", session=session.session_id if session else None, stats_gen=stats_gen, + scope=scope, columns=len(columns) if isinstance(columns, list) else None) as span: + try: + reply = _answer_stats_request(session, stats_gen, scope, started) + except Exception: + log.error("stats_request error session=%s: %s", session.session_id if session else None, + traceback.format_exc()) + reply = _aborted(stats_gen, scope, "error", session) + if reply["type"] == "stats_update": + span.set_attr(outcome="update", tier=reply["tier"]) + else: + span.set_attr(outcome=reply["reason"]) + return reply diff --git a/buckaroo/server/websocket_handler.py b/buckaroo/server/websocket_handler.py index 2232211c7..1b4ed6125 100644 --- a/buckaroo/server/websocket_handler.py +++ b/buckaroo/server/websocket_handler.py @@ -9,8 +9,9 @@ import tornado.websocket from buckaroo.pluggable_analysis_framework import perf_log -from buckaroo.server.data_loading import (handle_infinite_request, handle_infinite_request_buckaroo, handle_infinite_request_lazy, get_buckaroo_display_state) -from buckaroo.server.session import build_state_message +from buckaroo.server.data_loading import (handle_infinite_request, handle_infinite_request_buckaroo, handle_infinite_request_lazy) +from buckaroo.server.session import begin_stats_generation, dataflow_stats_tier +from buckaroo.server.stats_wire import (broadcast_state, build_state_message_for, handle_stats_request, parse_caps, refresh_session_snapshot) def _handle_infinite_request_xorq(xorq_dataflow, payload_args, search_string=""): @@ -38,6 +39,11 @@ def open(self, session_id): # highlight overlay below — never broadcast, never stored on the # session. self.search_string = "" + # Capabilities the client advertised on the URL (``?caps=a,b``). Recorded + # per connection because this method sends the first message before the + # client can say anything, and the other send sites push one shared + # snapshot (stats_wire.build_state_message_for reads it per client). + self.caps = parse_caps(self.get_query_argument("caps", "")) sessions = self.application.settings["sessions"] sessions.add_ws_client(session_id, self) @@ -45,7 +51,7 @@ def open(self, session_id): # search_string="" — fresh connection, no per-client typing yet. session = sessions.get(session_id) if session and (session.df is not None or session.ldf is not None or session.xorq_dataflow is not None): - self.write_message(json.dumps(build_state_message(session, search_string=self.search_string))) + self.write_message(json.dumps(build_state_message_for(session, self))) def on_message(self, message): try: @@ -59,6 +65,22 @@ def on_message(self, message): self._handle_infinite_request(msg.get("payload_args", {})) elif msg_type == "buckaroo_state_change": self._handle_buckaroo_state_change(msg.get("new_state") or {}) + elif msg_type == "stats_request": + self._handle_stats_request(msg) + + def _handle_stats_request(self, msg): + """Answer a client's ``stats_request`` with a ``stats_update`` or a + ``stats_aborted``. Synchronous, like ``infinite_request``: the request + runs the whole stats computation in this call (see ``stats_wire``). + + This branch is its own async context, so the session's telemetry sink + is bound here for the ``stats.request`` span and the stats spans under + it.""" + sessions = self.application.settings["sessions"] + session = sessions.get(self.session_id) + with perf_log.telemetry_context(self.session_id, session.tele_sink if session else None): + reply = handle_stats_request(session, msg) + self.write_message(json.dumps(reply)) def _handle_buckaroo_state_change(self, new_state): sessions = self.application.settings["sessions"] @@ -99,6 +121,12 @@ def _handle_buckaroo_state_change(self, new_state): log.debug("buckaroo_state_change no-op session=%s — skipping rebroadcast", self.session_id) return + # A deferred session whose stats were completed is at the full tier; + # the change reruns the cascade at the schema tier again (13 ms + # against the full stats), and the stats follow as requests. + if session.stats_delivery == "deferred": + dataflow.stats_tier = dataflow_stats_tier(session.stats_tier, session.stats_delivery) + # Propagate changes to the dataflow (mirrors BuckarooWidgetBase._buckaroo_state) if old_state.get("post_processing") != new_state.get("post_processing"): dataflow.post_processing_method = new_state.get("post_processing", "") @@ -108,39 +136,20 @@ def _handle_buckaroo_state_change(self, new_state): dataflow.quick_command_args = new_state.get("quick_command_args", {}) # Re-extract state from the dataflow — same helper works for both - # ServerDataflow and XorqServerDataflow (verified by probe). - buckaroo_state = get_buckaroo_display_state(dataflow) - session.df_display_args = buckaroo_state["df_display_args"] - session.df_data_dict = buckaroo_state["df_data_dict"] - session.df_meta = buckaroo_state["df_meta"] + # ServerDataflow and XorqServerDataflow (verified by probe). The + # state the stats describe has changed, so the generation moves on. + refresh_session_snapshot(session, dataflow) + begin_stats_generation(session) # Strip search_string before snapshotting onto the session — it # belongs to this client only (#851), so a future client that # connects shouldn't inherit it via build_state_message. session.buckaroo_state = {k: v for k, v in new_state.items() if k != "search_string"} - session.buckaroo_options = buckaroo_state["buckaroo_options"] - session.command_config = buckaroo_state["command_config"] - - # Re-apply component_config so theme settings survive state changes - if session.component_config and session.df_display_args: - for key in session.df_display_args: - dvc = session.df_display_args[key].get("df_viewer_config") - if dvc is not None: - dvc["component_config"] = { - **dvc.get("component_config", {}), - **session.component_config, - } # Broadcast updated state to all connected clients. Each # client gets its own search_string re-injected so a # dataflow rebuild from one tab doesn't silently clear the # search box on another (or on the typing client itself). - for client in list(session.ws_clients): - try: - msg = build_state_message(session, - search_string=getattr(client, "search_string", "")) - client.write_message(json.dumps(msg)) - except Exception: - session.ws_clients.discard(client) + broadcast_state(session) except Exception: tb = traceback.format_exc() log.error("buckaroo_state_change error session=%s: %s", self.session_id, tb) @@ -168,6 +177,14 @@ def _send_highlight_overlay(self, session): if not session.df_display_args: return term = self.search_string + # Build the message first: for a client without the stats_update + # capability it completes the session's stats, which replaces + # session.df_display_args, and the overlay is a copy of that. + # Pass self.search_string so the overlay's buckaroo_state + # round-trips the typed term back to this client (Codex P1 on + # #854 — without it the JS clears the search box on every + # keystroke). + msg = build_state_message_for(session, self) overlay = copy.deepcopy(session.df_display_args) for dva in overlay.values(): dvc = (dva or {}).get("df_viewer_config") or {} @@ -180,11 +197,6 @@ def _send_highlight_overlay(self, session): else: disp.pop("highlight_phrase", None) - # Pass self.search_string so the overlay's buckaroo_state - # round-trips the typed term back to this client (Codex P1 on - # #854 — without it the JS clears the search box on every - # keystroke). - msg = build_state_message(session, search_string=self.search_string) msg["df_display_args"] = overlay try: self.write_message(json.dumps(msg)) From dde83c168612e22e8676708297a9fae532143817 Mon Sep 17 00:00:00 2001 From: Paddy Mullen Date: Sun, 4 Oct 2026 02:25:14 -0400 Subject: [PATCH 12/18] test(stats): characterize process_df and the xorq histogram cache keys (rows-first s4) Both pass on the current code. process_df is written out as the process_column loop it is, and each xorq histogram query's snapshot-cache key is compared with the key of a query built in the test, so the refactor into resumable units that follows can be checked against something that does not go through the units. Co-Authored-By: Claude Sonnet 5.5 --- tests/unit/test_paf_v2.py | 67 ++++++++++++++++++++++++++++++++ tests/unit/test_xorq_stats_v2.py | 57 +++++++++++++++++++++++++++ 2 files changed, 124 insertions(+) diff --git a/tests/unit/test_paf_v2.py b/tests/unit/test_paf_v2.py index 6ec2c252c..a2978e836 100644 --- a/tests/unit/test_paf_v2.py +++ b/tests/unit/test_paf_v2.py @@ -736,3 +736,70 @@ def test_stat_output_values(self): assert col_stats['length'] == 4 assert isinstance(col_stats['distinct_per'], float) assert col_stats['distinct_per'] > 0 + + +# ============================================================================ +# Characterization: process_df is process_column run over each column +# ============================================================================ + +def _mixed_frames(): + """The same mixed-dtype columns as a pandas and a polars frame. Every column + has a different count for each of its values: polars does not keep the order + of equal counts, so ties would make two runs differ.""" + pl = pytest.importorskip("polars") + data = { + 'ints': [3, 1, 2, 3, 3, 1], + 'floats': [0.5, 0.5, 0.5, 2.5, 2.5, None], + 'words': ['x', 'y', 'x', 'z', 'x', 'y'], + 'flags': [True, False, True, True, False, True]} + return {'pandas': pd.DataFrame(data), 'polars': pl.DataFrame(data)} + + +def _stat_lists(): + from buckaroo.customizations.pd_stats_v2 import PD_ANALYSIS_V2 + from buckaroo.customizations.pl_stats_v2 import PL_ANALYSIS_V2 + return {'pandas': PD_ANALYSIS_V2, 'polars': PL_ANALYSIS_V2} + + +class TestProcessDfCharacterization: + """What ``process_df`` returns, written as the ``process_column`` loop it is + today, so a refactor of it into resumable units can be checked against + something that does not go through the units.""" + + @staticmethod + def _by_hand(pipeline, df, skip=()): + from buckaroo.df_util import old_col_new_col + expected = {} + for orig, rewritten in old_col_new_col(df): + if orig in skip: + expected[rewritten] = {'orig_col_name': orig, 'rewritten_col_name': rewritten} + continue + ser = df[orig] + stats, _errors = pipeline.process_column(column_name=rewritten, column_dtype=ser.dtype, raw_series=ser, + sampled_series=ser, raw_dataframe=df, + initial_stats={'orig_col_name': orig, 'rewritten_col_name': rewritten}) + expected[rewritten] = stats + return expected + + @pytest.mark.parametrize('backend', ['pandas', 'polars']) + @pytest.mark.parametrize('skip', [(), ('floats',)]) + def test_process_df_equals_the_column_loop(self, backend, skip): + df = _mixed_frames()[backend] + pipeline = StatPipeline(_stat_lists()[backend], unit_test=False) + result, errors = pipeline.process_df(df, skip_columns=set(skip)) + expected = self._by_hand(pipeline, df, skip) + assert errors == [] + assert list(result) == list(expected), "columns must keep the frame's order" + for col in expected: + assert repr(result[col]) == repr(expected[col]), col + + def test_errors_are_collected_per_column(self): + @stat() + def fails_on_words(ser: RawSeries) -> int: + if isinstance(ser.iloc[0], str): + raise ValueError('words') + return 0 + + pipeline = StatPipeline([length, fails_on_words], unit_test=False) + _result, errors = pipeline.process_df(_mixed_frames()['pandas']) + assert [(e.column, e.stat_key) for e in errors] == [('c', 'fails_on_words')] diff --git a/tests/unit/test_xorq_stats_v2.py b/tests/unit/test_xorq_stats_v2.py index 3565ef385..f1552cf14 100644 --- a/tests/unit/test_xorq_stats_v2.py +++ b/tests/unit/test_xorq_stats_v2.py @@ -876,3 +876,60 @@ def test_cached_stats_match_uncached(self, tmp_path): for col in plain: for key in ("length", "min", "max", "mean", "null_count", "distinct_count"): assert cold[col].get(key) == plain[col].get(key), (col, key) + + +class _RecordingCache: + """Wraps a snapshot cache and records the key of every query the pipeline + asks it about.""" + + def __init__(self, inner): + self.inner = inner + self.storage = inner.storage + self.keys = [] + + def calc_key(self, query): + key = self.inner.calc_key(query) + self.keys.append(key) + return key + + +class TestHistogramCacheKeys: + """The snapshot-cache key of each per-column histogram query, which warm + loads in another process depend on. The expected keys come from queries + built here in the shape ``histogram`` documents, not from the pipeline, so + a change to how the pipeline builds or schedules them is caught.""" + + @staticmethod + def _scan(tmp_path): + df = pd.DataFrame({ + "num": [float(i) for i in range(12)], "word": list("aabbccddeeff"), "flag": [True, False] * 6}) + path = tmp_path / "keys.parquet" + df.to_parquet(path) + return df, xo.deferred_read_parquet(str(path)) + + @staticmethod + def _numeric_histogram_query(expr, col, lo, hi): + bucket = ((expr[col].cast("float64") - lo) / (hi - lo) * 10).cast("int64").clip(lower=0, upper=9) + return (expr.mutate(__bucket=bucket).group_by("__bucket").aggregate(__count=lambda t: t.count()) + .order_by("__bucket")) + + @staticmethod + def _categorical_histogram_query(expr, col): + return (expr.group_by(col).aggregate(__count=lambda t: t.count()).order_by(xo.desc("__count")).limit(10)) + + def test_each_histogram_query_keeps_its_snapshot_key(self, tmp_path): + df, table = self._scan(tmp_path) + cache = xo.ParquetSnapshotCache.from_kwargs(source=xo.connect(), base_path=str(tmp_path / "cache")) + recorder = _RecordingCache(cache) + + XorqStatPipeline(XORQ_STATS_V2, unit_test=False, cache_storage=recorder).process_table(table) + + expected = { + "num": cache.calc_key(self._numeric_histogram_query(table, "num", float(df.num.min()), float(df.num.max()))), + "word": cache.calc_key(self._categorical_histogram_query(table, "word")), + "flag": cache.calc_key(self._categorical_histogram_query(table, "flag"))} + for col, key in expected.items(): + assert key in recorder.keys, f"the {col} histogram query is no longer keyed as it was" + # The batch aggregate and one histogram query per column, nothing else. + assert len(recorder.keys) == 1 + len(expected) + assert len(set(recorder.keys)) == len(recorder.keys) From 667ae34daf72b3cf9cdfba08059088b714c7f6dd Mon Sep 17 00:00:00 2001 From: Paddy Mullen Date: Sun, 4 Oct 2026 02:34:40 -0400 Subject: [PATCH 13/18] test(stats): failing tests for resumable stat units and the StatRun store (rows-first s4) Each stats class gets plan(state) and run(unit, acc): the fragments of the planned units, assembled through assemble_merged_sd, must equal the full-stats merged_sd on pandas, polars and xorq with init_sd, cleaning and overrides; a column group returns only its columns; skip_stat_columns columns get no unit; the xorq batch can be split by column chunk behind a default-off flag and is refused for anything but a plain parquet scan; the histogram cache keys stay put. A StatRun held on the session, keyed by (stats_gen, scope), keeps an append-only fragment list and accumulator, runs a unit only when asked, is dropped when stats_gen changes, and each connection keeps its own cursor. test_the_batch_is_one_query_by_default passes on the current code: it pins the default the flag must not change. Co-Authored-By: Claude Sonnet 5.5 --- .../dataflow/scoped_summary_stats_test.py | 33 ++ tests/unit/server/test_data_loading_polars.py | 213 +++++++++++- tests/unit/server/test_load_expr.py | 204 ++++++++++++ tests/unit/skip_columns_test.py | 60 ++++ tests/unit/test_paf_v2.py | 112 +++++++ tests/unit/test_xorq_stats_v2.py | 315 ++++++++++++++++++ 6 files changed, 936 insertions(+), 1 deletion(-) diff --git a/tests/unit/dataflow/scoped_summary_stats_test.py b/tests/unit/dataflow/scoped_summary_stats_test.py index a8b0a5339..40d6f0f07 100644 --- a/tests/unit/dataflow/scoped_summary_stats_test.py +++ b/tests/unit/dataflow/scoped_summary_stats_test.py @@ -26,6 +26,7 @@ from buckaroo.dataflow.autocleaning import (AutocleaningConfig, PandasAutocleaning) from buckaroo.dataflow.dataflow import CustomizableDataflow, StylingAnalysis from buckaroo.dataflow.sd_cache import split_chain_by_scope +from buckaroo.pluggable_analysis_framework.col_analysis import ColAnalysis from buckaroo.pluggable_analysis_framework.stat_func import stat @@ -316,3 +317,35 @@ def key_sets(): # with an input would add keys to that input. assert key_sets() == before assert inputs['chains'] == chains_before + + +class _OverridingPostProcessing(ColAnalysis): + """A post-processing method whose sd overrides a stat of qty (keyed by the + original column name) and adds one of its own.""" + provides_defaults = {} + post_processing_method = "override_post" + + @classmethod + def post_process_df(cls, df): + return [df, {"qty": {"mean": 99.5, "from_post": 1}}] + + +def _run_units(stats): + """Run every unit of a stats object built with ``run=False``, one at a time. + Returns the accumulator and the fragments in the order they came back.""" + state = stats.state + acc = stats.new_accumulator(state) + return acc, [stats.run(unit, acc) for unit in stats.plan(state)] + + +def _scope_sds_by_units(dataflow): + """Each scope's summary dict, built from the fragments of its planned units + where the dataflow built it by running the stats class whole.""" + from buckaroo.pluggable_analysis_framework.stat_units import merge_fragments, rewrite_sd + sds = {} + for scope in ("raw", "clean", "filt"): + scope_df = dataflow.processed_df if scope == "filt" else dataflow._compute_scope_df(scope) + stats = dataflow.build_stats(scope_df, run=False) + _acc, fragments = _run_units(stats) + sds[scope] = rewrite_sd(merge_fragments(fragments), stats.state.data) + return sds diff --git a/tests/unit/server/test_data_loading_polars.py b/tests/unit/server/test_data_loading_polars.py index 6ed441642..cd3500a21 100644 --- a/tests/unit/server/test_data_loading_polars.py +++ b/tests/unit/server/test_data_loading_polars.py @@ -15,17 +15,21 @@ the display state full stats give, without reading a value. """ import base64 +import dataclasses import datetime import decimal +import inspect import io import json import os import tempfile +import threading import pandas as pd import polars as pl import pyarrow.parquet as pq import pytest +from tornado.ioloop import IOLoop from buckaroo.dataflow.dataflow import assemble_merged_sd from buckaroo.dataflow.sd_cache import split_chain_by_scope @@ -36,7 +40,7 @@ from buckaroo.server.data_loading import ServerDataflow, handle_infinite_request_buckaroo from buckaroo.server.data_loading_polars import (PolarsServerDataflow, create_polars_dataflow, handle_infinite_request_buckaroo_polars, load_file_polars) from buckaroo.styling_helpers import float_, obj_ -from tests.unit.dataflow.scoped_summary_stats_test import _scope_inputs +from tests.unit.dataflow.scoped_summary_stats_test import (_OverridingPostProcessing, _run_units, _scope_inputs, _scope_sds_by_units) def _payload(start=0, end=100): @@ -424,3 +428,210 @@ def frame(): types[tier] = {v["orig_col_name"]: v["_type"] for v in sd.values()} assert types["full"] == {"words": "string", "mixed": "obj"} assert types["schema"] == {"words": "string", "mixed": "string"} + + +# --------------------------------------------------------------------------- +# Resumable stat units on the pandas and polars server dataflows (rows-first s4) +# --------------------------------------------------------------------------- + + +class TestStatUnits: + """``plan`` and ``run`` on the stats class a dataflow builds, in place of + the all-at-once constructor.""" + + def test_build_stats_without_running_computes_nothing(self, backend, monkeypatch): + dataflow = _build_dataflow(backend) + frames = _spy_stat_pipeline(monkeypatch) + stats = dataflow.build_stats(dataflow.processed_df, run=False) + stats.plan(stats.state) + assert frames == [] + assert stats.sdf == {} and stats.errs == {} + + def test_the_default_run_goes_through_the_units(self, backend, monkeypatch): + """Every stat the constructor computes comes from a unit, so the + inline path and a caller running units one at a time cannot differ.""" + unit_columns, column_calls = [], [] + original_run, original_column = StatPipeline.run, StatPipeline.process_column + + def spy_run(self, unit, acc): + unit_columns.append(unit.columns) + return original_run(self, unit, acc) + + def spy_column(self, *args, **kwargs): + column_calls.append(args or kwargs) + return original_column(self, *args, **kwargs) + + monkeypatch.setattr(StatPipeline, "run", spy_run) + monkeypatch.setattr(StatPipeline, "process_column", spy_column) + _three_scope_dataflow(backend) + assert {c for cols in unit_columns for c in cols} == {"price", "qty", "category"} + assert len(unit_columns) == len(column_calls), "a column's stats ran outside a unit" + + def test_units_assembled_equal_merged_sd_with_init_sd_cleaning_and_overrides(self, backend): + """The fragments of each scope's units, assembled by + ``assemble_merged_sd`` with ``init_sd``, a cleaning op (pandas), a + search filter and a post-processing override, are the dataflow's own + ``merged_sd``. init_sd is keyed by the original name; price is "a".""" + dataflow = _three_scope_dataflow(backend, _UNTIED_DATA, init_sd={ + "price": {"displayer_args": {"displayer": "string", "max_length": 99}, "init_only": 1}}, + column_config_overrides={"category": {"displayer_args": {"displayer": "string", "max_length": 5000}}}) + dataflow.add_analysis(_OverridingPostProcessing) + dataflow.post_processing_method = "override_post" + sd = dataflow.merged_sd + assert sd["a"]["init_only"] == 1 and sd["b"]["from_post"] == 1 and sd["b"]["mean"] == 99.5 + assert "filtered_length" in sd["a"], "the filter layer must be active" + if backend == "pandas": + assert "cleaned_length" in sd["a"], "the cleaning layer must be active" + + sds = _scope_sds_by_units(dataflow) + inputs = _scope_inputs(dataflow) + assembled = assemble_merged_sd(init_sd=inputs["init_sd"], cleaned_sd=inputs["cleaned_sd"], raw_sd=sds["raw"], + processed_sd=inputs["processed_sd"], processed_df=inputs["processed_df"], chains=inputs["chains"], + clean_sd=sds["clean"], filt_sd=sds["filt"]) + + assert _as_json(assembled) == _as_json(dataflow.merged_sd) + + def test_a_column_group_request_returns_only_those_columns(self, backend): + dataflow = _build_dataflow(backend, _frame(backend, _UNTIED_DATA)) + stats = dataflow.build_stats(dataflow.processed_df, run=False) + for names in (("qty", "category"), ("b", "c")): # original or rewritten names + state = dataclasses.replace(stats.state, columns=names) + acc = stats.new_accumulator(state) + fragments = [stats.run(unit, acc) for unit in stats.plan(state)] + assert [list(f) for f in fragments] == [["qty"], ["category"]] + assert list(acc.sd()) == ["qty", "category"] + + def test_skip_stat_columns_get_no_unit(self, backend): + dataflow = _build_dataflow(backend, _frame(backend, _UNTIED_DATA), skip_stat_columns=["qty"], + init_sd={"qty": {"_type": "float"}}) + stats = dataflow.build_stats(dataflow.processed_df, run=False) + assert stats.state.skip_columns == {"qty"} + assert [u.columns for u in stats.plan(stats.state)] == [("price",), ("category",)] + acc, fragments = _run_units(stats) + assert all("qty" not in f for f in fragments) + assert set(acc.sd()) == {"price", "qty", "category"} + + +def _stat_run(backend, gen=4): + from buckaroo.server.stat_run import StatRun + dataflow = _build_dataflow(backend, _frame(backend, _UNTIED_DATA)) + return StatRun(gen, "raw", dataflow.build_stats(dataflow.processed_df, run=False)), dataflow + + +class TestStatRun: + """The run a session holds for one generation and scope: the planned + units, the fragments they produced in the order they finished, and the + accumulator the units read. It runs a unit only when asked.""" + + def test_a_run_is_keyed_by_generation_and_scope_and_plans_its_units(self, backend): + run, _dataflow = _stat_run(backend) + assert run.key == (4, "raw") + assert run.status == "pending" and run.remaining == 3 + assert [u.columns for u in run.units] == [("price",), ("qty",), ("category",)] + assert run.fragments == [] and run.ran == [] + + def test_fragments_are_appended_in_completion_order(self, backend): + run, _dataflow = _stat_run(backend) + while True: + before = list(run.fragments) + fragment = run.run_next(prefer=("category",)) + if fragment is None: + break + assert run.fragments[:-1] == before, "nothing already in the list changes" + assert run.fragments[-1] is fragment + assert [list(f) for f in run.fragments] == [["category"], ["price"], ["qty"]] + assert len(run.ran) == 3 and len(set(run.ran)) == 3 + + def test_prefer_runs_the_units_that_name_those_columns_first(self, backend): + run, _dataflow = _stat_run(backend) + assert run.next_unit().columns == ("price",) + assert run.next_unit(prefer=("c",)).columns == ("category",), "a rewritten name works too" + assert run.next_unit(prefer=("nothing",)).columns == ("price",) + + def test_the_run_is_complete_when_no_unit_is_left(self, backend): + run, _dataflow = _stat_run(backend) + while run.run_next() is not None: + pass + assert (run.status, run.remaining) == ("complete", 0) + assert run.next_unit() is None and run.run_next() is None + assert len(run.fragments) == 3, "asking again runs nothing" + + def test_the_accumulator_holds_what_the_fragments_hold(self, backend): + from buckaroo.pluggable_analysis_framework.stat_units import merge_fragments + run, _dataflow = _stat_run(backend) + while run.run_next() is not None: + pass + assert merge_fragments(run.fragments) == run.acc.sd() + + def test_raw_sd_is_the_full_stats_summary_sd(self, backend): + run, dataflow = _stat_run(backend) + while run.run_next() is not None: + pass + assert _as_json(run.raw_sd()) == _as_json(dataflow.summary_sd) + + def test_cursors_read_one_list_each_at_its_own_pace(self, backend): + from buckaroo.pluggable_analysis_framework.stat_units import merge_fragments + from buckaroo.server.stat_run import StatCursor + run, _dataflow = _stat_run(backend) + a, b = StatCursor(), StatCursor() + run.run_next() + run.run_next() + first_a = a.take(run) + assert len(first_a) == 2 and a.take(run) == [] and a.caught_up(run) + run.run_next() + assert not a.caught_up(run) and b.position == 0, "one cursor's reads move no other" + all_b = b.take(run) + rest_a = a.take(run) + assert (len(all_b), len(rest_a)) == (3, 1) + assert merge_fragments(first_a + rest_a) == merge_fragments(all_b) + assert len(set(run.ran)) == 3, "no unit ran twice" + + def test_a_cursor_starts_again_on_another_run(self, backend): + from buckaroo.server.stat_run import StatCursor + first, _dataflow = _stat_run(backend, gen=1) + second, _dataflow = _stat_run(backend, gen=2) + first.run_next() + second.run_next() + second.run_next() + cursor = StatCursor() + assert len(cursor.take(first)) == 1 + assert len(cursor.take(second)) == 2, "a position in one run means nothing in another" + + def test_a_run_has_no_thread_timer_or_callback(self, backend, monkeypatch): + scheduled = [] + monkeypatch.setattr(threading.Thread, "start", lambda self, *a, **k: scheduled.append("thread")) + for name in ("add_callback", "call_later", "call_at", "add_timeout"): + monkeypatch.setattr(IOLoop, name, lambda self, *a, _n=name, **k: scheduled.append(_n)) + threads_before = threading.active_count() + run, _dataflow = _stat_run(backend) + while run.run_next() is not None: + pass + assert scheduled == [] and threading.active_count() == threads_before + assert not [name for name, value in vars(run).items() + if isinstance(value, (threading.Thread, threading.Timer)) or inspect.isroutine(value)] + + def test_a_unit_that_raises_fails_the_run(self): + from buckaroo.pluggable_analysis_framework.stat_units import StatAccumulator, StatState, StatUnit + from buckaroo.server.stat_run import StatRun + + class _FailingStats: + state = StatState(data=None) + + def plan(self, state): + return [StatUnit("one", ("x",), "column"), StatUnit("two", ("y",), "column")] + + def new_accumulator(self, state): + return StatAccumulator(state, columns=["x", "y"]) + + def run(self, unit, acc): + if unit.id == "two": + raise RuntimeError("boom") + return {"x": {"v": 1}} + + run = StatRun(1, "raw", _FailingStats()) + run.run_next() + with pytest.raises(RuntimeError, match="boom"): + run.run_next() + assert run.status == "error" and isinstance(run.error, RuntimeError) + assert run.run_next() is None, "a failed run is not retried" + assert run.fragments == [{"x": {"v": 1}}] diff --git a/tests/unit/server/test_load_expr.py b/tests/unit/server/test_load_expr.py index 5ab61e063..e222b56b2 100644 --- a/tests/unit/server/test_load_expr.py +++ b/tests/unit/server/test_load_expr.py @@ -37,6 +37,8 @@ from buckaroo.server import telemetry, xorq_loading # noqa: E402 from buckaroo.server.app import make_app as _make_app # noqa: E402 from buckaroo.server.websocket_handler import DataStreamHandler # noqa: E402 +from tests.unit.dataflow.scoped_summary_stats_test import ( # noqa: E402 + _OverridingPostProcessing, _scope_inputs, _scope_sds_by_units) pytestmark = pytest.mark.skipif( sys.platform == "win32", @@ -1913,6 +1915,208 @@ async def test_inline_sessions_send_no_df_meta_stats(self): self.assertIn("histogram_bins", _rows_by_stat(frame["df_data_dict"]["all_stats"])) + @tornado.testing.gen_test + async def test_a_stat_run_is_started_on_the_session_for_the_current_generation(self): + from buckaroo.server.stats_wire import start_stat_run + await self._load("sw-run", stats_delivery="deferred") + session = self._session("sw-run") + self.assertEqual(session.stat_runs, {}) + + with _count_stat_queries() as queries: + run = start_stat_run(session) + self.assertIs(start_stat_run(session), run, "one run per generation and scope") + + self.assertEqual(queries, [], "planning a run runs no query") + self.assertEqual(run.key, (session.stats_gen, "raw")) + self.assertEqual(session.stat_runs, {run.key: run}) + self.assertEqual(run.remaining, 4, "the scalar batch and one histogram unit for each of three columns") + self.assertEqual(run.units[0].id, "batch") + + @tornado.testing.gen_test + async def test_a_run_ends_with_the_summary_sd_an_inline_session_computes(self): + from buckaroo.server.stats_wire import start_stat_run + await self._load("sw-run-sd", stats_delivery="deferred") + session = self._session("sw-run-sd") + run = start_stat_run(session) + while run.run_next() is not None: + pass + inline = xorq_loading.XorqServerDataflow(session.expr, skip_main_serial=True) + self.assertEqual(run.status, "complete") + self.assertEqual(_as_json(run.raw_sd()), _as_json(inline.summary_sd)) + + @tornado.testing.gen_test + async def test_a_dataflow_field_change_drops_the_run_and_a_search_term_change_keeps_it(self): + from buckaroo.server.stats_wire import start_stat_run + await self._load("sw-drop", stats_delivery="deferred") + ws, _first = await self._connect("sw-drop", caps="stats_update") + session = self._session("sw-drop") + run = start_stat_run(session) + + ws.write_message(_state_change(search_string="alp")) + await _read_json(ws) # this client's highlight overlay + self.assertEqual(session.stat_runs, {run.key: run}, "a typed search term touches no stats") + + ws.write_message(_state_change(post_processing="first_three")) + await _read_json(ws) + self.assertEqual(session.stat_runs, {}) + self.assertEqual(session.stats_gen, run.stats_gen + 1) + again = start_stat_run(session) + self.assertEqual(again.key, (run.stats_gen + 1, "raw")) + self.assertIsNot(again, run) + + @tornado.testing.gen_test + async def test_reload_expr_drops_the_run(self): + from buckaroo.server.stats_wire import start_stat_run + await self._load("sw-drop-reload", stats_delivery="deferred") + session = self._session("sw-drop-reload") + start_stat_run(session) + resp = await _post(self.get_http_port(), "/reload_expr/sw-drop-reload", {}) + self.assertEqual(resp.code, 200, resp.body) + self.assertEqual(session.stat_runs, {}) + + @tornado.testing.gen_test + async def test_each_connection_keeps_its_own_cursor_into_the_run(self): + from buckaroo.pluggable_analysis_framework.stat_units import merge_fragments + from buckaroo.server.stats_wire import start_stat_run + await self._load("sw-cursor", stats_delivery="deferred") + await self._connect("sw-cursor", caps="stats_update") + await self._connect("sw-cursor", caps="stats_update") + session = self._session("sw-cursor") + cursor_a, cursor_b = [handler.stats_cursor for handler in session.ws_clients] + self.assertIsNot(cursor_a, cursor_b) + self.assertEqual((cursor_a.position, cursor_b.position), (0, 0)) + + run = start_stat_run(session) + run.run_next() + run.run_next() + first = cursor_a.take(run) + self.assertEqual((len(first), cursor_b.position), (2, 0)) + while run.run_next() is not None: + pass + rest, everything = cursor_a.take(run), cursor_b.take(run) + + self.assertEqual(merge_fragments(first + rest), merge_fragments(everything)) + self.assertEqual(len(set(run.ran)), len(run.units), "no unit ran twice") + + +def _wide_parquet_expr(tmp_path, name="wide"): + """Twelve rows, six columns, read from parquet. Each string column has a + different count per value, so a GROUP BY over it returns one order.""" + path = tmp_path / f"{name}.parquet" + pd.DataFrame({ + "n0": [float(i) for i in range(12)], "n1": [i * 3 for i in range(12)], + "n2": [float(i % 7) + 0.5 for i in range(12)], "word": list("aaaaabbbbccd"), + "label": list("xxxxxyyyyzzw"), "flag": [True] * 7 + [False] * 5}).to_parquet(path) + return xo.deferred_read_parquet(str(path)) + + +def _batch_queries(queries): + return [q for q in queries if type(q.op()).__name__ == "Aggregate"] + + +class TestStatUnits: + """Resumable units on the xorq server dataflow (rows-first s4): the + all-at-once constructor and a caller running ``plan``/``run`` one unit at a + time compute the same stats.""" + + def test_build_stats_without_running_issues_no_stat_query(self): + dataflow = _build_dataflow() + with _count_stat_queries() as queries: + stats = dataflow.build_stats(dataflow.processed_df, run=False) + stats.plan(stats.state) + stats.new_accumulator(stats.state) + assert queries == [] + assert stats.sdf == {} and stats.errs == {} + + def test_the_default_run_goes_through_the_units(self, monkeypatch): + """Every stat query of the constructor is issued from inside a unit.""" + depth, inside, outside, units = [0], [], [], [] + original_run, original_execute = XorqStatPipeline.run, XorqStatPipeline._execute + + def spy_run(self, unit, acc): + depth[0] += 1 + units.append(unit.id) + try: + return original_run(self, unit, acc) + finally: + depth[0] -= 1 + + def spy_execute(self, query): + (inside if depth[0] else outside).append(query) + return original_execute(self, query) + + monkeypatch.setattr(XorqStatPipeline, "run", spy_run) + monkeypatch.setattr(XorqStatPipeline, "_execute", spy_execute) + _three_scope_dataflow() + assert inside and outside == [] + assert units[0] == "batch" and "histogram:price" in units + + def test_units_assembled_equal_merged_sd_with_init_sd_cleaning_and_overrides(self): + from buckaroo.dataflow.dataflow import assemble_merged_sd + dataflow = _three_scope_dataflow(init_sd={ + "price": {"displayer_args": {"displayer": "string", "max_length": 99}, "init_only": 1}}, + column_config_overrides={"category": {"displayer_args": {"displayer": "string", "max_length": 5000}}}) + dataflow.add_analysis(_OverridingPostProcessing) + dataflow.post_processing_method = "override_post" + sd = dataflow.merged_sd + assert sd["a"]["init_only"] == 1 and sd["b"]["from_post"] == 1 and sd["b"]["mean"] == 99.5 + assert "cleaned_length" in sd["a"] and "filtered_length" in sd["a"], "both layers must be active" + + sds = _scope_sds_by_units(dataflow) + inputs = _scope_inputs(dataflow) + assembled = assemble_merged_sd(init_sd=inputs["init_sd"], cleaned_sd=inputs["cleaned_sd"], raw_sd=sds["raw"], + processed_sd=inputs["processed_sd"], processed_df=inputs["processed_df"], chains=inputs["chains"], + clean_sd=sds["clean"], filt_sd=sds["filt"]) + + assert _as_json(assembled) == _as_json(dataflow.merged_sd) + + def test_a_column_group_request_returns_only_those_columns(self): + import dataclasses + dataflow = _build_dataflow() + stats = dataflow.build_stats(dataflow.processed_df, run=False) + for names in (("qty", "category"), ("b", "c")): # original or rewritten names + state = dataclasses.replace(stats.state, columns=names) + acc = stats.new_accumulator(state) + fragments = [stats.run(unit, acc) for unit in stats.plan(state)] + assert [u.id for u in stats.plan(state)] == ["batch", "histogram:qty", "histogram:category"] + assert all(set(f) <= {"qty", "category"} for f in fragments) + assert list(acc.sd()) == ["qty", "category"] + + def test_skip_stat_columns_get_no_unit(self): + dataflow = _build_dataflow(skip_stat_columns=["qty"], init_sd={"qty": {"_type": "float"}}) + stats = dataflow.build_stats(dataflow.processed_df, run=False) + assert stats.state.skip_columns == {"qty"} + units = stats.plan(stats.state) + assert [u.id for u in units] == ["batch", "histogram:price", "histogram:category"] + assert all("qty" not in unit.columns for unit in units) + + def test_the_batch_is_one_query_by_default(self, tmp_path): + expr = _wide_parquet_expr(tmp_path) + with _count_stat_queries() as queries: + xorq_loading.XorqServerDataflow(expr, skip_main_serial=True) + assert len(_batch_queries(queries)) == 1 + + def test_stat_chunk_cells_splits_the_batch_of_a_parquet_scan(self, tmp_path): + expr = _wide_parquet_expr(tmp_path) + whole = xorq_loading.XorqServerDataflow(expr, skip_main_serial=True) + with _count_stat_queries() as queries: + # 36 cells at 12 rows is three columns per chunk: two batch queries. + split = xorq_loading.XorqServerDataflow(expr, skip_main_serial=True, stat_chunk_cells=36) + assert len(_batch_queries(queries)) == 2 + assert _as_json(split.merged_sd) == _as_json(whole.merged_sd) + + def test_the_split_is_refused_for_a_join(self, tmp_path): + left = _wide_parquet_expr(tmp_path, "left") + pd.DataFrame({"n0": [float(i) for i in range(12)], "extra": list(range(12))}).to_parquet( + tmp_path / "right.parquet") + joined = left.join(xo.deferred_read_parquet(str(tmp_path / "right.parquet")), "n0") + with _count_stat_queries() as queries: + split = xorq_loading.XorqServerDataflow(joined, skip_main_serial=True, stat_chunk_cells=12) + assert len(_batch_queries(queries)) == 1, "each chunk of a join would run the whole join again" + whole = xorq_loading.XorqServerDataflow(joined, skip_main_serial=True) + assert _as_json(split.merged_sd) == _as_json(whole.merged_sd) + + class TestReloadExpr(tornado.testing.AsyncHTTPTestCase): def get_app(self): return make_app() diff --git a/tests/unit/skip_columns_test.py b/tests/unit/skip_columns_test.py index bf0121be5..77ad8895a 100644 --- a/tests/unit/skip_columns_test.py +++ b/tests/unit/skip_columns_test.py @@ -61,6 +61,66 @@ def test_xorq_skip_columns_not_computed(tmp_path): assert "mean" not in sd["b"] +# --------------------------------------------------------------------------- +# Resumable units (rows-first s4): a skipped column gets no unit, so running +# every unit never touches it, and it still ends up in the result. +# --------------------------------------------------------------------------- + + +def _run_units(pipe, state): + acc = pipe.new_accumulator(state) + return acc, [pipe.run(unit, acc) for unit in pipe.plan(state)] + + +def test_pandas_skipped_column_gets_no_unit(): + from buckaroo.pluggable_analysis_framework.stat_units import StatState + df = pd.DataFrame({"a": [1, 2, 3], "b": [4, 5, 6], "c": [7, 8, 9]}) + pipe = StatPipeline(PD_ANALYSIS_V2, unit_test=False) + state = StatState(df, skip_columns=frozenset({"b"})) + assert [u.columns for u in pipe.plan(state)] == [("a",), ("c",)] + acc, fragments = _run_units(pipe, state) + assert [list(f) for f in fragments] == [["a"], ["c"]] + assert list(acc.sd()) == ["a", "b", "c"], "the skipped column keeps its place in the result" + assert acc.sd()["b"] == {"orig_col_name": "b", "rewritten_col_name": "b"} + assert pipe.process_df(df, skip_columns={"b"})[0]["b"] == acc.sd()["b"] + + +def test_polars_skipped_column_gets_no_unit(): + pl = pytest.importorskip("polars") + from buckaroo.customizations.pl_stats_v2 import PL_ANALYSIS_V2 + from buckaroo.pluggable_analysis_framework.stat_units import StatState + df = pl.DataFrame({"a": [1, 2, 3], "b": [4, 5, 6], "c": [7, 8, 9]}) + pipe = StatPipeline(PL_ANALYSIS_V2, unit_test=False) + state = StatState(df, skip_columns=frozenset({"b"})) + assert [u.columns for u in pipe.plan(state)] == [("a",), ("c",)] + acc, _fragments = _run_units(pipe, state) + assert set(acc.sd()["b"]) <= {"orig_col_name", "rewritten_col_name"} + + +def test_xorq_skipped_column_gets_no_unit(tmp_path): + xo = pytest.importorskip("xorq.api") + from buckaroo.pluggable_analysis_framework.xorq_stat_pipeline import XorqStatPipeline + from buckaroo.pluggable_analysis_framework.stat_units import StatState + from buckaroo.customizations.xorq_stats_v2 import XORQ_STATS_V2 + + p = tmp_path / "t.parquet" + pd.DataFrame({"a": [1, 2, 3], "b": [4, 5, 6]}).to_parquet(p) + expr = xo.deferred_read_parquet(str(p)) + + pipe = XorqStatPipeline(XORQ_STATS_V2) + state = StatState(expr, skip_columns=frozenset({"b"})) + units = pipe.plan(state) + assert all("b" not in unit.columns for unit in units) + assert [u.id for u in units] == ["batch", "histogram:a"] + acc, fragments = _run_units(pipe, state) + assert all("b" not in fragment for fragment in fragments) + sd = acc.sd() + # The column keeps what the batch gives every column, and no stat of its own. + assert sd["b"]["length"] == 3 and sd["b"]["dtype"] == "int64" + assert "mean" not in sd["b"] and "histogram" not in sd["b"] + assert sd == pipe.process_table(expr, skip_columns={"b"})[0] + + # --------------------------------------------------------------------------- # Notebook widget constructors must forward skip_stat_columns to the dataflow # (it was only wired into CustomizableDataflow + the server /load handler, not diff --git a/tests/unit/test_paf_v2.py b/tests/unit/test_paf_v2.py index a2978e836..8909bd7c3 100644 --- a/tests/unit/test_paf_v2.py +++ b/tests/unit/test_paf_v2.py @@ -803,3 +803,115 @@ def fails_on_words(ser: RawSeries) -> int: pipeline = StatPipeline([length, fails_on_words], unit_test=False) _result, errors = pipeline.process_df(_mixed_frames()['pandas']) assert [(e.column, e.stat_key) for e in errors] == [('c', 'fails_on_words')] + + +# ============================================================================ +# Resumable units: plan(state) and run(unit, acc) (rows-first s4) +# ============================================================================ + +class TestStatUnits: + """``plan(state)`` lists the units of a run, one per column, and + ``run(unit, acc)`` runs one and returns a ``{orig_col: {stat: value}}`` + fragment. ``process_df`` is these two run over every unit, so a caller that + runs them one at a time gets what it would have got all at once.""" + + @staticmethod + def _df(): + # Original names differ from the rewritten a, b, c. + return pd.DataFrame({'amount': [1, 2, 3, 1], 'word': ['x', 'y', 'x', 'x'], 'score': [0.5, 1.5, None, 2.5]}) + + @staticmethod + def _pipeline(): + return StatPipeline([length, null_count, distinct_count, distinct_per, nan_per], unit_test=False) + + @staticmethod + def _run_units(pipeline, state): + acc = pipeline.new_accumulator(state) + fragments = [pipeline.run(unit, acc) for unit in pipeline.plan(state)] + return acc, fragments + + def test_plan_is_one_unit_per_column_in_column_order(self): + from buckaroo.pluggable_analysis_framework.stat_units import StatState + units = self._pipeline().plan(StatState(self._df())) + assert [u.columns for u in units] == [('amount',), ('word',), ('score',)] + assert len({u.id for u in units}) == 3 + assert {u.phase for u in units} == {'column'} + assert all(u.after == () for u in units), "no column's stats read another column's" + + def test_plan_and_a_new_accumulator_compute_nothing(self, monkeypatch): + from buckaroo.pluggable_analysis_framework.stat_units import StatState + calls = [] + monkeypatch.setattr(StatPipeline, 'process_column', lambda self, *a, **k: calls.append(a)) + pipeline, state = self._pipeline(), StatState(self._df()) + pipeline.plan(state) + pipeline.new_accumulator(state) + assert calls == [] + + def test_run_returns_the_fragment_of_its_unit_keyed_by_original_name(self): + from buckaroo.pluggable_analysis_framework.stat_units import StatState + pipeline, state = self._pipeline(), StatState(self._df()) + acc = pipeline.new_accumulator(state) + fragment = pipeline.run(pipeline.plan(state)[1], acc) + assert list(fragment) == ['word'] + assert fragment['word']['length'] == 4 + assert fragment['word']['distinct_count'] == 2 + assert (fragment['word']['orig_col_name'], fragment['word']['rewritten_col_name']) == ('word', 'b') + assert list(acc.sd()) == ['word'], "a column with no unit run yet has no entry" + + def test_the_union_of_fragments_is_what_process_df_returns(self): + from buckaroo.pluggable_analysis_framework.stat_units import StatState, merge_fragments, rewrite_sd + df, pipeline = self._df(), self._pipeline() + acc, fragments = self._run_units(pipeline, StatState(df)) + expected, errors = pipeline.process_df(df) + assert rewrite_sd(merge_fragments(fragments), df) == expected + assert list(acc.sd()) == ['amount', 'word', 'score'] + assert [(e.column, e.stat_key) for e in acc.errors] == [(e.column, e.stat_key) for e in errors] + + def test_an_error_in_a_unit_is_recorded_on_the_accumulator(self): + from buckaroo.pluggable_analysis_framework.stat_units import StatState + + @stat() + def fails_on_words(ser: RawSeries) -> int: + if isinstance(ser.iloc[0], str): + raise ValueError('words') + return 0 + + pipeline = StatPipeline([length, fails_on_words], unit_test=False) + acc, fragments = self._run_units(pipeline, StatState(self._df())) + assert [(e.column, e.stat_key) for e in acc.errors] == [('b', 'fails_on_words')] + assert fragments[1]['word']['fails_on_words'] is None + + @pytest.mark.parametrize('names', [('word',), ('b',)]) + def test_a_column_group_returns_only_those_columns(self, names): + """A group is named by original or rewritten name, since the client + only knows the rewritten ones.""" + from buckaroo.pluggable_analysis_framework.stat_units import StatState + pipeline = self._pipeline() + state = StatState(self._df(), columns=names) + assert [u.columns for u in pipeline.plan(state)] == [('word',)] + acc, fragments = self._run_units(pipeline, state) + assert [list(f) for f in fragments] == [['word']] + assert list(acc.sd()) == ['word'] + + @pytest.mark.parametrize('names', [('score',), ('c',)]) + def test_priority_columns_are_planned_first(self, names): + from buckaroo.pluggable_analysis_framework.stat_units import StatState + units = self._pipeline().plan(StatState(self._df(), priority=names)) + assert [u.columns[0] for u in units] == ['score', 'amount', 'word'] + + def test_an_empty_frame_has_no_units(self): + from buckaroo.pluggable_analysis_framework.stat_units import StatState + pipeline, state = self._pipeline(), StatState(pd.DataFrame({'x': pd.Series([], dtype='int64')})) + assert pipeline.plan(state) == [] + assert pipeline.new_accumulator(state).sd() == {} + + @pytest.mark.parametrize('backend', ['pandas', 'polars']) + def test_real_stats_run_unit_by_unit_equal_process_df(self, backend): + from buckaroo.pluggable_analysis_framework.stat_units import StatState, merge_fragments, rewrite_sd + df = _mixed_frames()[backend] + pipeline = StatPipeline(_stat_lists()[backend], unit_test=False) + acc, fragments = self._run_units(pipeline, StatState(df)) + expected, errors = pipeline.process_df(df) + assert len(fragments) == len(df.columns) + assert repr(rewrite_sd(merge_fragments(fragments), df)) == repr(expected) + assert errors == [] and acc.errors == [] diff --git a/tests/unit/test_xorq_stats_v2.py b/tests/unit/test_xorq_stats_v2.py index f1552cf14..221d0fb4f 100644 --- a/tests/unit/test_xorq_stats_v2.py +++ b/tests/unit/test_xorq_stats_v2.py @@ -933,3 +933,318 @@ def test_each_histogram_query_keeps_its_snapshot_key(self, tmp_path): # The batch aggregate and one histogram query per column, nothing else. assert len(recorder.keys) == 1 + len(expected) assert len(set(recorder.keys)) == len(recorder.keys) + + +# ============================================================ +# Resumable units: plan(state) and run(unit, acc) (rows-first s4) +# ============================================================ + +_STATS_LOGGER = "buckaroo.pluggable_analysis_framework.xorq_stat_pipeline" + + +def _stable(sd): + """An sd as comparable JSON. Equal counts come back from a GROUP BY in no + fixed order, so the histogram lists are sorted.""" + import json + + def norm(value, key=None): + if isinstance(value, dict): + return {k: norm(v, k) for k, v in value.items()} + if isinstance(value, list): + items = [norm(v) for v in value] + return sorted(items, key=lambda item: json.dumps(item, sort_keys=True)) if key == "histogram" else items + return value + + return json.dumps(norm(sd), sort_keys=True, default=str) + + +class _QuerySpy: + """Records every query ``XorqStatPipeline`` executes while it is active. + ``kinds()`` tells the batch aggregate, whose outermost op is an Aggregate, + from a histogram query, which is ordered or limited on top of one.""" + + def __init__(self, monkeypatch): + self.queries = [] + original = XorqStatPipeline._execute + + def spy(pipeline, query): + self.queries.append(query) + return original(pipeline, query) + + monkeypatch.setattr(XorqStatPipeline, "_execute", spy) + + def kinds(self): + return ["batch" if type(q.op()).__name__ == "Aggregate" else "histogram" for q in self.queries] + + +def _run_units(pipeline, state): + acc = pipeline.new_accumulator(state) + return acc, [(unit, pipeline.run(unit, acc)) for unit in pipeline.plan(state)] + + +class TestStatUnits: + """``plan(state)`` is the scalar batch first, then one histogram GROUP BY + per column; ``run(unit, acc)`` runs one unit and returns the + ``{column: {stat: value}}`` fragment it produced. ``process_table`` is + these run in order.""" + + _ALL = ("ints", "floats", "strs", "bools") + + def test_the_scalar_batch_is_first_then_one_histogram_unit_per_column(self): + from buckaroo.pluggable_analysis_framework.stat_units import StatState + units = XorqStatPipeline(XORQ_STATS_V2, unit_test=False).plan(StatState(_make_table())) + assert [u.id for u in units] == ["batch", *(f"histogram:{c}" for c in self._ALL)] + assert units[0].columns == self._ALL + assert [u.columns for u in units[1:]] == [(c,) for c in self._ALL] + assert [u.phase for u in units] == ["batch", *["histogram"] * 4] + assert units[0].after == () + assert all(u.after == ("batch",) for u in units[1:]) + + @pytest.mark.parametrize("visible", [("strs", "bools"), ("c", "d")]) + def test_visible_columns_get_their_histogram_units_first(self, visible): + """Visible columns are named by original or rewritten name. The batch + stays first, since every histogram reads what it computes.""" + from buckaroo.pluggable_analysis_framework.stat_units import StatState + units = XorqStatPipeline(XORQ_STATS_V2, unit_test=False).plan(StatState(_make_table(), priority=visible)) + assert [u.id for u in units] == [ + "batch", "histogram:strs", "histogram:bools", "histogram:ints", "histogram:floats"] + + def test_plan_and_a_new_accumulator_issue_no_query(self, monkeypatch): + from buckaroo.pluggable_analysis_framework.stat_units import StatState + spy = _QuerySpy(monkeypatch) + pipeline, state = XorqStatPipeline(XORQ_STATS_V2, unit_test=False), StatState(_make_table()) + pipeline.plan(state) + pipeline.new_accumulator(state) + assert spy.queries == [] + + def test_the_batch_is_one_query_and_each_histogram_unit_one_more(self, monkeypatch): + from buckaroo.pluggable_analysis_framework.stat_units import StatState + spy = _QuerySpy(monkeypatch) + pipeline = XorqStatPipeline(XORQ_STATS_V2, unit_test=False) + state = StatState(_make_table()) + acc = pipeline.new_accumulator(state) + for unit in pipeline.plan(state): + before = len(spy.queries) + pipeline.run(unit, acc) + assert len(spy.queries) - before == 1, unit.id + assert spy.kinds() == ["batch", "histogram", "histogram", "histogram", "histogram"] + + def test_the_batch_fragment_has_the_scalar_stats_and_the_bins_but_no_histogram(self): + """``histogram_bins`` is a pure function of five scalars, so it comes + with the batch and ``color_map`` needs no histogram query.""" + from buckaroo.pluggable_analysis_framework.stat_units import StatState + _acc, run = _run_units(XorqStatPipeline(XORQ_STATS_V2, unit_test=False), StatState(_make_table())) + batch_unit, batch = run[0] + assert set(batch) == set(self._ALL) + assert {"length", "null_count", "min", "max", "mean", "std", "median", "distinct_count", "non_null_count", + "nan_per", "distinct_per", "histogram_bins", "_type"} <= set(batch["ints"]) + assert batch["ints"]["histogram_bins"][0] == 1.0 and len(batch["ints"]["histogram_bins"]) == 11 + assert all("histogram" not in stats for stats in batch.values()) + assert "mean" not in batch["strs"] + + def test_a_histogram_fragment_holds_the_histogram_of_its_column_only(self): + from buckaroo.pluggable_analysis_framework.stat_units import StatState + _acc, run = _run_units(XorqStatPipeline(XORQ_STATS_V2, unit_test=False), StatState(_make_table())) + for unit, fragment in run[1:]: + assert list(fragment) == list(unit.columns) + assert set(fragment[unit.columns[0]]) == {"histogram"} + + def test_the_union_of_fragments_is_process_table(self): + from buckaroo.pluggable_analysis_framework.stat_units import StatState, merge_fragments + table = _make_table_with_nulls() + pipeline = XorqStatPipeline(XORQ_STATS_V2, unit_test=False) + acc, run = _run_units(pipeline, StatState(table)) + expected, errors = pipeline.process_table(table) + assert _stable(merge_fragments(f for _u, f in run)) == _stable(expected) + assert _stable(acc.sd()) == _stable(expected) + assert errors == [] and acc.errors == [] + + def test_a_failing_stat_is_recorded_once_and_the_run_goes_on(self): + from buckaroo.pluggable_analysis_framework.stat_units import StatState + + @stat() + def will_fail(col: XorqColumn) -> int: + raise RuntimeError("intentional") + + pipeline = XorqStatPipeline([*XORQ_STATS_V2, will_fail], unit_test=False) + acc, run = _run_units(pipeline, StatState(_make_table())) + assert sorted(e.column for e in acc.errors if e.stat_key == "will_fail") == sorted(self._ALL) + assert len(acc.errors) == len({(e.column, e.stat_key) for e in acc.errors}), "an error is reported once" + assert run[0][1]["ints"]["will_fail"] is None + assert acc.sd()["ints"]["histogram"], "later units still ran" + + @pytest.mark.parametrize("names", [("ints", "strs"), ("a", "c")]) + def test_a_column_group_returns_only_those_columns(self, names, monkeypatch): + from buckaroo.pluggable_analysis_framework.stat_units import StatState + spy = _QuerySpy(monkeypatch) + pipeline = XorqStatPipeline(XORQ_STATS_V2, unit_test=False) + state = StatState(_make_table(), columns=names) + assert [u.id for u in pipeline.plan(state)] == ["batch", "histogram:ints", "histogram:strs"] + acc, run = _run_units(pipeline, state) + assert all(set(fragment) <= {"ints", "strs"} for _u, fragment in run) + assert list(acc.sd()) == ["ints", "strs"] + assert spy.kinds() == ["batch", "histogram", "histogram"] + full, _errors = pipeline.process_table(_make_table()) + assert _stable(acc.sd()) == _stable({c: full[c] for c in ("ints", "strs")}) + + +# ============================================================ +# Splitting the batch by column chunk (rows-first s4) +# ============================================================ + +_SIX_COLUMNS = ("n0", "n1", "n2", "word", "label", "flag") +_SIX_ROWS = 12 + + +def _six_column_frame(): + """Twelve rows, six columns. Every string column has a different count for + each value, so a GROUP BY over it returns the same order every time.""" + return pd.DataFrame({ + "n0": [float(i) for i in range(_SIX_ROWS)], "n1": [i * 3 for i in range(_SIX_ROWS)], + "n2": [float(i % 7) + 0.5 for i in range(_SIX_ROWS)], "word": list("aaaaabbbbccd"), + "label": list("xxxxxyyyyzzw"), "flag": [True] * 7 + [False] * 5}) + + +def _six_column_scan(tmp_path): + path = tmp_path / "six.parquet" + _six_column_frame().to_parquet(path) + return xo.deferred_read_parquet(str(path)) + + +def _chunked(chunk_cells): + return XorqStatPipeline(XORQ_STATS_V2, unit_test=False, chunk_cells=chunk_cells) + + +def _batch_units(units): + return [u for u in units if u.phase == "batch"] + + +class TestColumnChunkSplit: + """The batch aggregate split into one aggregate per chunk of columns, sized + by cells (rows x columns). It is off unless the pipeline is given + ``chunk_cells``, and it applies only to a plain parquet scan, because each + chunk of a join, an aggregate or a CSV read runs the whole plan again.""" + + def test_the_batch_is_not_split_by_default(self, tmp_path): + from buckaroo.pluggable_analysis_framework.stat_units import StatState + units = XorqStatPipeline(XORQ_STATS_V2, unit_test=False).plan( + StatState(_six_column_scan(tmp_path), rows=_SIX_ROWS)) + assert [u.id for u in _batch_units(units)] == ["batch"] + + def test_chunks_are_sized_by_cells(self, tmp_path): + from buckaroo.pluggable_analysis_framework.stat_units import StatState + table = _six_column_scan(tmp_path) + # 36 cells at 12 rows is three columns per chunk, 24 is two, and a + # chunk always holds at least one column. + for cells, size in ((36, 3), (24, 2), (1, 1)): + batches = _batch_units(_chunked(cells).plan(StatState(table, rows=_SIX_ROWS))) + assert [len(u.columns) for u in batches] == [size] * (6 // size), cells + assert [c for u in batches for c in u.columns] == list(_SIX_COLUMNS) + assert [u.id for u in _batch_units(_chunked(36).plan(StatState(table, rows=_SIX_ROWS)))] == [ + "batch:0", "batch:1"] + + def test_a_chunk_that_fits_every_column_is_the_single_batch(self, tmp_path): + from buckaroo.pluggable_analysis_framework.stat_units import StatState + units = _chunked(10_000).plan(StatState(_six_column_scan(tmp_path), rows=_SIX_ROWS)) + assert [u.id for u in _batch_units(units)] == ["batch"] + + def test_batches_come_first_and_each_histogram_follows_its_chunk(self, tmp_path): + from buckaroo.pluggable_analysis_framework.stat_units import StatState + units = _chunked(36).plan(StatState(_six_column_scan(tmp_path), rows=_SIX_ROWS)) + assert [u.phase for u in units] == ["batch"] * 2 + ["histogram"] * 6 + chunk_of = {c: u.id for u in units[:2] for c in u.columns} + for unit in units[2:]: + assert unit.after == (chunk_of[unit.columns[0]],) + + def test_visible_columns_fill_the_first_chunk(self, tmp_path): + from buckaroo.pluggable_analysis_framework.stat_units import StatState + units = _chunked(36).plan(StatState(_six_column_scan(tmp_path), rows=_SIX_ROWS, priority=("flag", "n2"))) + assert _batch_units(units)[0].columns == ("flag", "n2", "n0") + assert [u.columns[0] for u in units if u.phase == "histogram"][:2] == ["flag", "n2"] + + def test_an_unknown_row_count_keeps_the_single_batch(self, tmp_path): + """Cells need the row count, and asking for it is a query.""" + from buckaroo.pluggable_analysis_framework.stat_units import StatState + units = _chunked(36).plan(StatState(_six_column_scan(tmp_path))) + assert [u.id for u in _batch_units(units)] == ["batch"] + + def test_the_split_equals_the_single_batch(self, tmp_path, monkeypatch): + from buckaroo.pluggable_analysis_framework.stat_units import StatState + table = _six_column_scan(tmp_path) + single, errors = XorqStatPipeline(XORQ_STATS_V2, unit_test=False).process_table(table) + spy = _QuerySpy(monkeypatch) + acc, _run = _run_units(_chunked(36), StatState(table, rows=_SIX_ROWS)) + assert acc.errors == [] and errors == [] + assert _stable(acc.sd()) == _stable(single) + assert spy.kinds() == ["batch"] * 2 + ["histogram"] * 6 + + def test_process_table_takes_the_row_count_and_splits(self, tmp_path, monkeypatch): + table = _six_column_scan(tmp_path) + single, _errors = XorqStatPipeline(XORQ_STATS_V2, unit_test=False).process_table(table) + spy = _QuerySpy(monkeypatch) + split, errors = _chunked(24).process_table(table, rows=_SIX_ROWS) + assert errors == [] + assert _stable(split) == _stable(single) + assert spy.kinds().count("batch") == 3 + + def test_the_split_leaves_the_histogram_cache_keys_alone(self, tmp_path): + """A histogram query does not depend on how the batch was cut, so a + snapshot written before the split is still read after it.""" + table = _six_column_scan(tmp_path) + + def keys(name, **kwargs): + cache = xo.ParquetSnapshotCache.from_kwargs(source=xo.connect(), base_path=str(tmp_path / name)) + recorder = _RecordingCache(cache) + XorqStatPipeline(XORQ_STATS_V2, unit_test=False, cache_storage=recorder, **kwargs).process_table( + table, rows=_SIX_ROWS) + return recorder.keys + + single, split = keys("single"), keys("split", chunk_cells=36) + assert len(single) == 1 + 6 and len(split) == 2 + 6 + assert split[2:] == single[1:], "the six histogram queries keep their keys and their order" + assert not set(split[:2]) & set(single), "each chunk is a different batch query" + + def test_a_plain_parquet_scan_and_a_column_projection_of_it_may_be_split(self, tmp_path): + table = _six_column_scan(tmp_path) + assert XorqStatPipeline.chunk_refusal(table) is None + assert XorqStatPipeline.chunk_refusal(table.select("n0", "word", "flag")) is None + + @staticmethod + def _sources(tmp_path): + """Expressions that are not a plain parquet scan, by kind.""" + frame = _six_column_frame() + frame.to_parquet(tmp_path / "left.parquet") + frame[["n0", "n1"]].rename(columns={"n1": "other"}).to_parquet(tmp_path / "right.parquet") + frame.to_csv(tmp_path / "six.csv", index=False) + left = xo.deferred_read_parquet(str(tmp_path / "left.parquet")) + right = xo.deferred_read_parquet(str(tmp_path / "right.parquet")) + outer = left.join(right, "n0", how="outer") + return { + "join": left.join(right, "n0"), + "aggregate": left.group_by("word").aggregate(total=left.n0.sum(), peak=left.n1.max()), + "diff": outer.mutate(delta=outer.other - outer.n1), + "csv": xo.deferred_read_csv(str(tmp_path / "six.csv")), + "filter": left.filter(left.n0 > 2), + "derived column": left.mutate(twice=left.n0 * 2), + "cached": left.cache(cache=xo.ParquetSnapshotCache.from_kwargs( + source=xo.connect(), base_path=str(tmp_path / "snapshots"))), + "memtable": xo.memtable(frame), + "backend table": xo.connect().read_parquet(str(tmp_path / "left.parquet"))} + + @pytest.mark.parametrize("kind", ["join", "aggregate", "diff", "csv", "filter", "derived column", "cached", + "memtable", "backend table"]) + def test_the_split_is_refused_for_anything_but_a_plain_parquet_scan(self, kind, tmp_path, monkeypatch, caplog): + from buckaroo.pluggable_analysis_framework.stat_units import StatState + table = self._sources(tmp_path)[kind] + assert XorqStatPipeline.chunk_refusal(table), kind + + # Refused means the one batch it always was, and the same stats. + with caplog.at_level(logging.INFO, logger=_STATS_LOGGER): + units = _chunked(1).plan(StatState(table, rows=_SIX_ROWS)) + assert [u.id for u in _batch_units(units)] == ["batch"] + assert any("refused" in r.getMessage() for r in caplog.records), "a refusal must be logged" + spy = _QuerySpy(monkeypatch) + split, _errors = _chunked(1).process_table(table, rows=_SIX_ROWS) + single, _errors = XorqStatPipeline(XORQ_STATS_V2, unit_test=False).process_table(table) + assert spy.kinds().count("batch") == 2 # one per run: the refused split and the single batch + assert _stable(split) == _stable(single) From 799583f7f28a84db88cb66a2dbb3fd04d493e6b7 Mon Sep 17 00:00:00 2001 From: Paddy Mullen Date: Sun, 4 Oct 2026 02:51:32 -0400 Subject: [PATCH 14/18] feat(stats): resumable stat units and the StatRun store (rows-first s4) StatPipeline and XorqStatPipeline get plan(state), new_accumulator(state) and run(unit, acc); process_df and process_table are those three in order, so the inline path and a caller running units one at a time compute the same stats. pandas and polars plan one unit per non-skipped column. xorq plans the scalar batch (which also yields histogram_bins), then one histogram query per column, visible columns first. The batch can be cut into column chunks sized by cells when a host sets stat_chunk_cells, and only for a plain parquet scan: any other source falls back to the single batch and logs why. A StatRun, kept on the session by (stats_gen, scope), holds the planned units, an append-only fragment list and the accumulator. It has no thread, timer or callback, begin_stats_generation drops it, and each WebSocket connection keeps its own StatCursor into the list. CustomizableDataflow.build_stats is the one place a stats class is built, with run=False for plan/run callers. Two of the new tests are corrected here: the default-path test now ignores the PERVERSE_DF self-check units, and the xorq comparisons round floats to nine digits because a parallel aggregate adds its partial sums in varying order. The new-module imports in the tests move to module level. Co-Authored-By: Claude Sonnet 5.5 --- buckaroo/dataflow/dataflow.py | 29 +- .../df_stats_v2.py | 39 +- .../stat_pipeline.py | 80 +++-- .../stat_units.py | 181 ++++++++++ .../xorq_stat_pipeline.py | 333 +++++++++++++----- buckaroo/server/session.py | 9 +- buckaroo/server/stat_run.py | 118 +++++++ buckaroo/server/stats_wire.py | 19 + buckaroo/server/websocket_handler.py | 4 + buckaroo/server/xorq_loading.py | 8 +- buckaroo/xorq_buckaroo.py | 44 ++- .../dataflow/scoped_summary_stats_test.py | 2 +- tests/unit/server/test_data_loading_polars.py | 30 +- tests/unit/server/test_load_expr.py | 12 +- tests/unit/skip_columns_test.py | 4 +- tests/unit/test_paf_v2.py | 10 +- tests/unit/test_xorq_stats_v2.py | 24 +- 17 files changed, 742 insertions(+), 204 deletions(-) create mode 100644 buckaroo/pluggable_analysis_framework/stat_units.py create mode 100644 buckaroo/server/stat_run.py diff --git a/buckaroo/dataflow/dataflow.py b/buckaroo/dataflow/dataflow.py index 68b7bd416..1a61b3bc1 100644 --- a/buckaroo/dataflow/dataflow.py +++ b/buckaroo/dataflow/dataflow.py @@ -11,6 +11,7 @@ from ..serialization_utils import pd_to_obj, sd_to_parquet_b64, project_sd from buckaroo.pluggable_analysis_framework.utils import (filter_analysis) from buckaroo.pluggable_analysis_framework.df_stats_v2 import DfStatsV2 +from buckaroo.pluggable_analysis_framework.stat_units import Fragment, StatAccumulator, StatState, StatUnit from .autocleaning import SentinelAutocleaning from .dataflow_extras import (exception_protect, Sampling) from .styling_core import ( @@ -139,15 +140,22 @@ class DfStats(Protocol): ``DfStatsV2`` (pandas), ``PlDfStatsV2`` (polars) and ``XorqDfStatsV2`` share this surface but no base class. The frame argument is ``Any`` because each executor accepts only its own backend's frame type. + + ``run=False`` builds one without computing anything, for a caller that runs + ``plan(state)`` and ``run(unit, acc)`` itself (see ``stat_units``). """ sdf: TAny errs: TAny ap: TAny + state: StatState def __init__(self, df: TAny, col_analysis_objs: TAny, /, operating_df_name: TAny = ..., - debug: bool = ..., skip_columns: TAny = ...) -> None: ... + debug: bool = ..., skip_columns: TAny = ..., run: bool = ...) -> None: ... @classmethod def verify_analysis_objects(cls, col_analysis_objs: TAny, /) -> None: ... def add_analysis(self, a_obj: TAny, /) -> None: ... + def plan(self, state: Optional[StatState] = ...) -> List[StatUnit]: ... + def new_accumulator(self, state: Optional[StatState] = ...) -> StatAccumulator: ... + def run(self, unit: StatUnit, acc: StatAccumulator) -> Fragment: ... class DataFlow(ABCDataflow[DataFrameT], Generic[DataFrameT]): @@ -756,16 +764,25 @@ def _build_error_dataframe(self, e) -> pd.DataFrame: ### start summary stats block + def build_stats(self, processed_df: DataFrameT, run: bool = True) -> DfStats: + """The stats executor for ``processed_df``, built the one way every + summary-stats run builds it. With ``run=False`` it computes nothing yet + and the caller drives its ``plan`` and ``run``.""" + # ``run`` is passed only when it is False, so a DFStatsClass written + # before it existed still builds on the default path. + extra = {} if run else {'run': False} + return self.DFStatsClass( + processed_df, + self.analysis_klasses, + self.df_name, debug=self.debug, + skip_columns=getattr(self, 'skip_stat_columns', None), **extra) + #TAny closer to some error type @override def _get_summary_sd(self, processed_df: DataFrameT) -> Tuple[SDType, ErrDict]: if self.stats_tier == "schema": return self._get_schema_sd(processed_df), {} - stats = self.DFStatsClass( - processed_df, - self.analysis_klasses, - self.df_name, debug=self.debug, - skip_columns=getattr(self, 'skip_stat_columns', None)) + stats = self.build_stats(processed_df) sdf = stats.sdf if stats.errs: if self.debug: diff --git a/buckaroo/pluggable_analysis_framework/df_stats_v2.py b/buckaroo/pluggable_analysis_framework/df_stats_v2.py index 4c9fcadf1..83036a5ef 100644 --- a/buckaroo/pluggable_analysis_framework/df_stats_v2.py +++ b/buckaroo/pluggable_analysis_framework/df_stats_v2.py @@ -20,15 +20,17 @@ from .col_analysis import AObjs, ColAnalysis from .stat_pipeline import StatPipeline, errors_to_errdict +from .stat_units import StatState, UnitStats from .utils import FAST_SUMMARY_WHEN_GREATER from .safe_summary_df import output_full_reproduce -class DfStatsV2: +class DfStatsV2(UnitStats): """Tie a StatPipeline to a DataFrame, exposing ``.sdf`` and ``.errs``. Used by DataFlow, autocleaning, and other consumers as the pandas - summary-stats executor. + summary-stats executor. ``run=False`` builds it without computing anything, + for a caller that drives ``plan`` / ``run`` itself. """ ap_class = StatPipeline @@ -39,19 +41,23 @@ def verify_analysis_objects(cls, col_analysis_objs: AObjs) -> None: cls.ap_class(col_analysis_objs) def __init__(self, df_stats_df: pd.DataFrame, col_analysis_objs: AObjs, operating_df_name: str = None, - debug: bool = False, skip_columns=None) -> None: + debug: bool = False, skip_columns=None, run: bool = True) -> None: self.df = self.get_operating_df(df_stats_df, force_full_eval=False) self.col_order = self.df.columns self.ap = self.ap_class(col_analysis_objs) self.operating_df_name = operating_df_name self.debug = debug + self.state = StatState(self.df, frozenset(skip_columns or ())) - self.sdf, errors = self.ap.process_df(self.df, self.debug, skip_columns=skip_columns) - self.errs = errors_to_errdict(errors) + self.sdf = {} + self.errs = {} self.stat_errors = [] + if run: + self.sdf, errors = self.ap.process_df(self.df, self.debug, skip_columns=skip_columns) + self.errs = errors_to_errdict(errors) - if self.errs: - output_full_reproduce(self.errs, self.sdf, operating_df_name) + if self.errs: + output_full_reproduce(self.errs, self.sdf, operating_df_name) def get_operating_df(self, df: pd.DataFrame, force_full_eval: bool) -> pd.DataFrame: """Downsample large DataFrames for performance.""" @@ -80,8 +86,9 @@ def add_analysis(self, a_obj: Type[ColAnalysis]) -> None: self.ap.print_errors(errors + self.stat_errors) -class PlDfStatsV2: - """Polars summary-stats executor. Uses StatPipeline with @stat polars functions.""" +class PlDfStatsV2(UnitStats): + """Polars summary-stats executor. Uses StatPipeline with @stat polars functions. + ``run=False`` builds it without computing anything, as for ``DfStatsV2``.""" @classmethod def verify_analysis_objects(cls, objs): @@ -93,14 +100,18 @@ def get_operating_df(self, df): return df.sample(n=min(50_000, rows), seed=42) return df - def __init__(self, df, col_analysis_objs, operating_df_name=None, debug=False, skip_columns=None): + def __init__(self, df, col_analysis_objs, operating_df_name=None, debug=False, skip_columns=None, run=True): self.df = self.get_operating_df(df) self.ap = StatPipeline(col_analysis_objs, unit_test=False) - self.sdf, errors = self.ap.process_df(self.df, debug, skip_columns=skip_columns) - self.errs = errors_to_errdict(errors) + self.state = StatState(self.df, frozenset(skip_columns or ())) + self.sdf = {} + self.errs = {} self.stat_errors = [] - if self.errs: - output_full_reproduce(self.errs, self.sdf, operating_df_name) + if run: + self.sdf, errors = self.ap.process_df(self.df, debug, skip_columns=skip_columns) + self.errs = errors_to_errdict(errors) + if self.errs: + output_full_reproduce(self.errs, self.sdf, operating_df_name) def add_analysis(self, a_obj): """Add a new analysis class interactively.""" diff --git a/buckaroo/pluggable_analysis_framework/stat_pipeline.py b/buckaroo/pluggable_analysis_framework/stat_pipeline.py index f121624e0..f26e1de98 100644 --- a/buckaroo/pluggable_analysis_framework/stat_pipeline.py +++ b/buckaroo/pluggable_analysis_framework/stat_pipeline.py @@ -9,7 +9,7 @@ from __future__ import annotations import time -from typing import Any, Callable, Dict, List, Optional, Tuple +from typing import Any, Callable, Dict, List, Optional, Sequence, Tuple import pandas as pd @@ -20,6 +20,7 @@ from .col_analysis import ColAnalysis, ErrDict, SDType from .stat_func import (StatFunc, RawSeries, SampledSeries, RawDataFrame, XorqExpr, XorqExecute, RAW_MARKER_TYPES, MISSING, collect_stat_funcs) from .stat_result import Ok, Err, UpstreamError, StatError, StatResult, resolve_accumulator +from .stat_units import Fragment, StatAccumulator, StatState, StatUnit, UnitPipeline, columns_in_scope, prioritized from .typed_dag import build_typed_dag, build_column_dag, DAGConfigError from .utils import PERVERSE_DF @@ -173,7 +174,7 @@ def _execute_stat_func(sf: StatFunc, accumulator: Dict[str, StatResult], column_ inputs=kwargs.copy()) -class StatPipeline: +class StatPipeline(UnitPipeline): """Top-level orchestrator for the pluggable analysis framework. Accepts a mix of: @@ -277,9 +278,56 @@ def process_column(self, column_name: str, column_dtype, raw_series=None, sample return resolve_accumulator(accumulator, column_name, col_key_to_func) + @staticmethod + def _skipped(state: StatState, pairs: Sequence[Tuple[Any, str]]) -> set: + """The columns of ``pairs`` named in ``state.skip_columns``, by original + or rewritten name.""" + return {orig for orig, rewritten in pairs + if orig in state.skip_columns or rewritten in state.skip_columns} + + def plan(self, state: StatState) -> List[StatUnit]: + """One unit per column that is not skipped, in column order (or the + order ``state.priority`` asks for). A column's stats read only its own + column, so no unit waits for another. Nothing is computed.""" + if len(state.data) == 0: + return [] + pairs = columns_in_scope(state) + skip = self._skipped(state, pairs) + active = prioritized([pair for pair in pairs if pair[0] not in skip], state.priority) + return [StatUnit(id=f"column:{orig}", columns=(orig,), phase="column") for orig, _rewritten in active] + + def new_accumulator(self, state: StatState) -> StatAccumulator: + """The accumulator for a run of ``state``. A skipped column is entered + with its two names, as ``process_df`` always did, and a frame with no + rows has no entries at all.""" + if self.record_timings: + self.timings = [] + if len(state.data) == 0: + return StatAccumulator(state) + pairs = columns_in_scope(state) + acc = StatAccumulator(state, columns=[orig for orig, _rewritten in pairs], rewritten=dict(pairs)) + for orig in self._skipped(state, pairs): + acc.results[orig] = {'orig_col_name': orig, 'rewritten_col_name': acc.rewritten[orig]} + return acc + + def run(self, unit: StatUnit, acc: StatAccumulator) -> Fragment: + """Run one column's stats through the stat DAG and return them as + ``{orig_col: stats}``. The errors go on ``acc.errors``.""" + df = acc.state.data + fragment: Fragment = {} + for orig_col_name in unit.columns: + rewritten_col_name = acc.rewritten[orig_col_name] + ser = df[orig_col_name] + col_result, col_errors = self.process_column(column_name=rewritten_col_name, column_dtype=ser.dtype, + raw_series=ser, sampled_series=ser, raw_dataframe=df, + initial_stats={'orig_col_name': orig_col_name, 'rewritten_col_name': rewritten_col_name}) + fragment[orig_col_name] = col_result + acc.record({orig_col_name: col_result}, col_errors) + return fragment + def process_df(self, df: pd.DataFrame, debug: bool = False, skip_columns=None) -> Tuple[SDType, List[StatError]]: - """Process all columns of a DataFrame. + """Process all columns of a DataFrame: the planned units, run in order. ``skip_columns`` names columns whose summary stats are supplied externally (e.g. via ``init_sd`` — reused from a source dataframe in a @@ -293,35 +341,15 @@ def process_df(self, df: pd.DataFrame, debug: bool = False, if len(df) == 0: return {}, [] - if self.record_timings: - self.timings = [] - - skip = set(skip_columns or ()) - summary: SDType = {} - all_errors: List[StatError] = [] - - for orig_col_name, rewritten_col_name in old_col_new_col(df): - if orig_col_name in skip or rewritten_col_name in skip: - # Provided externally — keep the column, skip computation. - summary[rewritten_col_name] = { - 'orig_col_name': orig_col_name, 'rewritten_col_name': rewritten_col_name} - continue - ser = df[orig_col_name] - col_dtype = ser.dtype - - col_result, col_errors = self.process_column(column_name=rewritten_col_name, column_dtype=col_dtype, - raw_series=ser, sampled_series=ser, raw_dataframe=df, - initial_stats={'orig_col_name': orig_col_name, 'rewritten_col_name': rewritten_col_name}) - - summary[rewritten_col_name] = col_result - all_errors.extend(col_errors) + acc = self.run_all(StatState(df, frozenset(skip_columns or ()))) + summary: SDType = {acc.rewritten[orig]: stats for orig, stats in acc.sd().items()} if self.record_timings and perf_log.enabled() and not self._suppress_perf_summary: rec = perf_log.PerfRecorder(label=f"stats rows={len(df)} cols={len(summary)}") rec.extend_timings("pandas/polars", self.timings) rec.summary() - return summary, all_errors + return summary, acc.errors def unit_test(self) -> Tuple[bool, List[StatError]]: """Test the pipeline against PERVERSE_DF.""" diff --git a/buckaroo/pluggable_analysis_framework/stat_units.py b/buckaroo/pluggable_analysis_framework/stat_units.py new file mode 100644 index 000000000..ff3c5dd24 --- /dev/null +++ b/buckaroo/pluggable_analysis_framework/stat_units.py @@ -0,0 +1,181 @@ +"""Resumable stat units (rows-first s4). + +A stat run is a list of units a caller can run one at a time. ``plan(state)`` +lists them without computing anything, ``new_accumulator(state)`` makes the +object they share, and ``run(unit, acc)`` runs one and returns its *fragment*, +a ``{orig_col: {stat: value}}`` dict holding the stats that unit produced. The +pipelines' ``process_df`` and ``process_table`` are these three called in order, +so running a whole table and running its units one by one cannot differ. + +A column named in ``StatState.skip_columns`` (its stats come from ``init_sd``) +gets no unit. It still has an entry in ``acc.sd()``, with whatever structural +keys the pipeline gives a skipped column, because ``process_df`` and +``process_table`` always returned one. +""" +from __future__ import annotations + +from dataclasses import dataclass +from typing import Any, Dict, Iterable, Iterator, List, Optional, Sequence, Tuple + +from buckaroo.df_util import old_col_new_col + +from .col_analysis import SDType +from .stat_result import StatError + +# {orig_col: {stat: value}}: the stats one unit produced, for the columns it covers. +Fragment = Dict[Any, Dict[str, Any]] + + +@dataclass(frozen=True) +class StatUnit: + """One bounded piece of a stat run. + + ``id`` is unique within a plan. ``columns`` are original column names. + ``phase`` says what kind of work it is, and ``cost`` is a plain string + class (``column``, ``scan`` or ``query``) a scheduler can read. ``after`` + names the units whose results this one reads, which must have run first. + """ + id: str + columns: Tuple[Any, ...] + phase: str + after: Tuple[str, ...] = () + cost: str = "column" + + +@dataclass(frozen=True, eq=False) +class StatState: + """What a stat run describes: the frame or expression, and how to cut it. + + ``skip_columns`` get no unit. ``columns`` restricts the run to a group (a + column-group request), and ``priority`` names columns to plan first. Both + take original or rewritten (``a, b, c``) names, since a client only knows + the rewritten ones. ``rows`` is the row count when the caller already has + it, which the xorq column-chunk split needs and which is never queried for. + """ + data: Any + skip_columns: frozenset = frozenset() + columns: Optional[Tuple[Any, ...]] = None + priority: Tuple[Any, ...] = () + rows: Optional[int] = None + + +def columns_in_scope(state: StatState) -> List[Tuple[Any, str]]: + """``(orig, rewritten)`` for each column the run covers, in data order. The + rewritten name is the column's position in the whole of ``state.data``, + whatever group is asked for.""" + pairs = old_col_new_col(state.data) + if state.columns is None: + return list(pairs) + wanted = set(state.columns) + return [(orig, rewritten) for orig, rewritten in pairs if orig in wanted or rewritten in wanted] + + +def prioritized(pairs: Sequence[Tuple[Any, str]], priority: Sequence[Any]) -> List[Tuple[Any, str]]: + """``pairs`` with the columns named in ``priority`` first, in that order, + and the rest in their own order.""" + rank: Dict[Any, int] = {} + for position, name in enumerate(priority): + rank.setdefault(name, position) + + def key(pair: Tuple[Any, str]) -> int: + ranks = [rank[name] for name in pair if name in rank] + return min(ranks) if ranks else len(rank) + + return sorted(pairs, key=key) + + +class StatAccumulator: + """What the units of one run share: the state, the stats recorded so far + for each column (``results``, keyed by original name) and the errors the + units reported. + + ``record`` is how a unit's fragment lands here. ``sd()`` is the result so + far in data order. A column no unit has run for is absent from it, unless + it is a skipped column, whose entry is there from the start. + """ + + def __init__(self, state: StatState, columns: Sequence[Any] = (), rewritten: Optional[Dict[Any, str]] = None): + self.state = state + self.order: List[Any] = list(columns) + self.rewritten: Dict[Any, str] = dict(rewritten or {}) + self.results: Dict[Any, Dict[str, Any]] = {} + self.errors: List[StatError] = [] + + def record(self, fragment: Fragment, errors: Iterable[StatError] = ()) -> None: + for col, stats in fragment.items(): + self.results.setdefault(col, {}).update(stats) + self.errors.extend(errors) + + def sd(self) -> Dict[Any, Dict[str, Any]]: + return {col: self.results[col] for col in self.order if col in self.results} + + +def merge_fragments(fragments: Iterable[Fragment]) -> Fragment: + """The union of fragments: each column's stats from every fragment that + has it. Neither the fragments nor their column dicts are changed.""" + merged: Fragment = {} + for fragment in fragments: + for col, stats in fragment.items(): + merged.setdefault(col, {}).update(stats) + return merged + + +def rewrite_sd(results: Dict[Any, Dict[str, Any]], data: Any) -> SDType: + """``results`` (keyed by original column name) as the summary dict the + dataflow carries: keyed by the rewritten name, with ``orig_col_name`` and + ``rewritten_col_name`` set from ``data``'s columns. Empty when ``results`` + is, as ``process_df`` is for an empty frame.""" + if not results: + return {} + rewritten: SDType = {} + for orig_col, rewritten_col in old_col_new_col(data): + col_meta = dict(results.get(orig_col, {})) + col_meta['orig_col_name'] = orig_col + col_meta['rewritten_col_name'] = rewritten_col + rewritten[rewritten_col] = col_meta + return rewritten + + +class UnitPipeline: + """The loop ``StatPipeline`` and ``XorqStatPipeline`` share: plan, then run + each unit. A subclass implements ``plan``, ``new_accumulator`` and ``run``.""" + + def plan(self, state: StatState) -> List[StatUnit]: + raise NotImplementedError + + def new_accumulator(self, state: StatState) -> StatAccumulator: + raise NotImplementedError + + def run(self, unit: StatUnit, acc: StatAccumulator) -> Fragment: + raise NotImplementedError + + def iter_units(self, state: StatState, acc: StatAccumulator) -> Iterator[Tuple[StatUnit, Fragment]]: + """Plan the run and run its units in order, yielding each with its + fragment as it finishes. Nothing runs until the generator is advanced, + so a caller can stop between units.""" + for unit in self.plan(state): + yield unit, self.run(unit, acc) + + def run_all(self, state: StatState) -> StatAccumulator: + """Every unit, in order. What ``process_df`` and ``process_table`` are.""" + acc = self.new_accumulator(state) + for _unit, _fragment in self.iter_units(state, acc): + pass + return acc + + +class UnitStats: + """``plan``, ``new_accumulator`` and ``run`` for a stats class (``DfStatsV2`` + and its siblings) that holds a pipeline in ``ap`` and the state it analyzes + in ``state``. The state defaults to the stats class's own.""" + ap: Any + state: StatState + + def plan(self, state: Optional[StatState] = None) -> List[StatUnit]: + return self.ap.plan(self.state if state is None else state) + + def new_accumulator(self, state: Optional[StatState] = None) -> StatAccumulator: + return self.ap.new_accumulator(self.state if state is None else state) + + def run(self, unit: StatUnit, acc: StatAccumulator) -> Fragment: + return self.ap.run(unit, acc) diff --git a/buckaroo/pluggable_analysis_framework/xorq_stat_pipeline.py b/buckaroo/pluggable_analysis_framework/xorq_stat_pipeline.py index 37eff339f..93ab518f6 100644 --- a/buckaroo/pluggable_analysis_framework/xorq_stat_pipeline.py +++ b/buckaroo/pluggable_analysis_framework/xorq_stat_pipeline.py @@ -1,14 +1,17 @@ """Xorq-backed stat pipeline for the v2 framework. -Two-phase execution: +Two-phase execution, each phase a kind of resumable unit (see ``stat_units``): 1. Batch aggregate — every @stat with an XorqColumn parameter contributes one ibis scalar expression. All such expressions across all columns are folded into a single ``table.aggregate(...)`` query and executed - once. - 2. Per-column post-batch — computed stats (deps only on other stats) and - XorqExpr-param stats (e.g. histograms that need their own query) - run through the standard typed-DAG executor with results written into - the per-column accumulator. + once. This is the ``batch`` unit, and it also computes the stats that + are pure functions of its results (``histogram_bins``, ``nan_per``...). + Behind a host's opt-in the batch can be cut into one aggregate per chunk + of columns (``chunk_cells``), for a plain parquet scan only. + 2. Per-column post-batch — XorqExpr-param stats (e.g. histograms that need + their own query), and any stat reading one, run through the standard + typed-DAG executor with results written into the per-column + accumulator. One ``histogram:`` unit per column. Errors are captured into ``StatError`` via the standard Ok/Err mechanism; nothing is silently swallowed. Construction validates the DAG up front and @@ -23,7 +26,7 @@ import os import time from contextlib import nullcontext -from typing import Any, Dict, List, Tuple +from typing import Any, Dict, List, Optional, Sequence, Tuple import pandas as pd @@ -33,6 +36,8 @@ from .stat_func import XorqColumn, XorqExpr, XorqExecute, RAW_MARKER_TYPES, StatFunc from .stat_pipeline import _execute_stat_func, _normalize_inputs, errors_to_errdict from .stat_result import Err, Ok, StatError, StatResult, resolve_accumulator +from .stat_units import (Fragment, StatAccumulator, StatState, StatUnit, UnitPipeline, UnitStats, columns_in_scope, + prioritized) from .typed_dag import build_column_dag, build_typed_dag from .utils import PERVERSE_DF @@ -41,10 +46,12 @@ try: import xorq.api as xo + from xorq.expr.relations import Read + from xorq.vendor.ibis.expr import operations as ops HAS_XORQ = True except ImportError: - xo = None + xo = Read = ops = None HAS_XORQ = False log = logging.getLogger(__name__) @@ -100,7 +107,29 @@ def _is_batch_func(sf: StatFunc) -> bool: return True -class XorqStatPipeline: +class XorqAccumulator(StatAccumulator): + """``StatAccumulator`` for a xorq run. ``working`` holds each column's + Ok/Err accumulator, which the units read and write; ``emitted`` the keys of + it already returned in a fragment, so each stat is returned and any error in + it reported once; ``phases`` caches how each dtype's stats split (see + ``XorqStatPipeline._column_phases``). ``sd()`` is every column's stats so + far in accumulator order, a failed stat as None.""" + + def __init__(self, state: StatState, columns: Sequence[Any] = (), rewritten: Optional[Dict[Any, str]] = None, + skipped: Sequence[Any] = ()): + super().__init__(state, columns, rewritten) + self.skipped: Tuple[Any, ...] = tuple(skipped) + self.working: Dict[Any, Dict[str, StatResult]] = {} + self.emitted: Dict[Any, set] = {} + self.phases: Dict[Any, Any] = {} + self.schema: Any = None + + def sd(self) -> Dict[Any, Dict[str, Any]]: + return {col: {key: (result.value if isinstance(result, Ok) else None if isinstance(result, Err) else result) + for key, result in self.working[col].items()} for col in self.order} + + +class XorqStatPipeline(UnitPipeline): """v2 stat pipeline for ``ibis.Table`` inputs. Accepts the same kinds of inputs as ``StatPipeline``: @@ -110,7 +139,13 @@ class XorqStatPipeline: - ``ColAnalysis`` subclasses (via v1 adapter) Use ``process_table(table)`` to run the pipeline; returns - ``(SDType, List[StatError])``. + ``(SDType, List[StatError])``. ``plan(state)``, ``new_accumulator(state)`` + and ``run(unit, acc)`` are the same run one unit at a time. + + ``chunk_cells`` (off when ``None``) cuts the batch aggregate into one + aggregate per chunk of columns holding about that many cells, when the + caller gives the row count and the source is a plain parquet scan + (``chunk_refusal``). """ # Keys that the pipeline pre-populates per column. Listed as external @@ -123,7 +158,7 @@ class XorqStatPipeline: "distinct_count"}) def __init__(self, stat_funcs: list, backend: Any = None, unit_test: bool = True, - cache_storage=None): + cache_storage=None, chunk_cells: Optional[int] = None): if not HAS_XORQ: raise ImportError( "xorq is required for XorqStatPipeline. " @@ -138,12 +173,13 @@ def __init__(self, stat_funcs: list, backend: Any = None, unit_test: bool = True self._original_inputs = list(stat_funcs) self.backend = backend self.cache_storage = cache_storage + self.chunk_cells = chunk_cells - # Per-run snapshot-cache counters, (re)initialised in process_table. + # Per-run snapshot-cache counters, (re)initialised in new_accumulator. # Set here so the attribute always exists (e.g. for the unit_test() # run kicked off below, which disables the cache). self._cache_stats = _new_cache_run_stats() - # Per-run perf recorder, (re)initialised in process_table when the + # Per-run perf recorder, (re)initialised in new_accumulator when the # BUCKAROO_PERF toggle is on; None otherwise. self._perf = None # Set during the unit_test() DAG self-check so its PERVERSE_DF run @@ -324,19 +360,21 @@ def _span(self, label, **fields): return nullcontext() return perf_log.perf_span(label, **fields) - def process_table(self, table, skip_columns=None) -> Tuple[SDType, List[StatError]]: + def process_table(self, table, skip_columns=None, rows=None) -> Tuple[SDType, List[StatError]]: # Each per-column query runs directly against ``table`` (the source # expression). When a snapshot cache is set, queries are keyed against # that source so the content-addressed key is stable across processes; # ``_execute_cached`` serves a hit from the snapshot parquet or executes - # and writes one on a miss. - self._cache_stats = _new_cache_run_stats() - self._perf = (perf_log.PerfRecorder() - if perf_log.enabled() and not self._suppress_perf_summary else None) + # and writes one on a miss. ``rows`` is the row count when the caller + # has it; only the column-chunk split reads it. + state = StatState(table, frozenset(skip_columns or ()), rows=rows) + acc = self.new_accumulator(state) _t0 = time.perf_counter() with self._span("stat.xorq.total") as span: try: - return self._process_table_impl(table, skip_columns=skip_columns) + for _unit, _fragment in self.iter_units(state, acc): + pass + return acc.sd(), acc.errors finally: self._cache_stats["secs"] = round(time.perf_counter() - _t0, 4) self._log_cache_stats(span) @@ -347,27 +385,159 @@ def process_table(self, table, skip_columns=None) -> Tuple[SDType, List[StatErro f"misses={self._cache_stats['misses']}") self._perf.summary() - def _process_table_impl(self, table, skip_columns=None) -> Tuple[SDType, List[StatError]]: + @staticmethod + def chunk_refusal(table) -> Optional[str]: + """Why ``table`` may not have its batch split by column chunk, or + ``None`` when it may. + + Each chunk is its own aggregate over the source. On a parquet scan a + chunk reads only its columns, so the chunks cost about what one batch + does (0.90-1.16x at 10-12M rows) and each needs less memory. On + anything else a chunk runs the plan again: a join or an aggregate + executes in full per chunk, and a CSV read reparses the file (2.1-2.2x). + So the split is allowed for a parquet ``Read``, or a projection of + plain columns of one, and refused for everything else, including an + expression this does not recognise. A refusal can only keep the single + batch; it never turns the split on. + """ + op = table.op() + if isinstance(op, ops.Project): + if not all(isinstance(value, ops.Field) for value in op.values.values()): + return "it computes columns instead of reading them" + op = op.parent + if not isinstance(op, Read): + return f"not a parquet scan ({type(op).__name__})" + if op.method_name != "read_parquet": + return f"not a parquet scan ({op.method_name})" + return None + + def _batch_chunks(self, state: StatState, columns: List[Any]) -> List[List[Any]]: + """``columns`` cut into the chunks the batch runs in: one, unless the + split is on, the row count is known and the source allows it.""" + if not self.chunk_cells: + return [columns] + if state.rows is None: + log.debug("xorq stat column chunking needs the row count; running one batch") + return [columns] + refusal = self.chunk_refusal(state.data) + if refusal is not None: + log.info("xorq stat column chunking refused: %s", refusal) + return [columns] + per_chunk = max(1, self.chunk_cells // max(state.rows, 1)) + return [columns[i:i + per_chunk] for i in range(0, len(columns), per_chunk)] or [columns] + + def _column_phases(self, cache: Dict[Any, Any], dtype) -> Tuple[List[StatFunc], List[StatFunc], Dict[str, StatFunc]]: + """A column's stat funcs split into those that need no further query + once the batch has run (``pure``) and those that run one or read the + result of one that does (``query``), each in dependency order, plus the + stat key -> func map errors are reported against.""" + if dtype not in cache: + col_funcs = build_column_dag(self.all_stat_funcs, dtype, external_keys=self.EXTERNAL_KEYS) + pure: List[StatFunc] = [] + query: List[StatFunc] = [] + query_keys: set = set() + for sf in col_funcs: + if any(r.type in (XorqExpr, XorqExecute) or r.name in query_keys for r in sf.requires): + query.append(sf) + query_keys.update(sk.name for sk in sf.provides) + else: + pure.append(sf) + key_to_func = {sk.name: sf for sf in col_funcs for sk in sf.provides} + cache[dtype] = (pure, query, key_to_func) + return cache[dtype] + + def plan(self, state: StatState) -> List[StatUnit]: + """The scalar batch first, then one histogram GROUP BY per column that + has one, columns named in ``state.priority`` first. The batch is cut + into chunks only as ``_batch_chunks`` allows, and every chunk precedes + every histogram, since a histogram reads what its chunk computes. + Skipped columns get no unit. Nothing is computed: no query is sent.""" + table = state.data + pairs = columns_in_scope(state) + if not pairs: + return [] schema = table.schema() - columns = list(table.columns) - # Columns whose stats are supplied externally (via init_sd) keep their - # structural metadata (name/dtype/length) but get no stat expressions - # built — so the column's data is never scanned. - skip = set(skip_columns or ()) - - # Pre-populate every column accumulator with the externally-provided - # keys. ``length`` is filled in by the batch query below. ``min`` / - # ``max`` start as None so dependents (histogram) don't cascade- - # exclude on non-numeric columns; ``min`` / ``max`` overwrite for - # numeric cols. ``distinct_count`` likewise starts as None so float - # columns (where the stat is column_filtered out) keep their - # dependents (histogram, histogram_bins, distinct_per) runnable. - accumulators: Dict[str, Dict[str, StatResult]] = {} - for col in columns: - accumulators[col] = {"orig_col_name": Ok(col), "rewritten_col_name": Ok(col), "dtype": Ok(str(schema[col])), + # xorq has always matched skip_columns on the column's own name only. + skip = {orig for orig, _rewritten in pairs if orig in state.skip_columns} + active = [orig for orig, _rewritten in prioritized([p for p in pairs if p[0] not in skip], state.priority)] + chunks = self._batch_chunks(state, active) + units: List[StatUnit] = [] + batch_of: Dict[Any, str] = {} + for i, chunk in enumerate(chunks): + unit_id = "batch" if len(chunks) == 1 else f"batch:{i}" + units.append(StatUnit(id=unit_id, columns=tuple(chunk), phase="batch", cost="scan")) + batch_of.update({col: unit_id for col in chunk}) + phases: Dict[Any, Any] = {} + for col in active: + if self._column_phases(phases, schema[col])[1]: + units.append(StatUnit(id=f"histogram:{col}", columns=(col,), phase="histogram", + after=(batch_of[col],), cost="query")) + return units + + def new_accumulator(self, state: StatState) -> "XorqAccumulator": + """The accumulator for a run of ``state``, and the start of the run's + counters: the snapshot-cache hit/miss counts and the perf recorder. + + Pre-populate every column accumulator with the externally-provided + keys. ``length`` is filled in by the batch query. ``min`` / ``max`` + start as None so dependents (histogram) don't cascade-exclude on + non-numeric columns; ``min`` / ``max`` overwrite for numeric cols. + ``distinct_count`` likewise starts as None so float columns (where the + stat is column_filtered out) keep their dependents (histogram, + histogram_bins, distinct_per) runnable. A skipped column keeps these + (its stats are supplied via init_sd and its data is never scanned).""" + self._cache_stats = _new_cache_run_stats() + self._perf = (perf_log.PerfRecorder() + if perf_log.enabled() and not self._suppress_perf_summary else None) + schema = state.data.schema() + pairs = columns_in_scope(state) + acc = XorqAccumulator(state, columns=[orig for orig, _rewritten in pairs], rewritten=dict(pairs), + skipped=[orig for orig, _rewritten in pairs if orig in state.skip_columns]) + for col in acc.order: + acc.working[col] = {"orig_col_name": Ok(col), "rewritten_col_name": Ok(col), "dtype": Ok(str(schema[col])), "length": Ok(0), "min": Ok(None), "max": Ok(None), "distinct_count": Ok(None)} - - # ---- Phase 1: batch aggregate ---------------------------------- + acc.emitted[col] = set() + acc.schema = schema + return acc + + def run(self, unit: StatUnit, acc: StatAccumulator) -> Fragment: + """Run one unit and return the stats it produced, ``{col: {stat: value}}`` + for its columns. A stat that fails is recorded on ``acc.errors`` once + and the run goes on.""" + assert isinstance(acc, XorqAccumulator) + if unit.phase == "batch": + return self._run_batch(unit, acc) + return self._run_histogram(unit, acc) + + def _run_funcs(self, funcs: List[StatFunc], col: Any, acc: "XorqAccumulator") -> None: + """Run ``funcs`` (in dependency order) for one column against its + accumulator, skipping stats whose results are already there (typically + the batch-phase stats).""" + col_accum = acc.working[col] + for sf in funcs: + if sf.provides and all(sk.name in col_accum for sk in sf.provides): + continue + if self._perf is not None: + t0 = time.perf_counter() + _execute_stat_func(sf, col_accum, col, raw_series=None, sampled_series=None, raw_dataframe=None, + xorq_expr=acc.state.data, xorq_execute=self._execute) + if self._perf is not None: + self._perf.record("xorq/per-column", col, sf.name, time.perf_counter() - t0) + + def _emit(self, acc: "XorqAccumulator", col: Any, key_to_func: Dict[str, StatFunc]) -> Dict[str, Any]: + """The stats of ``col`` that no earlier unit has returned, as plain + values (a failed stat as None), recorded on ``acc`` with their errors.""" + working = acc.working[col] + fresh = {key: result for key, result in working.items() if key not in acc.emitted[col]} + plain, errors = resolve_accumulator(fresh, col, key_to_func) + acc.emitted[col].update(fresh) + acc.record({col: plain}, errors) + return plain + + def _run_batch(self, unit: StatUnit, acc: "XorqAccumulator") -> Fragment: + table = acc.state.data + schema = acc.schema + columns = list(unit.columns) # ``length`` is a table-level scalar (same value for every column), # so it goes in once as ``__total_length__`` rather than as N # per-column expressions. @@ -378,8 +548,6 @@ def _process_table_impl(self, table, skip_columns=None) -> Tuple[SDType, List[St continue xorq_col_param = next(r.name for r in sf.requires if r.type is XorqColumn) for col in columns: - if col in skip: - continue col_dtype = schema[col] if sf.column_filter is not None and not sf.column_filter(col_dtype): continue @@ -387,7 +555,7 @@ def _process_table_impl(self, table, skip_columns=None) -> Tuple[SDType, List[St expr = sf.func(**{xorq_col_param: table[col]}) except Exception as e: for sk in sf.provides: - accumulators[col][sk.name] = Err(error=e, stat_func_name=sf.name, column_name=col, + acc.working[col][sk.name] = Err(error=e, stat_func_name=sf.name, column_name=col, inputs={"col": col}) continue if expr is None: @@ -397,7 +565,7 @@ def _process_table_impl(self, table, skip_columns=None) -> Tuple[SDType, List[St expr = expr.name(f"{col}|{stat_name}") except Exception as e: for sk in sf.provides: - accumulators[col][sk.name] = Err(error=e, stat_func_name=sf.name, column_name=col, + acc.working[col][sk.name] = Err(error=e, stat_func_name=sf.name, column_name=col, inputs={"col": col}) continue batch_items.append((col, sf, expr)) @@ -413,54 +581,40 @@ def _process_table_impl(self, table, skip_columns=None) -> Tuple[SDType, List[St # length stays at the prepopulated 0 so consumers still see something. for col, sf, _ in batch_items: for sk in sf.provides: - accumulators[col][sk.name] = Err(error=e, stat_func_name=sf.name, column_name=col, inputs={}) + acc.working[col][sk.name] = Err(error=e, stat_func_name=sf.name, column_name=col, inputs={}) else: total_length = _to_python_scalar(result_df[TOTAL_LENGTH_KEY].iloc[0]) if total_length is None: total_length = 0 - for col in columns: - accumulators[col]["length"] = Ok(total_length) + # A skipped column has no unit but has always been given the length. + for col in [*columns, *acc.skipped]: + acc.working[col]["length"] = Ok(total_length) for col, sf, _ in batch_items: stat_name = sf.provides[0].name col_stat = f"{col}|{stat_name}" if col_stat in result_df.columns: raw_val = result_df[col_stat].iloc[0] - accumulators[col][stat_name] = Ok(_to_python_scalar(raw_val)) + acc.working[col][stat_name] = Ok(_to_python_scalar(raw_val)) else: - accumulators[col][stat_name] = Err(error=KeyError( + acc.working[col][stat_name] = Err(error=KeyError( f"missing aggregate column {col_stat!r} in result"), stat_func_name=sf.name, column_name=col, inputs={}) - # ---- Phase 2: per-column post-batch ---------------------------- - all_errors: List[StatError] = [] - summary: SDType = {} - + # The stats that are pure functions of the ones just computed, so that + # histogram_bins (and with it color_map) needs no histogram query. + fragment: Fragment = {} for col in columns: - col_accum = accumulators[col] - col_dtype = schema[col] - col_funcs = build_column_dag(self.all_stat_funcs, col_dtype, external_keys=self.EXTERNAL_KEYS) - - for sf in col_funcs if col not in skip else []: - # Skip stats whose results are already in the accumulator - # (typically the batch-phase stats). - if sf.provides and all(sk.name in col_accum for sk in sf.provides): - continue - if self._perf is not None: - t0 = time.perf_counter() - _execute_stat_func(sf, col_accum, col, raw_series=None, sampled_series=None, raw_dataframe=None, - xorq_expr=table, xorq_execute=self._execute) - if self._perf is not None: - self._perf.record("xorq/per-column", col, sf.name, time.perf_counter() - t0) - - col_key_to_func: Dict[str, StatFunc] = {} - for sf in col_funcs: - for sk in sf.provides: - col_key_to_func[sk.name] = sf - - plain, errors = resolve_accumulator(col_accum, col, col_key_to_func) - summary[col] = plain - all_errors.extend(errors) - - return summary, all_errors + pure, _query, key_to_func = self._column_phases(acc.phases, schema[col]) + self._run_funcs(pure, col, acc) + fragment[col] = self._emit(acc, col, key_to_func) + return fragment + + def _run_histogram(self, unit: StatUnit, acc: "XorqAccumulator") -> Fragment: + fragment: Fragment = {} + for col in unit.columns: + _pure, query, key_to_func = self._column_phases(acc.phases, acc.schema[col]) + self._run_funcs(query, col, acc) + fragment[col] = self._emit(acc, col, key_to_func) + return fragment def add_stat(self, stat_func_or_class) -> Tuple[bool, List[StatError]]: """Add a stat function or ColAnalysis class interactively. @@ -501,7 +655,7 @@ def add_stat(self, stat_func_or_class) -> Tuple[bool, List[StatError]]: return True, [] -class XorqDfStatsV2: +class XorqDfStatsV2(UnitStats): """Stats wrapper for xorq table inputs. Mirrors the ``DfStatsV2`` / ``PlDfStatsV2`` surface (``.sdf``, ``.errs``, @@ -516,7 +670,9 @@ class XorqDfStatsV2: Stats execute through ``XorqStatPipeline`` — a single batched ``table.aggregate(...)`` query plus per-column histogram queries — pushing computation to the backend instead of materialising the - entire table. + entire table. ``run=False`` builds the wrapper without running anything, + for a caller that drives ``plan`` / ``run`` itself. ``chunk_cells`` and + ``rows`` are the column-chunk split's size and the row count it needs. """ @classmethod @@ -526,21 +682,26 @@ def verify_analysis_objects(cls, objs): XorqStatPipeline(objs, unit_test=False) def __init__(self, table, col_analysis_objs, operating_df_name=None, debug=False, - cache_storage=None, skip_columns=None): + cache_storage=None, skip_columns=None, chunk_cells=None, rows=None, run=True): self.table = table # Skip the unit_test PERVERSE_DF run on each widget construction — # it doubles the SQL query count (issue #709). The DAG-validation # cost is already paid by verify_analysis_objects on first set up # and by the test suite. Mirrors PlDfStatsV2. self.ap = XorqStatPipeline(col_analysis_objs, unit_test=False, - cache_storage=cache_storage) + cache_storage=cache_storage, chunk_cells=chunk_cells) self.operating_df_name = operating_df_name self.debug = debug - self.sdf, errors = self.ap.process_table(self.table, skip_columns=skip_columns) - self.errs = errors_to_errdict(errors) + self.rows = rows + self.state = StatState(table, frozenset(skip_columns or ()), rows=rows) + self.sdf: SDType = {} + self.errs = {} self.stat_errors = [] - if self.errs: - output_full_reproduce(self.errs, self.sdf, operating_df_name) + if run: + self.sdf, errors = self.ap.process_table(self.table, skip_columns=skip_columns, rows=rows) + self.errs = errors_to_errdict(errors) + if self.errs: + output_full_reproduce(self.errs, self.sdf, operating_df_name) def cache_run_stats(self) -> dict: """The summary-stat cache outcome of the last pipeline run (#943) — @@ -556,7 +717,7 @@ def add_analysis(self, a_obj): so DataFlow.add_analysis works against a xorq-backed stats wrapper. """ passed, errors = self.ap.add_stat(a_obj) - self.sdf, self.stat_errors = self.ap.process_table(self.table) + self.sdf, self.stat_errors = self.ap.process_table(self.table, rows=self.rows) self.errs = errors_to_errdict(self.stat_errors) if not passed: print("DAG validation failed") diff --git a/buckaroo/server/session.py b/buckaroo/server/session.py index dc09c377b..754ddc1ce 100644 --- a/buckaroo/server/session.py +++ b/buckaroo/server/session.py @@ -90,6 +90,11 @@ class SessionState: stats_gen: int = 0 stats_status: str = "complete" stats_reason: Optional[str] = None + # The ``StatRun``s (stat_run.py) of the current generation, keyed by + # ``(stats_gen, scope)``: the planned units, the fragments they have + # produced and the accumulator they read. Dropped whenever the generation + # changes, since a run describes one state of the dataflow. + stat_runs: dict = field(default_factory=dict) # Companion telemetry sink (#943): a fire-and-forget POST callable, built # once from the /load_expr payload's telemetry_url on the IOLoop (where # make_http_sink captures AsyncHTTPClient/IOLoop.current()). Stored here so @@ -130,10 +135,12 @@ def touch(self) -> None: def begin_stats_generation(session: "SessionState") -> None: """Start a new stats generation: the session's dataflow state has changed, so stats for the previous one no longer describe it. The status restarts - from the session's policy pair.""" + from the session's policy pair, and the stat runs of the old generation are + dropped.""" session.stats_gen += 1 session.stats_status, session.stats_reason = initial_stats_status( session.stats_tier, session.stats_delivery) + session.stat_runs.clear() def stats_meta(session: "SessionState") -> Optional[dict]: diff --git a/buckaroo/server/stat_run.py b/buckaroo/server/stat_run.py new file mode 100644 index 000000000..d775f2a0c --- /dev/null +++ b/buckaroo/server/stat_run.py @@ -0,0 +1,118 @@ +"""The stat run a session keeps for one stats generation (rows-first s4). + +A ``StatRun`` holds the planned units of one ``(stats_gen, scope)``, the +fragments they produced in the order they finished, and the accumulator the +units read. It runs a unit only when its caller asks: it has no thread, timer or +callback, so there is nothing to cancel, and the session drops it when the +generation changes (``begin_stats_generation``). + +The fragment list is append-only and shared. A ``StatCursor`` is one client's +position in it, held by that client's WebSocket handler, so work is done once +and every client reads all of it at its own pace. +""" +import itertools +import time +from typing import Any, List, Optional, Sequence, Tuple + +from buckaroo.pluggable_analysis_framework.stat_units import Fragment, StatUnit, rewrite_sd + +_run_ids = itertools.count(1) + + +class StatRun: + """``stats`` is a stats class (``DfStatsV2`` and its siblings) built with + ``run=False``: the run plans it once and runs its units through it. + + ``status`` is ``pending`` until the last unit has run, then ``complete``; + ``error`` once a unit raised, 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 + still finishes. + """ + + def __init__(self, stats_gen: int, scope: str, stats: Any): + self.id = next(_run_ids) + self.stats_gen = stats_gen + self.scope = scope + self.stats = stats + self.units: List[StatUnit] = stats.plan(stats.state) + self.acc = stats.new_accumulator(stats.state) + self.fragments: List[Fragment] = [] + self.ran: List[str] = [] + self.status = "pending" if self.units else "complete" + self.error: Optional[Exception] = None + self.created_at = time.time() + + @property + def key(self) -> Tuple[int, str]: + return (self.stats_gen, self.scope) + + @property + def remaining(self) -> int: + """The number of units not yet run.""" + return len(self.units) - len(self.ran) + + def next_unit(self, prefer: Sequence[Any] = ()) -> Optional[StatUnit]: + """The unit ``run_next`` would run: the first whose prerequisites have + run, or the first of those that covers a column named in ``prefer`` + (by original or rewritten name), so the columns a client is looking at + come first. ``None`` when no unit is left.""" + ran = set(self.ran) + ready = [unit for unit in self.units if unit.id not in ran and all(a in ran for a in unit.after)] + if prefer: + wanted = set(prefer) + for unit in ready: + if any(col in wanted or self.acc.rewritten.get(col) in wanted for col in unit.columns): + return unit + return ready[0] if ready else None + + def run_next(self, prefer: Sequence[Any] = ()) -> Optional[Fragment]: + """Run one unit and return its fragment, which is also appended to + ``fragments``. ``None`` when the run is complete or failed. A unit that + raises fails the run, and the exception propagates once.""" + if self.status != "pending": + return None + unit = self.next_unit(prefer) + if unit is None: + self.status = "complete" + return None + try: + fragment = self.stats.run(unit, self.acc) + except Exception as exc: + self.status, self.error = "error", exc + raise + self.ran.append(unit.id) + self.fragments.append(fragment) + if self.remaining == 0: + self.status = "complete" + return fragment + + def raw_sd(self) -> dict: + """The run's stats as the summary dict the dataflow carries, keyed by + rewritten column name: what ``summary_sd`` holds after a whole run.""" + return rewrite_sd(self.acc.sd(), self.acc.state.data) + + +class StatCursor: + """One client's position in a ``StatRun``'s fragment list. It follows one + run at a time and starts again at the beginning of any other, since a + position in one run means nothing in another.""" + + def __init__(self) -> None: + self.run_id: Optional[int] = None + self.position = 0 + + def _follow(self, run: StatRun) -> None: + if self.run_id != run.id: + self.run_id, self.position = run.id, 0 + + def take(self, run: StatRun) -> List[Fragment]: + """The fragments of ``run`` this cursor has not returned yet, and move + past them.""" + self._follow(run) + unseen = run.fragments[self.position:] + self.position += len(unseen) + return unseen + + def caught_up(self, run: StatRun) -> bool: + """Whether the cursor has returned every fragment ``run`` has now.""" + return (self.position if self.run_id == run.id else 0) >= len(run.fragments) diff --git a/buckaroo/server/stats_wire.py b/buckaroo/server/stats_wire.py index 7aeecf2d9..b019af1fd 100644 --- a/buckaroo/server/stats_wire.py +++ b/buckaroo/server/stats_wire.py @@ -31,6 +31,7 @@ from buckaroo.pluggable_analysis_framework import perf_log from buckaroo.server.data_loading import get_buckaroo_display_state from buckaroo.server.session import SessionState, build_state_message +from buckaroo.server.stat_run import StatRun log = logging.getLogger("buckaroo.server.stats_wire") @@ -83,6 +84,24 @@ def refresh_session_snapshot(session: SessionState, dataflow: Any) -> None: dvc["component_config"] = {**dvc.get("component_config", {}), **session.component_config} +def start_stat_run(session: SessionState, scope: str = "raw") -> Optional[StatRun]: + """The ``StatRun`` of the session's current generation for ``scope``: the one + it holds, or a new one planned from the dataflow and stored on the session. + Planning sends no query. The run analyzes the frame ``assign_full_stats`` + does, so running every unit gives the stats that call computes whole. + ``None`` when there is no dataflow to plan from or the scope is not one a + request can name.""" + dataflow = session_dataflow(session) + if dataflow is None or dataflow.processed_df is None or scope not in STATS_SCOPES: + return None + key = (session.stats_gen, scope) + run = session.stat_runs.get(key) + if run is None: + run = StatRun(session.stats_gen, scope, dataflow.build_stats(dataflow.processed_df, run=False)) + session.stat_runs[key] = run + return run + + def assign_full_stats(dataflow: Any) -> None: """Take a schema-tier dataflow to the full tier: compute the current state's full stats (or find them in ``summary_stats_cache``, where an diff --git a/buckaroo/server/websocket_handler.py b/buckaroo/server/websocket_handler.py index 1b4ed6125..7a5289917 100644 --- a/buckaroo/server/websocket_handler.py +++ b/buckaroo/server/websocket_handler.py @@ -11,6 +11,7 @@ from buckaroo.pluggable_analysis_framework import perf_log from buckaroo.server.data_loading import (handle_infinite_request, handle_infinite_request_buckaroo, handle_infinite_request_lazy) from buckaroo.server.session import begin_stats_generation, dataflow_stats_tier +from buckaroo.server.stat_run import StatCursor from buckaroo.server.stats_wire import (broadcast_state, build_state_message_for, handle_stats_request, parse_caps, refresh_session_snapshot) @@ -44,6 +45,9 @@ def open(self, session_id): # client can say anything, and the other send sites push one shared # snapshot (stats_wire.build_state_message_for reads it per client). self.caps = parse_caps(self.get_query_argument("caps", "")) + # This client's position in the session's StatRun fragment list (the + # run itself is shared); it ends with the connection. + self.stats_cursor = StatCursor() sessions = self.application.settings["sessions"] sessions.add_ws_client(session_id, self) diff --git a/buckaroo/server/xorq_loading.py b/buckaroo/server/xorq_loading.py index 60a9226cf..f60c9cc0e 100644 --- a/buckaroo/server/xorq_loading.py +++ b/buckaroo/server/xorq_loading.py @@ -64,6 +64,11 @@ class XorqServerDataflow(XorqDataflow): and project-authored post-processing classes discovered under ``/post_processing/*.py``; the built-in xorq stats are kept first so collisions resolve to the built-in. + + ``stat_chunk_cells`` is the host's request (off when ``None``) to run the + stats batch aggregate in chunks of about that many cells instead of one + query; the stats pipeline still refuses it for any source that is not a + plain parquet scan (``XorqStatPipeline.chunk_refusal``). """ sampling_klass = XorqInfiniteSampling @@ -72,13 +77,14 @@ class XorqServerDataflow(XorqDataflow): DFStatsClass = XorqDfStatsV2 analysis_klasses = _XORQ_ANALYSIS_KLASSES - def __init__(self, expr, *args, extra_klasses=None, cache_storage_path=None, **kwargs): + def __init__(self, expr, *args, extra_klasses=None, cache_storage_path=None, stat_chunk_cells=None, **kwargs): if extra_klasses: # Per-instance override — class-level _XORQ_ANALYSIS_KLASSES is # left untouched so other sessions / direct widget usage don't # inherit one project's stats. self.analysis_klasses = list(_XORQ_ANALYSIS_KLASSES) + list(extra_klasses) self.cache_storage = _make_cache_storage(cache_storage_path) + self.stat_chunk_cells = stat_chunk_cells super().__init__(expr, *args, **kwargs) diff --git a/buckaroo/xorq_buckaroo.py b/buckaroo/xorq_buckaroo.py index e1aad2315..569cb8a06 100644 --- a/buckaroo/xorq_buckaroo.py +++ b/buckaroo/xorq_buckaroo.py @@ -14,7 +14,7 @@ import traceback import weakref from io import BytesIO -from typing import TYPE_CHECKING, Any, cast +from typing import TYPE_CHECKING, Any, Optional, cast import pandas as pd @@ -29,11 +29,12 @@ from .customizations.xorq_autoclean_conf import NoCleaningConfXorq from .customizations.xorq_stats_v2 import XORQ_STATS_V2, schema_stats from .dataflow.autocleaning import PandasAutocleaning -from .dataflow.dataflow import CustomizableDataflow +from .dataflow.dataflow import CustomizableDataflow, DfStats from .dataflow.dataflow_extras import Sampling from .df_util import old_col_new_col from .pluggable_analysis_framework import perf_log from .pluggable_analysis_framework.col_analysis import ColAnalysis +from .pluggable_analysis_framework.stat_units import rewrite_sd from .pluggable_analysis_framework.xorq_stat_pipeline import XorqDfStatsV2 from .serialization_utils import pd_to_obj, to_parquet, send_infinite_resp @@ -145,8 +146,14 @@ class XorqDataflow(CustomizableDataflow["XorqExpr | pd.DataFrame"]): 2. ``_get_summary_sd`` re-keys the summary dict from original column names (what ``XorqStatPipeline`` produces) to the rewritten ``a, b, c`` names that ``pd_to_obj`` and the styling layer expect. + + ``stat_chunk_cells`` (off when ``None``) is a host's request to cut the + stats batch aggregate into chunks of that many cells; the stats pipeline + still refuses it for any source that is not a plain parquet scan. """ + stat_chunk_cells: Optional[int] = None + def populate_df_meta(self) -> None: if self.processed_df is None: self.df_meta = { @@ -175,15 +182,8 @@ def _get_summary_sd(self, processed_df: "XorqExpr | pd.DataFrame"): return empty, {} if self.stats_tier == "schema": return self._get_schema_sd(processed_df), {} - cache_storage = getattr(self, 'cache_storage', None) - # The owning widget injects XorqDfStatsV2 as DFStatsClass (via its - # InnerDataFlow subclass); the cast exposes cache_run_stats() below. - stats_klass = cast("type[XorqDfStatsV2]", self.DFStatsClass) with perf_log.perf_span("firstpull.summary_stats") as span: - stats = stats_klass( - processed_df, self.analysis_klasses, self.df_name, - debug=self.debug, cache_storage=cache_storage, - skip_columns=getattr(self, 'skip_stat_columns', None)) + stats = cast("XorqDfStatsV2", self.build_stats(processed_df)) # Attach the summary-stat cache signal (#944) to the span so a # telemetry consumer learns whether the stats were cached — the one # signal only the server observes (#943). Carries the write side too @@ -207,13 +207,23 @@ def _get_summary_sd(self, processed_df: "XorqExpr | pd.DataFrame"): errs = stats.errs else: errs = {} - rewritten = {} - for orig_col, rewritten_col in old_col_new_col(processed_df): - col_meta = dict(sdf.get(orig_col, {})) - col_meta['orig_col_name'] = orig_col - col_meta['rewritten_col_name'] = rewritten_col - rewritten[rewritten_col] = col_meta - return rewritten, errs + return rewrite_sd(sdf, processed_df), errs + + def build_stats(self, processed_df: "XorqExpr | pd.DataFrame", run: bool = True) -> DfStats: + """The stats executor for an expression: as the base class builds it, + with the snapshot cache and the column-chunk split. The split needs the + row count, which is the cached count ``populate_df_meta`` already took, + and is looked up only when the split is on.""" + # The owning widget injects XorqDfStatsV2 as DFStatsClass (via its + # InnerDataFlow subclass), whose constructor takes the extra keywords + # below, so it is typed as itself here and handed back as the protocol. + stats_klass = cast("type[XorqDfStatsV2]", self.DFStatsClass) + chunk_cells = self.stat_chunk_cells + return cast("DfStats", stats_klass( + processed_df, self.analysis_klasses, self.df_name, + debug=self.debug, cache_storage=getattr(self, 'cache_storage', None), + skip_columns=getattr(self, 'skip_stat_columns', None), chunk_cells=chunk_cells, + rows=_expr_count(processed_df) if chunk_cells else None, run=run)) def _get_schema_sd(self, processed_df: "XorqExpr | pd.DataFrame") -> dict: """Identity, dtype, typing flags and row count for every column of an diff --git a/tests/unit/dataflow/scoped_summary_stats_test.py b/tests/unit/dataflow/scoped_summary_stats_test.py index 40d6f0f07..bf908c6a4 100644 --- a/tests/unit/dataflow/scoped_summary_stats_test.py +++ b/tests/unit/dataflow/scoped_summary_stats_test.py @@ -28,6 +28,7 @@ from buckaroo.dataflow.sd_cache import split_chain_by_scope from buckaroo.pluggable_analysis_framework.col_analysis import ColAnalysis from buckaroo.pluggable_analysis_framework.stat_func import stat +from buckaroo.pluggable_analysis_framework.stat_units import merge_fragments, rewrite_sd _AddOrigResult = TypedDict('_AddOrigResult', {'cleaning_ops': Any, 'add_orig': Any}) @@ -341,7 +342,6 @@ def _run_units(stats): def _scope_sds_by_units(dataflow): """Each scope's summary dict, built from the fragments of its planned units where the dataflow built it by running the stats class whole.""" - from buckaroo.pluggable_analysis_framework.stat_units import merge_fragments, rewrite_sd sds = {} for scope in ("raw", "clean", "filt"): scope_df = dataflow.processed_df if scope == "filt" else dataflow._compute_scope_df(scope) diff --git a/tests/unit/server/test_data_loading_polars.py b/tests/unit/server/test_data_loading_polars.py index cd3500a21..0e53572b0 100644 --- a/tests/unit/server/test_data_loading_polars.py +++ b/tests/unit/server/test_data_loading_polars.py @@ -36,9 +36,11 @@ from buckaroo.jlisp.lisp_utils import s as lisp_sym from buckaroo.pluggable_analysis_framework.col_analysis import ColAnalysis from buckaroo.pluggable_analysis_framework.stat_pipeline import StatPipeline +from buckaroo.pluggable_analysis_framework.stat_units import StatAccumulator, StatState, StatUnit, merge_fragments from buckaroo.pluggable_analysis_framework.utils import PERVERSE_DF from buckaroo.server.data_loading import ServerDataflow, handle_infinite_request_buckaroo from buckaroo.server.data_loading_polars import (PolarsServerDataflow, create_polars_dataflow, handle_infinite_request_buckaroo_polars, load_file_polars) +from buckaroo.server.stat_run import StatCursor, StatRun from buckaroo.styling_helpers import float_, obj_ from tests.unit.dataflow.scoped_summary_stats_test import (_OverridingPostProcessing, _run_units, _scope_inputs, _scope_sds_by_units) @@ -448,24 +450,31 @@ def test_build_stats_without_running_computes_nothing(self, backend, monkeypatch assert stats.sdf == {} and stats.errs == {} def test_the_default_run_goes_through_the_units(self, backend, monkeypatch): - """Every stat the constructor computes comes from a unit, so the - inline path and a caller running units one at a time cannot differ.""" - unit_columns, column_calls = [], [] + """Every column the constructor computes is computed inside a unit, so + the inline path and a caller running units one at a time cannot + differ. (The DAG self-check runs ``PERVERSE_DF`` through the same + units, which has no part of the loaded data in it.)""" + depth, inside, outside, unit_columns = [0], [], [], [] original_run, original_column = StatPipeline.run, StatPipeline.process_column def spy_run(self, unit, acc): - unit_columns.append(unit.columns) - return original_run(self, unit, acc) + depth[0] += 1 + if list(acc.state.data.columns) != list(PERVERSE_DF.columns): + unit_columns.append(unit.columns) + try: + return original_run(self, unit, acc) + finally: + depth[0] -= 1 def spy_column(self, *args, **kwargs): - column_calls.append(args or kwargs) + (inside if depth[0] else outside).append(kwargs.get("column_name")) return original_column(self, *args, **kwargs) monkeypatch.setattr(StatPipeline, "run", spy_run) monkeypatch.setattr(StatPipeline, "process_column", spy_column) _three_scope_dataflow(backend) assert {c for cols in unit_columns for c in cols} == {"price", "qty", "category"} - assert len(unit_columns) == len(column_calls), "a column's stats ran outside a unit" + assert inside and outside == [], "a column's stats ran outside a unit" def test_units_assembled_equal_merged_sd_with_init_sd_cleaning_and_overrides(self, backend): """The fragments of each scope's units, assembled by @@ -513,7 +522,6 @@ def test_skip_stat_columns_get_no_unit(self, backend): def _stat_run(backend, gen=4): - from buckaroo.server.stat_run import StatRun dataflow = _build_dataflow(backend, _frame(backend, _UNTIED_DATA)) return StatRun(gen, "raw", dataflow.build_stats(dataflow.processed_df, run=False)), dataflow @@ -557,7 +565,6 @@ def test_the_run_is_complete_when_no_unit_is_left(self, backend): assert len(run.fragments) == 3, "asking again runs nothing" def test_the_accumulator_holds_what_the_fragments_hold(self, backend): - from buckaroo.pluggable_analysis_framework.stat_units import merge_fragments run, _dataflow = _stat_run(backend) while run.run_next() is not None: pass @@ -570,8 +577,6 @@ def test_raw_sd_is_the_full_stats_summary_sd(self, backend): assert _as_json(run.raw_sd()) == _as_json(dataflow.summary_sd) def test_cursors_read_one_list_each_at_its_own_pace(self, backend): - from buckaroo.pluggable_analysis_framework.stat_units import merge_fragments - from buckaroo.server.stat_run import StatCursor run, _dataflow = _stat_run(backend) a, b = StatCursor(), StatCursor() run.run_next() @@ -587,7 +592,6 @@ def test_cursors_read_one_list_each_at_its_own_pace(self, backend): assert len(set(run.ran)) == 3, "no unit ran twice" def test_a_cursor_starts_again_on_another_run(self, backend): - from buckaroo.server.stat_run import StatCursor first, _dataflow = _stat_run(backend, gen=1) second, _dataflow = _stat_run(backend, gen=2) first.run_next() @@ -611,8 +615,6 @@ def test_a_run_has_no_thread_timer_or_callback(self, backend, monkeypatch): if isinstance(value, (threading.Thread, threading.Timer)) or inspect.isroutine(value)] def test_a_unit_that_raises_fails_the_run(self): - from buckaroo.pluggable_analysis_framework.stat_units import StatAccumulator, StatState, StatUnit - from buckaroo.server.stat_run import StatRun class _FailingStats: state = StatState(data=None) diff --git a/tests/unit/server/test_load_expr.py b/tests/unit/server/test_load_expr.py index e222b56b2..b2c1b2252 100644 --- a/tests/unit/server/test_load_expr.py +++ b/tests/unit/server/test_load_expr.py @@ -1,5 +1,6 @@ """End-to-end tests for POST /load_expr — server load path for XorqBuckarooInfiniteWidget over a xorq/ibis expression.""" +import dataclasses import datetime import io import json @@ -29,13 +30,16 @@ from xorq.common.utils.provenance_utils import read_parquet_provenance # noqa: E402 from xorq.expr.relations import CachedNode # noqa: E402 +from buckaroo.dataflow.dataflow import assemble_merged_sd # noqa: E402 from buckaroo.dataflow.sd_cache import split_chain_by_scope # noqa: E402 from buckaroo.jlisp.lisp_utils import s as lisp_sym # noqa: E402 from buckaroo.pluggable_analysis_framework.col_analysis import ColAnalysis # noqa: E402 +from buckaroo.pluggable_analysis_framework.stat_units import merge_fragments # noqa: E402 from buckaroo.pluggable_analysis_framework.xorq_stat_pipeline import XorqStatPipeline # noqa: E402 from buckaroo.serialization_utils import resolve_summary_stats_payload # noqa: E402 from buckaroo.server import telemetry, xorq_loading # noqa: E402 from buckaroo.server.app import make_app as _make_app # noqa: E402 +from buckaroo.server.stats_wire import start_stat_run # noqa: E402 from buckaroo.server.websocket_handler import DataStreamHandler # noqa: E402 from tests.unit.dataflow.scoped_summary_stats_test import ( # noqa: E402 _OverridingPostProcessing, _scope_inputs, _scope_sds_by_units) @@ -1917,7 +1921,6 @@ async def test_inline_sessions_send_no_df_meta_stats(self): @tornado.testing.gen_test async def test_a_stat_run_is_started_on_the_session_for_the_current_generation(self): - from buckaroo.server.stats_wire import start_stat_run await self._load("sw-run", stats_delivery="deferred") session = self._session("sw-run") self.assertEqual(session.stat_runs, {}) @@ -1934,7 +1937,6 @@ async def test_a_stat_run_is_started_on_the_session_for_the_current_generation(s @tornado.testing.gen_test async def test_a_run_ends_with_the_summary_sd_an_inline_session_computes(self): - from buckaroo.server.stats_wire import start_stat_run await self._load("sw-run-sd", stats_delivery="deferred") session = self._session("sw-run-sd") run = start_stat_run(session) @@ -1946,7 +1948,6 @@ async def test_a_run_ends_with_the_summary_sd_an_inline_session_computes(self): @tornado.testing.gen_test async def test_a_dataflow_field_change_drops_the_run_and_a_search_term_change_keeps_it(self): - from buckaroo.server.stats_wire import start_stat_run await self._load("sw-drop", stats_delivery="deferred") ws, _first = await self._connect("sw-drop", caps="stats_update") session = self._session("sw-drop") @@ -1966,7 +1967,6 @@ async def test_a_dataflow_field_change_drops_the_run_and_a_search_term_change_ke @tornado.testing.gen_test async def test_reload_expr_drops_the_run(self): - from buckaroo.server.stats_wire import start_stat_run await self._load("sw-drop-reload", stats_delivery="deferred") session = self._session("sw-drop-reload") start_stat_run(session) @@ -1976,8 +1976,6 @@ async def test_reload_expr_drops_the_run(self): @tornado.testing.gen_test async def test_each_connection_keeps_its_own_cursor_into_the_run(self): - from buckaroo.pluggable_analysis_framework.stat_units import merge_fragments - from buckaroo.server.stats_wire import start_stat_run await self._load("sw-cursor", stats_delivery="deferred") await self._connect("sw-cursor", caps="stats_update") await self._connect("sw-cursor", caps="stats_update") @@ -2052,7 +2050,6 @@ def spy_execute(self, query): assert units[0] == "batch" and "histogram:price" in units def test_units_assembled_equal_merged_sd_with_init_sd_cleaning_and_overrides(self): - from buckaroo.dataflow.dataflow import assemble_merged_sd dataflow = _three_scope_dataflow(init_sd={ "price": {"displayer_args": {"displayer": "string", "max_length": 99}, "init_only": 1}}, column_config_overrides={"category": {"displayer_args": {"displayer": "string", "max_length": 5000}}}) @@ -2071,7 +2068,6 @@ def test_units_assembled_equal_merged_sd_with_init_sd_cleaning_and_overrides(sel assert _as_json(assembled) == _as_json(dataflow.merged_sd) def test_a_column_group_request_returns_only_those_columns(self): - import dataclasses dataflow = _build_dataflow() stats = dataflow.build_stats(dataflow.processed_df, run=False) for names in (("qty", "category"), ("b", "c")): # original or rewritten names diff --git a/tests/unit/skip_columns_test.py b/tests/unit/skip_columns_test.py index 77ad8895a..1ffef9850 100644 --- a/tests/unit/skip_columns_test.py +++ b/tests/unit/skip_columns_test.py @@ -9,6 +9,7 @@ import pytest from buckaroo.pluggable_analysis_framework.stat_pipeline import StatPipeline +from buckaroo.pluggable_analysis_framework.stat_units import StatState from buckaroo.customizations.pd_stats_v2 import PD_ANALYSIS_V2 @@ -73,7 +74,6 @@ def _run_units(pipe, state): def test_pandas_skipped_column_gets_no_unit(): - from buckaroo.pluggable_analysis_framework.stat_units import StatState df = pd.DataFrame({"a": [1, 2, 3], "b": [4, 5, 6], "c": [7, 8, 9]}) pipe = StatPipeline(PD_ANALYSIS_V2, unit_test=False) state = StatState(df, skip_columns=frozenset({"b"})) @@ -88,7 +88,6 @@ def test_pandas_skipped_column_gets_no_unit(): def test_polars_skipped_column_gets_no_unit(): pl = pytest.importorskip("polars") from buckaroo.customizations.pl_stats_v2 import PL_ANALYSIS_V2 - from buckaroo.pluggable_analysis_framework.stat_units import StatState df = pl.DataFrame({"a": [1, 2, 3], "b": [4, 5, 6], "c": [7, 8, 9]}) pipe = StatPipeline(PL_ANALYSIS_V2, unit_test=False) state = StatState(df, skip_columns=frozenset({"b"})) @@ -100,7 +99,6 @@ def test_polars_skipped_column_gets_no_unit(): def test_xorq_skipped_column_gets_no_unit(tmp_path): xo = pytest.importorskip("xorq.api") from buckaroo.pluggable_analysis_framework.xorq_stat_pipeline import XorqStatPipeline - from buckaroo.pluggable_analysis_framework.stat_units import StatState from buckaroo.customizations.xorq_stats_v2 import XORQ_STATS_V2 p = tmp_path / "t.parquet" diff --git a/tests/unit/test_paf_v2.py b/tests/unit/test_paf_v2.py index 8909bd7c3..d63b3632a 100644 --- a/tests/unit/test_paf_v2.py +++ b/tests/unit/test_paf_v2.py @@ -14,6 +14,7 @@ from buckaroo.pluggable_analysis_framework.typed_dag import (build_typed_dag, build_column_dag, DAGConfigError) from buckaroo.pluggable_analysis_framework.column_filters import (is_numeric, is_string, is_temporal, is_boolean, any_of, not_) from buckaroo.pluggable_analysis_framework.stat_pipeline import (StatPipeline, _normalize_inputs, errors_to_errdict) +from buckaroo.pluggable_analysis_framework.stat_units import StatState, merge_fragments, rewrite_sd from buckaroo.pluggable_analysis_framework.col_analysis import ColAnalysis from buckaroo.pluggable_analysis_framework.utils import PERVERSE_DF @@ -831,7 +832,6 @@ def _run_units(pipeline, state): return acc, fragments def test_plan_is_one_unit_per_column_in_column_order(self): - from buckaroo.pluggable_analysis_framework.stat_units import StatState units = self._pipeline().plan(StatState(self._df())) assert [u.columns for u in units] == [('amount',), ('word',), ('score',)] assert len({u.id for u in units}) == 3 @@ -839,7 +839,6 @@ def test_plan_is_one_unit_per_column_in_column_order(self): assert all(u.after == () for u in units), "no column's stats read another column's" def test_plan_and_a_new_accumulator_compute_nothing(self, monkeypatch): - from buckaroo.pluggable_analysis_framework.stat_units import StatState calls = [] monkeypatch.setattr(StatPipeline, 'process_column', lambda self, *a, **k: calls.append(a)) pipeline, state = self._pipeline(), StatState(self._df()) @@ -848,7 +847,6 @@ def test_plan_and_a_new_accumulator_compute_nothing(self, monkeypatch): assert calls == [] def test_run_returns_the_fragment_of_its_unit_keyed_by_original_name(self): - from buckaroo.pluggable_analysis_framework.stat_units import StatState pipeline, state = self._pipeline(), StatState(self._df()) acc = pipeline.new_accumulator(state) fragment = pipeline.run(pipeline.plan(state)[1], acc) @@ -859,7 +857,6 @@ def test_run_returns_the_fragment_of_its_unit_keyed_by_original_name(self): assert list(acc.sd()) == ['word'], "a column with no unit run yet has no entry" def test_the_union_of_fragments_is_what_process_df_returns(self): - from buckaroo.pluggable_analysis_framework.stat_units import StatState, merge_fragments, rewrite_sd df, pipeline = self._df(), self._pipeline() acc, fragments = self._run_units(pipeline, StatState(df)) expected, errors = pipeline.process_df(df) @@ -868,7 +865,6 @@ def test_the_union_of_fragments_is_what_process_df_returns(self): assert [(e.column, e.stat_key) for e in acc.errors] == [(e.column, e.stat_key) for e in errors] def test_an_error_in_a_unit_is_recorded_on_the_accumulator(self): - from buckaroo.pluggable_analysis_framework.stat_units import StatState @stat() def fails_on_words(ser: RawSeries) -> int: @@ -885,7 +881,6 @@ def fails_on_words(ser: RawSeries) -> int: def test_a_column_group_returns_only_those_columns(self, names): """A group is named by original or rewritten name, since the client only knows the rewritten ones.""" - from buckaroo.pluggable_analysis_framework.stat_units import StatState pipeline = self._pipeline() state = StatState(self._df(), columns=names) assert [u.columns for u in pipeline.plan(state)] == [('word',)] @@ -895,19 +890,16 @@ def test_a_column_group_returns_only_those_columns(self, names): @pytest.mark.parametrize('names', [('score',), ('c',)]) def test_priority_columns_are_planned_first(self, names): - from buckaroo.pluggable_analysis_framework.stat_units import StatState units = self._pipeline().plan(StatState(self._df(), priority=names)) assert [u.columns[0] for u in units] == ['score', 'amount', 'word'] def test_an_empty_frame_has_no_units(self): - from buckaroo.pluggable_analysis_framework.stat_units import StatState pipeline, state = self._pipeline(), StatState(pd.DataFrame({'x': pd.Series([], dtype='int64')})) assert pipeline.plan(state) == [] assert pipeline.new_accumulator(state).sd() == {} @pytest.mark.parametrize('backend', ['pandas', 'polars']) def test_real_stats_run_unit_by_unit_equal_process_df(self, backend): - from buckaroo.pluggable_analysis_framework.stat_units import StatState, merge_fragments, rewrite_sd df = _mixed_frames()[backend] pipeline = StatPipeline(_stat_lists()[backend], unit_test=False) acc, fragments = self._run_units(pipeline, StatState(df)) diff --git a/tests/unit/test_xorq_stats_v2.py b/tests/unit/test_xorq_stats_v2.py index 221d0fb4f..2a32bbad5 100644 --- a/tests/unit/test_xorq_stats_v2.py +++ b/tests/unit/test_xorq_stats_v2.py @@ -19,6 +19,7 @@ XorqStatPipeline, XorqColumn) from buckaroo.pluggable_analysis_framework.stat_func import stat # noqa: E402 +from buckaroo.pluggable_analysis_framework.stat_units import StatState, merge_fragments # noqa: E402 from buckaroo.customizations.xorq_stats_v2 import ( # noqa: E402 XORQ_STATS_V2) @@ -944,7 +945,9 @@ def test_each_histogram_query_keeps_its_snapshot_key(self, tmp_path): def _stable(sd): """An sd as comparable JSON. Equal counts come back from a GROUP BY in no - fixed order, so the histogram lists are sorted.""" + fixed order, so the histogram lists are sorted, and the partial sums of a + parallel aggregate add in an order that changes the last digit of a float, + so floats are compared to nine significant digits.""" import json def norm(value, key=None): @@ -953,6 +956,8 @@ def norm(value, key=None): if isinstance(value, list): items = [norm(v) for v in value] return sorted(items, key=lambda item: json.dumps(item, sort_keys=True)) if key == "histogram" else items + if isinstance(value, float): + return float(f"{value:.9g}") return value return json.dumps(norm(sd), sort_keys=True, default=str) @@ -991,7 +996,6 @@ class TestStatUnits: _ALL = ("ints", "floats", "strs", "bools") def test_the_scalar_batch_is_first_then_one_histogram_unit_per_column(self): - from buckaroo.pluggable_analysis_framework.stat_units import StatState units = XorqStatPipeline(XORQ_STATS_V2, unit_test=False).plan(StatState(_make_table())) assert [u.id for u in units] == ["batch", *(f"histogram:{c}" for c in self._ALL)] assert units[0].columns == self._ALL @@ -1004,13 +1008,11 @@ def test_the_scalar_batch_is_first_then_one_histogram_unit_per_column(self): def test_visible_columns_get_their_histogram_units_first(self, visible): """Visible columns are named by original or rewritten name. The batch stays first, since every histogram reads what it computes.""" - from buckaroo.pluggable_analysis_framework.stat_units import StatState units = XorqStatPipeline(XORQ_STATS_V2, unit_test=False).plan(StatState(_make_table(), priority=visible)) assert [u.id for u in units] == [ "batch", "histogram:strs", "histogram:bools", "histogram:ints", "histogram:floats"] def test_plan_and_a_new_accumulator_issue_no_query(self, monkeypatch): - from buckaroo.pluggable_analysis_framework.stat_units import StatState spy = _QuerySpy(monkeypatch) pipeline, state = XorqStatPipeline(XORQ_STATS_V2, unit_test=False), StatState(_make_table()) pipeline.plan(state) @@ -1018,7 +1020,6 @@ def test_plan_and_a_new_accumulator_issue_no_query(self, monkeypatch): assert spy.queries == [] def test_the_batch_is_one_query_and_each_histogram_unit_one_more(self, monkeypatch): - from buckaroo.pluggable_analysis_framework.stat_units import StatState spy = _QuerySpy(monkeypatch) pipeline = XorqStatPipeline(XORQ_STATS_V2, unit_test=False) state = StatState(_make_table()) @@ -1032,7 +1033,6 @@ def test_the_batch_is_one_query_and_each_histogram_unit_one_more(self, monkeypat def test_the_batch_fragment_has_the_scalar_stats_and_the_bins_but_no_histogram(self): """``histogram_bins`` is a pure function of five scalars, so it comes with the batch and ``color_map`` needs no histogram query.""" - from buckaroo.pluggable_analysis_framework.stat_units import StatState _acc, run = _run_units(XorqStatPipeline(XORQ_STATS_V2, unit_test=False), StatState(_make_table())) batch_unit, batch = run[0] assert set(batch) == set(self._ALL) @@ -1043,14 +1043,12 @@ def test_the_batch_fragment_has_the_scalar_stats_and_the_bins_but_no_histogram(s assert "mean" not in batch["strs"] def test_a_histogram_fragment_holds_the_histogram_of_its_column_only(self): - from buckaroo.pluggable_analysis_framework.stat_units import StatState _acc, run = _run_units(XorqStatPipeline(XORQ_STATS_V2, unit_test=False), StatState(_make_table())) for unit, fragment in run[1:]: assert list(fragment) == list(unit.columns) assert set(fragment[unit.columns[0]]) == {"histogram"} def test_the_union_of_fragments_is_process_table(self): - from buckaroo.pluggable_analysis_framework.stat_units import StatState, merge_fragments table = _make_table_with_nulls() pipeline = XorqStatPipeline(XORQ_STATS_V2, unit_test=False) acc, run = _run_units(pipeline, StatState(table)) @@ -1060,7 +1058,6 @@ def test_the_union_of_fragments_is_process_table(self): assert errors == [] and acc.errors == [] def test_a_failing_stat_is_recorded_once_and_the_run_goes_on(self): - from buckaroo.pluggable_analysis_framework.stat_units import StatState @stat() def will_fail(col: XorqColumn) -> int: @@ -1075,7 +1072,6 @@ def will_fail(col: XorqColumn) -> int: @pytest.mark.parametrize("names", [("ints", "strs"), ("a", "c")]) def test_a_column_group_returns_only_those_columns(self, names, monkeypatch): - from buckaroo.pluggable_analysis_framework.stat_units import StatState spy = _QuerySpy(monkeypatch) pipeline = XorqStatPipeline(XORQ_STATS_V2, unit_test=False) state = StatState(_make_table(), columns=names) @@ -1126,13 +1122,11 @@ class TestColumnChunkSplit: chunk of a join, an aggregate or a CSV read runs the whole plan again.""" def test_the_batch_is_not_split_by_default(self, tmp_path): - from buckaroo.pluggable_analysis_framework.stat_units import StatState units = XorqStatPipeline(XORQ_STATS_V2, unit_test=False).plan( StatState(_six_column_scan(tmp_path), rows=_SIX_ROWS)) assert [u.id for u in _batch_units(units)] == ["batch"] def test_chunks_are_sized_by_cells(self, tmp_path): - from buckaroo.pluggable_analysis_framework.stat_units import StatState table = _six_column_scan(tmp_path) # 36 cells at 12 rows is three columns per chunk, 24 is two, and a # chunk always holds at least one column. @@ -1144,12 +1138,10 @@ def test_chunks_are_sized_by_cells(self, tmp_path): "batch:0", "batch:1"] def test_a_chunk_that_fits_every_column_is_the_single_batch(self, tmp_path): - from buckaroo.pluggable_analysis_framework.stat_units import StatState units = _chunked(10_000).plan(StatState(_six_column_scan(tmp_path), rows=_SIX_ROWS)) assert [u.id for u in _batch_units(units)] == ["batch"] def test_batches_come_first_and_each_histogram_follows_its_chunk(self, tmp_path): - from buckaroo.pluggable_analysis_framework.stat_units import StatState units = _chunked(36).plan(StatState(_six_column_scan(tmp_path), rows=_SIX_ROWS)) assert [u.phase for u in units] == ["batch"] * 2 + ["histogram"] * 6 chunk_of = {c: u.id for u in units[:2] for c in u.columns} @@ -1157,19 +1149,16 @@ def test_batches_come_first_and_each_histogram_follows_its_chunk(self, tmp_path) assert unit.after == (chunk_of[unit.columns[0]],) def test_visible_columns_fill_the_first_chunk(self, tmp_path): - from buckaroo.pluggable_analysis_framework.stat_units import StatState units = _chunked(36).plan(StatState(_six_column_scan(tmp_path), rows=_SIX_ROWS, priority=("flag", "n2"))) assert _batch_units(units)[0].columns == ("flag", "n2", "n0") assert [u.columns[0] for u in units if u.phase == "histogram"][:2] == ["flag", "n2"] def test_an_unknown_row_count_keeps_the_single_batch(self, tmp_path): """Cells need the row count, and asking for it is a query.""" - from buckaroo.pluggable_analysis_framework.stat_units import StatState units = _chunked(36).plan(StatState(_six_column_scan(tmp_path))) assert [u.id for u in _batch_units(units)] == ["batch"] def test_the_split_equals_the_single_batch(self, tmp_path, monkeypatch): - from buckaroo.pluggable_analysis_framework.stat_units import StatState table = _six_column_scan(tmp_path) single, errors = XorqStatPipeline(XORQ_STATS_V2, unit_test=False).process_table(table) spy = _QuerySpy(monkeypatch) @@ -1234,7 +1223,6 @@ def _sources(tmp_path): @pytest.mark.parametrize("kind", ["join", "aggregate", "diff", "csv", "filter", "derived column", "cached", "memtable", "backend table"]) def test_the_split_is_refused_for_anything_but_a_plain_parquet_scan(self, kind, tmp_path, monkeypatch, caplog): - from buckaroo.pluggable_analysis_framework.stat_units import StatState table = self._sources(tmp_path)[kind] assert XorqStatPipeline.chunk_refusal(table), kind From 746ae8c5dd47d92ae85a60193da74d4e9fb9babc Mon Sep 17 00:00:00 2001 From: Paddy Mullen Date: Sun, 4 Oct 2026 03:25:31 -0400 Subject: [PATCH 15/18] test(stats): failing tests for column names read in two namespaces at once (rows-first s4) A column group, the priority columns and StatRun's prefer hint accept original or rewritten (a, b, c) names and read each name in both. When an original name equals another column's rewritten name, a group returns the columns of both, priority ranks both first, and prefer picks the unit of whichever comes first in the frame. On a frame whose columns are c, b, a (rewritten a, b, c): - StatState(df, columns=('c',)) plans the units of c and a. - priority=('a',) plans c before a. - StatRun.next_unit(prefer=('a',)) picks the unit of c. The new tests also pin the namespace argument the fix adds, so a client that holds the rewritten names can ask for exactly those: StatState(namespace=) and StatRun.next_unit/run_next(namespace=), with "any" (the default), "original" and "rewritten". Co-Authored-By: Claude Sonnet 5.5 --- tests/unit/server/test_data_loading_polars.py | 41 ++++++++++++ tests/unit/test_paf_v2.py | 62 +++++++++++++++++++ tests/unit/test_xorq_stats_v2.py | 34 ++++++++++ 3 files changed, 137 insertions(+) diff --git a/tests/unit/server/test_data_loading_polars.py b/tests/unit/server/test_data_loading_polars.py index 0e53572b0..c636b6cda 100644 --- a/tests/unit/server/test_data_loading_polars.py +++ b/tests/unit/server/test_data_loading_polars.py @@ -526,6 +526,19 @@ def _stat_run(backend, gen=4): return StatRun(gen, "raw", dataflow.build_stats(dataflow.processed_df, run=False)), dataflow +# Original names that are also rewritten names of other columns: the rewritten +# names are c -> a, b -> b, a -> c. +_PERMUTED_DATA = { + "c": [12.5, 12.5, 12.5, 18.9, 18.9, 1e9], + "b": [1, 1, 1, 2, 2, 3], + "a": ["a", "a", "a", "b", "b", "c"]} + + +def _permuted_stat_run(backend, gen=4): + dataflow = _build_dataflow(backend, _frame(backend, _PERMUTED_DATA)) + return StatRun(gen, "raw", dataflow.build_stats(dataflow.processed_df, run=False)), dataflow + + class TestStatRun: """The run a session holds for one generation and scope: the planned units, the fragments they produced in the order they finished, and the @@ -556,6 +569,34 @@ def test_prefer_runs_the_units_that_name_those_columns_first(self, backend): assert run.next_unit(prefer=("c",)).columns == ("category",), "a rewritten name works too" assert run.next_unit(prefer=("nothing",)).columns == ("price",) + def test_prefer_reads_a_name_that_is_an_original_name_as_that_column(self, backend): + """"a" is the original name of the last column and the rewritten name + of the first. Read once, as the original name, it picks the last.""" + run, _dataflow = _permuted_stat_run(backend) + assert [u.columns for u in run.units] == [("c",), ("b",), ("a",)] + assert run.next_unit(prefer=("a",)).columns == ("a",) + assert run.next_unit(prefer=("b",)).columns == ("b",) + assert run.next_unit(prefer=("c",)).columns == ("c",) + + def test_prefer_in_the_rewritten_namespace_reads_a_client_s_names(self, backend): + """A client knows only the rewritten names: its "c" is the column + whose original name is "a".""" + run, _dataflow = _permuted_stat_run(backend) + assert run.next_unit(prefer=("c",), namespace="rewritten").columns == ("a",) + assert run.next_unit(prefer=("a",), namespace="rewritten").columns == ("c",) + assert run.next_unit(prefer=("zzz",), namespace="rewritten").columns == ("c",) + + def test_run_next_runs_the_unit_prefer_picks_in_its_namespace(self, backend): + run, _dataflow = _permuted_stat_run(backend) + assert list(run.run_next(prefer=("c",), namespace="rewritten")) == ["a"] + assert list(run.run_next(prefer=("c",))) == ["c"] + assert run.ran == ["column:a", "column:c"] + + def test_prefer_with_an_unknown_namespace_is_an_error(self, backend): + run, _dataflow = _permuted_stat_run(backend) + with pytest.raises(ValueError, match="namespace"): + run.next_unit(prefer=("c",), namespace="both") + def test_the_run_is_complete_when_no_unit_is_left(self, backend): run, _dataflow = _stat_run(backend) while run.run_next() is not None: diff --git a/tests/unit/test_paf_v2.py b/tests/unit/test_paf_v2.py index d63b3632a..e60b4e892 100644 --- a/tests/unit/test_paf_v2.py +++ b/tests/unit/test_paf_v2.py @@ -907,3 +907,65 @@ def test_real_stats_run_unit_by_unit_equal_process_df(self, backend): assert len(fragments) == len(df.columns) assert repr(rewrite_sd(merge_fragments(fragments), df)) == repr(expected) assert errors == [] and acc.errors == [] + + # A name can be an original column name and, at the same time, the + # rewritten name of another column. ``permuted`` has that: its rewritten + # names are c -> a, b -> b, a -> c. + @staticmethod + def _permuted(): + return pd.DataFrame({'c': [1, 2, 3, 1], 'b': ['x', 'y', 'x', 'x'], 'a': [0.5, 1.5, None, 2.5]}) + + @pytest.mark.parametrize('columns, expected', [ + (('c',), ['c']), + (('a',), ['a']), + (('a', 'c'), ['c', 'a']), + (('b',), ['b']), + ]) + def test_a_name_that_is_an_original_name_picks_that_column_only(self, columns, expected): + """'c' is the original name of one column and the rewritten name of + another. Read once, as the original name, it asks for one column.""" + state = StatState(self._permuted(), columns=columns) + assert [u.columns[0] for u in self._pipeline().plan(state)] == expected + + def test_a_rewritten_name_that_is_no_original_name_still_picks_its_column(self): + df = self._permuted().rename(columns={'b': 'word'}) + state = StatState(df, columns=('b',)) + assert [u.columns for u in self._pipeline().plan(state)] == [('word',)] + + @pytest.mark.parametrize('columns, expected', [ + (('c',), ['a']), + (('a',), ['c']), + (('b',), ['b']), + (('a', 'c'), ['c', 'a']), + ]) + def test_the_rewritten_namespace_reads_every_name_as_a_rewritten_name(self, columns, expected): + state = StatState(self._permuted(), columns=columns, namespace='rewritten') + pipeline = self._pipeline() + assert [u.columns[0] for u in pipeline.plan(state)] == expected + acc, fragments = self._run_units(pipeline, state) + assert list(acc.sd()) == expected and [list(f) for f in fragments] == [[c] for c in expected] + + def test_the_original_namespace_ignores_a_rewritten_name(self): + df = self._permuted().rename(columns={'b': 'word'}) + state = StatState(df, columns=('b', 'c'), namespace='original') + assert [u.columns for u in self._pipeline().plan(state)] == [('c',)] + + def test_a_group_outside_the_frame_is_empty(self): + assert self._pipeline().plan(StatState(self._permuted(), columns=('zzz',))) == [] + + def test_priority_reads_a_name_that_is_an_original_name_as_that_column(self): + state = StatState(self._permuted(), priority=('a',)) + assert [u.columns[0] for u in self._pipeline().plan(state)] == ['a', 'c', 'b'] + + @pytest.mark.parametrize('namespace, priority, expected', [ + ('original', ('a',), ['a', 'c', 'b']), + ('rewritten', ('c',), ['a', 'c', 'b']), + ('rewritten', ('a',), ['c', 'b', 'a']), + ]) + def test_priority_reads_each_name_once_in_its_namespace(self, namespace, priority, expected): + state = StatState(self._permuted(), priority=priority, namespace=namespace) + assert [u.columns[0] for u in self._pipeline().plan(state)] == expected + + def test_an_unknown_namespace_is_an_error(self): + with pytest.raises(ValueError, match='namespace'): + self._pipeline().plan(StatState(self._permuted(), columns=('c',), namespace='both')) diff --git a/tests/unit/test_xorq_stats_v2.py b/tests/unit/test_xorq_stats_v2.py index 2a32bbad5..3e420cec9 100644 --- a/tests/unit/test_xorq_stats_v2.py +++ b/tests/unit/test_xorq_stats_v2.py @@ -1083,6 +1083,40 @@ def test_a_column_group_returns_only_those_columns(self, names, monkeypatch): full, _errors = pipeline.process_table(_make_table()) assert _stable(acc.sd()) == _stable({c: full[c] for c in ("ints", "strs")}) + @staticmethod + def _permuted_table(): + """'c' is the original name of one column and the rewritten name of + another: the rewritten names are c -> a, b -> b, a -> c.""" + return xo.memtable(pd.DataFrame({"c": [1, 2, 3, 4, 5, 6, 7], "b": [7, 6, 5, 4, 3, 2, 1], + "a": [1, 3, 5, 7, 9, 11, 13]})) + + @pytest.mark.parametrize("columns, expected", [(("c",), ["c"]), (("a",), ["a"]), (("a", "c"), ["c", "a"])]) + def test_a_group_name_that_is_an_original_name_picks_that_column_only(self, columns, expected): + state = StatState(self._permuted_table(), columns=columns) + units = XorqStatPipeline(XORQ_STATS_V2, unit_test=False).plan(state) + assert [u.id for u in units] == ["batch", *(f"histogram:{c}" for c in expected)] + assert units[0].columns == tuple(expected) + + @pytest.mark.parametrize("columns, namespace, expected", [ + (("c",), "rewritten", ["a"]), + (("a",), "rewritten", ["c"]), + (("c",), "original", ["c"]), + ]) + def test_a_group_reads_each_name_in_its_namespace(self, columns, namespace, expected): + state = StatState(self._permuted_table(), columns=columns, namespace=namespace) + units = XorqStatPipeline(XORQ_STATS_V2, unit_test=False).plan(state) + assert [u.id for u in units] == ["batch", *(f"histogram:{c}" for c in expected)] + assert units[0].columns == tuple(expected) + + def test_priority_reads_a_name_that_is_an_original_name_as_that_column(self): + units = XorqStatPipeline(XORQ_STATS_V2, unit_test=False).plan(StatState(self._permuted_table(), priority=("a",))) + assert [u.id for u in units] == ["batch", "histogram:a", "histogram:c", "histogram:b"] + + def test_priority_reads_each_name_in_the_rewritten_namespace(self): + state = StatState(self._permuted_table(), priority=("c",), namespace="rewritten") + units = XorqStatPipeline(XORQ_STATS_V2, unit_test=False).plan(state) + assert [u.id for u in units] == ["batch", "histogram:a", "histogram:c", "histogram:b"] + # ============================================================ # Splitting the batch by column chunk (rows-first s4) From fa05adba91065fac5a7f8e59b675eededc4665b9 Mon Sep 17 00:00:00 2001 From: Paddy Mullen Date: Sun, 4 Oct 2026 03:33:20 -0400 Subject: [PATCH 16/18] fix(stats): read each requested column name in one namespace (rows-first s4) A column group, StatState.priority and StatRun's prefer hint matched a 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: a group returned columns that were not asked for, and prefer ran the unit of whichever column came first in the frame. resolve_names reads each name once. In "any", the default and what the code did for names that do not collide, 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. StatState(namespace=) and StatRun.next_unit/run_next(namespace=) also take "original" and "rewritten", so a caller that holds the client's rewritten names says so and gets exactly those columns. 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 now takes the state, since it needs the frame's columns and the namespace to read the priority names. Co-Authored-By: Claude Sonnet 5.5 --- .../stat_pipeline.py | 2 +- .../stat_units.py | 72 ++++++++++++++----- .../xorq_stat_pipeline.py | 2 +- buckaroo/server/stat_run.py | 26 ++++--- 4 files changed, 74 insertions(+), 28 deletions(-) diff --git a/buckaroo/pluggable_analysis_framework/stat_pipeline.py b/buckaroo/pluggable_analysis_framework/stat_pipeline.py index f26e1de98..812553ae4 100644 --- a/buckaroo/pluggable_analysis_framework/stat_pipeline.py +++ b/buckaroo/pluggable_analysis_framework/stat_pipeline.py @@ -293,7 +293,7 @@ def plan(self, state: StatState) -> List[StatUnit]: return [] pairs = columns_in_scope(state) skip = self._skipped(state, pairs) - active = prioritized([pair for pair in pairs if pair[0] not in skip], state.priority) + active = prioritized(state, [pair for pair in pairs if pair[0] not in skip]) return [StatUnit(id=f"column:{orig}", columns=(orig,), phase="column") for orig, _rewritten in active] def new_accumulator(self, state: StatState) -> StatAccumulator: diff --git a/buckaroo/pluggable_analysis_framework/stat_units.py b/buckaroo/pluggable_analysis_framework/stat_units.py index ff3c5dd24..c5a02c1c4 100644 --- a/buckaroo/pluggable_analysis_framework/stat_units.py +++ b/buckaroo/pluggable_analysis_framework/stat_units.py @@ -42,21 +42,62 @@ class (``column``, ``scan`` or ``query``) a scheduler can read. ``after`` cost: str = "column" +# How a caller writes the column names it hands in (``StatState.columns`` and +# ``priority``, and ``StatRun``'s ``prefer``): ``original`` names only, +# ``rewritten`` (``a, b, c``) names only, or ``any``. +NAMESPACES = ("any", "original", "rewritten") + + +def check_namespace(namespace: str) -> None: + if namespace not in NAMESPACES: + raise ValueError(f"namespace must be one of {NAMESPACES}, not {namespace!r}") + + +def resolve_names(pairs: Sequence[Tuple[Any, str]], names: Iterable[Any], namespace: str = "any") -> List[Any]: + """The original names of the columns that ``names`` pick out of ``pairs``, + the ``(orig, rewritten)`` of every column of the frame, in the order given. + A name that matches no column is ignored. + + Each name is read once, so it picks one column. In ``any``, 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. That makes ``any`` right for a caller that + writes original names. A caller that holds the rewritten names, as a client + does, must say ``rewritten``: when one of its names is also another column's + original name (a frame with columns ``c, b, a`` has a column named ``c`` + and a column rewritten to ``c``), ``any`` would read it as the original. + """ + check_namespace(namespace) + originals = {orig for orig, _rewritten in pairs} + by_rewritten = {rewritten: orig for orig, rewritten in pairs} + picked: List[Any] = [] + for name in names: + if namespace != "rewritten" and name in originals: + picked.append(name) + elif namespace != "original" and name in by_rewritten: + picked.append(by_rewritten[name]) + return picked + + @dataclass(frozen=True, eq=False) class StatState: """What a stat run describes: the frame or expression, and how to cut it. ``skip_columns`` get no unit. ``columns`` restricts the run to a group (a column-group request), and ``priority`` names columns to plan first. Both - take original or rewritten (``a, b, c``) names, since a client only knows - the rewritten ones. ``rows`` is the row count when the caller already has - it, which the xorq column-chunk split needs and which is never queried for. + are written in ``namespace`` (see ``resolve_names``): a client that only + knows the rewritten ``a, b, c`` names says ``rewritten``. ``rows`` is the + row count when the caller already has it, which the xorq column-chunk split + needs and which is never queried for. """ data: Any skip_columns: frozenset = frozenset() columns: Optional[Tuple[Any, ...]] = None priority: Tuple[Any, ...] = () rows: Optional[int] = None + namespace: str = "any" + + def __post_init__(self) -> None: + check_namespace(self.namespace) def columns_in_scope(state: StatState) -> List[Tuple[Any, str]]: @@ -66,22 +107,21 @@ def columns_in_scope(state: StatState) -> List[Tuple[Any, str]]: pairs = old_col_new_col(state.data) if state.columns is None: return list(pairs) - wanted = set(state.columns) - return [(orig, rewritten) for orig, rewritten in pairs if orig in wanted or rewritten in wanted] + wanted = set(resolve_names(pairs, state.columns, state.namespace)) + return [pair for pair in pairs if pair[0] in wanted] -def prioritized(pairs: Sequence[Tuple[Any, str]], priority: Sequence[Any]) -> List[Tuple[Any, str]]: - """``pairs`` with the columns named in ``priority`` first, in that order, - and the rest in their own order.""" +def prioritized(state: StatState, pairs: Sequence[Tuple[Any, str]]) -> List[Tuple[Any, str]]: + """``pairs`` with the columns named in ``state.priority`` first, in that + order, and the rest in their own order. The names are read against every + column of ``state.data``, not only ``pairs``, so a name picks the same + column whatever group is asked for.""" + if not state.priority: + return list(pairs) rank: Dict[Any, int] = {} - for position, name in enumerate(priority): - rank.setdefault(name, position) - - def key(pair: Tuple[Any, str]) -> int: - ranks = [rank[name] for name in pair if name in rank] - return min(ranks) if ranks else len(rank) - - return sorted(pairs, key=key) + for position, orig in enumerate(resolve_names(old_col_new_col(state.data), state.priority, state.namespace)): + rank.setdefault(orig, position) + return sorted(pairs, key=lambda pair: rank.get(pair[0], len(rank))) class StatAccumulator: diff --git a/buckaroo/pluggable_analysis_framework/xorq_stat_pipeline.py b/buckaroo/pluggable_analysis_framework/xorq_stat_pipeline.py index 93ab518f6..9601e2e03 100644 --- a/buckaroo/pluggable_analysis_framework/xorq_stat_pipeline.py +++ b/buckaroo/pluggable_analysis_framework/xorq_stat_pipeline.py @@ -459,7 +459,7 @@ def plan(self, state: StatState) -> List[StatUnit]: schema = table.schema() # xorq has always matched skip_columns on the column's own name only. skip = {orig for orig, _rewritten in pairs if orig in state.skip_columns} - active = [orig for orig, _rewritten in prioritized([p for p in pairs if p[0] not in skip], state.priority)] + active = [orig for orig, _rewritten in prioritized(state, [p for p in pairs if p[0] not in skip])] chunks = self._batch_chunks(state, active) units: List[StatUnit] = [] batch_of: Dict[Any, str] = {} diff --git a/buckaroo/server/stat_run.py b/buckaroo/server/stat_run.py index d775f2a0c..e2431297e 100644 --- a/buckaroo/server/stat_run.py +++ b/buckaroo/server/stat_run.py @@ -14,7 +14,8 @@ import time from typing import Any, List, Optional, Sequence, Tuple -from buckaroo.pluggable_analysis_framework.stat_units import Fragment, StatUnit, rewrite_sd +from buckaroo.df_util import old_col_new_col +from buckaroo.pluggable_analysis_framework.stat_units import Fragment, StatUnit, resolve_names, rewrite_sd _run_ids = itertools.count(1) @@ -41,6 +42,7 @@ def __init__(self, stats_gen: int, scope: str, stats: Any): self.status = "pending" if self.units else "complete" self.error: Optional[Exception] = None self.created_at = time.time() + self._frame_columns: Optional[List[Tuple[Any, str]]] = None @property def key(self) -> Tuple[int, str]: @@ -51,27 +53,31 @@ def remaining(self) -> int: """The number of units not yet run.""" return len(self.units) - len(self.ran) - def next_unit(self, prefer: Sequence[Any] = ()) -> Optional[StatUnit]: + def next_unit(self, prefer: Sequence[Any] = (), namespace: str = "any") -> Optional[StatUnit]: """The unit ``run_next`` would run: the first whose prerequisites have - run, or the first of those that covers a column named in ``prefer`` - (by original or rewritten name), so the columns a client is looking at - come first. ``None`` when no unit is left.""" + run, or the first of those that covers a column named in ``prefer``, so + the columns a client is looking at come first. The names are written in + ``namespace`` (see ``resolve_names``): a client that holds the rewritten + ``a, b, c`` names says ``rewritten``. ``None`` when no unit is left.""" ran = set(self.ran) ready = [unit for unit in self.units if unit.id not in ran and all(a in ran for a in unit.after)] if prefer: - wanted = set(prefer) + if self._frame_columns is None: + self._frame_columns = old_col_new_col(self.acc.state.data) + wanted = set(resolve_names(self._frame_columns, prefer, namespace)) for unit in ready: - if any(col in wanted or self.acc.rewritten.get(col) in wanted for col in unit.columns): + if any(col in wanted for col in unit.columns): return unit return ready[0] if ready else None - def run_next(self, prefer: Sequence[Any] = ()) -> Optional[Fragment]: + def run_next(self, prefer: Sequence[Any] = (), namespace: str = "any") -> Optional[Fragment]: """Run one unit and return its fragment, which is also appended to ``fragments``. ``None`` when the run is complete or failed. A unit that - raises fails the run, and the exception propagates once.""" + raises fails the run, and the exception propagates once. ``prefer`` and + ``namespace`` are as in ``next_unit``.""" if self.status != "pending": return None - unit = self.next_unit(prefer) + unit = self.next_unit(prefer, namespace) if unit is None: self.status = "complete" return None From 2a73c8ed958292fd148aa148190df41e424a9e1d Mon Sep 17 00:00:00 2001 From: Paddy Mullen Date: Sun, 4 Oct 2026 04:03:17 -0400 Subject: [PATCH 17/18] test(server): failing tests for incremental stats requests, per-client cursors and the final assignment (rows-first s5) Covers the request branch that runs units for a time budget (at least one, one cold unit per budget), the per-handler cursors over a shared StatRun, the final assignment that frees the run and carries the rebuilt df_display_args, a legacy client finishing a run in progress, a concurrent infinite_request between units, the end-to-end order of first initial_state, rows and stats, /reload_expr mid-run, and the stats.unit and firstpull.stats_total spans. Co-Authored-By: Claude Sonnet 5.5 --- tests/unit/server/test_data_loading_polars.py | 228 +++++++++ tests/unit/server/test_load_expr.py | 481 +++++++++++++++++- 2 files changed, 708 insertions(+), 1 deletion(-) diff --git a/tests/unit/server/test_data_loading_polars.py b/tests/unit/server/test_data_loading_polars.py index c636b6cda..089ad0743 100644 --- a/tests/unit/server/test_data_loading_polars.py +++ b/tests/unit/server/test_data_loading_polars.py @@ -34,10 +34,13 @@ from buckaroo.dataflow.dataflow import assemble_merged_sd from buckaroo.dataflow.sd_cache import split_chain_by_scope from buckaroo.jlisp.lisp_utils import s as lisp_sym +from buckaroo.pluggable_analysis_framework import perf_log from buckaroo.pluggable_analysis_framework.col_analysis import ColAnalysis from buckaroo.pluggable_analysis_framework.stat_pipeline import StatPipeline from buckaroo.pluggable_analysis_framework.stat_units import StatAccumulator, StatState, StatUnit, merge_fragments from buckaroo.pluggable_analysis_framework.utils import PERVERSE_DF +from buckaroo.serialization_utils import resolve_summary_stats_payload +from buckaroo.server import stats_wire from buckaroo.server.data_loading import ServerDataflow, handle_infinite_request_buckaroo from buckaroo.server.data_loading_polars import (PolarsServerDataflow, create_polars_dataflow, handle_infinite_request_buckaroo_polars, load_file_polars) from buckaroo.server.stat_run import StatCursor, StatRun @@ -678,3 +681,228 @@ def run(self, unit, acc): assert run.status == "error" and isinstance(run.error, RuntimeError) assert run.run_next() is None, "a failed run is not retried" assert run.fragments == [{"x": {"v": 1}}] + + +class _FakeClock: + """A clock a test moves by hand, so a budget can be spent without waiting.""" + + def __init__(self): + self.now = 0.0 + + def __call__(self): + return self.now + + +class _CostedStats: + """A stats class whose units run for real and cost what the test says on a + fake clock: ``cost`` is the default and ``costs`` maps a unit id to its own.""" + + def __init__(self, stats, clock, cost, costs=None): + self._stats, self._clock, self._cost, self._costs = stats, clock, cost, costs or {} + self.state = stats.state + + def plan(self, state): + return self._stats.plan(state) + + def new_accumulator(self, state): + return self._stats.new_accumulator(state) + + def run(self, unit, acc): + fragment = self._stats.run(unit, acc) + self._clock.now += self._costs.get(unit.id, self._cost) + return fragment + + +def _costed_run(backend, clock, cost, costs=None, gen=4): + run, dataflow = _stat_run(backend, gen) + run.stats = _CostedStats(run.stats, clock, cost, costs) + return run, dataflow + + +def _stat_rows(payload): + """The decoded ``all_stats`` rows of a payload, keyed by stat name.""" + return {row["index"]: row for row in resolve_summary_stats_payload(payload)} + + +class TestRunUnits: + """Running a ``StatRun``'s units for one request (rows-first s5): as many as + fit in the request's time budget, and at least one.""" + + def test_units_that_cost_little_run_together_within_the_budget(self, backend): + clock = _FakeClock() + run, _dataflow = _costed_run(backend, clock, cost=0.01) + assert stats_wire.run_units(run, 0.075, clock=clock) == 3 + assert run.status == "complete" and run.remaining == 0 + + def test_a_request_stops_once_the_budget_is_spent(self, backend): + clock = _FakeClock() + run, _dataflow = _costed_run(backend, clock, cost=0.04) + assert stats_wire.run_units(run, 0.075, clock=clock) == 2 + assert run.status == "pending" and run.remaining == 1 + + def test_a_unit_that_costs_more_than_the_budget_is_the_only_one_of_its_request(self, backend): + clock = _FakeClock() + run, _dataflow = _costed_run(backend, clock, cost=0.2) + assert [stats_wire.run_units(run, 0.075, clock=clock) for _ in range(3)] == [1, 1, 1] + assert run.status == "complete" + + def test_a_cold_unit_after_warm_ones_ends_the_request(self, backend): + """A snapshot-cache hit costs milliseconds and a miss is a query: the + units that hit run together, and the first miss is the last of its + request.""" + clock = _FakeClock() + run, _dataflow = _costed_run(backend, clock, cost=0.005, costs={"column:category": 0.5}) + assert stats_wire.run_units(run, 0.075, prefer=("c",), clock=clock) == 1 + assert run.ran == ["column:category"] + assert stats_wire.run_units(run, 0.075, clock=clock) == 2 + assert run.status == "complete" + + def test_one_unit_runs_whatever_the_budget(self, backend): + clock = _FakeClock() + run, _dataflow = _costed_run(backend, clock, cost=0.0) + assert stats_wire.run_units(run, 0, clock=clock) == 1 + assert run.remaining == 2 + + def test_no_budget_runs_every_unit(self, backend): + clock = _FakeClock() + run, _dataflow = _costed_run(backend, clock, cost=10.0) + assert stats_wire.run_units(run, None, clock=clock) == 3 + assert run.status == "complete" + + def test_a_complete_run_runs_nothing(self, backend): + clock = _FakeClock() + run, _dataflow = _costed_run(backend, clock, cost=0.0) + stats_wire.run_units(run, None, clock=clock) + assert stats_wire.run_units(run, 0.075, clock=clock) == 0 + assert len(run.fragments) == 3 + + def test_the_columns_hint_is_read_in_the_client_s_rewritten_names(self, backend): + """The grid knows only a, b, c: "c" is the third column, category.""" + clock = _FakeClock() + run, _dataflow = _costed_run(backend, clock, cost=1.0) + stats_wire.run_units(run, 0, prefer=["c", "a"], clock=clock) + stats_wire.run_units(run, 0, prefer=["c", "a"], clock=clock) + assert run.ran == ["column:category", "column:price"] + + def test_a_unit_that_raises_fails_the_run_and_the_error_propagates(self, backend): + clock = _FakeClock() + run, _dataflow = _costed_run(backend, clock, cost=0.0) + + def boom(unit, acc): + raise RuntimeError("unit failed") + + run.stats.run = boom + with pytest.raises(RuntimeError, match="unit failed"): + stats_wire.run_units(run, 0.075, clock=clock) + assert run.status == "error" + + def test_each_unit_is_a_stats_unit_span_on_the_bound_sink(self, backend): + clock = _FakeClock() + run, _dataflow = _costed_run(backend, clock, cost=0.0) + spans = [] + with perf_log.telemetry_context("run-units", spans.append): + stats_wire.run_units(run, None, clock=clock, session_id="run-units") + units = [r for r in spans if r["name"] == "stats.unit"] + assert [r["attrs"]["unit"] for r in units] == run.ran + for record, unit in zip(units, run.units): + assert record["trace"] == "run-units" + attrs = record["attrs"] + assert (attrs["session"], attrs["stats_gen"]) == ("run-units", 4) + assert (attrs["phase"], attrs["cost"], attrs["columns"]) == (unit.phase, unit.cost, len(unit.columns)) + + def test_the_run_adds_up_the_time_its_units_took(self, backend): + run, _dataflow = _stat_run(backend) + assert run.elapsed_s == 0 + run.run_next() + first = run.elapsed_s + assert first > 0 + run.run_next() + assert run.elapsed_s > first + + def test_a_clean_run_has_no_errors(self, backend): + run, _dataflow = _stat_run(backend) + stats_wire.run_units(run, None) + assert run.errs() == {} + + +class TestPartialPayload: + """The ``all_stats`` payload of the fragments a client has not seen: the + stats of the state, assembled the way ``merged_sd`` is, for the columns + those fragments cover.""" + + def test_the_payload_of_every_fragment_is_the_dataflow_s_all_stats(self, backend): + run, dataflow = _stat_run(backend) + stats_wire.run_units(run, None) + payload = stats_wire.partial_payload(dataflow, run, run.fragments) + assert (payload["format"], payload["layout"]) == ("parquet_b64", "wide") + assert _stat_rows(payload) == _stat_rows(dataflow.df_data_dict["all_stats"]) + + def test_the_payload_holds_only_the_columns_the_fragments_cover(self, backend): + run, _dataflow = _stat_run(backend) + fragment = run.run_next(prefer=("category",)) + payload = stats_wire.partial_payload(_dataflow, run, [fragment]) + columns = {name for row in _stat_rows(payload).values() for name in row if name != "index"} + assert columns == {"c"}, "category is the third column, rewritten to c" + + def test_a_search_filter_keys_the_run_s_stats_as_filtered(self, backend): + """With a search active the run analyzes the filtered frame, so its + stats are the ``filtered_*`` keys of ``merged_sd`` and the bare keys + stay the unfiltered scope's.""" + dataflow = _three_scope_dataflow(backend, _UNTIED_DATA) + run = StatRun(4, "raw", dataflow.build_stats(dataflow.processed_df, run=False)) + stats_wire.run_units(run, None) + assert _as_json(dataflow._assemble_merged_sd(run.raw_sd())) == _as_json(dataflow.merged_sd) + + def test_init_sd_overrides_win_over_the_run_s_stats_as_in_merged_sd(self, backend): + dataflow = _build_dataflow(backend, _frame(backend, _UNTIED_DATA), + init_sd={"qty": {"mean": 99.5, "from_init": 1}}) + run = StatRun(4, "raw", dataflow.build_stats(dataflow.processed_df, run=False)) + stats_wire.run_units(run, None) + assembled = dataflow._assemble_merged_sd(run.raw_sd()) + assert _as_json(assembled) == _as_json(dataflow.merged_sd) + assert assembled["b"]["from_init"] == 1 + + +class TestAssembledMergedSd: + """``_assemble_merged_sd`` is the body of the ``merged_sd`` observer.""" + + def test_it_is_merged_sd_with_no_sd_standing_in(self, backend): + dataflow = _three_scope_dataflow(backend, _UNTIED_DATA) + assert _as_json(dataflow._assemble_merged_sd()) == _as_json(dataflow.merged_sd) + + def test_a_running_sd_stands_in_for_the_scopes_that_share_the_filt_chain(self, backend): + dataflow = _build_dataflow(backend, _frame(backend, _UNTIED_DATA), stats_tier="schema") + full = _build_dataflow(backend, _frame(backend, _UNTIED_DATA)) + assembled = dataflow._assemble_merged_sd(full.summary_sd) + assert _as_json(assembled) == _as_json(full.merged_sd) + assert "mean" not in dataflow.merged_sd["a"], "the dataflow itself is still at the schema tier" + + +class TestHighlightedDisplayArgs: + def test_the_term_is_set_on_string_columns_of_a_copy(self, backend): + dataflow = _build_dataflow(backend, _frame(backend, _UNTIED_DATA)) + original = json.dumps(dataflow.df_display_args, sort_keys=True, default=str) + highlighted = stats_wire.highlighted_display_args(dataflow.df_display_args, "ca") + assert json.dumps(dataflow.df_display_args, sort_keys=True, default=str) == original + phrases = [col["displayer_args"]["highlight_phrase"] + for col in highlighted["main"]["df_viewer_config"]["column_config"] + if col.get("displayer_args", {}).get("highlight_phrase")] + assert phrases == [["ca"]] + + def test_an_empty_term_removes_the_highlight(self, backend): + dataflow = _build_dataflow(backend, _frame(backend, _UNTIED_DATA)) + highlighted = stats_wire.highlighted_display_args(dataflow.df_display_args, "ca") + cleared = stats_wire.highlighted_display_args(highlighted, "") + assert json.dumps(cleared, sort_keys=True, default=str) == json.dumps( + dataflow.df_display_args, sort_keys=True, default=str) + + +class TestDisplayArgsHash: + def test_equal_content_hashes_equal_whatever_the_key_order(self): + assert stats_wire.display_args_hash({"a": 1, "b": [1, 2]}) == stats_wire.display_args_hash({"b": [1, 2], "a": 1}) + + def test_a_changed_value_changes_the_hash(self): + assert stats_wire.display_args_hash({"a": 1}) != stats_wire.display_args_hash({"a": 2}) + + def test_nan_hashes_equal_to_nan(self): + assert stats_wire.display_args_hash({"a": float("nan")}) == stats_wire.display_args_hash({"a": float("nan")}) diff --git a/tests/unit/server/test_load_expr.py b/tests/unit/server/test_load_expr.py index b2c1b2252..7b186fcb7 100644 --- a/tests/unit/server/test_load_expr.py +++ b/tests/unit/server/test_load_expr.py @@ -8,6 +8,7 @@ import shutil import sys import tempfile +import time from contextlib import contextmanager from pathlib import Path from types import SimpleNamespace @@ -37,8 +38,9 @@ from buckaroo.pluggable_analysis_framework.stat_units import merge_fragments # noqa: E402 from buckaroo.pluggable_analysis_framework.xorq_stat_pipeline import XorqStatPipeline # noqa: E402 from buckaroo.serialization_utils import resolve_summary_stats_payload # noqa: E402 -from buckaroo.server import telemetry, xorq_loading # noqa: E402 +from buckaroo.server import stats_wire, telemetry, xorq_loading # noqa: E402 from buckaroo.server.app import make_app as _make_app # noqa: E402 +from buckaroo.server.stat_run import StatRun # noqa: E402 from buckaroo.server.stats_wire import start_stat_run # noqa: E402 from buckaroo.server.websocket_handler import DataStreamHandler # noqa: E402 from tests.unit.dataflow.scoped_summary_stats_test import ( # noqa: E402 @@ -1440,6 +1442,43 @@ def spy(pipeline, query): yield queries +def _merge_stat_payloads(payloads): + """What a client holds after key-merging ``stats_update`` payloads into + ``all_stats``, as ``StatsChannel`` does: a cell a payload carries replaces + the held one, and a null (the wide pivot's fill for a stat a column did not + carry) never replaces a value.""" + merged = {} + for payload in payloads: + for stat, row in _rows_by_stat(payload).items(): + cells = merged.setdefault(stat, {}) + for column, value in row.items(): + if column == "index" or (value is None and cells.get(column) is not None): + continue + cells[column] = value + return merged + + +class _SlowStats: + """A stats class whose units run for real, each after ``delay`` seconds, as + a slow query would. ``fail_on`` names a unit that raises instead.""" + + def __init__(self, stats, delay=0.0, fail_on=None): + self._stats, self._delay, self._fail_on = stats, delay, fail_on + self.state = stats.state + + def plan(self, state): + return self._stats.plan(state) + + def new_accumulator(self, state): + return self._stats.new_accumulator(state) + + def run(self, unit, acc): + time.sleep(self._delay) + if unit.id == self._fail_on: + raise RuntimeError(f"{unit.id} failed") + return self._stats.run(unit, acc) + + class TestStatsWire(tornado.testing.AsyncHTTPTestCase): """``stats_request``, ``stats_update`` and ``stats_aborted`` on a deferred ``/load_expr`` session, the ``stats_gen`` counter and ``df_meta.stats``, and @@ -1996,6 +2035,446 @@ async def test_each_connection_keeps_its_own_cursor_into_the_run(self): self.assertEqual(merge_fragments(first + rest), merge_fragments(everything)) self.assertEqual(len(set(run.ran)), len(run.units), "no unit ran twice") + def _budget(self, seconds): + """Set the time budget of an incremental request for the rest of the + test. Zero runs exactly one unit per request, whatever it costs.""" + patcher = patch.object(stats_wire, "STATS_BUDGET_S", seconds, create=True) + patcher.start() + self.addCleanup(patcher.stop) + + async def _pull_stats(self, ws, gen, limit=12, **fields): + """Send incremental ``stats_request``s on ``ws`` until the final reply + (or an abort) and return every reply.""" + replies = [] + for _ in range(limit): + ws.write_message(_stats_request(gen, incremental=True, **fields)) + reply = await _read_json(ws) + replies.append(reply) + if reply["type"] != "stats_update" or reply["final"]: + break + return replies + + def _wrapped_run(self, sid, **kwargs): + """Give the session a StatRun for its current generation whose units run + for real but through ``_SlowStats``.""" + session = self._session(sid) + dataflow = session.xorq_dataflow + run = StatRun(session.stats_gen, "raw", + _SlowStats(dataflow.build_stats(dataflow.processed_df, run=False), **kwargs)) + session.stat_runs[run.key] = run + return run + + async def _infinite_window(self, ws, end=5): + ws.write_message(json.dumps({"type": "infinite_request", + "payload_args": {"start": 0, "end": end, "sourceName": "default", "origEnd": end}})) + resp = await _read_json(ws) + pq.read_table(io.BytesIO(await ws.read_message())) + return resp + + @tornado.testing.gen_test + async def test_incremental_requests_run_one_unit_each_and_end_with_the_complete_state(self): + inline = await self._inline_frame("sw-inc-inline") + await self._load("sw-inc", stats_delivery="deferred") + ws, first = await self._connect("sw-inc", caps="stats_update") + gen = self._stats(first)["gen"] + self._budget(0) + + with _count_stat_queries() as queries: + replies = await self._pull_stats(ws, gen) + + self.assertEqual([r["type"] for r in replies], ["stats_update"] * 4) + self.assertEqual([r["final"] for r in replies], [False, False, False, True]) + self.assertEqual([r["remaining"] for r in replies], [3, 2, 1, 0]) + self.assertEqual(len(queries), 4, "one query for each unit, none repeated") + for reply in replies: + self.assertEqual((reply["stats_gen"], reply["scope"], reply["tier"]), (gen, "raw", "full")) + self.assertIsInstance(reply["elapsed_ms"], (int, float)) + self.assertNotIn("histogram", _rows_by_stat(replies[0]["payload"]), "the scalar batch comes first") + self.assertIn("histogram", _rows_by_stat(replies[1]["payload"])) + self.assertEqual(_as_json(_merge_stat_payloads(r["payload"] for r in replies)), + _as_json(_merge_stat_payloads([inline["df_data_dict"]["all_stats"]]))) + + @tornado.testing.gen_test + async def test_the_last_unit_assigns_the_stats_and_frees_the_run(self): + await self._load("sw-assign", stats_delivery="deferred") + ws, first = await self._connect("sw-assign", caps="stats_update") + gen = self._stats(first)["gen"] + session = self._session("sw-assign") + dataflow = session.xorq_dataflow + key = dataflow._scope_cache_key(split_chain_by_scope(dataflow.operations)["filt"], tier="full") + self._budget(0) + + ws.write_message(_stats_request(gen, incremental=True)) + await _read_json(ws) + self.assertEqual(len(session.stat_runs[(gen, "raw")].fragments), 1) + self.assertEqual((session.stats_status, dataflow.stats_tier), ("pending", "schema")) + self.assertNotIn(key, dataflow.summary_stats_cache, "a partial run writes nothing into the dataflow") + + await self._pull_stats(ws, gen) + self.assertEqual(session.stat_runs, {}, "the accumulator is freed with the run") + self.assertEqual((session.stats_status, dataflow.stats_tier), ("complete", "full")) + self.assertIn(key, dataflow.summary_stats_cache) + self.assertIn("histogram_bins", _rows_by_stat(session.df_data_dict["all_stats"]), + "the session snapshot is refreshed in the same step") + self.assertEqual(_as_json(dataflow.summary_sd), _as_json(dataflow.summary_stats_cache[key])) + + @tornado.testing.gen_test + async def test_a_long_budget_runs_every_unit_in_one_request(self): + await self._load("sw-warm-budget", stats_delivery="deferred") + ws, first = await self._connect("sw-warm-budget", caps="stats_update") + self._budget(3600) + with _count_stat_queries() as queries: + replies = await self._pull_stats(ws, self._stats(first)["gen"]) + self.assertEqual(len(replies), 1) + self.assertEqual((replies[0]["final"], replies[0]["remaining"]), (True, 0)) + self.assertEqual(len(queries), 4) + + @tornado.testing.gen_test + async def test_the_columns_hint_orders_the_units_of_a_request(self): + await self._load("sw-hint", stats_delivery="deferred") + ws, first = await self._connect("sw-hint", caps="stats_update") + gen = self._stats(first)["gen"] + self._budget(0) + ws.write_message(_stats_request(gen, incremental=True, columns=["c"])) + await _read_json(ws) + ws.write_message(_stats_request(gen, incremental=True, columns=["c"])) + reply = await _read_json(ws) + run = self._session("sw-hint").stat_runs[(gen, "raw")] + self.assertEqual(run.ran, ["batch", "histogram:category"], + "the batch every histogram reads goes first, then the unit of the column named, c being category") + columns = {name for row in _rows_by_stat(reply["payload"]).values() for name in row if name != "index"} + self.assertEqual(columns, {"c"}) + + @tornado.testing.gen_test + async def test_a_stale_gen_mid_run_gets_stats_aborted_and_the_next_generation_starts_over(self): + await self._load("sw-mid-stale", stats_delivery="deferred") + ws, first = await self._connect("sw-mid-stale", caps="stats_update") + gen = self._stats(first)["gen"] + self._budget(0) + ws.write_message(_stats_request(gen, incremental=True)) + self.assertFalse((await _read_json(ws))["final"]) + + ws.write_message(_state_change(quick_command_args={"search": ["a"]})) + self._assert_pending(await _read_json(ws), gen + 1) + with _count_stat_queries() as queries: + ws.write_message(_stats_request(gen, incremental=True)) + aborted = await _read_json(ws) + self.assertEqual((aborted["type"], aborted["reason"], aborted["current_gen"]), + ("stats_aborted", "stale", gen + 1)) + self.assertEqual(queries, [], "a stale request runs no unit") + + replies = await self._pull_stats(ws, gen + 1) + self.assertEqual([r["remaining"] for r in replies], [3, 2, 1, 0], "the new generation runs every unit again") + + @tornado.testing.gen_test + async def test_client_b_opening_after_client_a_finished_gets_the_complete_all_stats(self): + inline = await self._inline_frame("sw-after-inline") + await self._load("sw-after", stats_delivery="deferred") + a, a_open = await self._connect("sw-after", caps="stats_update") + gen = self._stats(a_open)["gen"] + self._budget(0) + self.assertEqual(len(await self._pull_stats(a, gen)), 4) + + with _count_stat_queries() as queries: + b, b_open = await self._connect("sw-after", caps="stats_update") + self._assert_complete(b_open, gen, inline) + self.assertEqual(queries, [], "the stats A's run produced are not computed again for B") + + b.write_message(_stats_request(gen, incremental=True)) + reply = await _read_json(b) + self.assertEqual((reply["type"], reply["final"], reply["remaining"]), ("stats_update", True, 0)) + self.assertEqual(queries, []) + + @tornado.testing.gen_test + async def test_two_interleaved_clients_each_end_with_full_stats_and_run_no_unit_twice(self): + inline = await self._inline_frame("sw-inter-inline") + await self._load("sw-inter", stats_delivery="deferred") + a, a_open = await self._connect("sw-inter", caps="stats_update") + b, _ = await self._connect("sw-inter", caps="stats_update") + gen = self._stats(a_open)["gen"] + self._budget(0) + received = {"a": [], "b": []} + ran = [] + + with _count_stat_queries() as queries: + for _ in range(4): + for name, ws in (("a", a), ("b", b)): + ws.write_message(_stats_request(gen, incremental=True)) + received[name].append(await _read_json(ws)) + ran.append(len(queries)) + + self.assertEqual(ran, [1, 1, 2, 2, 3, 3, 4, 4], + "each unit ran once, for whichever client asked first, and the other was caught up without a query") + full = _as_json(_merge_stat_payloads([inline["df_data_dict"]["all_stats"]])) + for name in ("a", "b"): + replies = received[name] + self.assertTrue(all(r["type"] == "stats_update" for r in replies)) + self.assertTrue(replies[-1]["final"]) + self.assertEqual(_as_json(_merge_stat_payloads(r["payload"] for r in replies)), full, name) + + @tornado.testing.gen_test + async def test_the_final_reply_carries_the_rebuilt_display_args_when_its_hash_differs(self): + await self._load("sw-dargs", stats_delivery="deferred") + a, a_open = await self._connect("sw-dargs", caps="stats_update") + b, _ = await self._connect("sw-dargs", caps="stats_update") + gen = self._stats(a_open)["gen"] + session = self._session("sw-dargs") + self._budget(0) + + replies = await self._pull_stats(a, gen) + self.assertEqual(len(replies), 4) + for reply in replies[:-1]: + self.assertNotIn("df_display_args", reply) + rebuilt = replies[-1]["df_display_args"] + self.assertEqual(_as_json(rebuilt), _as_json(session.df_display_args)) + self.assertNotEqual(_as_json(rebuilt), _as_json(a_open["df_display_args"]), + "full stats change the float column's minWidth") + + # B asked for nothing while A ran, and still holds the schema tier's config. + b.write_message(_stats_request(gen, incremental=True)) + late = await _read_json(b) + self.assertEqual((late["final"], late["remaining"]), (True, 0)) + self.assertEqual(_as_json(late["df_display_args"]), _as_json(session.df_display_args)) + b.write_message(_stats_request(gen, incremental=True)) + self.assertNotIn("df_display_args", await _read_json(b), "B holds it now") + + @tornado.testing.gen_test + async def test_the_final_reply_carries_no_display_args_when_they_did_not_change(self): + build = str(xo.build_expr(xo.memtable({"qty": [1, 2, 1, 3, 2], "category": ["a", "b", "a", "c", "b"]}, + name="t2"), builds_dir=self.builds_root)) + await self._load("sw-same-args", build_dir=build, stats_delivery="deferred") + ws, first = await self._connect("sw-same-args", caps="stats_update") + self._budget(0) + replies = await self._pull_stats(ws, self._stats(first)["gen"]) + self.assertEqual(len(replies), 3, "the scalar batch and one histogram for each of two columns") + self.assertTrue(replies[-1]["final"]) + self.assertNotIn("df_display_args", replies[-1], "int and string columns are styled the same at both tiers") + + @tornado.testing.gen_test + async def test_the_final_display_args_keep_the_clients_search_highlight(self): + await self._load("sw-highlight", stats_delivery="deferred") + ws, first = await self._connect("sw-highlight", caps="stats_update") + ws.write_message(_state_change(search_string="ca")) + self._assert_pending(await _read_json(ws), self._stats(first)["gen"]) + self._budget(0) + replies = await self._pull_stats(ws, self._stats(first)["gen"]) + self.assertEqual(len(replies), 4) + columns = replies[-1]["df_display_args"]["main"]["df_viewer_config"]["column_config"] + self.assertTrue([cc for cc in columns if cc.get("displayer_args", {}).get("highlight_phrase") == ["ca"]], + "rebuilt config replaces the one with the highlight, so it must carry it") + + @tornado.testing.gen_test + async def test_a_filtered_state_ends_as_complete_as_the_inline_session_s(self): + search = {"search": ["a"]} + inline = await self._inline_frame("sw-filt-inline", quick_command_args=search) + await self._load("sw-filt", stats_delivery="deferred") + ws, first = await self._connect("sw-filt", caps="stats_update") + ws.write_message(_state_change(quick_command_args=search)) + gen = self._stats(await _read_json(ws))["gen"] + self._budget(0) + replies = await self._pull_stats(ws, gen) + self.assertEqual(len(replies), 4) + _, late = await self._connect("sw-filt", caps="stats_update") + self._assert_complete(late, gen, inline) + + @tornado.testing.gen_test + async def test_a_legacy_client_connecting_mid_run_gets_a_complete_state_from_the_remaining_units(self): + inline = await self._inline_frame("sw-mid-inline") + await self._load("sw-mid", stats_delivery="deferred") + a, a_open = await self._connect("sw-mid", caps="stats_update") + gen = self._stats(a_open)["gen"] + self._budget(0) + + with _count_stat_queries() as queries: + for _ in range(2): + a.write_message(_stats_request(gen, incremental=True)) + self.assertFalse((await _read_json(a))["final"]) + self.assertEqual(len(queries), 2) + _, legacy = await self._connect("sw-mid") + self.assertEqual(len(queries), 4, "the legacy client's frame ran the two units still to run") + a.write_message(_stats_request(gen, incremental=True)) + reply = await _read_json(a) + self.assertEqual(len(queries), 4, "A is answered from the stats the legacy client's frame completed") + self._assert_complete(legacy, gen, inline) + self.assertEqual((reply["final"], reply["remaining"]), (True, 0)) + self.assertEqual(_rows_by_stat(reply["payload"]), _rows_by_stat(inline["df_data_dict"]["all_stats"])) + + @tornado.testing.gen_test + async def test_a_request_with_no_incremental_flag_finishes_a_run_that_has_progress(self): + inline = await self._inline_frame("sw-finish-inline") + await self._load("sw-finish", stats_delivery="deferred") + ws, first = await self._connect("sw-finish", caps="stats_update") + gen = self._stats(first)["gen"] + self._budget(0) + with _count_stat_queries() as queries: + for _ in range(2): + ws.write_message(_stats_request(gen, incremental=True)) + self.assertFalse((await _read_json(ws))["final"]) + ws.write_message(_stats_request(gen)) + update = await _read_json(ws) + self.assertEqual((update["type"], update["final"], update["remaining"]), ("stats_update", True, 0)) + self.assertEqual(len(queries), 4, "the whole-run request ran the two units left, not all four") + self.assertEqual(_rows_by_stat(update["payload"]), _rows_by_stat(inline["df_data_dict"]["all_stats"])) + + @tornado.testing.gen_test + async def test_a_concurrent_infinite_request_is_served_between_units_within_the_stall_bound(self): + await self._load("sw-stall", stats_delivery="deferred") + a, a_open = await self._connect("sw-stall", caps="stats_update") + b, _ = await self._connect("sw-stall", caps="stats_update") + gen = self._stats(a_open)["gen"] + self._wrapped_run("sw-stall", delay=0.04) + + a.write_message(_stats_request(gen, incremental=True)) + started = time.perf_counter() + rows = await self._infinite_window(b) + rows_latency = time.perf_counter() - started + first = await _read_json(a) + + self.assertEqual((rows["type"], rows["length"]), ("infinite_resp", 5)) + self.assertFalse(first["final"], "rows came back while the stats were still running") + self.assertEqual(self._session("sw-stall").stats_status, "pending") + replies = [first] + await self._pull_stats(a, gen) + self.assertTrue(replies[-1]["final"]) + # Plan 1's proposed ceiling on the longest single request. + for reply in replies: + self.assertLess(reply["elapsed_ms"], 250) + self.assertLess(rows_latency, 0.25) + + @tornado.testing.gen_test + async def test_the_first_initial_state_precedes_any_data_stat_query_and_rows_precede_the_stats(self): + with _count_stat_queries() as queries: + await self._load("sw-first", stats_delivery="deferred") + ws, first = await self._connect("sw-first", caps="stats_update") + gen = self._stats(first)["gen"] + self._assert_pending(first, gen) + self.assertEqual((await self._infinite_window(ws))["length"], 5) + self.assertEqual(queries, [], "the first initial_state and the first rows precede every data stat query") + + self._budget(0) + ws.write_message(_stats_request(gen, incremental=True)) + reply = await _read_json(ws) + self.assertEqual((reply["type"], reply["final"]), ("stats_update", False)) + self.assertEqual(len(queries), 1) + self.assertEqual((await self._infinite_window(ws))["type"], "infinite_resp") + self.assertEqual(self._session("sw-first").stats_status, "pending", "rows are served before stats complete") + + @tornado.testing.gen_test + async def test_reload_expr_mid_run_bumps_the_generation_and_drops_the_run(self): + await self._load("sw-reload-run", stats_delivery="deferred") + ws, first = await self._connect("sw-reload-run", caps="stats_update") + gen = self._stats(first)["gen"] + session = self._session("sw-reload-run") + self._budget(0) + ws.write_message(_stats_request(gen, incremental=True)) + self.assertFalse((await _read_json(ws))["final"]) + self.assertEqual(list(session.stat_runs), [(gen, "raw")]) + + resp = await _post(self.get_http_port(), "/reload_expr/sw-reload-run", {}) + self.assertEqual(resp.code, 200, resp.body) + self._assert_pending(await _read_json(ws), gen + 1) + self.assertEqual(session.stat_runs, {}) + + with _count_stat_queries() as queries: + ws.write_message(_stats_request(gen, incremental=True)) + aborted = await _read_json(ws) + self.assertEqual((aborted["reason"], aborted["current_gen"]), ("stale", gen + 1)) + self.assertEqual(queries, []) + ws.write_message(_stats_request(gen + 1, incremental=True)) + reply = await _read_json(ws) + self.assertEqual((reply["remaining"], len(queries)), (3, 1), "the new generation starts from its first unit") + + @tornado.testing.gen_test + async def test_stats_unit_and_completion_spans_are_emitted_on_the_session_sink(self): + captured: list = [] + with patch.object(telemetry, "make_http_sink", lambda url, **kw: captured.append): + await self._load("sw-spans", stats_delivery="deferred", + telemetry_url="http://companion.invalid/internal/telemetry") + ws, first = await self._connect("sw-spans", caps="stats_update") + gen = self._stats(first)["gen"] + self._budget(0) + self.assertEqual(len(await self._pull_stats(ws, gen)), 4) + + units = [r for r in captured if r["name"] == "stats.unit"] + self.assertEqual([r["attrs"]["unit"] for r in units], + ["batch", "histogram:price", "histogram:qty", "histogram:category"]) + for record in units: + self.assertEqual(record["trace"], "sw-spans") + self.assertEqual((record["attrs"]["session"], record["attrs"]["stats_gen"]), ("sw-spans", gen)) + self.assertEqual([r["attrs"]["phase"] for r in units], ["batch", "histogram", "histogram", "histogram"]) + totals = [r for r in captured if r["name"] == "firstpull.stats_total"] + self.assertEqual(len(totals), 1, "one completion span for the run, not one for each request") + self.assertEqual((totals[0]["attrs"]["stats_gen"], totals[0]["attrs"]["units"]), (gen, 4)) + self.assertGreater(totals[0]["attrs"]["run_secs"], 0) + requests = [r for r in captured if r["name"] == "stats.request"] + self.assertEqual([r["attrs"]["outcome"] for r in requests], ["update"] * 4) + self.assertEqual([r["attrs"]["units"] for r in requests], [1, 1, 1, 1]) + + @tornado.testing.gen_test + async def test_the_completion_span_reports_the_snapshot_cache_outcome_of_the_units(self): + captured: list = [] + cache_root = tempfile.mkdtemp() + self.addCleanup(shutil.rmtree, cache_root, ignore_errors=True) + with patch.object(telemetry, "make_http_sink", lambda url, **kw: captured.append): + await self._load("sw-cache", stats_delivery="deferred", cache_storage_path=cache_root, + telemetry_url="http://companion.invalid/internal/telemetry") + ws, first = await self._connect("sw-cache", caps="stats_update") + gen = self._stats(first)["gen"] + self._budget(3600) + self.assertEqual(len(await self._pull_stats(ws, gen)), 1) + # A reload builds a new dataflow with an empty in-process cache, so the + # same state's units now read their snapshots. + resp = await _post(self.get_http_port(), "/reload_expr/sw-cache", {}) + self.assertEqual(resp.code, 200, resp.body) + await _read_json(ws) + self.assertEqual(len(await self._pull_stats(ws, gen + 1)), 1) + + cold, warm = [r["attrs"] for r in captured if r["name"] == "firstpull.stats_total"] + self.assertEqual((cold["cache_status"], cold["cache_hits"], cold["cache_misses"]), ("miss", 0, 4)) + self.assertEqual((warm["cache_status"], warm["cache_hits"], warm["cache_misses"]), ("hit", 4, 0)) + + @tornado.testing.gen_test + async def test_a_unit_that_raises_fails_the_generation_until_the_next(self): + await self._load("sw-unit-fail", stats_delivery="deferred") + ws, first = await self._connect("sw-unit-fail", caps="stats_update") + gen = self._stats(first)["gen"] + session = self._session("sw-unit-fail") + self._wrapped_run("sw-unit-fail", fail_on="histogram:price") + self._budget(0) + + ws.write_message(_stats_request(gen, incremental=True)) + self.assertFalse((await _read_json(ws))["final"]) + ws.write_message(_stats_request(gen, incremental=True)) + aborted = await _read_json(ws) + self.assertEqual((aborted["type"], aborted["stats_gen"], aborted["reason"]), ("stats_aborted", gen, "error")) + self.assertEqual((session.stats_status, session.stats_reason), ("error", "stats_failed")) + self.assertEqual(session.stat_runs, {}) + + with _count_stat_queries() as queries: + ws.write_message(_stats_request(gen, incremental=True)) + self.assertEqual((await _read_json(ws))["reason"], "error") + self.assertEqual(queries, [], "a failed generation is not retried by every request") + _, legacy = await self._connect("sw-unit-fail") + self.assertEqual(list(_rows_by_stat(legacy["df_data_dict"]["all_stats"])), ["dtype"]) + + ws.write_message(_state_change(quick_command_args={"search": ["a"]})) + self._assert_pending(await _read_json(ws), gen + 1) + self.assertEqual(len(await self._pull_stats(ws, gen + 1)), 4) + + @tornado.testing.gen_test + async def test_a_non_boolean_incremental_flag_is_aborted_as_a_bad_request(self): + await self._load("sw-bad-flag", stats_delivery="deferred") + ws, first = await self._connect("sw-bad-flag", caps="stats_update") + gen = self._stats(first)["gen"] + with _count_stat_queries() as queries: + ws.write_message(_stats_request(gen, incremental="yes")) + aborted = await _read_json(ws) + self.assertEqual((aborted["type"], aborted["reason"], aborted["stats_gen"]), + ("stats_aborted", "bad_request", gen)) + self.assertEqual(queries, [], "a malformed request must not fall back to the whole run") + self._budget(0) + self.assertFalse((await self._pull_stats(ws, gen, limit=1))[0]["final"], + "the session still serves a good request") + def _wide_parquet_expr(tmp_path, name="wide"): """Twelve rows, six columns, read from parquet. Each string column has a From ad26fdf464651a14ab75e0650a11171ba0e597e9 Mon Sep 17 00:00:00 2001 From: Paddy Mullen Date: Sun, 4 Oct 2026 04:13:18 -0400 Subject: [PATCH 18/18] feat(stats): stats_request runs units for a time budget with per-client cursors and the final assignment (rows-first s5) A stats_request with incremental: true runs the units of the session's StatRun for about STATS_BUDGET_S (at least one) and answers with the fragments the asking client has not seen, read through the handler's own StatCursor. The request that runs the last unit does the final assignment in complete_stats: the full-tier summary_stats_cache entry, summary_sd, the session snapshot and the status in one step, the run freed, and the final reply carries the rebuilt df_display_args when its digest differs from the one the client holds. Without the field a request is still the whole run, now finishing the run's remaining units when a client has made progress. A legacy client's complete frame does the same. Co-Authored-By: Claude Sonnet 5.5 --- buckaroo/dataflow/dataflow.py | 46 ++- buckaroo/server/stat_run.py | 13 +- buckaroo/server/stats_wire.py | 316 ++++++++++++++++-- buckaroo/server/websocket_handler.py | 25 +- tests/unit/server/test_data_loading_polars.py | 17 +- tests/unit/server/test_load_expr.py | 57 +++- 6 files changed, 400 insertions(+), 74 deletions(-) diff --git a/buckaroo/dataflow/dataflow.py b/buckaroo/dataflow/dataflow.py index 1a61b3bc1..ea1650d9d 100644 --- a/buckaroo/dataflow/dataflow.py +++ b/buckaroo/dataflow/dataflow.py @@ -565,21 +565,31 @@ def setup_options_from_analysis(self): @observe('summary_sd', 'processed_result', 'clean_sd_key', 'filt_sd_key') @exception_protect('merged_sd-protector') def _merged_sd(self, change): - # Bare keys come from the raw scope's SD (computed on - # sampled_df). ``cleaned_*`` keys are layered on top from the - # clean scope's SD when cleaning is active; ``filtered_*`` keys - # are layered on top from the filt scope's SD when a search - # filter is active. Scope SDs are read from the keyed cache - # (#783) — the cache observer ``_populate_sd_cache`` is what - # computes and stores them; this observer just assembles the - # wire shape #777's `?key` JS consumes. - # # clean_sd_key and filt_sd_key are in the observed set so this # fires after ``_populate_sd_cache`` has updated the pointers # (which it always does, even on a pure cache hit) — guarantees - # the cache lookups below see the right keys for the current - # state. - + # the cache lookups in ``_assemble_merged_sd`` see the right keys + # for the current state. + self.merged_sd = self._assemble_merged_sd() + + def _assemble_merged_sd(self, running_sd: Optional[SDType] = None) -> SDType: + """The wire ``merged_sd`` for the current state. + + Bare keys come from the raw scope's SD (computed on + sampled_df). ``cleaned_*`` keys are layered on top from the + clean scope's SD when cleaning is active; ``filtered_*`` keys + are layered on top from the filt scope's SD when a search + filter is active. Scope SDs are read from the keyed cache + (#783) — the cache observer ``_populate_sd_cache`` is what + computes and stores them; this only assembles the wire shape + #777's `?key` JS consumes. + + ``running_sd`` is an SD still being computed for the filt scope's + frame (a ``StatRun``'s accumulated results). It stands in for every + scope whose chain is the filt chain, which are the scopes the cache + holds one entry for, so a caller can assemble the wire shape of + stats that are not in the cache yet. + """ # Resolve scope SDs. Falls back to summary_sd / cleaned_sd # for pre-cache-population states (initial startup, the brief # window before _populate_sd_cache has fired). @@ -594,9 +604,17 @@ def _merged_sd(self, change): if filt_sd is None: filt_sd = self.summary_sd or {} - self.merged_sd = assemble_merged_sd( + chains = split_chain_by_scope(self.operations) + if running_sd is not None: + if chains['raw'] == chains['filt']: + raw_sd = running_sd + if chains['clean'] == chains['filt']: + clean_sd = running_sd + filt_sd = running_sd + + return assemble_merged_sd( self.init_sd, self.cleaned_sd, raw_sd, self.processed_sd, self.processed_df, - split_chain_by_scope(self.operations), clean_sd=clean_sd, filt_sd=filt_sd) + chains, clean_sd=clean_sd, filt_sd=filt_sd) def _compute_scope_df(self, scope: str): """Return the df that scope's SD should be computed against. diff --git a/buckaroo/server/stat_run.py b/buckaroo/server/stat_run.py index e2431297e..82a13fba7 100644 --- a/buckaroo/server/stat_run.py +++ b/buckaroo/server/stat_run.py @@ -15,6 +15,7 @@ from typing import Any, List, Optional, Sequence, Tuple from buckaroo.df_util import old_col_new_col +from buckaroo.pluggable_analysis_framework.stat_pipeline import errors_to_errdict from buckaroo.pluggable_analysis_framework.stat_units import Fragment, StatUnit, resolve_names, rewrite_sd _run_ids = itertools.count(1) @@ -27,7 +28,8 @@ class StatRun: ``status`` is ``pending`` until the last unit has run, then ``complete``; ``error`` once a unit raised, 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 - still finishes. + still finishes. ``elapsed_s`` is the time the units have taken, summed over + every request that ran one. """ def __init__(self, stats_gen: int, scope: str, stats: Any): @@ -42,6 +44,7 @@ def __init__(self, stats_gen: int, scope: str, stats: Any): self.status = "pending" if self.units else "complete" self.error: Optional[Exception] = None self.created_at = time.time() + self.elapsed_s = 0.0 self._frame_columns: Optional[List[Tuple[Any, str]]] = None @property @@ -81,11 +84,14 @@ def run_next(self, prefer: Sequence[Any] = (), namespace: str = "any") -> Option if unit is None: self.status = "complete" return None + started = time.perf_counter() try: fragment = self.stats.run(unit, self.acc) except Exception as exc: self.status, self.error = "error", exc raise + finally: + self.elapsed_s += time.perf_counter() - started self.ran.append(unit.id) self.fragments.append(fragment) if self.remaining == 0: @@ -97,6 +103,11 @@ def raw_sd(self) -> dict: rewritten column name: what ``summary_sd`` holds after a whole run.""" return rewrite_sd(self.acc.sd(), self.acc.state.data) + def errs(self) -> dict: + """The stats that failed inside a unit, as the ``{(col, stat): (error, + None)}`` dict the dataflow's ``errs`` holds.""" + return errors_to_errdict(self.acc.errors) + class StatCursor: """One client's position in a ``StatRun``'s fragment list. It follows one diff --git a/buckaroo/server/stats_wire.py b/buckaroo/server/stats_wire.py index b019af1fd..950daf29a 100644 --- a/buckaroo/server/stats_wire.py +++ b/buckaroo/server/stats_wire.py @@ -1,4 +1,4 @@ -"""Server side of the stats wire protocol (rows-first s3). +"""Server side of the stats wire protocol (rows-first s3 and s5). A deferred session (``stats_delivery="deferred"``) publishes its dataflow at the schema tier and leaves the rest of the stats to later. Two kinds of client @@ -17,21 +17,32 @@ which calls it per client), because the session holds one shared snapshot and the client is known only to the handler that owns the connection. -The request here is the whole run: one synchronous call that computes the full -stats and applies the final assignment (``assign_full_stats``, -``refresh_session_snapshot``). Resumable units generalize it later. +A ``stats_request`` takes one of two shapes. Without ``incremental`` it is the +whole run: one synchronous call that runs every unit still to run, applies the +final assignment and answers with the complete ``all_stats``. With +``incremental: true`` it runs the units of the session's ``StatRun`` for about +``STATS_BUDGET_S`` (at least one, so a cold unit is the only one of its +request) and answers with the fragments this client has not seen. The request +that runs the last unit applies the final assignment (``complete_stats``): the +full-tier cache entry, ``summary_sd``, the session snapshot and the status, +in one step. A client that follows the run reads the same list through its own +cursor, so work is done once whoever asks. """ +import copy +import hashlib import json import logging import time import traceback -from typing import Any, Optional +from typing import Any, Callable, Dict, Optional, Sequence from buckaroo.dataflow.sd_cache import split_chain_by_scope +from buckaroo.df_util import old_col_new_col from buckaroo.pluggable_analysis_framework import perf_log +from buckaroo.pluggable_analysis_framework.stat_units import Fragment from buckaroo.server.data_loading import get_buckaroo_display_state from buckaroo.server.session import SessionState, build_state_message -from buckaroo.server.stat_run import StatRun +from buckaroo.server.stat_run import StatCursor, StatRun log = logging.getLogger("buckaroo.server.stats_wire") @@ -47,6 +58,13 @@ # generation. STATS_SCOPES = ("raw",) +# The time an incremental ``stats_request`` may spend running units, in +# seconds. A request always runs one unit and stops starting new ones once this +# is spent, so a snapshot-cache hit (milliseconds) shares its request with the +# others and a query is alone in its own. It cannot cut a unit short: the +# longest unit is what a waiting ``infinite_request`` stalls behind. +STATS_BUDGET_S = 0.075 + def parse_caps(raw: str) -> frozenset: """The capabilities in a ``?caps=a,b`` value.""" @@ -102,11 +120,88 @@ def start_stat_run(session: SessionState, scope: str = "raw") -> Optional[StatRu return run -def assign_full_stats(dataflow: Any) -> None: - """Take a schema-tier dataflow to the full tier: compute the current - state's full stats (or find them in ``summary_stats_cache``, where an - earlier visit to the same state left them), write them under the full-tier - key, then assign ``summary_sd``. Assigning runs the cascade, which fills the +def display_args_hash(display_args: Any) -> str: + """A digest of a ``df_display_args`` value, equal for equal content whatever + the key order. JSON, because that is how the value reaches a client, and it + reads a NaN the same each time.""" + return hashlib.sha1(json.dumps(display_args, sort_keys=True, default=str).encode()).hexdigest() + + +def highlighted_display_args(display_args: dict, term: str) -> dict: + """A copy of ``df_display_args`` with ``term`` set as the highlight phrase + of every string column's displayer, or with the highlight removed when the + term is empty (#851). A copy, so the shared session snapshot is never + changed.""" + overlay = copy.deepcopy(display_args) + for dva in overlay.values(): + dvc = (dva or {}).get("df_viewer_config") or {} + for col in dvc.get("column_config", []) or []: + disp = col.get("displayer_args") + if not isinstance(disp, dict) or disp.get("displayer") != "string": + continue + if term: + disp["highlight_phrase"] = [term] + else: + disp.pop("highlight_phrase", None) + return overlay + + +def run_units(run: StatRun, budget_s: Optional[float], prefer: Sequence[str] = (), + clock: Callable[[], float] = time.perf_counter, session_id: Optional[str] = None) -> int: + """Run units of ``run`` for one request, and return how many. + + At least one runs while the run is pending. Another starts only while the + time spent is under ``budget_s``, so units that cost milliseconds (their + queries are snapshot-cache hits) run together and one that is a query ends + its request. ``budget_s=None`` runs every unit still to run. ``prefer`` are + column names as a client writes them (the rewritten ``a, b, c``), whose + units go first. A unit that raises fails the run and the exception + propagates. Each unit is a ``stats.unit`` span.""" + started = clock() + ran = 0 + while run.status == "pending": + if ran and budget_s is not None and clock() - started >= budget_s: + break + unit = run.next_unit(prefer, "rewritten") + if unit is None: + run.run_next() # no unit is left: this marks the run complete + break + with perf_log.perf_span("stats.unit", session=session_id, stats_gen=run.stats_gen, unit=unit.id, + phase=unit.phase, cost=unit.cost, columns=len(unit.columns)): + run.run_next(prefer, "rewritten") + ran += 1 + return ran + + +def partial_payload(dataflow: Any, run: StatRun, fragments: Sequence[Fragment]) -> Any: + """The ``all_stats`` payload (an inline wide ``DFEnvelope``) of the columns + ``fragments`` cover: their stats so far, assembled as ``merged_sd`` is, so + ``init_sd``, a processing step's sd and the ``cleaned_*`` and ``filtered_*`` + layers apply as they will in the complete state. A client merges it into + its ``all_stats`` key by key, so a stat sent again with a later fragment + only replaces itself.""" + rewritten = dict(old_col_new_col(run.acc.state.data)) + touched = {rewritten[col] for fragment in fragments for col in fragment if col in rewritten} + merged = dataflow._assemble_merged_sd(run.raw_sd()) + return dataflow._sd_to_jsondf({col: stats for col, stats in merged.items() if col in touched}) + + +def _full_stats_key(dataflow: Any) -> Any: + """The ``summary_stats_cache`` key of the current state's full-tier sd.""" + return dataflow._scope_cache_key(split_chain_by_scope(dataflow.operations)["filt"], tier="full") + + +def _cached_full_sd(dataflow: Any) -> Any: + """The full-tier sd of the current state if an earlier visit to it left one.""" + return dataflow.summary_stats_cache.get(_full_stats_key(dataflow)) + + +def assign_full_stats(dataflow: Any, computed: Optional[tuple] = None) -> None: + """Take a schema-tier dataflow to the full tier: use the current state's + full stats (found in ``summary_stats_cache``, where an earlier visit to the + same state left them, or ``computed``, the ``(sd, errs)`` of a finished + ``StatRun``, or computed here whole), write them under the full-tier key, + then assign ``summary_sd``. Assigning runs the cascade, which fills the other scopes' entries it still lacks and rebuilds ``merged_sd``, ``df_data_dict`` and ``df_display_args``. @@ -115,14 +210,13 @@ def assign_full_stats(dataflow: Any) -> None: plain attribute, flipped before the compute so ``_get_summary_sd`` runs the full pipeline, and put back if anything raises.""" tier = dataflow.stats_tier - filt_chain = split_chain_by_scope(dataflow.operations)["filt"] - key = dataflow._scope_cache_key(filt_chain, tier="full") + key = _full_stats_key(dataflow) dataflow.stats_tier = "full" try: sd = dataflow.summary_stats_cache.get(key) errs = {} if sd is None: - sd, errs = dataflow._get_summary_sd(dataflow.processed_df) + sd, errs = computed if computed is not None else dataflow._get_summary_sd(dataflow.processed_df) dataflow.summary_stats_cache = {**dataflow.summary_stats_cache, key: sd} dataflow.summary_sd = sd dataflow.errs = errs @@ -131,12 +225,48 @@ def assign_full_stats(dataflow: Any) -> None: raise +def _run_summary(dataflow: Any, run: StatRun) -> tuple: + """The ``(sd, errs)`` of a finished run, as ``_get_summary_sd`` returns them + (it raises on a failed stat in debug mode, and so does this).""" + errs = run.errs() + if errs and dataflow.debug: + raise Exception("Error executing analysis") + return run.raw_sd(), errs + + +def _run_span_attrs(run: StatRun) -> dict: + """What the completion span says about a run that spanned requests: the + units it took and their summed time, and for a backend with a snapshot + cache (xorq) the outcome ``firstpull.summary_stats`` carries for a whole + run (#943).""" + attrs: Dict[str, Any] = {"units": len(run.units), "run_secs": round(run.elapsed_s, 4)} + cache_run_stats = getattr(run.stats, "cache_run_stats", None) + if cache_run_stats is not None: + cs = cache_run_stats() + attrs.update(cache_status=cs["status"], cache_hits=cs["hits"], cache_misses=cs["misses"], + cache_secs=attrs["run_secs"], cache_snapshots=cs["snapshots"], cache_bytes=cs["bytes"], + cache_write_errors=cs["write_errors"]) + return attrs + + +def _fail_stats(session: SessionState) -> None: + """A stats run failed: that is the session's state for this generation, so + no request retries it and the next generation starts clean.""" + session.stats_status, session.stats_reason = "error", "stats_failed" + session.stat_runs.clear() + + def complete_stats(session: SessionState) -> bool: - """Run the stats a pending session is missing, in one synchronous call, and - publish them: the final assignment (``assign_full_stats``), then the session - snapshot refreshed and the status set to ``complete`` in the same step, so - ``all_stats`` and ``df_meta.stats`` cannot disagree. Returns whether the - session is complete. + """Run the stats a pending session is missing and publish them: the final + assignment (``assign_full_stats``), then the session snapshot refreshed and + the status set to ``complete`` in the same step, so ``all_stats`` and + ``df_meta.stats`` cannot disagree. Returns whether the session is complete. + + When the generation has a ``StatRun``, its remaining units run (none if it + is finished) and its results are what is assigned, so a unit that a client + already ran is not run again. With no run the whole computation is one + ``_get_summary_sd`` call, which runs the same units in the same order. The + run is freed with the assignment. A failure is the session's state for this generation (``error``, reason ``stats_failed``) and is not retried by the next request; the next @@ -146,18 +276,26 @@ def complete_stats(session: SessionState) -> bool: dataflow = session_dataflow(session) if session.stats_status != "pending" or dataflow is None: return False + run_key = (session.stats_gen, "raw") with ( perf_log.telemetry_context(session.session_id, session.tele_sink), - perf_log.perf_span("firstpull.stats_total", session=session.session_id, stats_gen=session.stats_gen), + perf_log.perf_span("firstpull.stats_total", session=session.session_id, stats_gen=session.stats_gen) as span, ): try: - assign_full_stats(dataflow) + run = session.stat_runs.get(run_key) + if run is not None and _cached_full_sd(dataflow) is None: + run_units(run, None, session_id=session.session_id) + span.set_attr(**_run_span_attrs(run)) + assign_full_stats(dataflow, computed=_run_summary(dataflow, run)) + else: + assign_full_stats(dataflow) refresh_session_snapshot(session, dataflow) except Exception: log.error("stats run failed session=%s stats_gen=%s: %s", session.session_id, session.stats_gen, traceback.format_exc()) - session.stats_status, session.stats_reason = "error", "stats_failed" + _fail_stats(session) return False + session.stat_runs.pop(run_key, None) session.stats_status, session.stats_reason = "complete", None return True @@ -169,10 +307,17 @@ def build_state_message_for(session: SessionState, client: Any, metadata: Option session gets its missing stats run first, so its message is complete; a capable client gets the session snapshot as it is (stats-free while the session is pending) and pulls the rest. The search term is the recipient's - own (#851).""" + own (#851). + + A capable client is also told apart from the display config it holds: the + digest of the one a pending frame carried is kept on it, so the final + ``stats_update`` can say whether the config it would send differs.""" if (session.stats_delivery == "deferred" and session.stats_status == "pending" and not client_has_cap(client, STATS_UPDATE_CAP)): complete_stats(session) + if client_has_cap(client, STATS_UPDATE_CAP): + client.display_args_hash = (display_args_hash(session.df_display_args) + if session.stats_status == "pending" else None) return build_state_message(session, metadata=metadata, search_string=getattr(client, "search_string", "")) @@ -204,7 +349,80 @@ def _aborted(stats_gen: Any, scope: Any, reason: str, session: Optional[SessionS return msg -def _answer_stats_request(session: Optional[SessionState], stats_gen: Any, scope: Any, started: float) -> dict: +def _elapsed_ms(started: float) -> float: + return round((time.perf_counter() - started) * 1000, 1) + + +def _plan_run(session: SessionState) -> Optional[StatRun]: + """The session's ``StatRun`` for the current generation, or ``None`` when it + cannot be planned (there is no frame, or the dataflow's stats class cannot + plan the one it has: a post-processor that failed leaves an error frame). + The whole-run path then decides what the stats are.""" + try: + return start_stat_run(session) + except Exception: + log.warning("stat units not planned session=%s stats_gen=%s: %s", session.session_id, session.stats_gen, + traceback.format_exc()) + return None + + +def _serve_units(session: SessionState, client: Any, prefer: Sequence[str], stats_gen: Any, scope: Any, + started: float, info: dict) -> Optional[dict]: + """The reply to an incremental request on a pending session, when the + answer is the fragments this client has not seen: a partial ``stats_update`` + (``final`` false), or ``stats_aborted`` when a unit failed. ``None`` when + the request is not answered that way: the stats are cached or cannot be + planned (the whole-run path takes them), or the run is now finished and the + caller publishes it. + + A client that is behind the run is caught up and no unit runs for it, so + work is done once and a client that asks late reads it all from its own + cursor.""" + dataflow = session_dataflow(session) + if _cached_full_sd(dataflow) is not None: + return None + run = _plan_run(session) + if run is None: + return None + cursor = getattr(client, "stats_cursor", None) or StatCursor() + if cursor.caught_up(run): + try: + info["units"] = run_units(run, STATS_BUDGET_S, prefer, session_id=session.session_id) + except Exception: + log.error("stat unit failed session=%s stats_gen=%s: %s", session.session_id, session.stats_gen, + traceback.format_exc()) + _fail_stats(session) + return _aborted(stats_gen, scope, "error", session) + fragments = cursor.take(run) + if run.status != "pending": + return None + return {"type": "stats_update", "stats_gen": stats_gen, "scope": scope, "tier": session.stats_tier, + "final": False, "remaining": run.remaining, "payload": partial_payload(dataflow, run, fragments), + "elapsed_ms": _elapsed_ms(started)} + + +def _rebuilt_display_args(session: SessionState, client: Any) -> Optional[dict]: + """The session's display config, when it differs from the one this client + holds: the stats-derived parts of a config (a float column's ``minWidth``) + change when the stats arrive. The client's own highlight is kept on it, + since the config replaces the one that carries it. ``None`` when the client + holds the current config, or there is no record of one to compare with.""" + held = getattr(client, "display_args_hash", None) + if held is None: + return None + client.display_args_hash = None + if display_args_hash(session.df_display_args) == held: + return None + term = getattr(client, "search_string", "") + return highlighted_display_args(session.df_display_args, term) if term else session.df_display_args + + +def _answer_stats_request(session: Optional[SessionState], msg: dict, client: Any, started: float, + info: dict) -> dict: + stats_gen, scope = msg.get("stats_gen"), msg.get("scope", "raw") + columns = msg.get("columns") + prefer = [c for c in columns if isinstance(c, str)] if isinstance(columns, list) else [] + incremental = msg.get("incremental") dataflow = session_dataflow(session) if session is None or dataflow is None: return _aborted(stats_gen, scope, "no_data") @@ -212,22 +430,47 @@ def _answer_stats_request(session: Optional[SessionState], stats_gen: Any, scope return _aborted(stats_gen, scope, "stale", session) if scope not in STATS_SCOPES: return _aborted(stats_gen, scope, "unsupported_scope", session) + if incremental is not None and not isinstance(incremental, bool): + return _aborted(stats_gen, scope, "bad_request", session) if session.stats_status == "not_computed": return _aborted(stats_gen, scope, "not_requestable", session) - if session.stats_status == "error" or (session.stats_status == "pending" and not complete_stats(session)): + if session.stats_status == "pending" and incremental: + reply = _serve_units(session, client, prefer, stats_gen, scope, started, info) + if reply is not None: + return reply + if session.stats_status == "pending": + complete_stats(session) + if session.stats_status != "complete": return _aborted(stats_gen, scope, "error", session) # Complete: from here the answer is the dataflow's own all_stats, with no # query, whether this request ran the stats or an earlier one did. - return {"type": "stats_update", "stats_gen": stats_gen, "scope": scope, "tier": dataflow.stats_tier, - "final": True, "payload": session.df_data_dict["all_stats"], - "elapsed_ms": round((time.perf_counter() - started) * 1000, 1)} + reply = {"type": "stats_update", "stats_gen": stats_gen, "scope": scope, "tier": dataflow.stats_tier, + "final": True, "remaining": 0, "payload": session.df_data_dict["all_stats"], + "elapsed_ms": _elapsed_ms(started)} + display_args = _rebuilt_display_args(session, client) + if display_args is not None: + reply["df_display_args"] = display_args + return reply -def handle_stats_request(session: Optional[SessionState], msg: dict) -> dict: - """Answer a ``stats_request {stats_gen, scope, columns?}``: a - ``stats_update`` carrying the complete ``all_stats`` as an inline wide - ``DFEnvelope`` (self-contained, so binary pairing stays single-slot), or a - ``stats_aborted``. ``columns`` is a hint a whole-run request ignores. +def handle_stats_request(session: Optional[SessionState], msg: dict, client: Any = None) -> dict: + """Answer a ``stats_request {stats_gen, scope, columns?, incremental?}``: a + ``stats_update`` or a ``stats_aborted``. ``client`` is the handler that + received it, which holds the connection's cursor into the run and the + display config it was last sent. + + Without ``incremental`` the request is the whole run: every unit still to + run, then a ``stats_update`` that is ``final`` and carries the complete + ``all_stats`` as an inline wide ``DFEnvelope`` (self-contained, so binary + pairing stays single-slot). With ``incremental: true`` it is a time-boxed + step (see ``_serve_units``): a ``stats_update`` with ``final`` false + carries the stats of the columns this client has not seen, ``remaining`` + counts the units left, and the reply that follows the last unit is the + ``final`` one, with the complete ``all_stats`` and, when the stats changed + it, the rebuilt ``df_display_args``. ``columns`` are the grid's column + names (``a, b, c``), a hint for which units go first, and a whole-run + request ignores it. A malformed request is ``stats_aborted`` with reason + ``bad_request``, and is not run as a whole run. The ``stats.request`` span records the request and its ``outcome`` (``update`` or the abort reason), which is where updates sent, requests @@ -235,16 +478,19 @@ def handle_stats_request(session: Optional[SessionState], msg: dict) -> dict: telemetry sink around this call.""" stats_gen, scope, columns = msg.get("stats_gen"), msg.get("scope", "raw"), msg.get("columns") started = time.perf_counter() + info: dict = {} with perf_log.perf_span("stats.request", session=session.session_id if session else None, stats_gen=stats_gen, scope=scope, columns=len(columns) if isinstance(columns, list) else None) as span: try: - reply = _answer_stats_request(session, stats_gen, scope, started) + reply = _answer_stats_request(session, msg, client, started, info) except Exception: log.error("stats_request error session=%s: %s", session.session_id if session else None, traceback.format_exc()) reply = _aborted(stats_gen, scope, "error", session) if reply["type"] == "stats_update": - span.set_attr(outcome="update", tier=reply["tier"]) + span.set_attr(outcome="update", tier=reply["tier"], final=reply["final"], remaining=reply["remaining"]) else: span.set_attr(outcome=reply["reason"]) + if "units" in info: + span.set_attr(units=info["units"]) return reply diff --git a/buckaroo/server/websocket_handler.py b/buckaroo/server/websocket_handler.py index 7a5289917..ec394f599 100644 --- a/buckaroo/server/websocket_handler.py +++ b/buckaroo/server/websocket_handler.py @@ -1,4 +1,3 @@ -import copy import json import logging import os @@ -12,7 +11,7 @@ from buckaroo.server.data_loading import (handle_infinite_request, handle_infinite_request_buckaroo, handle_infinite_request_lazy) from buckaroo.server.session import begin_stats_generation, dataflow_stats_tier from buckaroo.server.stat_run import StatCursor -from buckaroo.server.stats_wire import (broadcast_state, build_state_message_for, handle_stats_request, parse_caps, refresh_session_snapshot) +from buckaroo.server.stats_wire import (broadcast_state, build_state_message_for, handle_stats_request, highlighted_display_args, parse_caps, refresh_session_snapshot) def _handle_infinite_request_xorq(xorq_dataflow, payload_args, search_string=""): @@ -48,6 +47,10 @@ def open(self, session_id): # This client's position in the session's StatRun fragment list (the # run itself is shared); it ends with the connection. self.stats_cursor = StatCursor() + # Digest of the df_display_args a stats-free frame last sent this + # client, kept so the final stats_update can say whether the rebuilt + # config differs (stats_wire.build_state_message_for sets it). + self.display_args_hash = None sessions = self.application.settings["sessions"] sessions.add_ws_client(session_id, self) @@ -75,7 +78,9 @@ def on_message(self, message): def _handle_stats_request(self, msg): """Answer a client's ``stats_request`` with a ``stats_update`` or a ``stats_aborted``. Synchronous, like ``infinite_request``: the request - runs the whole stats computation in this call (see ``stats_wire``). + runs its units in this call, a time-boxed few for an incremental one and + every unit left for the others (see ``stats_wire``), so a waiting + ``infinite_request`` is served between requests. This branch is its own async context, so the session's telemetry sink is bound here for the ``stats.request`` span and the stats spans under @@ -83,7 +88,7 @@ def _handle_stats_request(self, msg): sessions = self.application.settings["sessions"] session = sessions.get(self.session_id) with perf_log.telemetry_context(self.session_id, session.tele_sink if session else None): - reply = handle_stats_request(session, msg) + reply = handle_stats_request(session, msg, self) self.write_message(json.dumps(reply)) def _handle_buckaroo_state_change(self, new_state): @@ -189,17 +194,7 @@ def _send_highlight_overlay(self, session): # #854 — without it the JS clears the search box on every # keystroke). msg = build_state_message_for(session, self) - overlay = copy.deepcopy(session.df_display_args) - for dva in overlay.values(): - dvc = (dva or {}).get("df_viewer_config") or {} - for col in dvc.get("column_config", []) or []: - disp = col.get("displayer_args") - if not isinstance(disp, dict) or disp.get("displayer") != "string": - continue - if term: - disp["highlight_phrase"] = [term] - else: - disp.pop("highlight_phrase", None) + overlay = highlighted_display_args(session.df_display_args, term) msg["df_display_args"] = overlay try: diff --git a/tests/unit/server/test_data_loading_polars.py b/tests/unit/server/test_data_loading_polars.py index 089ad0743..ed4bcfd16 100644 --- a/tests/unit/server/test_data_loading_polars.py +++ b/tests/unit/server/test_data_loading_polars.py @@ -719,6 +719,12 @@ def _costed_run(backend, clock, cost, costs=None, gen=4): return run, dataflow +def _without_memory_usage(sd): + """``sd`` without the ``memory_usage`` stats, which pandas reads off the + index and so depend on what ran on the frame before.""" + return {col: {k: v for k, v in stats.items() if not k.endswith("memory_usage")} for col, stats in sd.items()} + + def _stat_rows(payload): """The decoded ``all_stats`` rows of a payload, keyed by stat name.""" return {row["index"]: row for row in resolve_summary_stats_payload(payload)} @@ -780,9 +786,9 @@ def test_the_columns_hint_is_read_in_the_client_s_rewritten_names(self, backend) """The grid knows only a, b, c: "c" is the third column, category.""" clock = _FakeClock() run, _dataflow = _costed_run(backend, clock, cost=1.0) - stats_wire.run_units(run, 0, prefer=["c", "a"], clock=clock) - stats_wire.run_units(run, 0, prefer=["c", "a"], clock=clock) - assert run.ran == ["column:category", "column:price"] + stats_wire.run_units(run, 0, prefer=["c"], clock=clock) + stats_wire.run_units(run, 0, prefer=["b"], clock=clock) + assert run.ran == ["column:category", "column:qty"] def test_a_unit_that_raises_fails_the_run_and_the_error_propagates(self, backend): clock = _FakeClock() @@ -841,7 +847,7 @@ def test_the_payload_holds_only_the_columns_the_fragments_cover(self, backend): run, _dataflow = _stat_run(backend) fragment = run.run_next(prefer=("category",)) payload = stats_wire.partial_payload(_dataflow, run, [fragment]) - columns = {name for row in _stat_rows(payload).values() for name in row if name != "index"} + columns = {name for row in _stat_rows(payload).values() for name in row if name not in ("index", "level_0")} assert columns == {"c"}, "category is the third column, rewritten to c" def test_a_search_filter_keys_the_run_s_stats_as_filtered(self, backend): @@ -851,7 +857,8 @@ def test_a_search_filter_keys_the_run_s_stats_as_filtered(self, backend): dataflow = _three_scope_dataflow(backend, _UNTIED_DATA) run = StatRun(4, "raw", dataflow.build_stats(dataflow.processed_df, run=False)) stats_wire.run_units(run, None) - assert _as_json(dataflow._assemble_merged_sd(run.raw_sd())) == _as_json(dataflow.merged_sd) + assert _as_json(_without_memory_usage(dataflow._assemble_merged_sd(run.raw_sd()))) == _as_json( + _without_memory_usage(dataflow.merged_sd)) def test_init_sd_overrides_win_over_the_run_s_stats_as_in_merged_sd(self, backend): dataflow = _build_dataflow(backend, _frame(backend, _UNTIED_DATA), diff --git a/tests/unit/server/test_load_expr.py b/tests/unit/server/test_load_expr.py index 7b186fcb7..813955234 100644 --- a/tests/unit/server/test_load_expr.py +++ b/tests/unit/server/test_load_expr.py @@ -1452,7 +1452,7 @@ def _merge_stat_payloads(payloads): for stat, row in _rows_by_stat(payload).items(): cells = merged.setdefault(stat, {}) for column, value in row.items(): - if column == "index" or (value is None and cells.get(column) is not None): + if column in ("index", "level_0") or (value is None and cells.get(column) is not None): continue cells[column] = value return merged @@ -2142,7 +2142,8 @@ async def test_the_columns_hint_orders_the_units_of_a_request(self): run = self._session("sw-hint").stat_runs[(gen, "raw")] self.assertEqual(run.ran, ["batch", "histogram:category"], "the batch every histogram reads goes first, then the unit of the column named, c being category") - columns = {name for row in _rows_by_stat(reply["payload"]).values() for name in row if name != "index"} + columns = {name for row in _rows_by_stat(reply["payload"]).values() for name in row + if name not in ("index", "level_0")} self.assertEqual(columns, {"c"}) @tornado.testing.gen_test @@ -2453,8 +2454,8 @@ async def test_a_unit_that_raises_fails_the_generation_until_the_next(self): ws.write_message(_stats_request(gen, incremental=True)) self.assertEqual((await _read_json(ws))["reason"], "error") self.assertEqual(queries, [], "a failed generation is not retried by every request") - _, legacy = await self._connect("sw-unit-fail") - self.assertEqual(list(_rows_by_stat(legacy["df_data_dict"]["all_stats"])), ["dtype"]) + self.assertEqual(list(_rows_by_stat(session.df_data_dict["all_stats"])), ["dtype"], + "the snapshot keeps the schema tier") ws.write_message(_state_change(quick_command_args={"search": ["a"]})) self._assert_pending(await _read_json(ws), gen + 1) @@ -2475,6 +2476,54 @@ async def test_a_non_boolean_incremental_flag_is_aborted_as_a_bad_request(self): self.assertFalse((await self._pull_stats(ws, gen, limit=1))[0]["final"], "the session still serves a good request") + @tornado.testing.gen_test + async def test_an_incremental_request_for_a_state_with_cached_stats_runs_no_unit(self): + await self._load("sw-inc-cached", stats_delivery="deferred") + ws, first = await self._connect("sw-inc-cached", caps="stats_update") + gen = self._stats(first)["gen"] + self._budget(0) + self.assertEqual(len(await self._pull_stats(ws, gen)), 4) + ws.write_message(_state_change(quick_command_args={"search": ["a"]})) + self._assert_pending(await _read_json(ws), gen + 1) + self.assertEqual(len(await self._pull_stats(ws, gen + 1)), 4) + + ws.write_message(_state_change(quick_command_args={})) + self._assert_pending(await _read_json(ws), gen + 2) + with _count_stat_queries() as queries: + replies = await self._pull_stats(ws, gen + 2) + self.assertEqual([(r["final"], r["remaining"]) for r in replies], [(True, 0)]) + self.assertEqual(queries, [], "the first state's stats are in summary_stats_cache") + self.assertEqual(self._session("sw-inc-cached").stat_runs, {}) + + @tornado.testing.gen_test + async def test_an_incremental_request_that_cannot_plan_units_takes_the_whole_run_path(self): + """A post-processor that fails leaves an error frame, which the xorq + stats class cannot plan units for. The request is still answered, as the + whole run is.""" + with open(os.path.join(self.project_root, "post_processing", "boom.py"), "w") as f: + f.write("def process(expr):\n raise ValueError('boom')\n") + await self._load("sw-no-plan", stats_delivery="deferred") + ws, first = await self._connect("sw-no-plan", caps="stats_update") + gen = self._stats(first)["gen"] + ws.write_message(_state_change(post_processing="boom")) + pushed = await _read_json(ws) + self.assertEqual((self._stats(pushed)["status"], self._stats(pushed)["gen"]), ("pending", gen + 1)) + self._budget(0) + with self.assertLogs("buckaroo.server.stats_wire", level="WARNING") as logs: + replies = await self._pull_stats(ws, gen + 1, limit=3) + self.assertIn("stat units not planned", logs.output[0]) + self.assertEqual([(r["type"], r["final"]) for r in replies], [("stats_update", True)]) + self.assertEqual(self._session("sw-no-plan").stats_status, "complete") + + @tornado.testing.gen_test + async def test_a_columns_hint_that_is_not_a_list_of_names_is_only_a_hint(self): + await self._load("sw-junk-hint", stats_delivery="deferred") + ws, first = await self._connect("sw-junk-hint", caps="stats_update") + self._budget(0) + replies = await self._pull_stats(ws, self._stats(first)["gen"], columns=[["c"], 5, None, "c"]) + self.assertEqual([r["type"] for r in replies], ["stats_update"] * 4) + self.assertEqual(self._session("sw-junk-hint").stat_runs, {}) + def _wide_parquet_expr(tmp_path, name="wide"): """Twelve rows, six columns, read from parquet. Each string column has a