From a9419af0c1d11f8137b61d0f8246cdc484954268 Mon Sep 17 00:00:00 2001 From: Paddy Mullen Date: Sat, 3 Oct 2026 22:51:32 -0400 Subject: [PATCH 01/39] 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 c1391701449dbb838d4d8e89e9f574de845bf08e Mon Sep 17 00:00:00 2001 From: Paddy Mullen Date: Sat, 3 Oct 2026 22:53:59 -0400 Subject: [PATCH 02/39] test(server): failing tests for the stats policy, routing and probes (rows-first p31) New tests/unit/server/test_stats_policy.py for buckaroo/server/stats_policy.py, which does not exist yet: - resolve_stats_policy: a table of (backend, source_kind, rows, cols, host_tier) against (tier_target, auto_request, requestable, reason), including the boundaries of each threshold and the three tallyman entries over 70 s. - The ceiling holds for host_tier="full", for a reload that re-resolves against a larger row count, and for a force request, and every tier a result lists as requestable is one a force would be granted. - Probes add no data query: the parquet footer probe reads a small share of a counting file object and decodes no row group, the dtype probe never collects a LazyFrame or executes a xorq expression, and a known xorq count is an input. - route_polars_entry returns "xorq" above R and "eager" at or below it. - Threshold overrides from BUCKAROO_* environment variables. Nothing imports the module yet. The tests error at fixture setup on the missing module until the implementation lands. Co-Authored-By: Claude Sonnet 5.5 --- tests/unit/server/test_stats_policy.py | 421 +++++++++++++++++++++++++ 1 file changed, 421 insertions(+) create mode 100644 tests/unit/server/test_stats_policy.py diff --git a/tests/unit/server/test_stats_policy.py b/tests/unit/server/test_stats_policy.py new file mode 100644 index 000000000..6e2abf98c --- /dev/null +++ b/tests/unit/server/test_stats_policy.py @@ -0,0 +1,421 @@ +"""Unit tests for ``buckaroo.server.stats_policy`` (rows-first p31). + +Everything here is pure Python: the policy is a function of numbers, the +probes read a schema or a parquet footer, and nothing needs a server. The +module is not wired into any handler, session or dataflow in this phase. + +Thresholds are provisional, so the table tests pass an explicit +``StatsLimits`` instead of reading the module defaults. The few tests that +use the defaults assert only facts the plan requires of any calibration +(the 78M-row entry cannot reach ``full``). +""" +import importlib +import importlib.util +import io +import json +import logging + +import numpy as np +import pandas as pd +import polars as pl +import pyarrow as pa +import pyarrow.parquet as pq +import pytest + +ENV_FIELDS = {"BUCKAROO_STATS_FULL_AUTO_ROWS": "full_auto_rows", + "BUCKAROO_STATS_SCALAR_AUTO_CELLS": "scalar_auto_cells", "BUCKAROO_STATS_CEILING_FULL_ROWS": "ceiling_full_rows", + "BUCKAROO_STATS_CEILING_SCALAR_CELLS": "ceiling_scalar_cells", "BUCKAROO_POLARS_ROUTE_ROWS": "polars_route_rows"} + +TIER_RANK = {"schema": 0, "scalar": 1, "full": 2} + + +@pytest.fixture +def sp(): + spec = importlib.util.find_spec("buckaroo.server.stats_policy") + assert spec is not None, "buckaroo.server.stats_policy is not importable" + return importlib.import_module("buckaroo.server.stats_policy") + + +@pytest.fixture(autouse=True) +def clean_env(monkeypatch): + """A developer's own BUCKAROO_STATS_* overrides must not leak in.""" + for name in ENV_FIELDS: + monkeypatch.delenv(name, raising=False) + + +@pytest.fixture +def lim(sp): + """The plan's proposed starting values, pinned so a recalibration of the + module defaults does not move these tables.""" + return sp.StatsLimits( + full_auto_rows=10_000_000, + scalar_auto_cells=500_000_000, + ceiling_full_rows=50_000_000, + ceiling_scalar_cells=None, + polars_route_rows=10_000_000) + + +def _resolve(sp, lim, backend, source_kind, rows, cols, host_tier=None, bytes=None): + return sp.resolve_stats_policy( + backend, source_kind, rows, cols, bytes=bytes, host_tier=host_tier, limits=lim) + + +# (backend, source_kind, rows, cols, host_tier) -> (tier_target, auto_request, +# requestable, reason). The row counts of 42.3M, 54.1M and 78.0M and the 44 +# columns are the three tallyman entries over 70 s; 10.8M x 43 is parking_2017. +RESOLVE_TABLE = [ + # Eager backends resolve to full and ignore the host tier. + ("pandas", "memory", 52_814, 27, None, ("full", True, [], None)), + ("pandas", "parquet", 1_000_000, 43, None, ("full", True, [], None)), + ("polars", "parquet", 10_800_000, 43, None, ("full", True, [], None)), + ("polars", "parquet", 10_800_000, 43, "schema", ("full", True, [], None)), + ("pandas", "csv", 52_814, 27, "scalar", ("full", True, [], None)), + # xorq sized by the policy alone. + ("xorq", "parquet", 0, 0, None, ("full", True, [], None)), + ("xorq", "parquet", 52_814, 27, None, ("full", True, [], None)), + ("xorq", "parquet", 10_000_000, 43, None, ("full", True, [], None)), + ("xorq", "parquet", 10_000_001, 43, None, ("scalar", True, ["full"], "size")), + ("xorq", "parquet", 10_800_000, 43, None, ("scalar", True, ["full"], "size")), + ("xorq", "parquet", 11_363_636, 44, None, ("scalar", True, ["full"], "size")), + ("xorq", "parquet", 11_363_637, 44, None, ("schema", False, ["scalar", "full"], "size")), + ("xorq", "csv", 42_300_000, 44, None, ("schema", False, ["scalar", "full"], "size")), + ("xorq", "csv", 50_000_000, 44, None, ("schema", False, ["scalar", "full"], "size")), + ("xorq", "csv", 54_100_000, 44, None, ("schema", False, ["scalar"], "size")), + ("xorq", "csv", 78_000_000, 44, None, ("schema", False, ["scalar"], "size")), + ("xorq", "csv", 78_000_000, 44, "auto", ("schema", False, ["scalar"], "size")), + # The host lowers freely. + ("xorq", "parquet", 1_000_000, 43, "schema", ("schema", False, ["scalar", "full"], "host")), + ("xorq", "parquet", 1_000_000, 43, "scalar", ("scalar", True, ["full"], "host")), + ("xorq", "parquet", 1_000_000, 43, "full", ("full", True, [], None)), + # The host raises only as far as the ceiling. + ("xorq", "parquet", 12_000_000, 43, "full", ("full", True, [], None)), + ("xorq", "csv", 42_300_000, 44, "full", ("full", True, [], None)), + ("xorq", "csv", 50_000_000, 44, "full", ("full", True, [], None)), + ("xorq", "csv", 50_000_001, 44, "full", ("scalar", True, [], "ceiling")), + ("xorq", "csv", 78_000_000, 44, "full", ("scalar", True, [], "ceiling")), + ("xorq", "csv", 78_000_000, 44, "scalar", ("scalar", True, [], "host")), + ("xorq", "csv", 78_000_000, 44, "schema", ("schema", False, ["scalar"], "host")), +] + + +def _table_id(case): + backend, _kind, rows, cols, host_tier, _expected = case + return f"{backend}-{rows}x{cols}-{host_tier}" + + +class TestResolveStatsPolicy: + @pytest.mark.parametrize( + "backend, source_kind, rows, cols, host_tier, expected", RESOLVE_TABLE, + ids=[_table_id(c) for c in RESOLVE_TABLE]) + def test_table(self, sp, lim, backend, source_kind, rows, cols, host_tier, expected): + out = _resolve(sp, lim, backend, source_kind, rows, cols, host_tier) + got = (out["tier_target"], out["auto_request"], out["requestable"], out["reason"]) + assert got == expected + + def test_result_has_exactly_the_planned_keys(self, sp, lim): + out = _resolve(sp, lim, "xorq", "parquet", 12_400_000, 43) + assert set(out) == {"tier_target", "auto_request", "requestable", "reason", "estimate"} + + def test_result_is_json_serialisable(self, sp, lim): + """It becomes df_meta.stats, so it has to survive json.dumps.""" + out = _resolve(sp, lim, "xorq", "csv", 78_000_000, 44, "full") + assert json.loads(json.dumps(out)) == out + + def test_estimate_reports_the_inputs(self, sp, lim): + out = _resolve(sp, lim, "xorq", "parquet", 12_400_000, 43) + assert out["estimate"] == {"rows": 12_400_000, "cols": 43} + + def test_estimate_includes_bytes_when_given(self, sp, lim): + out = _resolve(sp, lim, "xorq", "parquet", 12_400_000, 43, bytes=987_654_321) + assert out["estimate"] == {"rows": 12_400_000, "cols": 43, "bytes": 987_654_321} + + def test_numpy_integer_counts_are_accepted(self, sp, lim): + out = _resolve(sp, lim, "xorq", "parquet", np.int64(11_000_000), np.int64(43)) + assert out["tier_target"] == "scalar" + assert out["estimate"] == {"rows": 11_000_000, "cols": 43} + + @pytest.mark.parametrize("rows, cols", [(-1, 3), (3, -1), (None, 3), (3, None), (1.5, 3)]) + def test_bad_counts_are_rejected(self, sp, lim, rows, cols): + with pytest.raises((TypeError, ValueError)): + _resolve(sp, lim, "xorq", "parquet", rows, cols) + + def test_unknown_host_tier_is_rejected(self, sp, lim): + with pytest.raises(ValueError): + _resolve(sp, lim, "xorq", "parquet", 1_000, 3, "everything") + + def test_unknown_host_tier_is_rejected_for_eager_backends_too(self, sp, lim): + with pytest.raises(ValueError): + _resolve(sp, lim, "pandas", "memory", 1_000, 3, "everything") + + def test_defaults_refuse_full_for_the_78m_row_entry(self, sp): + """No calibration may let a host request full on the 215 s entry.""" + out = sp.resolve_stats_policy("xorq", "csv", 78_000_000, 44, host_tier="full") + assert TIER_RANK[out["tier_target"]] < TIER_RANK["full"] + assert out["reason"] == "ceiling" + assert "full" not in out["requestable"] + + def test_defaults_let_the_control_recover_parking_2017(self, sp): + """10.8M x 43 columns took 3.7 s of stats and was fine, so a default + that withholds it must offer full on demand.""" + out = sp.resolve_stats_policy("xorq", "parquet", 10_800_000, 43) + assert out["tier_target"] == "full" or "full" in out["requestable"] + + +class TestCeiling: + """The ceiling is applied inside the function, so every caller gets it: + load, /reload_expr and a stats_request {force} all end in this call.""" + + HOST_TIERS = [None, "auto", "schema", "scalar", "full"] + + def test_host_full_above_the_ceiling_resolves_lower(self, sp, lim): + out = _resolve(sp, lim, "xorq", "csv", 78_000_000, 44, "full") + assert out["tier_target"] == "scalar" + assert out["reason"] == "ceiling" + + def test_reload_re_resolves_against_the_new_row_count(self, sp, lim): + """/reload_expr keeps the stored host tier and re-resolves. An entry + that fit under the ceiling and then grew past it loses full.""" + before = _resolve(sp, lim, "xorq", "parquet", 40_000_000, 44, "full") + assert before["tier_target"] == "full" + after = _resolve(sp, lim, "xorq", "parquet", 78_000_000, 44, "full") + assert after["tier_target"] == "scalar" + assert after["reason"] == "ceiling" + + def test_force_cannot_exceed_the_ceiling(self, sp, lim): + """A force request passes the requested tier as host_tier.""" + loaded = _resolve(sp, lim, "xorq", "csv", 78_000_000, 44) + assert loaded["tier_target"] == "schema" + assert "full" not in loaded["requestable"] + forced = _resolve(sp, lim, "xorq", "csv", 78_000_000, 44, "full") + assert TIER_RANK[forced["tier_target"]] < TIER_RANK["full"] + assert forced["reason"] == "ceiling" + + def test_ceiling_applies_to_the_automatic_choice_too(self, sp): + """A server whose ceiling sits below its own auto threshold is + clamped, with no host tier involved.""" + lim = sp.StatsLimits(full_auto_rows=10_000_000, ceiling_full_rows=5_000_000) + out = sp.resolve_stats_policy("xorq", "parquet", 7_000_000, 10, limits=lim) + assert out["tier_target"] == "scalar" + assert out["reason"] == "ceiling" + + def test_scalar_ceiling_can_leave_nothing_requestable(self, sp): + lim = sp.StatsLimits(ceiling_full_rows=50_000_000, ceiling_scalar_cells=1_000_000_000) + out = sp.resolve_stats_policy("xorq", "csv", 78_000_000, 44, host_tier="full", limits=lim) + assert out["tier_target"] == "schema" + assert out["reason"] == "ceiling" + assert out["auto_request"] is False + assert out["requestable"] == [] + + def test_scalar_ceiling_does_not_touch_smaller_entries(self, sp): + lim = sp.StatsLimits(ceiling_full_rows=50_000_000, ceiling_scalar_cells=1_000_000_000) + out = sp.resolve_stats_policy("xorq", "parquet", 1_000_000, 44, host_tier="full", limits=lim) + assert out["tier_target"] == "full" + + @pytest.mark.parametrize("rows", + [0, 1_000, 10_000_000, 11_000_000, 50_000_000, 50_000_001, 78_000_000, 500_000_000]) + @pytest.mark.parametrize("cols", [1, 44]) + @pytest.mark.parametrize("host_tier", HOST_TIERS) + def test_full_never_resolves_above_the_ceiling(self, sp, lim, rows, cols, host_tier): + out = _resolve(sp, lim, "xorq", "parquet", rows, cols, host_tier) + if rows > lim.ceiling_full_rows: + assert out["tier_target"] != "full" + assert "full" not in out["requestable"] + + @pytest.mark.parametrize("rows", [1_000, 10_800_000, 42_300_000, 50_000_001, 78_000_000]) + def test_requestable_is_exactly_what_a_force_would_grant(self, sp, lim, rows): + """Every tier the result offers is granted when asked for, and every + tier above the target that it does not offer is clipped.""" + out = _resolve(sp, lim, "xorq", "parquet", rows, 44) + for tier in ("scalar", "full"): + if TIER_RANK[tier] <= TIER_RANK[out["tier_target"]]: + continue + forced = _resolve(sp, lim, "xorq", "parquet", rows, 44, tier) + if tier in out["requestable"]: + assert forced["tier_target"] == tier + else: + assert forced["tier_target"] != tier + assert forced["reason"] == "ceiling" + + +class TestLimitsFromEnv: + def test_defaults_when_the_environment_is_empty(self, sp): + assert sp.StatsLimits.from_env() == sp.StatsLimits() + + @pytest.mark.parametrize("name, field", ENV_FIELDS.items()) + def test_each_variable_overrides_its_field(self, sp, monkeypatch, name, field): + monkeypatch.setenv(name, "12345") + assert getattr(sp.StatsLimits.from_env(), field) == 12345 + + def test_underscored_integers_are_accepted(self, sp, monkeypatch): + monkeypatch.setenv("BUCKAROO_STATS_FULL_AUTO_ROWS", "2_000_000") + assert sp.StatsLimits.from_env().full_auto_rows == 2_000_000 + + @pytest.mark.parametrize("name, field", ENV_FIELDS.items()) + @pytest.mark.parametrize("bad", ["lots", "-5", "1.5"]) + def test_invalid_value_falls_back_to_the_default_and_warns(self, sp, monkeypatch, caplog, name, field, bad): + monkeypatch.setenv(name, bad) + with caplog.at_level(logging.WARNING): + limits = sp.StatsLimits.from_env() + assert getattr(limits, field) == getattr(sp.StatsLimits(), field) + assert name in caplog.text + + def test_empty_value_means_unset(self, sp, monkeypatch): + monkeypatch.setenv("BUCKAROO_STATS_CEILING_SCALAR_CELLS", "") + assert sp.StatsLimits.from_env().ceiling_scalar_cells is None + + def test_resolve_reads_the_environment_when_no_limits_are_passed(self, sp, monkeypatch): + monkeypatch.setenv("BUCKAROO_STATS_CEILING_FULL_ROWS", "1000") + out = sp.resolve_stats_policy("xorq", "parquet", 5_000, 10, host_tier="full") + assert out["tier_target"] == "scalar" + assert out["reason"] == "ceiling" + + def test_the_environment_is_read_on_each_call(self, sp, monkeypatch): + args = ("xorq", "parquet", 5_000, 10) + assert sp.resolve_stats_policy(*args, host_tier="full")["tier_target"] == "full" + monkeypatch.setenv("BUCKAROO_STATS_CEILING_FULL_ROWS", "1000") + assert sp.resolve_stats_policy(*args, host_tier="full")["tier_target"] == "scalar" + + +class _CountingFile(io.BytesIO): + """A binary file that counts the bytes read out of it.""" + + def __init__(self, data): + super().__init__(data) + self.bytes_read = 0 + + def read(self, size=-1): + out = super().read(size) + self.bytes_read += len(out) + return out + + def read1(self, size=-1): + out = super().read1(size) + self.bytes_read += len(out) + return out + + def readinto(self, b): + n = super().readinto(b) + self.bytes_read += n + return n + + +def _parquet_bytes(rows, row_group_size): + """Incompressible float columns, so the file is mostly row-group data.""" + rng = np.random.default_rng(0) + table = pa.table({f"c{i}": rng.random(rows) for i in range(4)}) + buf = io.BytesIO() + pq.write_table(table, buf, row_group_size=row_group_size, compression="none") + return buf.getvalue() + + +class TestProbes: + def test_parquet_rows_reads_the_footer_and_no_row_group(self, sp): + data = _parquet_bytes(200_000, 50_000) + source = _CountingFile(data) + assert sp.probe_parquet_rows(source) == 200_000 + assert len(data) > 5_000_000 + assert source.bytes_read * 20 < len(data) + + def test_parquet_rows_from_a_path_decodes_no_data(self, sp, tmp_path, monkeypatch): + path = tmp_path / "t.parquet" + path.write_bytes(_parquet_bytes(120_000, 40_000)) + + def boom(*args, **kwargs): + raise AssertionError("the footer probe decoded row-group data") + + monkeypatch.setattr(pq.ParquetFile, "read", boom) + monkeypatch.setattr(pq.ParquetFile, "read_row_group", boom) + monkeypatch.setattr(pq.ParquetFile, "iter_batches", boom) + monkeypatch.setattr(pq, "read_table", boom) + assert sp.probe_parquet_rows(str(path)) == 120_000 + assert sp.probe_parquet_rows(path) == 120_000 + + def test_parquet_rows_of_an_empty_file(self, sp): + buf = io.BytesIO() + pq.write_table(pa.table({"a": pa.array([], type=pa.int64())}), buf) + assert sp.probe_parquet_rows(io.BytesIO(buf.getvalue())) == 0 + + def test_dtypes_of_a_pandas_frame(self, sp): + df = pd.DataFrame({"a": [1, 2], "b": ["x", "y"]}) + assert sp.probe_dtypes(df) == {"a": "int64", "b": str(df["b"].dtype)} + + def test_dtypes_of_a_polars_frame(self, sp): + df = pl.DataFrame({"a": [1, 2], "b": ["x", "y"]}) + assert sp.probe_dtypes(df) == {"a": "Int64", "b": "String"} + + @pytest.mark.filterwarnings("error") + def test_dtypes_of_a_polars_scan_run_no_collect(self, sp, tmp_path, monkeypatch): + """Reading ``LazyFrame.dtypes`` resolves the plan and warns, so the + error filter also pins that the probe goes through ``collect_schema``.""" + path = tmp_path / "t.parquet" + pl.DataFrame({"a": [1, 2], "b": ["x", "y"]}).write_parquet(path) + lf = pl.scan_parquet(path) + + def boom(*args, **kwargs): + raise AssertionError("the dtype probe collected a LazyFrame") + + monkeypatch.setattr(pl.LazyFrame, "collect", boom) + assert sp.probe_dtypes(lf) == {"a": "Int64", "b": "String"} + + def test_dtypes_keep_column_order(self, sp): + df = pd.DataFrame({"z": [1], "a": [2], "m": [3]}) + assert list(sp.probe_dtypes(df)) == ["z", "a", "m"] + + def test_dtypes_of_something_without_a_schema_raise(self, sp): + with pytest.raises(TypeError): + sp.probe_dtypes(object()) + + def test_xorq_probes_and_a_known_count_execute_nothing(self, sp, monkeypatch): + """The count is an input: the caller passes the one load already took + (``_expr_count``) and the policy never runs one itself.""" + xo = pytest.importorskip("xorq.api") + from xorq.vendor.ibis.expr.types.core import Expr + + expr = xo.memtable({"a": [1, 2, 3], "b": ["x", "y", "z"]}, name="t") + + def boom(*args, **kwargs): + raise AssertionError("a probe or the policy executed a xorq query") + + for name in ("execute", "to_pyarrow", "to_pyarrow_batches", "to_pandas"): + if hasattr(Expr, name): + monkeypatch.setattr(Expr, name, boom) + + dtypes = sp.probe_dtypes(expr) + assert dtypes == {"a": "int64", "b": "string"} + out = sp.resolve_stats_policy("xorq", "expr", 78_000_000, len(dtypes), host_tier="full") + assert out["reason"] == "ceiling" + + +class TestRoutePolarsEntry: + def test_above_r_goes_to_xorq(self, sp, lim): + r = lim.polars_route_rows + assert sp.route_polars_entry(r + 1, 43, limits=lim) == "xorq" + assert sp.route_polars_entry(78_000_000, 44, limits=lim) == "xorq" + + def test_at_and_below_r_stay_eager(self, sp, lim): + r = lim.polars_route_rows + assert sp.route_polars_entry(r, 43, limits=lim) == "eager" + assert sp.route_polars_entry(52_814, 27, limits=lim) == "eager" + assert sp.route_polars_entry(0, 0, limits=lim) == "eager" + + def test_r_comes_from_the_limits_object(self, sp): + lim = sp.StatsLimits(polars_route_rows=1_000) + assert sp.route_polars_entry(1_001, 3, limits=lim) == "xorq" + assert sp.route_polars_entry(1_000, 3, limits=lim) == "eager" + + def test_r_comes_from_the_environment_by_default(self, sp, monkeypatch): + assert sp.route_polars_entry(5_000, 3) == "eager" + monkeypatch.setenv("BUCKAROO_POLARS_ROUTE_ROWS", "1000") + assert sp.route_polars_entry(5_000, 3) == "xorq" + + def test_defaults_route_the_78m_row_entry_to_xorq(self, sp): + assert sp.route_polars_entry(78_000_000, 44) == "xorq" + + def test_the_result_is_one_of_two_strings(self, sp, lim): + for rows in (0, 10_000_000, 10_000_001): + assert sp.route_polars_entry(rows, 43, limits=lim) in ("xorq", "eager") + + @pytest.mark.parametrize("rows, cols", [(-1, 3), (3, -1), (None, 3), (1.5, 3)]) + def test_bad_counts_are_rejected(self, sp, lim, rows, cols): + with pytest.raises((TypeError, ValueError)): + sp.route_polars_entry(rows, cols, limits=lim) From 8895d26c739f02299ce022270cdb9d052a9c6a68 Mon Sep 17 00:00:00 2001 From: Paddy Mullen Date: Sat, 3 Oct 2026 23:03:13 -0400 Subject: [PATCH 03/39] feat(server): stats policy, routing and probes, not wired (rows-first p31) New buckaroo/server/stats_policy.py, pure and imported by nothing yet: - resolve_stats_policy(backend, source_kind, rows, cols, bytes, host_tier, limits) returns tier_target, auto_request, requestable, reason and estimate over the tiers schema < scalar < full. A ceiling is computed inside the function, so every caller gets the lower of the requested and ceiling tiers with reason "ceiling". A host tier lowers freely and raises only to the ceiling. Eager pandas and polars resolve to full. - route_polars_entry(rows, cols) returns "xorq" above R rows and "eager" at or below it. - probe_dtypes reads a schema without collecting or executing anything, and probe_parquet_rows reads a parquet footer's num_rows without decoding data. A xorq count is an input to the policy, never computed by it. - The thresholds are provisional constants gathered in StatsLimits, each with a BUCKAROO_* environment override read on every call. Co-Authored-By: Claude Sonnet 5.5 --- buckaroo/server/stats_policy.py | 245 ++++++++++++++++++++++++++++++++ 1 file changed, 245 insertions(+) create mode 100644 buckaroo/server/stats_policy.py diff --git a/buckaroo/server/stats_policy.py b/buckaroo/server/stats_policy.py new file mode 100644 index 000000000..01884dfb5 --- /dev/null +++ b/buckaroo/server/stats_policy.py @@ -0,0 +1,245 @@ +"""Which stats tier an entry should reach, decided from its size. + +Pure functions. Nothing in the handlers, the session or the dataflow imports +this module yet; it exists so the policy can be tested without a backend +before it is wired in. + +Tiers, lowest to highest: + +* ``schema``: dtype, ``_type``, identity keys and the row count. Never + requested automatically. +* ``scalar``: adds ``null_count``, ``min``, ``max``, ``mean``, ``std``, + ``empty_count`` and ``histogram_bins``. No histogram query, + ``value_counts``, median or (on xorq) ``distinct_count``. +* ``full``: today's stats. + +:func:`resolve_stats_policy` picks the tier. A hard ceiling is applied inside +it, so the result is the lower of the requested tier and the ceiling tier +whoever asks (a load, ``/reload_expr`` or a ``stats_request`` with ``force``). +A host tier can lower the result freely and raise it only up to the ceiling. +Eager pandas and polars resolve to ``full``: their stats sample 50,000 rows, +so the cost does not grow with the file. + +:func:`route_polars_entry` is the separate size rule for polars: an entry +above R rows belongs on xorq's ``/load_expr``, not eager ``/load``. + +The thresholds are PROVISIONAL. They are proposals taken from repo constants +and a few tallyman telemetry entries, and calibration measurements are +expected to replace them. They live in the ``DEFAULT_*`` constants below, and each has an +environment override that is read on every call, not at import. A value that +is not a non-negative integer is ignored with a warning. + +* ``BUCKAROO_STATS_FULL_AUTO_ROWS`` (10,000,000 rows): ``full`` is chosen + automatically up to this many rows. +* ``BUCKAROO_STATS_SCALAR_AUTO_CELLS`` (500,000,000 cells): above the full + threshold, ``scalar`` is chosen automatically up to this many cells + (rows x columns), from the 0.40-0.54 s scalar batch on 464M cells. Above it + the entry stays at ``schema``. +* ``BUCKAROO_STATS_CEILING_FULL_ROWS`` (50,000,000 rows): hard ceiling on + ``full``. It has to refuse the 78M-row entry whose stats took 215 s; the + number itself has no other evidence behind it. +* ``BUCKAROO_STATS_CEILING_SCALAR_CELLS`` (unset, no ceiling): hard ceiling + on ``scalar``, in cells. Above it only ``schema`` is allowed. +* ``BUCKAROO_POLARS_ROUTE_ROWS`` (10,000,000 rows): R for + :func:`route_polars_entry`. A memory threshold, set apart from the stats + ones. +""" +from __future__ import annotations + +import logging +import operator +import os +from dataclasses import dataclass +from typing import Any + +import pyarrow.parquet as pq + +log = logging.getLogger("buckaroo.server.stats_policy") + +# Provisional values; see the module docstring. +DEFAULT_FULL_AUTO_ROWS: int = 10_000_000 +DEFAULT_SCALAR_AUTO_CELLS: int = 500_000_000 +DEFAULT_CEILING_FULL_ROWS: int = 50_000_000 +DEFAULT_CEILING_SCALAR_CELLS: int | None = None +DEFAULT_POLARS_ROUTE_ROWS: int = 10_000_000 + +# Environment variable -> StatsLimits field. +_ENV_FIELDS = {"BUCKAROO_STATS_FULL_AUTO_ROWS": "full_auto_rows", + "BUCKAROO_STATS_SCALAR_AUTO_CELLS": "scalar_auto_cells", "BUCKAROO_STATS_CEILING_FULL_ROWS": "ceiling_full_rows", + "BUCKAROO_STATS_CEILING_SCALAR_CELLS": "ceiling_scalar_cells", "BUCKAROO_POLARS_ROUTE_ROWS": "polars_route_rows"} + +# Lowest to highest. +TIERS = ("schema", "scalar", "full") + +# Backends whose stats cost does not grow with the file. +EAGER_BACKENDS = frozenset({"pandas", "polars"}) + + +def _env_int(name: str) -> int | None: + raw = os.environ.get(name, "").strip() + if not raw: + return None + try: + value = int(raw) + except ValueError: + value = -1 + if value < 0: + log.warning("%s=%r is not a non-negative integer; using the default", name, raw) + return None + return value + + +@dataclass(frozen=True) +class StatsLimits: + """The thresholds :func:`resolve_stats_policy` and + :func:`route_polars_entry` read. ``ceiling_scalar_cells=None`` means + ``scalar`` has no ceiling.""" + full_auto_rows: int = DEFAULT_FULL_AUTO_ROWS + scalar_auto_cells: int = DEFAULT_SCALAR_AUTO_CELLS + ceiling_full_rows: int = DEFAULT_CEILING_FULL_ROWS + ceiling_scalar_cells: int | None = DEFAULT_CEILING_SCALAR_CELLS + polars_route_rows: int = DEFAULT_POLARS_ROUTE_ROWS + + @classmethod + def from_env(cls) -> StatsLimits: + """The defaults with any ``BUCKAROO_*`` overrides applied.""" + overrides = {} + for name, field in _ENV_FIELDS.items(): + value = _env_int(name) + if value is not None: + overrides[field] = value + return cls(**overrides) + + +def _count(name: str, value: Any) -> int: + try: + count = operator.index(value) + except TypeError: + raise TypeError(f"{name} must be an integer, got {value!r}") from None + if count < 0: + raise ValueError(f"{name} must not be negative, got {count}") + return count + + +def _size_tier(rows: int, cells: int, limits: StatsLimits) -> str: + """The tier the server picks on its own for an entry of this size.""" + if rows <= limits.full_auto_rows: + return "full" + if cells <= limits.scalar_auto_cells: + return "scalar" + return "schema" + + +def _ceiling_tier(rows: int, cells: int, limits: StatsLimits) -> str: + """The highest tier any caller may reach for an entry of this size.""" + if limits.ceiling_scalar_cells is not None and cells > limits.ceiling_scalar_cells: + return "schema" + return "full" if rows <= limits.ceiling_full_rows else "scalar" + + +def resolve_stats_policy(backend: str, source_kind: str, rows: int, cols: int, bytes: int | None = None, + host_tier: str | None = None, limits: StatsLimits | None = None) -> dict[str, Any]: + """Resolve the stats tier for an entry. + + ``rows`` is the count load already took (the parquet footer's, or xorq's + cached ``_expr_count``); this function never counts anything. ``host_tier`` + is the request's ``stats_tier`` (``None`` or ``"auto"`` for the server's + own choice); a ``stats_request`` with ``force`` passes the tier it asks for + here. ``limits`` defaults to :meth:`StatsLimits.from_env`. + + Returns a JSON-serialisable dict: + + * ``tier_target``: the tier the session should reach. + * ``auto_request``: whether the client should request stats up to + ``tier_target`` on its own. False only for ``schema``. + * ``requestable``: tiers above ``tier_target`` that an on-demand request + may still be granted, in ascending order, up to the ceiling. + * ``reason``: ``None`` when ``tier_target`` is ``full``; otherwise + ``"ceiling"`` if the ceiling clipped the request, ``"host"`` if the host + asked for this tier, ``"size"`` if the server chose it. The ``"cost"`` + reason belongs to the cost guard and is never returned here. + * ``estimate``: the inputs the decision rests on (``rows``, ``cols`` and + ``bytes`` when given). + + Eager backends ignore ``host_tier``. ``source_kind`` and ``bytes`` are + accepted for the rules that will read them (a CSV size rule, and a bytes + threshold if calibration wants one); no threshold reads them yet. + """ + rows = _count("rows", rows) + cols = _count("cols", cols) + if host_tier == "auto": + host_tier = None + if host_tier is not None and host_tier not in TIERS: + raise ValueError(f"host_tier must be one of {TIERS} or 'auto', got {host_tier!r}") + if limits is None: + limits = StatsLimits.from_env() + + if backend in EAGER_BACKENDS: + requested = ceiling = "full" + from_host = False + else: + cells = rows * cols + ceiling = _ceiling_tier(rows, cells, limits) + from_host = host_tier is not None + requested = host_tier or _size_tier(rows, cells, limits) + + target = min(requested, ceiling, key=TIERS.index) + if target == "full": + reason = None + elif target != requested: + reason = "ceiling" + elif from_host: + reason = "host" + else: + reason = "size" + + requestable = list(TIERS[TIERS.index(target) + 1:TIERS.index(ceiling) + 1]) + estimate: dict[str, int] = {"rows": rows, "cols": cols} + if bytes is not None: + estimate["bytes"] = bytes + return {"tier_target": target, "auto_request": target != "schema", "requestable": requestable, "reason": reason, + "estimate": estimate} + + +def route_polars_entry(rows: int, cols: int, limits: StatsLimits | None = None) -> str: + """``"xorq"`` for a polars entry above R rows, ``"eager"`` otherwise. + + Eager polars holds about 6.7 GB at 10.8M rows and cannot hold the 78M-row + entry, so a host sends entries above R to ``/load_expr`` instead of + ``/load`` with ``backend="polars"``. R is ``limits.polars_route_rows``. + ``cols`` is accepted so the rule can become a cell budget without a + signature change; only ``rows`` decides today. + """ + rows = _count("rows", rows) + _count("cols", cols) + if limits is None: + limits = StatsLimits.from_env() + return "xorq" if rows > limits.polars_route_rows else "eager" + + +def probe_dtypes(obj: Any) -> dict[str, str]: + """Column name to dtype string, from the schema alone. + + Takes a pandas or polars frame (a ``LazyFrame`` is not collected) or a xorq + expression. No row is read and no query runs. + """ + # Polars first: reading ``LazyFrame.dtypes`` resolves the plan and warns. + if hasattr(obj, "collect_schema"): + pairs = obj.collect_schema().items() + elif callable(getattr(obj, "schema", None)): + pairs = obj.schema().items() + elif hasattr(getattr(obj, "dtypes", None), "items"): + pairs = obj.dtypes.items() + else: + raise TypeError(f"cannot read a schema from {type(obj).__name__}") + return {str(name): str(dtype) for name, dtype in pairs} + + +def probe_parquet_rows(source: Any) -> int: + """A parquet file's row count, from its footer. + + ``source`` is a path or a binary file object. Only the footer metadata is + read, so the cost does not grow with the file (0.4 ms on 4M rows). A + single file only; a dataset directory is not handled. + """ + return int(pq.read_metadata(source).num_rows) From 9b96220122c42bdac52b0f081a856885c3bb030e Mon Sep 17 00:00:00 2001 From: Paddy Mullen Date: Sat, 3 Oct 2026 22:56:48 -0400 Subject: [PATCH 04/39] 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 05/39] 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 f03532ca7c3c38645f3d5c72d9d3bca299d7e8af Mon Sep 17 00:00:00 2001 From: Paddy Mullen Date: Sat, 3 Oct 2026 23:17:32 -0400 Subject: [PATCH 06/39] test(server): failing tests for bytes and source_kind validation in resolve_stats_policy (rows-first p31) bytes and source_kind were documented as validated but are not: bytes=-1, 'lots', 1.5 and object() are accepted, a np.int64 is echoed unchanged so json.dumps of the result fails, a host tier passed positionally lands in the bytes slot, and source_kind=None or 5 is accepted. Co-Authored-By: Claude Sonnet 5.5 --- tests/unit/server/test_stats_policy.py | 23 +++++++++++++++++++++++ 1 file changed, 23 insertions(+) diff --git a/tests/unit/server/test_stats_policy.py b/tests/unit/server/test_stats_policy.py index 6e2abf98c..bb60bbacc 100644 --- a/tests/unit/server/test_stats_policy.py +++ b/tests/unit/server/test_stats_policy.py @@ -139,6 +139,29 @@ def test_bad_counts_are_rejected(self, sp, lim, rows, cols): with pytest.raises((TypeError, ValueError)): _resolve(sp, lim, "xorq", "parquet", rows, cols) + def test_numpy_integer_bytes_are_echoed_as_plain_ints(self, sp, lim): + """df.memory_usage(deep=True).sum() is a np.int64, and the result has + to survive json.dumps.""" + out = _resolve(sp, lim, "xorq", "parquet", 12_400_000, 43, bytes=np.int64(987_654_321)) + assert type(out["estimate"]["bytes"]) is int + assert json.loads(json.dumps(out)) == out + + @pytest.mark.parametrize("bad", [-1, "lots", 1.5, object()]) + def test_bad_bytes_are_rejected(self, sp, lim, bad): + with pytest.raises((TypeError, ValueError)): + _resolve(sp, lim, "xorq", "parquet", 12_400_000, 43, bytes=bad) + + def test_a_host_tier_in_the_bytes_slot_is_rejected(self, sp): + """bytes sits before host_tier, so a positional slip must not be taken + for a byte count.""" + with pytest.raises((TypeError, ValueError)): + sp.resolve_stats_policy("xorq", "csv", 12_000_000, 43, "full") + + @pytest.mark.parametrize("bad", [None, 5, b"csv"]) + def test_bad_source_kind_is_rejected(self, sp, lim, bad): + with pytest.raises(TypeError): + _resolve(sp, lim, "xorq", bad, 1_000, 3) + def test_unknown_host_tier_is_rejected(self, sp, lim): with pytest.raises(ValueError): _resolve(sp, lim, "xorq", "parquet", 1_000, 3, "everything") From 21cd73e16843527d71cbb03f60654535f7b431b8 Mon Sep 17 00:00:00 2001 From: Paddy Mullen Date: Sat, 3 Oct 2026 23:24:11 -0400 Subject: [PATCH 07/39] 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 a2cc7ce5600de08707815d3aaf01ddce62efa9f0 Mon Sep 17 00:00:00 2001 From: Paddy Mullen Date: Sat, 3 Oct 2026 23:32:57 -0400 Subject: [PATCH 08/39] fix(server): validate bytes and source_kind in resolve_stats_policy (rows-first p31) bytes goes through the same integer check as rows and cols, so a negative, float or non-numeric value raises, a numpy integer is echoed as a plain int and the result stays JSON-serialisable, and a host tier passed positionally into the bytes slot raises instead of being echoed. source_kind must be a str; its vocabulary stays open for the callers that will read it. Co-Authored-By: Claude Sonnet 5.5 --- buckaroo/server/stats_policy.py | 8 ++++++++ 1 file changed, 8 insertions(+) diff --git a/buckaroo/server/stats_policy.py b/buckaroo/server/stats_policy.py index 01884dfb5..a19dc7b38 100644 --- a/buckaroo/server/stats_policy.py +++ b/buckaroo/server/stats_policy.py @@ -164,9 +164,17 @@ def resolve_stats_policy(backend: str, source_kind: str, rows: int, cols: int, b Eager backends ignore ``host_tier``. ``source_kind`` and ``bytes`` are accepted for the rules that will read them (a CSV size rule, and a bytes threshold if calibration wants one); no threshold reads them yet. + ``source_kind`` must be a ``str``, its vocabulary being the caller's. + ``bytes`` is ``None`` or a non-negative integer, echoed as a plain ``int``. + A value of the wrong type raises ``TypeError`` and a negative count + ``ValueError``. """ + if not isinstance(source_kind, str): + raise TypeError(f"source_kind must be a str, got {source_kind!r}") rows = _count("rows", rows) cols = _count("cols", cols) + if bytes is not None: + bytes = _count("bytes", bytes) if host_tier == "auto": host_tier = None if host_tier is not None and host_tier not in TIERS: From 1c89cfd6a82f5b3ee2c2e366784e43a0a3ab028d Mon Sep 17 00:00:00 2001 From: Paddy Mullen Date: Sat, 3 Oct 2026 23:49:42 -0400 Subject: [PATCH 09/39] 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 10/39] 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 11/39] 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 12/39] 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 3fd6b5fac9777e03d0afcb97d2e4b30535461378 Mon Sep 17 00:00:00 2001 From: Paddy Mullen Date: Sun, 4 Oct 2026 01:14:21 -0400 Subject: [PATCH 13/39] test(server): failing tests for the calibrated stats policy thresholds (rows-first p31b) The boundary tables now run on the module defaults and are written against the phase-0 proposals: full auto up to 12M rows and 520M cells, scalar auto up to 1.0B cells, full refused above 25M rows or 1.0B cells (force included), a scalar ceiling of 4.0B cells, and polars routed to xorq above 8M rows. Each threshold has an equal, a one-below and a one-above row, in rows and in cells. New tests pin each default to its literal value and cover the two new environment overrides, BUCKAROO_STATS_FULL_AUTO_CELLS and BUCKAROO_STATS_CEILING_FULL_CELLS. They fail on the current constants (10M rows, 500M cells, 50M-row ceiling, no scalar ceiling, R of 10M) and on the missing full_auto_cells and ceiling_full_cells fields. Co-Authored-By: Claude Sonnet 5.5 --- tests/unit/server/test_stats_policy.py | 187 +++++++++++++++++++++---- 1 file changed, 161 insertions(+), 26 deletions(-) diff --git a/tests/unit/server/test_stats_policy.py b/tests/unit/server/test_stats_policy.py index bb60bbacc..9a032b7b3 100644 --- a/tests/unit/server/test_stats_policy.py +++ b/tests/unit/server/test_stats_policy.py @@ -1,14 +1,16 @@ -"""Unit tests for ``buckaroo.server.stats_policy`` (rows-first p31). +"""Unit tests for ``buckaroo.server.stats_policy`` (rows-first p31, p31b). Everything here is pure Python: the policy is a function of numbers, the probes read a schema or a parquet footer, and nothing needs a server. The module is not wired into any handler, session or dataflow in this phase. -Thresholds are provisional, so the table tests pass an explicit -``StatsLimits`` instead of reading the module defaults. The few tests that -use the defaults assert only facts the plan requires of any calibration -(the 78M-row entry cannot reach ``full``). +The thresholds are the provisional values proposed by the phase-0 +measurements (p31b). The boundary tables run on the module defaults, so a +change to a ``DEFAULT_*`` constant fails them and the table is updated +together with the constant. ``TestCalibratedDefaults`` pins each default to +its literal value. """ +import dataclasses import importlib import importlib.util import io @@ -22,10 +24,21 @@ import pyarrow.parquet as pq import pytest -ENV_FIELDS = {"BUCKAROO_STATS_FULL_AUTO_ROWS": "full_auto_rows", +ENV_FIELDS = {"BUCKAROO_STATS_FULL_AUTO_ROWS": "full_auto_rows", "BUCKAROO_STATS_FULL_AUTO_CELLS": "full_auto_cells", "BUCKAROO_STATS_SCALAR_AUTO_CELLS": "scalar_auto_cells", "BUCKAROO_STATS_CEILING_FULL_ROWS": "ceiling_full_rows", + "BUCKAROO_STATS_CEILING_FULL_CELLS": "ceiling_full_cells", "BUCKAROO_STATS_CEILING_SCALAR_CELLS": "ceiling_scalar_cells", "BUCKAROO_POLARS_ROUTE_ROWS": "polars_route_rows"} +# The phase-0 values (measurements-phase0.md, "Proposed values"), written out +# so a change to a module default has to change a test too. +CALIBRATED = {"full_auto_rows": 12_000_000, "full_auto_cells": 520_000_000, "scalar_auto_cells": 1_000_000_000, + "ceiling_full_rows": 25_000_000, "ceiling_full_cells": 1_000_000_000, "ceiling_scalar_cells": 4_000_000_000, + "polars_route_rows": 8_000_000} +CALIBRATED_CONSTANTS = {"DEFAULT_FULL_AUTO_ROWS": 12_000_000, "DEFAULT_FULL_AUTO_CELLS": 520_000_000, + "DEFAULT_SCALAR_AUTO_CELLS": 1_000_000_000, "DEFAULT_CEILING_FULL_ROWS": 25_000_000, + "DEFAULT_CEILING_FULL_CELLS": 1_000_000_000, "DEFAULT_CEILING_SCALAR_CELLS": 4_000_000_000, + "DEFAULT_POLARS_ROUTE_ROWS": 8_000_000} + TIER_RANK = {"schema": 0, "scalar": 1, "full": 2} @@ -45,14 +58,9 @@ def clean_env(monkeypatch): @pytest.fixture def lim(sp): - """The plan's proposed starting values, pinned so a recalibration of the - module defaults does not move these tables.""" - return sp.StatsLimits( - full_auto_rows=10_000_000, - scalar_auto_cells=500_000_000, - ceiling_full_rows=50_000_000, - ceiling_scalar_cells=None, - polars_route_rows=10_000_000) + """The module defaults. The tables are written against the calibrated + values, so moving a ``DEFAULT_*`` constant fails them.""" + return sp.StatsLimits() def _resolve(sp, lim, backend, source_kind, rows, cols, host_tier=None, bytes=None): @@ -63,38 +71,88 @@ def _resolve(sp, lim, backend, source_kind, rows, cols, host_tier=None, bytes=No # (backend, source_kind, rows, cols, host_tier) -> (tier_target, auto_request, # requestable, reason). The row counts of 42.3M, 54.1M and 78.0M and the 44 # columns are the three tallyman entries over 70 s; 10.8M x 43 is parking_2017. +# Each threshold has an equal, a one-below and a one-above row, in rows and, +# where the rule reads cells, in cells (a row count or a column count moves +# the product). The comments give the cell count where it decides the row. RESOLVE_TABLE = [ - # Eager backends resolve to full and ignore the host tier. + # Eager backends resolve to full and ignore the host tier and the cell bounds. ("pandas", "memory", 52_814, 27, None, ("full", True, [], None)), ("pandas", "parquet", 1_000_000, 43, None, ("full", True, [], None)), ("polars", "parquet", 10_800_000, 43, None, ("full", True, [], None)), ("polars", "parquet", 10_800_000, 43, "schema", ("full", True, [], None)), ("pandas", "csv", 52_814, 27, "scalar", ("full", True, [], None)), - # xorq sized by the policy alone. + ("pandas", "memory", 30_000_000, 100, None, ("full", True, [], None)), + # xorq sized by the policy alone: full up to 12M rows. ("xorq", "parquet", 0, 0, None, ("full", True, [], None)), ("xorq", "parquet", 52_814, 27, None, ("full", True, [], None)), - ("xorq", "parquet", 10_000_000, 43, None, ("full", True, [], None)), - ("xorq", "parquet", 10_000_001, 43, None, ("scalar", True, ["full"], "size")), - ("xorq", "parquet", 10_800_000, 43, None, ("scalar", True, ["full"], "size")), - ("xorq", "parquet", 11_363_636, 44, None, ("scalar", True, ["full"], "size")), - ("xorq", "parquet", 11_363_637, 44, None, ("schema", False, ["scalar", "full"], "size")), - ("xorq", "csv", 42_300_000, 44, None, ("schema", False, ["scalar", "full"], "size")), - ("xorq", "csv", 50_000_000, 44, None, ("schema", False, ["scalar", "full"], "size")), + ("xorq", "parquet", 10_800_000, 43, None, ("full", True, [], None)), + ("xorq", "parquet", 11_999_999, 43, None, ("full", True, [], None)), + ("xorq", "parquet", 12_000_000, 43, None, ("full", True, [], None)), + ("xorq", "parquet", 12_000_001, 43, None, ("scalar", True, ["full"], "size")), + ("xorq", "parquet", 12_000_000, 10, None, ("full", True, [], None)), + ("xorq", "parquet", 12_000_001, 10, None, ("scalar", True, ["full"], "size")), + # The 11.7M-row CSV entries: 515.8M cells at 44 columns, 527.5M at 45. + ("xorq", "csv", 11_721_603, 44, None, ("full", True, [], None)), + ("xorq", "csv", 11_721_603, 45, None, ("scalar", True, ["full"], "size")), + # full auto also stops at 520M cells, reached by rows or by columns. + ("xorq", "parquet", 9_999_999, 52, None, ("full", True, [], None)), + ("xorq", "parquet", 10_000_000, 52, None, ("full", True, [], None)), + ("xorq", "parquet", 10_000_001, 52, None, ("scalar", True, ["full"], "size")), + ("xorq", "parquet", 8_000_000, 65, None, ("full", True, [], None)), + ("xorq", "parquet", 8_000_000, 66, None, ("scalar", True, ["full"], "size")), + # scalar auto up to 1.0B cells; full stays requestable up to its own ceiling. + ("xorq", "parquet", 11_000_000, 90, None, ("scalar", True, ["full"], "size")), + ("xorq", "parquet", 11_000_000, 100, None, ("schema", False, ["scalar"], "size")), + ("xorq", "parquet", 19_999_999, 50, None, ("scalar", True, ["full"], "size")), + ("xorq", "parquet", 20_000_000, 50, None, ("scalar", True, ["full"], "size")), + ("xorq", "parquet", 20_000_001, 50, None, ("schema", False, ["scalar"], "size")), + ("xorq", "parquet", 22_727_272, 44, None, ("scalar", True, ["full"], "size")), + ("xorq", "parquet", 22_727_273, 44, None, ("schema", False, ["scalar"], "size")), + ("xorq", "parquet", 12_500_000, 80, None, ("scalar", True, ["full"], "size")), + ("xorq", "parquet", 12_500_000, 81, None, ("schema", False, ["scalar"], "size")), + # The three slow tallyman entries are above 1.8B cells and start at schema. + ("xorq", "csv", 42_300_000, 44, None, ("schema", False, ["scalar"], "size")), + ("xorq", "csv", 50_000_000, 44, None, ("schema", False, ["scalar"], "size")), ("xorq", "csv", 54_100_000, 44, None, ("schema", False, ["scalar"], "size")), ("xorq", "csv", 78_000_000, 44, None, ("schema", False, ["scalar"], "size")), ("xorq", "csv", 78_000_000, 44, "auto", ("schema", False, ["scalar"], "size")), + # Past the scalar ceiling (4.0B cells) nothing is requestable. + ("xorq", "csv", 100_000_000, 40, None, ("schema", False, ["scalar"], "size")), + ("xorq", "csv", 100_000_001, 40, None, ("schema", False, [], "size")), # The host lowers freely. ("xorq", "parquet", 1_000_000, 43, "schema", ("schema", False, ["scalar", "full"], "host")), ("xorq", "parquet", 1_000_000, 43, "scalar", ("scalar", True, ["full"], "host")), ("xorq", "parquet", 1_000_000, 43, "full", ("full", True, [], None)), - # The host raises only as far as the ceiling. + # The host raises only as far as the ceiling on full: 25M rows ... ("xorq", "parquet", 12_000_000, 43, "full", ("full", True, [], None)), - ("xorq", "csv", 42_300_000, 44, "full", ("full", True, [], None)), - ("xorq", "csv", 50_000_000, 44, "full", ("full", True, [], None)), + ("xorq", "parquet", 24_999_999, 10, "full", ("full", True, [], None)), + ("xorq", "parquet", 25_000_000, 10, "full", ("full", True, [], None)), + ("xorq", "parquet", 25_000_001, 10, "full", ("scalar", True, [], "ceiling")), + ("xorq", "parquet", 25_000_000, 10, None, ("scalar", True, ["full"], "size")), + ("xorq", "parquet", 25_000_001, 10, None, ("scalar", True, [], "size")), + # ... or 1.0B cells, reached by rows or by columns. + ("xorq", "parquet", 24_999_999, 40, "full", ("full", True, [], None)), + ("xorq", "parquet", 25_000_000, 40, "full", ("full", True, [], None)), + ("xorq", "parquet", 25_000_000, 41, "full", ("scalar", True, [], "ceiling")), + ("xorq", "parquet", 20_000_000, 50, "full", ("full", True, [], None)), + ("xorq", "parquet", 20_000_001, 50, "full", ("scalar", True, [], "ceiling")), + ("xorq", "parquet", 11_000_000, 90, "full", ("full", True, [], None)), + ("xorq", "parquet", 11_000_000, 91, "full", ("scalar", True, [], "ceiling")), + ("xorq", "csv", 42_300_000, 44, "full", ("scalar", True, [], "ceiling")), + ("xorq", "csv", 50_000_000, 44, "full", ("scalar", True, [], "ceiling")), ("xorq", "csv", 50_000_001, 44, "full", ("scalar", True, [], "ceiling")), ("xorq", "csv", 78_000_000, 44, "full", ("scalar", True, [], "ceiling")), ("xorq", "csv", 78_000_000, 44, "scalar", ("scalar", True, [], "host")), ("xorq", "csv", 78_000_000, 44, "schema", ("schema", False, ["scalar"], "host")), + # The ceiling on scalar is 4.0B cells, reached by rows or by columns. + ("xorq", "csv", 99_999_999, 40, "scalar", ("scalar", True, [], "host")), + ("xorq", "csv", 100_000_000, 40, "scalar", ("scalar", True, [], "host")), + ("xorq", "csv", 100_000_001, 40, "scalar", ("schema", False, [], "ceiling")), + ("xorq", "csv", 50_000_000, 80, "scalar", ("scalar", True, [], "host")), + ("xorq", "csv", 50_000_000, 81, "scalar", ("schema", False, [], "ceiling")), + ("xorq", "csv", 100_000_000, 40, "full", ("scalar", True, [], "ceiling")), + ("xorq", "csv", 100_000_001, 40, "full", ("schema", False, [], "ceiling")), + ("xorq", "csv", 100_000_001, 40, "schema", ("schema", False, [], "host")), ] @@ -184,6 +242,23 @@ def test_defaults_let_the_control_recover_parking_2017(self, sp): assert out["tier_target"] == "full" or "full" in out["requestable"] +class TestCalibratedDefaults: + """The module defaults are the phase-0 proposals, one literal per value.""" + + def test_limits_defaults(self, sp): + assert dataclasses.asdict(sp.StatsLimits()) == CALIBRATED + + def test_module_constants(self, sp): + assert {name: getattr(sp, name, None) for name in CALIBRATED_CONSTANTS} == CALIBRATED_CONSTANTS + + def test_from_env_with_a_clean_environment_is_the_calibrated_set(self, sp): + assert dataclasses.asdict(sp.StatsLimits.from_env()) == CALIBRATED + + def test_every_limit_has_an_environment_override(self, sp): + """A new limit without an override could not be tuned on a server.""" + assert set(ENV_FIELDS.values()) == {f.name for f in dataclasses.fields(sp.StatsLimits)} + + class TestCeiling: """The ceiling is applied inside the function, so every caller gets it: load, /reload_expr and a stats_request {force} all end in this call.""" @@ -260,6 +335,42 @@ def test_requestable_is_exactly_what_a_force_would_grant(self, sp, lim, rows): assert forced["reason"] == "ceiling" + @pytest.mark.parametrize("rows, cols", + [(20_000_001, 50), (25_000_000, 41), (11_000_000, 91), (24_999_999, 41), (42_300_000, 44), (78_000_000, 44)]) + @pytest.mark.parametrize("host_tier", HOST_TIERS) + def test_full_never_resolves_above_the_cell_ceiling(self, sp, lim, rows, cols, host_tier): + """1.0B cells, whatever the row count and the host asks for.""" + assert rows * cols > 1_000_000_000 + out = _resolve(sp, lim, "xorq", "parquet", rows, cols, host_tier) + assert out["tier_target"] != "full" + assert "full" not in out["requestable"] + + @pytest.mark.parametrize("rows, cols", + [(25_000_001, 10), (25_000_001, 1), (30_000_000, 5), (42_300_000, 4), (78_000_000, 1)]) + @pytest.mark.parametrize("host_tier", HOST_TIERS) + def test_full_never_resolves_above_25m_rows_however_few_the_cells(self, sp, lim, rows, cols, host_tier): + """The row ceiling holds on its own: these entries are under 1.0B cells.""" + assert rows * cols < 1_000_000_000 + out = _resolve(sp, lim, "xorq", "parquet", rows, cols, host_tier) + assert out["tier_target"] != "full" + assert "full" not in out["requestable"] + + @pytest.mark.parametrize("rows, cols", + [(100_000_001, 40), (50_000_000, 81), (1_000_000, 4_001), (500_000_000, 10)]) + @pytest.mark.parametrize("host_tier", HOST_TIERS) + def test_nothing_above_the_scalar_ceiling_resolves_past_schema(self, sp, lim, rows, cols, host_tier): + """4.0B cells, the placeholder scalar ceiling.""" + assert rows * cols > 4_000_000_000 + out = _resolve(sp, lim, "xorq", "parquet", rows, cols, host_tier) + assert out["tier_target"] == "schema" + assert out["requestable"] == [] + + def test_a_scalar_ceiling_is_on_by_default(self, sp): + out = sp.resolve_stats_policy("xorq", "csv", 100_000_001, 40, host_tier="scalar", limits=sp.StatsLimits()) + assert out["tier_target"] == "schema" + assert out["reason"] == "ceiling" + + class TestLimitsFromEnv: def test_defaults_when_the_environment_is_empty(self, sp): assert sp.StatsLimits.from_env() == sp.StatsLimits() @@ -292,6 +403,22 @@ def test_resolve_reads_the_environment_when_no_limits_are_passed(self, sp, monke assert out["tier_target"] == "scalar" assert out["reason"] == "ceiling" + def test_the_full_auto_cell_bound_comes_from_the_environment(self, sp, monkeypatch): + args = ("xorq", "parquet", 5_000, 10) + assert sp.resolve_stats_policy(*args)["tier_target"] == "full" + monkeypatch.setenv("BUCKAROO_STATS_FULL_AUTO_CELLS", "1000") + out = sp.resolve_stats_policy(*args) + assert out["tier_target"] == "scalar" + assert out["reason"] == "size" + + def test_the_full_ceiling_cell_bound_comes_from_the_environment(self, sp, monkeypatch): + args = ("xorq", "parquet", 5_000, 10) + assert sp.resolve_stats_policy(*args, host_tier="full")["tier_target"] == "full" + monkeypatch.setenv("BUCKAROO_STATS_CEILING_FULL_CELLS", "1000") + out = sp.resolve_stats_policy(*args, host_tier="full") + assert out["tier_target"] == "scalar" + assert out["reason"] == "ceiling" + def test_the_environment_is_read_on_each_call(self, sp, monkeypatch): args = ("xorq", "parquet", 5_000, 10) assert sp.resolve_stats_policy(*args, host_tier="full")["tier_target"] == "full" @@ -431,6 +558,14 @@ def test_r_comes_from_the_environment_by_default(self, sp, monkeypatch): monkeypatch.setenv("BUCKAROO_POLARS_ROUTE_ROWS", "1000") assert sp.route_polars_entry(5_000, 3) == "xorq" + @pytest.mark.parametrize("rows, expected", + [(0, "eager"), (7_999_999, "eager"), (8_000_000, "eager"), (8_000_001, "xorq"), (10_000_000, "xorq"), + (10_800_000, "xorq")]) + def test_default_r_is_8m_rows(self, sp, rows, expected): + """Equal, one below and one above R, on the module defaults; 10.8M + rows (parking_2017) was eager under the old 10M threshold.""" + assert sp.route_polars_entry(rows, 43) == expected + def test_defaults_route_the_78m_row_entry_to_xorq(self, sp): assert sp.route_polars_entry(78_000_000, 44) == "xorq" From d7ba7bd963efae445d3f8a235ed8178f191d8a04 Mon Sep 17 00:00:00 2001 From: Paddy Mullen Date: Sun, 4 Oct 2026 01:26:27 -0400 Subject: [PATCH 14/39] 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 3eef709755941bb04c7faeb3ef473bc5efd69e30 Mon Sep 17 00:00:00 2001 From: Paddy Mullen Date: Sun, 4 Oct 2026 01:29:01 -0400 Subject: [PATCH 15/39] feat(server): calibrate the stats policy thresholds from the phase-0 measurements (rows-first p31b) Full auto now needs at most 12M rows and 520M cells (was 10M rows), scalar auto goes to a 1.0B cell budget (was 500M), `full` is refused above 25M rows or 1.0B cells including for a forced request (was 50M rows), scalar gets a 4.0B cell ceiling by default (was unset; an extrapolation with no measurement behind it), and polars routes to xorq above 8M rows (was 10M; the figure assumes pre_limit False). Each DEFAULT_* constant carries the measurement it comes from. The two new bounds are full_auto_cells and ceiling_full_cells, appended to StatsLimits, with BUCKAROO_STATS_FULL_AUTO_CELLS and BUCKAROO_STATS_CEILING_FULL_CELLS as overrides under the same parsing rules. _size_tier and _ceiling_tier read them; signatures and the result shape are unchanged. Three existing tests that hard-coded the old boundaries are updated: the numpy-integer case (11M rows is now full), the reload case (a 40M-row entry is now refused full) and the empty-environment case for the scalar ceiling (an empty value now means the 4.0B default, not no ceiling). Co-Authored-By: Claude Sonnet 5.5 --- buckaroo/server/stats_policy.py | 102 +++++++++++++++++-------- tests/unit/server/test_stats_policy.py | 34 ++++++++- 2 files changed, 100 insertions(+), 36 deletions(-) diff --git a/buckaroo/server/stats_policy.py b/buckaroo/server/stats_policy.py index a19dc7b38..04d171bc2 100644 --- a/buckaroo/server/stats_policy.py +++ b/buckaroo/server/stats_policy.py @@ -23,26 +23,31 @@ :func:`route_polars_entry` is the separate size rule for polars: an entry above R rows belongs on xorq's ``/load_expr``, not eager ``/load``. -The thresholds are PROVISIONAL. They are proposals taken from repo constants -and a few tallyman telemetry entries, and calibration measurements are -expected to replace them. They live in the ``DEFAULT_*`` constants below, and each has an -environment override that is read on every call, not at import. A value that -is not a non-negative integer is ignored with a warning. - -* ``BUCKAROO_STATS_FULL_AUTO_ROWS`` (10,000,000 rows): ``full`` is chosen - automatically up to this many rows. -* ``BUCKAROO_STATS_SCALAR_AUTO_CELLS`` (500,000,000 cells): above the full - threshold, ``scalar`` is chosen automatically up to this many cells - (rows x columns), from the 0.40-0.54 s scalar batch on 464M cells. Above it - the entry stays at ``schema``. -* ``BUCKAROO_STATS_CEILING_FULL_ROWS`` (50,000,000 rows): hard ceiling on - ``full``. It has to refuse the 78M-row entry whose stats took 215 s; the - number itself has no other evidence behind it. -* ``BUCKAROO_STATS_CEILING_SCALAR_CELLS`` (unset, no ceiling): hard ceiling - on ``scalar``, in cells. Above it only ``schema`` is allowed. -* ``BUCKAROO_POLARS_ROUTE_ROWS`` (10,000,000 rows): R for +The thresholds are PROVISIONAL. They are the values the phase-0 measurements +proposed (one Apple M4 Pro, 115 tallyman telemetry loads, xorq stats on +parquet and CSV slices, eager polars RSS), and they rest on gaps: the +telemetry has nothing between 11.8M and 42.3M rows, nothing was measured on +the xorq side above 12M rows, and the scalar ceiling is an extrapolation. +They live in the ``DEFAULT_*`` constants below, each with the measurement it +comes from, and each has an environment override that is read on every call, +not at import. A value that is not a non-negative integer is ignored with a +warning. + +* ``BUCKAROO_STATS_FULL_AUTO_ROWS`` (12,000,000 rows) and + ``BUCKAROO_STATS_FULL_AUTO_CELLS`` (520,000,000 cells): ``full`` is chosen + automatically while both hold. +* ``BUCKAROO_STATS_SCALAR_AUTO_CELLS`` (1,000,000,000 cells): otherwise + ``scalar`` is chosen automatically up to this many cells (rows x columns). + Above it the entry stays at ``schema``. +* ``BUCKAROO_STATS_CEILING_FULL_ROWS`` (25,000,000 rows) and + ``BUCKAROO_STATS_CEILING_FULL_CELLS`` (1,000,000,000 cells): hard ceiling + on ``full``, refused above either bound, including for a forced request. +* ``BUCKAROO_STATS_CEILING_SCALAR_CELLS`` (4,000,000,000 cells): hard ceiling + on ``scalar``, in cells. Above it only ``schema`` is allowed. This number + is an extrapolation with no measurement behind it. +* ``BUCKAROO_POLARS_ROUTE_ROWS`` (8,000,000 rows): R for :func:`route_polars_entry`. A memory threshold, set apart from the stats - ones. + ones, and sized for ``pre_limit`` False (see the constant below). """ from __future__ import annotations @@ -56,17 +61,41 @@ log = logging.getLogger("buckaroo.server.stats_policy") -# Provisional values; see the module docstring. -DEFAULT_FULL_AUTO_ROWS: int = 10_000_000 -DEFAULT_SCALAR_AUTO_CELLS: int = 500_000_000 -DEFAULT_CEILING_FULL_ROWS: int = 50_000_000 -DEFAULT_CEILING_SCALAR_CELLS: int | None = None -DEFAULT_POLARS_ROUTE_ROWS: int = 10_000_000 +# Provisional values from the phase-0 measurements; see the module docstring. +# 105 of 115 telemetry loads are at or below 12M rows (stats p90 4.4 s). Every +# load above it is one of three CSV unions (42.3M, 54.1M, 78.0M rows) at 75, +# 114 and 215 s. +DEFAULT_FULL_AUTO_ROWS: int = 12_000_000 +# Parquet 12M x 43 (516M cells) full stats took 2.9 s over 8 repetitions. The +# bound only bites on an entry wider than 43 columns at 12M rows. +DEFAULT_FULL_AUTO_CELLS: int = 520_000_000 +# The scalar class costs 1.1-1.3 ms per Mcell on a parquet scan and 3-6 ms on a +# CSV union. At the worst telemetry batch rates (30-48 ms per Mcell, scalar +# being 0.44-0.71 of the batch) that is about 20-34 s at 1.0B cells. +DEFAULT_SCALAR_AUTO_CELLS: int = 1_000_000_000 +# The ceiling on full separates the slowest in-limit load (11.8M rows, 22 s) +# from the fastest one above it (42.3M rows, 75 s). 25M rows and 1.0B cells are +# near the geometric midpoint of that gap (22M rows, 1.0B cells), and it is a +# gap in the telemetry, so the true limit could sit anywhere inside it. +DEFAULT_CEILING_FULL_ROWS: int = 25_000_000 +DEFAULT_CEILING_FULL_CELLS: int = 1_000_000_000 +# An extrapolation with no measurement behind it. The only evidence above 1.8B +# cells is the 3.36B-cell entry whose batch took 162.5 s, and scalar was not +# measured there. ``None`` means no ceiling. +DEFAULT_CEILING_SCALAR_CELLS: int | None = 4_000_000_000 +# Eager polars RSS is 0.28 GB + 0.543 GB per Mrow, and with ``pre_limit`` False +# a sorted window peaks at 0.23 + 1.62 GB per Mrow (3.0x the frame). With a +# 32 GB budget, four files resident and one window in flight, 8M rows is 27 GB +# and 10M is 33.5 GB. That figure assumes ``pre_limit`` False; main's 1M +# ``pre_limit`` windows do not copy the frame and the resident-only bound would +# be about 14M rows. +DEFAULT_POLARS_ROUTE_ROWS: int = 8_000_000 # Environment variable -> StatsLimits field. _ENV_FIELDS = {"BUCKAROO_STATS_FULL_AUTO_ROWS": "full_auto_rows", "BUCKAROO_STATS_SCALAR_AUTO_CELLS": "scalar_auto_cells", "BUCKAROO_STATS_CEILING_FULL_ROWS": "ceiling_full_rows", - "BUCKAROO_STATS_CEILING_SCALAR_CELLS": "ceiling_scalar_cells", "BUCKAROO_POLARS_ROUTE_ROWS": "polars_route_rows"} + "BUCKAROO_STATS_CEILING_SCALAR_CELLS": "ceiling_scalar_cells", "BUCKAROO_POLARS_ROUTE_ROWS": "polars_route_rows", + "BUCKAROO_STATS_FULL_AUTO_CELLS": "full_auto_cells", "BUCKAROO_STATS_CEILING_FULL_CELLS": "ceiling_full_cells"} # Lowest to highest. TIERS = ("schema", "scalar", "full") @@ -99,6 +128,9 @@ class StatsLimits: ceiling_full_rows: int = DEFAULT_CEILING_FULL_ROWS ceiling_scalar_cells: int | None = DEFAULT_CEILING_SCALAR_CELLS polars_route_rows: int = DEFAULT_POLARS_ROUTE_ROWS + # Added after the first five fields, so a positional construction keeps its meaning. + full_auto_cells: int = DEFAULT_FULL_AUTO_CELLS + ceiling_full_cells: int = DEFAULT_CEILING_FULL_CELLS @classmethod def from_env(cls) -> StatsLimits: @@ -123,7 +155,7 @@ def _count(name: str, value: Any) -> int: def _size_tier(rows: int, cells: int, limits: StatsLimits) -> str: """The tier the server picks on its own for an entry of this size.""" - if rows <= limits.full_auto_rows: + if rows <= limits.full_auto_rows and cells <= limits.full_auto_cells: return "full" if cells <= limits.scalar_auto_cells: return "scalar" @@ -134,7 +166,9 @@ def _ceiling_tier(rows: int, cells: int, limits: StatsLimits) -> str: """The highest tier any caller may reach for an entry of this size.""" if limits.ceiling_scalar_cells is not None and cells > limits.ceiling_scalar_cells: return "schema" - return "full" if rows <= limits.ceiling_full_rows else "scalar" + if rows <= limits.ceiling_full_rows and cells <= limits.ceiling_full_cells: + return "full" + return "scalar" def resolve_stats_policy(backend: str, source_kind: str, rows: int, cols: int, bytes: int | None = None, @@ -212,11 +246,15 @@ def resolve_stats_policy(backend: str, source_kind: str, rows: int, cols: int, b def route_polars_entry(rows: int, cols: int, limits: StatsLimits | None = None) -> str: """``"xorq"`` for a polars entry above R rows, ``"eager"`` otherwise. - Eager polars holds about 6.7 GB at 10.8M rows and cannot hold the 78M-row + Eager polars holds about 6.8 GB at 12M rows and cannot hold the 78M-row entry, so a host sends entries above R to ``/load_expr`` instead of - ``/load`` with ``backend="polars"``. R is ``limits.polars_route_rows``. - ``cols`` is accepted so the rule can become a cell budget without a - signature change; only ``rows`` decides today. + ``/load`` with ``backend="polars"``. R is ``limits.polars_route_rows``, + 8M rows by default. That figure assumes ``PolarsServerSampling.pre_limit`` + is False, so a sorted window copies the frame (3.0x its size at the peak); + with the 1M ``pre_limit`` on main windows do not copy it and the same + memory budget would allow about 14M rows. ``cols`` is accepted so the rule + can become a cell budget without a signature change; only ``rows`` decides + today. """ rows = _count("rows", rows) _count("cols", cols) diff --git a/tests/unit/server/test_stats_policy.py b/tests/unit/server/test_stats_policy.py index 9a032b7b3..4e9bc7d62 100644 --- a/tests/unit/server/test_stats_policy.py +++ b/tests/unit/server/test_stats_policy.py @@ -188,9 +188,9 @@ def test_estimate_includes_bytes_when_given(self, sp, lim): assert out["estimate"] == {"rows": 12_400_000, "cols": 43, "bytes": 987_654_321} def test_numpy_integer_counts_are_accepted(self, sp, lim): - out = _resolve(sp, lim, "xorq", "parquet", np.int64(11_000_000), np.int64(43)) + out = _resolve(sp, lim, "xorq", "parquet", np.int64(13_000_000), np.int64(43)) assert out["tier_target"] == "scalar" - assert out["estimate"] == {"rows": 11_000_000, "cols": 43} + assert out["estimate"] == {"rows": 13_000_000, "cols": 43} @pytest.mark.parametrize("rows, cols", [(-1, 3), (3, -1), (None, 3), (3, None), (1.5, 3)]) def test_bad_counts_are_rejected(self, sp, lim, rows, cols): @@ -273,7 +273,7 @@ def test_host_full_above_the_ceiling_resolves_lower(self, sp, lim): def test_reload_re_resolves_against_the_new_row_count(self, sp, lim): """/reload_expr keeps the stored host tier and re-resolves. An entry that fit under the ceiling and then grew past it loses full.""" - before = _resolve(sp, lim, "xorq", "parquet", 40_000_000, 44, "full") + before = _resolve(sp, lim, "xorq", "parquet", 20_000_000, 44, "full") assert before["tier_target"] == "full" after = _resolve(sp, lim, "xorq", "parquet", 78_000_000, 44, "full") assert after["tier_target"] == "scalar" @@ -371,6 +371,25 @@ def test_a_scalar_ceiling_is_on_by_default(self, sp): assert out["reason"] == "ceiling" + @pytest.mark.parametrize("rows, cols", + [(12_000_001, 43), (10_000_001, 52), (11_000_000, 91), (20_000_000, 50), (20_000_001, 50), (22_727_273, 44), + (25_000_000, 40), (25_000_000, 41), (25_000_001, 10), (50_000_000, 81), (100_000_000, 40), + (100_000_001, 40)]) + def test_requestable_matches_a_force_at_the_row_and_cell_boundaries(self, sp, lim, rows, cols): + """With the cell bounds, the tiers a result offers are still exactly + the ones a force request is granted.""" + out = _resolve(sp, lim, "xorq", "parquet", rows, cols) + for tier in ("scalar", "full"): + if TIER_RANK[tier] <= TIER_RANK[out["tier_target"]]: + continue + forced = _resolve(sp, lim, "xorq", "parquet", rows, cols, tier) + if tier in out["requestable"]: + assert forced["tier_target"] == tier + else: + assert forced["tier_target"] != tier + assert forced["reason"] == "ceiling" + + class TestLimitsFromEnv: def test_defaults_when_the_environment_is_empty(self, sp): assert sp.StatsLimits.from_env() == sp.StatsLimits() @@ -395,7 +414,7 @@ def test_invalid_value_falls_back_to_the_default_and_warns(self, sp, monkeypatch def test_empty_value_means_unset(self, sp, monkeypatch): monkeypatch.setenv("BUCKAROO_STATS_CEILING_SCALAR_CELLS", "") - assert sp.StatsLimits.from_env().ceiling_scalar_cells is None + assert sp.StatsLimits.from_env().ceiling_scalar_cells == sp.StatsLimits().ceiling_scalar_cells def test_resolve_reads_the_environment_when_no_limits_are_passed(self, sp, monkeypatch): monkeypatch.setenv("BUCKAROO_STATS_CEILING_FULL_ROWS", "1000") @@ -419,6 +438,13 @@ def test_the_full_ceiling_cell_bound_comes_from_the_environment(self, sp, monkey assert out["tier_target"] == "scalar" assert out["reason"] == "ceiling" + def test_the_scalar_ceiling_comes_from_the_environment(self, sp, monkeypatch): + args = ("xorq", "parquet", 5_000, 10) + monkeypatch.setenv("BUCKAROO_STATS_CEILING_SCALAR_CELLS", "1000") + out = sp.resolve_stats_policy(*args, host_tier="scalar") + assert out["tier_target"] == "schema" + assert out["reason"] == "ceiling" + def test_the_environment_is_read_on_each_call(self, sp, monkeypatch): args = ("xorq", "parquet", 5_000, 10) assert sp.resolve_stats_policy(*args, host_tier="full")["tier_target"] == "full" From 6bd63ed7e14c6e31b33bbdf6c6910735e5e07ffe Mon Sep 17 00:00:00 2001 From: Paddy Mullen Date: Sun, 4 Oct 2026 01:40:09 -0400 Subject: [PATCH 16/39] 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 17/39] 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 18/39] 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 19/39] 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 20/39] 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 21/39] 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 22/39] 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 23/39] 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 24/39] 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 From 751ea324cc05279eb4d5122347a3aaebe7a46837 Mon Sep 17 00:00:00 2001 From: Paddy Mullen Date: Sun, 4 Oct 2026 05:12:35 -0400 Subject: [PATCH 25/39] test(server): failing tests for the stats policy on the wire and the handlers (rows-first p33) stats_tier takes auto, full, scalar or schema on /load_expr and /reload_expr and is stored with the pair. The policy resolves after the schema-tier dataflow and the count exist, is stored on the session, is reported in df_meta.stats (tier_target, reason, auto_request, requestable, estimate, omitted_keys, approx_keys, demand_columns, each with its documented default) and is applied at WebSocket open only for a client that sends ?caps=stats_update,stats_ondemand. The version-skew cases (no caps, stats_update only, both bits, an old server) and /load, which keeps resolving to full, are covered. Co-Authored-By: Claude Sonnet 5.5 --- tests/unit/server/test_load_expr.py | 402 +++++++++++++++++++++++++ tests/unit/server/test_server.py | 25 ++ tests/unit/server/test_stats_policy.py | 254 +++++++++++++++- 3 files changed, 679 insertions(+), 2 deletions(-) diff --git a/tests/unit/server/test_load_expr.py b/tests/unit/server/test_load_expr.py index 813955234..faa71435e 100644 --- a/tests/unit/server/test_load_expr.py +++ b/tests/unit/server/test_load_expr.py @@ -38,6 +38,7 @@ 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 session as session_mod # 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 @@ -2525,6 +2526,407 @@ async def test_a_columns_hint_that_is_not_a_list_of_names_is_only_a_hint(self): self.assertEqual(self._session("sw-junk-hint").stat_runs, {}) +# The table _build_stats_wire_dir builds: five rows, three columns. +_WIRE_ESTIMATE = {"rows": 5, "cols": 3} +_ONDEMAND = "stats_update,stats_ondemand" +# Limits that put the table in the scalar or the schema tier by size. +_SCALAR_BY_SIZE = {"full_auto_rows": 3, "scalar_auto_cells": 1_000} +_SCHEMA_BY_SIZE = {"full_auto_rows": 3, "scalar_auto_cells": 10} + + +@contextmanager +def _stats_limits(**limits): + """Server stats thresholds for the block, as ``BUCKAROO_STATS_*`` + environment overrides, which ``resolve_stats_policy`` reads on every call.""" + env = {f"BUCKAROO_STATS_{name.upper()}": str(value) for name, value in limits.items()} + with patch.dict(os.environ, env): + yield + + +# (host tier, limits, df_meta.stats an ondemand client sees, minus gen). The +# five-by-three table is full by default; each set of limits moves it down. +_POLICY_CASES = { + "auto, within every threshold": ("auto", {}, {"status": "pending", "tier": "schema", "tier_target": "full", + "requestable": [], "estimate": _WIRE_ESTIMATE}), + "auto, scalar by size": ("auto", {"full_auto_rows": 3}, {"status": "not_computed", "tier": "schema", + "reason": "size", "tier_target": "scalar", "estimate": _WIRE_ESTIMATE}), + "auto, schema by size": ("auto", {"full_auto_rows": 3, "scalar_auto_cells": 10}, {"status": "not_computed", + "tier": "schema", "reason": "size", "tier_target": "schema", "auto_request": False, + "requestable": ["scalar", "full"], "estimate": _WIRE_ESTIMATE}), + "full, over the ceiling": ("full", {"ceiling_full_rows": 3}, {"status": "not_computed", "tier": "schema", + "reason": "ceiling", "tier_target": "scalar", "requestable": [], "estimate": _WIRE_ESTIMATE}), + "scalar, named by the host": ("scalar", {}, {"status": "not_computed", "tier": "schema", "reason": "host", + "tier_target": "scalar", "estimate": _WIRE_ESTIMATE}), + "schema, named by the host": ("schema", {}, {"status": "not_computed", "tier": "schema", "reason": "host", + "tier_target": "schema", "auto_request": False, "requestable": ["scalar", "full"], + "estimate": _WIRE_ESTIMATE}), + "scalar, over the scalar ceiling": ("scalar", {"ceiling_scalar_cells": 10}, {"status": "not_computed", + "tier": "schema", "reason": "ceiling", "tier_target": "schema", "auto_request": False, "requestable": [], + "estimate": _WIRE_ESTIMATE}), +} + + +class TestStatsPolicyWire(tornado.testing.AsyncHTTPTestCase): + """The stats policy on a ``/load_expr`` session (rows-first p33): + ``stats_tier`` ``auto | full | scalar | schema`` stored with the pair, the + policy resolved after the schema-tier dataflow and the count exist, reported + in ``df_meta.stats`` and applied at WebSocket open for a client that + advertises ``stats_ondemand`` alone, with the version-skew cases.""" + + 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) + 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, limits=None, **body): + with _stats_limits(**(limits or {})): + 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 _reload(self, sid, limits=None, **body): + with _stats_limits(**(limits or {})): + resp = await _post(self.get_http_port(), f"/reload_expr/{sid}", body) + self.assertEqual(resp.code, 200, resp.body) + + async def _connect(self, sid, caps=None): + 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 + + def _schema_rows(self, frame): + return list(_rows_by_stat(frame["df_data_dict"]["all_stats"])) + + async def _inline_frame(self, sid): + """The complete message an inline session of the same build sends.""" + await self._load(sid) + return (await self._connect(sid))[1] + + def _assert_complete(self, frame, inline_frame): + stats = self._stats(frame) + self.assertEqual((stats["status"], stats["tier"]), ("complete", "full")) + self.assertIn("histogram_bins", _rows_by_stat(frame["df_data_dict"]["all_stats"])) + self.assertEqual(_comparable(frame), _comparable(inline_frame)) + + @tornado.testing.gen_test + async def test_every_stats_tier_value_is_accepted_and_stored(self): + for tier in ("auto", "full", "scalar", "schema"): + sid = f"pw-accept-{tier}" + await self._load(sid, stats_tier=tier, stats_delivery="deferred") + session = self._session(sid) + self.assertEqual((session.stats_tier, session.stats_delivery), (tier, "deferred")) + self.assertEqual(session.xorq_dataflow.stats_tier, "schema") + resp = await _post(self.get_http_port(), f"/reload_expr/{sid}", {"stats_tier": "auto"}) + self.assertEqual(resp.code, 200, resp.body) + self.assertEqual(session.stats_tier, "auto") + resp = await _post(self.get_http_port(), "/load_expr", + {"session": "pw-bad", "build_dir": self.build_path, "stats_tier": "none"}) + self.assertEqual((resp.code, json.loads(resp.body)["error_code"]), (400, "invalid_stats_tier")) + + @tornado.testing.gen_test + async def test_the_default_tier_stays_full(self): + await self._load("pw-default", stats_delivery="deferred") + session = self._session("pw-default") + self.assertEqual(session.stats_tier, "full") + self.assertEqual(session.stats_policy["tier_target"], "full") + + @tornado.testing.gen_test + async def test_an_inline_session_has_no_schema_dataflow_to_resolve_against(self): + """The constructor runs the stats of an inline session, so auto is full + there and the session sends the message it always has.""" + with _count_stat_queries() as queries: + await self._load("pw-auto-inline", limits={"full_auto_rows": 1, "ceiling_full_rows": 1}, + stats_tier="auto") + session = self._session("pw-auto-inline") + self.assertEqual((session.stats_tier, session.stats_delivery, session.stats_policy), ("auto", "inline", None)) + self.assertEqual(session.xorq_dataflow.stats_tier, "full") + self.assertTrue(queries, "an inline session runs its stats in the constructor") + _, frame = await self._connect("pw-auto-inline", caps=_ONDEMAND) + self.assertNotIn("stats", frame["df_meta"]) + + @tornado.testing.gen_test + async def test_the_policy_resolves_once_the_schema_dataflow_and_the_count_exist(self): + events, calls = [], [] + original_metadata = xorq_loading.get_xorq_metadata + original_resolve = stats_wire.resolve_stats_policy + + def metadata(*args, **kwargs): + events.append("dataflow and count") + return original_metadata(*args, **kwargs) + + def resolve(*args, **kwargs): + events.append("policy") + calls.append((args, kwargs)) + return original_resolve(*args, **kwargs) + + with ( + patch.object(xorq_loading, "get_xorq_metadata", metadata), + patch.object(stats_wire, "resolve_stats_policy", resolve), + _count_stat_queries() as queries, + ): + await self._load("pw-order", stats_tier="auto", stats_delivery="deferred") + self.assertEqual(events, ["dataflow and count", "policy"]) + (args, kwargs), = calls + self.assertEqual((args[0], args[2], args[3], kwargs["host_tier"]), ("xorq", 5, 3, "auto")) + self.assertEqual(queries, [], "resolving the policy runs no data stat query") + + @tornado.testing.gen_test + async def test_the_resolved_policy_is_stored_on_the_session(self): + limits = {"full_auto_rows": 3, "scalar_auto_cells": 10} + await self._load("pw-store", limits=limits, stats_tier="auto", stats_delivery="deferred") + session = self._session("pw-store") + self.assertEqual(session.stats_policy, {"tier_target": "schema", "auto_request": False, + "requestable": ["scalar", "full"], "reason": "size", "estimate": _WIRE_ESTIMATE}) + self.assertEqual((session.stats_status, session.stats_reason), ("not_computed", "size")) + + @tornado.testing.gen_test + async def test_df_meta_stats_carries_the_policy_fields_for_an_ondemand_client(self): + for index, (name, (tier, limits, expected)) in enumerate(_POLICY_CASES.items()): + sid = f"pw-fields-{index}" + await self._load(sid, limits=limits, stats_tier=tier, stats_delivery="deferred") + _, frame = await self._connect(sid, caps=_ONDEMAND) + stats = self._stats(frame) + self.assertEqual(stats, {**expected, "gen": stats["gen"]}, name) + self.assertEqual(self._schema_rows(frame), ["dtype"], f"{name}: the first message is the schema tier") + + @tornado.testing.gen_test + async def test_a_warm_repost_with_the_same_auto_tier_short_circuits(self): + body = {"session": "pw-warm", "build_dir": self.build_path, "stats_tier": "auto", "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": "pw-warm", "build_dir": self.build_path}) + self.assertEqual((same.code, omitted.code), (200, 200)) + session = self._session("pw-warm") + self.assertEqual((session.stats_tier, session.stats_delivery), ("auto", "deferred")) + + @tornado.testing.gen_test + async def test_a_changed_tier_rebuilds_and_resolves_again(self): + sid = "pw-change" + await self._load(sid, stats_tier="auto", stats_delivery="deferred") + self.assertEqual(self._session(sid).stats_policy["tier_target"], "full") + calls = [] + original = xorq_loading.load_expr_build_dir + + 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 self._load(sid, stats_tier="scalar", stats_delivery="deferred") + self.assertEqual(len(calls), 1, "a changed tier must rebuild") + self.assertEqual(self._session(sid).stats_policy["tier_target"], "scalar") + await self._load(sid, stats_tier="scalar", stats_delivery="deferred") + self.assertEqual(len(calls), 1, "the same tier takes the warm exit") + + @tornado.testing.gen_test + async def test_reload_expr_resolves_the_policy_again(self): + sid = "pw-reload" + await self._load(sid, stats_tier="auto", stats_delivery="deferred") + ws, first = await self._connect(sid, caps=_ONDEMAND) + gen = self._stats(first)["gen"] + self.assertEqual(self._stats(first)["tier_target"], "full") + + await self._reload(sid, limits={"full_auto_rows": 3}) + pushed = self._stats(await _read_json(ws)) + self.assertEqual((pushed["status"], pushed["reason"], pushed["tier_target"]), + ("not_computed", "size", "scalar")) + self.assertGreater(pushed["gen"], gen) + self.assertEqual(self._session(sid).stats_policy["tier_target"], "scalar") + + await self._reload(sid, stats_tier="schema") + pushed = self._stats(await _read_json(ws)) + self.assertEqual((pushed["reason"], pushed["tier_target"], pushed["auto_request"]), ("host", "schema", False)) + self.assertEqual(self._session(sid).stats_tier, "schema") + + await self._reload(sid, stats_tier="auto") + pushed = self._stats(await _read_json(ws)) + self.assertEqual((pushed["status"], pushed["tier_target"]), ("pending", "full")) + + @tornado.testing.gen_test + async def test_a_reload_that_omits_the_tier_keeps_the_session_s_and_resolves_it_against_the_stored_count(self): + """The expression is reused, so its count is the one the load took + (``session.metadata``), and the reload runs no second count.""" + sid = "pw-reload-keep" + await self._load(sid, stats_tier="auto", stats_delivery="deferred") + session = self._session(sid) + session.metadata = {**session.metadata, "rows": 20_000_000} + await self._reload(sid) + self.assertEqual(session.stats_tier, "auto") + self.assertEqual((session.stats_policy["tier_target"], session.stats_policy["estimate"]), + ("scalar", {"rows": 20_000_000, "cols": 3})) + + @tornado.testing.gen_test + async def test_the_policy_survives_a_dataflow_field_change(self): + await self._load("pw-keep", limits=_SCALAR_BY_SIZE, stats_tier="auto", stats_delivery="deferred") + ws, first = await self._connect("pw-keep", caps=_ONDEMAND) + gen = self._stats(first)["gen"] + ws.write_message(_state_change(quick_command_args={"search": ["a"]})) + stats = self._stats(await _read_json(ws)) + self.assertEqual((stats["status"], stats["reason"], stats["tier_target"], stats["gen"]), + ("not_computed", "size", "scalar", gen + 1)) + + # -- version skew ------------------------------------------------------- + + async def _policy_session(self, sid, **limits): + """A deferred ``auto`` session that resolved to the schema tier by size.""" + await self._load(sid, limits=limits or _SCHEMA_BY_SIZE, stats_tier="auto", stats_delivery="deferred") + + @tornado.testing.gen_test + async def test_no_caps_gets_a_complete_state_at_connect(self): + inline = await self._inline_frame("pw-skew-inline") + await self._policy_session("pw-skew-none") + _, legacy = await self._connect("pw-skew-none") + self._assert_complete(legacy, inline) + self.assertEqual(self._session("pw-skew-none").stats_status, "complete") + + @tornado.testing.gen_test + async def test_stats_ondemand_without_stats_update_is_a_client_with_no_caps(self): + inline = await self._inline_frame("pw-skew-inline-2") + await self._policy_session("pw-skew-odonly") + _, frame = await self._connect("pw-skew-odonly", caps="stats_ondemand") + self._assert_complete(frame, inline) + + @tornado.testing.gen_test + async def test_stats_update_only_is_served_as_a_session_headed_for_full(self): + await self._policy_session("pw-skew-update") + ws, first = await self._connect("pw-skew-update", caps="stats_update") + stats = self._stats(first) + self.assertEqual(stats, {"status": "pending", "tier": "schema", "gen": stats["gen"]}) + self.assertEqual(self._schema_rows(first), ["dtype"]) + + ws.write_message(_stats_request(stats["gen"])) + update = await _read_json(ws) + self.assertEqual((update["type"], update["final"], update["tier"]), ("stats_update", True, "full")) + self.assertIn("histogram_bins", _rows_by_stat(update["payload"])) + self.assertEqual(self._session("pw-skew-update").stats_status, "complete") + + @tornado.testing.gen_test + async def test_stats_update_only_pulls_units_for_a_policy_session(self): + patcher = patch.object(stats_wire, "STATS_BUDGET_S", 0) + patcher.start() + self.addCleanup(patcher.stop) + await self._policy_session("pw-skew-units") + ws, first = await self._connect("pw-skew-units", caps="stats_update") + gen = self._stats(first)["gen"] + replies = [] + for _ in range(12): + ws.write_message(_stats_request(gen, incremental=True)) + replies.append(await _read_json(ws)) + if replies[-1]["final"]: + break + 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(self._session("pw-skew-units").stats_status, "complete") + + @tornado.testing.gen_test + async def test_both_bits_get_the_schema_tier_not_computed_and_run_no_stat_query(self): + await self._policy_session("pw-skew-both") + with _count_stat_queries() as queries: + ws, first = await self._connect("pw-skew-both", caps=_ONDEMAND) + stats = self._stats(first) + self.assertEqual((stats["status"], stats["tier"], stats["reason"], stats["tier_target"]), + ("not_computed", "schema", "size", "schema")) + self.assertEqual((stats["auto_request"], stats["requestable"]), (False, ["scalar", "full"])) + self.assertEqual(self._schema_rows(first), ["dtype"]) + self.assertEqual(queries, []) + # The server enforces it: nothing is requestable until the request path serves a tier. + ws.write_message(_stats_request(stats["gen"])) + aborted = await _read_json(ws) + self.assertEqual((aborted["type"], aborted["reason"]), ("stats_aborted", "not_requestable")) + self.assertEqual(self._session("pw-skew-both").stats_status, "not_computed") + + @tornado.testing.gen_test + async def test_the_three_client_kinds_see_one_session_each_their_own_way(self): + inline = await self._inline_frame("pw-skew-inline-3") + await self._policy_session("pw-skew-mixed") + _, ondemand = await self._connect("pw-skew-mixed", caps=_ONDEMAND) + _, update = await self._connect("pw-skew-mixed", caps="stats_update") + self.assertEqual(self._stats(ondemand)["status"], "not_computed") + self.assertEqual(self._stats(update)["status"], "pending") + _, legacy = await self._connect("pw-skew-mixed") + self._assert_complete(legacy, inline) + + @tornado.testing.gen_test + async def test_an_ondemand_client_is_sent_its_frame_before_a_legacy_client_completes_the_session(self): + await self._policy_session("pw-skew-order", **_SCALAR_BY_SIZE) + a, a_open = await self._connect("pw-skew-order", caps=_ONDEMAND) + gen = self._stats(a_open)["gen"] + b, b_open = await self._connect("pw-skew-order") + self.assertEqual(self._stats(b_open)["status"], "complete") + + a.write_message(_state_change(quick_command_args={"search": ["a"]})) + a_frame, b_frame = await _read_json(a), await _read_json(b) + self.assertEqual((self._stats(a_frame)["status"], self._stats(a_frame)["gen"]), ("not_computed", gen + 1)) + self.assertEqual((self._stats(b_frame)["status"], self._stats(b_frame)["gen"]), ("complete", gen + 1)) + + @tornado.testing.gen_test + async def test_a_tier_the_host_named_reaches_every_client_as_the_host_chose_it(self): + await self._load("pw-skew-host", stats_tier="scalar", stats_delivery="deferred") + _, legacy = await self._connect("pw-skew-host") + self.assertEqual(self._schema_rows(legacy), ["dtype"]) + self.assertEqual(self._stats(legacy)["status"], "not_computed") + ws, update = await self._connect("pw-skew-host", caps="stats_update") + self.assertEqual(self._stats(update), {"status": "not_computed", "tier": "schema", + "gen": self._stats(update)["gen"], "reason": "host"}) + ws.write_message(_stats_request(self._stats(update)["gen"])) + self.assertEqual((await _read_json(ws))["reason"], "not_requestable") + _, ondemand = await self._connect("pw-skew-host", caps=_ONDEMAND) + self.assertEqual(self._stats(ondemand)["tier_target"], "scalar") + + @tornado.testing.gen_test + async def test_a_client_of_an_old_server_reads_a_frame_with_no_stats_as_complete(self): + """An old server sends no ``df_meta.stats`` and no policy field; a new + client falls back to the documented defaults and waits for nothing.""" + await self._load("pw-old") + _, frame = await self._connect("pw-old", caps=_ONDEMAND) + stats = session_mod.stats_with_defaults(frame["df_meta"]) + self.assertEqual((stats["status"], stats["tier"], stats["tier_target"]), ("complete", "full", "full")) + self.assertEqual((stats["auto_request"], stats["requestable"]), (True, ["full"])) + self.assertEqual(stats["demand_columns"], []) + + @tornado.testing.gen_test + async def test_an_explicit_full_changes_no_frame(self): + """With ``stats_tier`` full the message is the one sent without the + field, and the policy is recorded but not reported.""" + for delivery in ("inline", "deferred"): + await self._load(f"pw-same-{delivery}-a", stats_delivery=delivery) + await self._load(f"pw-same-{delivery}-b", stats_tier="full", stats_delivery=delivery) + # A client with no caps completes a deferred session, so it goes last. + for caps in ("stats_update", _ONDEMAND, None): + _, before = await self._connect(f"pw-same-{delivery}-a", caps=caps) + _, after = await self._connect(f"pw-same-{delivery}-b", caps=caps) + self.assertEqual(_as_json(before["df_meta"]), _as_json(after["df_meta"]), (delivery, caps)) + self.assertEqual(_comparable(before), _comparable(after), (delivery, caps)) + self.assertEqual(list(before), list(after)) + self.assertEqual(self._session("pw-same-deferred-b").stats_policy["tier_target"], "full") + self.assertIsNone(self._session("pw-same-inline-b").stats_policy) + + 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.""" diff --git a/tests/unit/server/test_server.py b/tests/unit/server/test_server.py index b1b8f9032..5d4c1f87d 100644 --- a/tests/unit/server/test_server.py +++ b/tests/unit/server/test_server.py @@ -11,6 +11,7 @@ import tornado.websocket from buckaroo.server.app import make_app as _make_app +from buckaroo.server.session import build_state_message # Temp file cleanup fails on Windows due to file locking (WinError 32) pytestmark = pytest.mark.skipif(sys.platform == "win32", reason="Temp file locking prevents cleanup on Windows") @@ -242,6 +243,30 @@ def test_load_buckaroo_without_optional_configs_keeps_defaults(self): finally: os.unlink(f.name) + def test_load_does_not_take_stats_tier_and_resolves_to_full(self): + """/load serves the eager backends, whose stats cost does not grow with + the file, so they resolve to full: the request fields are not read and + a policy left on the session by an earlier /load_expr is dropped + (rows-first p33).""" + with tempfile.NamedTemporaryFile(suffix=".csv", delete=False) as f: + _write_test_csv(f.name) + try: + body = {"session": "stats-tier-1", "path": f.name, "mode": "buckaroo", "stats_tier": "schema", + "stats_delivery": "deferred"} + for _ in range(2): + resp = self.fetch("/load", method="POST", body=json.dumps(body), + headers={"Content-Type": "application/json"}) + self.assertEqual(resp.code, 200) + session = self._app.settings["sessions"].get("stats-tier-1") + self.assertEqual((session.stats_tier, session.stats_delivery), ("full", "inline")) + self.assertIsNone(session.stats_policy) + self.assertEqual(session.stats_status, "complete") + self.assertNotIn("stats", build_state_message(session)["df_meta"]) + # A /load_expr that came first would have left a policy behind. + session.stats_policy = {"tier_target": "schema", "reason": "size"} + finally: + os.unlink(f.name) + class TestSessionPage(tornado.testing.AsyncHTTPTestCase): def get_app(self): diff --git a/tests/unit/server/test_stats_policy.py b/tests/unit/server/test_stats_policy.py index 4e9bc7d62..a16f0ca6e 100644 --- a/tests/unit/server/test_stats_policy.py +++ b/tests/unit/server/test_stats_policy.py @@ -1,8 +1,10 @@ """Unit tests for ``buckaroo.server.stats_policy`` (rows-first p31, p31b). Everything here is pure Python: the policy is a function of numbers, the -probes read a schema or a parquet footer, and nothing needs a server. The -module is not wired into any handler, session or dataflow in this phase. +probes read a schema or a parquet footer, and nothing needs a server. The last +sections (rows-first p33) test how the session carries the result and how +``df_meta.stats`` reports it, still without a server; the HTTP and WebSocket +side is in test_load_expr.py. The thresholds are the provisional values proposed by the phase-0 measurements (p31b). The boundary tables run on the module defaults, so a @@ -16,6 +18,7 @@ import io import json import logging +from types import SimpleNamespace import numpy as np import pandas as pd @@ -24,6 +27,9 @@ import pyarrow.parquet as pq import pytest +from buckaroo.server import session as session_mod +from buckaroo.server import stats_wire + ENV_FIELDS = {"BUCKAROO_STATS_FULL_AUTO_ROWS": "full_auto_rows", "BUCKAROO_STATS_FULL_AUTO_CELLS": "full_auto_cells", "BUCKAROO_STATS_SCALAR_AUTO_CELLS": "scalar_auto_cells", "BUCKAROO_STATS_CEILING_FULL_ROWS": "ceiling_full_rows", "BUCKAROO_STATS_CEILING_FULL_CELLS": "ceiling_full_cells", @@ -603,3 +609,247 @@ def test_the_result_is_one_of_two_strings(self, sp, lim): def test_bad_counts_are_rejected(self, sp, lim, rows, cols): with pytest.raises((TypeError, ValueError)): sp.route_polars_entry(rows, cols, limits=lim) + + +# --------------------------------------------------------------------------- +# The policy on the session and the wire (rows-first p33) +# --------------------------------------------------------------------------- + +# The table the wire tests load: five rows and three columns. +ESTIMATE = {"rows": 5, "cols": 3} +ONDEMAND_CAPS = frozenset({"stats_update", "stats_ondemand"}) +UPDATE_CAPS = frozenset({"stats_update"}) +LEGACY_CAPS = frozenset() + +# Limits that put a five-by-three table in each tier. +SCALAR_BY_SIZE = dict(full_auto_rows=3, scalar_auto_cells=1_000) +SCHEMA_BY_SIZE = dict(full_auto_rows=3, scalar_auto_cells=10) +FULL_OVER_CEILING = dict(ceiling_full_rows=3) +SCALAR_OVER_CEILING = dict(ceiling_scalar_cells=10) + + +def _session(sp, stats_tier="auto", delivery="deferred", limits=None, resolve=True): + """A buckaroo-mode session after a load that resolved its policy and began + its first stats generation, the way ``/load_expr`` leaves one.""" + session = session_mod.SessionState(session_id="s", path="p") + session.mode = "buckaroo" + session.df_meta = {"total_rows": 5} + session.stats_tier, session.stats_delivery = stats_tier, delivery + session.stats_policy = ( + sp.resolve_stats_policy("xorq", "xorq_build", 5, 3, host_tier=stats_tier, + limits=sp.StatsLimits(**(limits or {}))) + if resolve else None) + session_mod.begin_stats_generation(session) + return session + + +def _client(caps): + return SimpleNamespace(caps=caps, search_string="") + + +# (host tier, delivery, limits) -> what an ``stats_ondemand`` client is told. +# A field equal to its documented default is left out, and ``tier_target`` and +# ``estimate`` are always there once a policy is reported. +POLICY_FRAMES = [ + ("auto", "deferred", {}, {"status": "pending", "tier": "schema", "gen": 1, "tier_target": "full", + "requestable": [], "estimate": ESTIMATE}), + ("auto", "deferred", SCALAR_BY_SIZE, {"status": "not_computed", "tier": "schema", "gen": 1, "reason": "size", + "tier_target": "scalar", "estimate": ESTIMATE}), + ("auto", "deferred", SCHEMA_BY_SIZE, {"status": "not_computed", "tier": "schema", "gen": 1, "reason": "size", + "tier_target": "schema", "auto_request": False, "requestable": ["scalar", "full"], "estimate": ESTIMATE}), + ("full", "deferred", FULL_OVER_CEILING, {"status": "not_computed", "tier": "schema", "gen": 1, + "reason": "ceiling", "tier_target": "scalar", "requestable": [], "estimate": ESTIMATE}), + ("scalar", "deferred", {}, {"status": "not_computed", "tier": "schema", "gen": 1, "reason": "host", + "tier_target": "scalar", "estimate": ESTIMATE}), + ("schema", "deferred", {}, {"status": "not_computed", "tier": "schema", "gen": 1, "reason": "host", + "tier_target": "schema", "auto_request": False, "requestable": ["scalar", "full"], "estimate": ESTIMATE}), + ("scalar", "inline", SCALAR_OVER_CEILING, {"status": "not_computed", "tier": "schema", "gen": 1, + "reason": "ceiling", "tier_target": "schema", "auto_request": False, "requestable": [], "estimate": ESTIMATE}), + ("schema", "inline", {}, {"status": "not_computed", "tier": "schema", "gen": 1, "reason": "host", + "tier_target": "schema", "auto_request": False, "requestable": ["scalar", "full"], "estimate": ESTIMATE}), +] + + +class TestStatsFieldDefaults: + """What a client assumes for a ``df_meta.stats`` field the server leaves out.""" + + def test_a_frame_with_no_stats_is_complete(self): + """An old server sends no ``df_meta.stats``, and a new client reads it as complete.""" + assert session_mod.stats_with_defaults({"total_rows": 3}) == {"status": "complete", "tier": "full", + "tier_target": "full", "gen": None, "reason": None, "auto_request": True, "requestable": ["full"], + "estimate": None, "omitted_keys": [], "approx_keys": [], "demand_columns": []} + + def test_a_pending_frame_is_headed_for_full(self): + out = session_mod.stats_with_defaults({"stats": {"status": "pending", "tier": "schema", "gen": 3}}) + assert (out["status"], out["tier"], out["tier_target"], out["gen"]) == ("pending", "schema", "full", 3) + assert (out["auto_request"], out["requestable"]) == (True, ["full"]) + + def test_a_not_computed_frame_without_a_target_targets_the_tier_it_has(self): + out = session_mod.stats_with_defaults({"stats": {"status": "not_computed", "tier": "schema", "gen": 2, + "reason": "host"}}) + assert (out["tier_target"], out["reason"]) == ("schema", "host") + + def test_fields_the_frame_carries_win(self): + stats = {"status": "not_computed", "tier": "schema", "gen": 2, "reason": "size", "tier_target": "scalar", + "auto_request": False, "requestable": ["scalar", "full"], "estimate": ESTIMATE, "omitted_keys": ["a"], + "approx_keys": ["b"], "demand_columns": ["c"]} + out = session_mod.stats_with_defaults({"stats": stats}) + assert {key: out[key] for key in stats} == stats + + def test_the_defaults_are_not_shared_between_calls(self): + first = session_mod.stats_with_defaults({}) + first["requestable"].append("scalar") + first["omitted_keys"].append("x") + second = session_mod.stats_with_defaults({}) + assert (second["requestable"], second["omitted_keys"]) == (["full"], []) + + def test_the_documented_defaults(self): + assert session_mod.STATS_FIELD_DEFAULTS == {"auto_request": True, "requestable": ["full"], "omitted_keys": [], + "approx_keys": [], "demand_columns": []} + assert session_mod.STATS_TIER_REQUESTS == ("auto", "full", "scalar", "schema") + assert session_mod.STATS_REASONS == ("size", "host", "cost", "ceiling") + + +class TestInitialStatsStatus: + @pytest.mark.parametrize("stats_tier, delivery, policy, expected", [ + ("full", "inline", None, ("complete", None)), + ("full", "deferred", None, ("pending", None)), + ("schema", "inline", None, ("not_computed", "host")), + ("scalar", "deferred", None, ("not_computed", "host")), + # auto resolves against a schema-tier dataflow, which an inline session never has. + ("auto", "inline", None, ("complete", None)), + ("auto", "deferred", None, ("pending", None)), + ("auto", "deferred", {"tier_target": "full", "reason": None}, ("pending", None)), + ("auto", "deferred", {"tier_target": "scalar", "reason": "size"}, ("not_computed", "size")), + ("auto", "deferred", {"tier_target": "schema", "reason": "size"}, ("not_computed", "size")), + ("full", "deferred", {"tier_target": "scalar", "reason": "ceiling"}, ("not_computed", "ceiling")), + ("schema", "deferred", {"tier_target": "schema", "reason": "host"}, ("not_computed", "host")), + ("scalar", "inline", {"tier_target": "schema", "reason": "ceiling"}, ("not_computed", "ceiling"))]) + def test_the_status_follows_the_pair_and_the_policy(self, stats_tier, delivery, policy, expected): + assert session_mod.initial_stats_status(stats_tier, delivery, policy) == expected + + @pytest.mark.parametrize("stats_tier, delivery, expected", [ + ("full", "inline", "full"), ("auto", "inline", "full"), ("full", "deferred", "schema"), + ("auto", "deferred", "schema"), ("scalar", "inline", "schema"), ("schema", "inline", "schema"), + ("scalar", "deferred", "schema")]) + def test_the_dataflow_is_built_at_the_schema_tier_unless_stats_run_inline(self, stats_tier, delivery, expected): + """No scalar-tier units exist yet, so a scalar target builds at the schema tier.""" + assert session_mod.dataflow_stats_tier(stats_tier, delivery) == expected + + +class TestStatsMeta: + @pytest.mark.parametrize("stats_tier, delivery, limits, expected", POLICY_FRAMES) + def test_a_stats_ondemand_client_is_told_the_policy(self, sp, stats_tier, delivery, limits, expected): + session = _session(sp, stats_tier, delivery, limits) + assert session_mod.stats_meta(session) == expected + message = session_mod.build_state_message(session) + assert message["df_meta"] == {"total_rows": 5, "stats": expected} + + @pytest.mark.parametrize("stats_tier, limits", [("auto", SCALAR_BY_SIZE), ("auto", SCHEMA_BY_SIZE), + ("full", FULL_OVER_CEILING)]) + def test_a_client_without_stats_ondemand_is_told_the_stats_are_on_their_way(self, sp, stats_tier, limits): + """A target the server chose below full is applied only for a client that + can take it; the others are served as a deferred session headed for full.""" + session = _session(sp, stats_tier, "deferred", limits) + pending = {"status": "pending", "tier": "schema", "gen": 1} + assert session_mod.stats_meta(session, ondemand=False) == pending + assert session_mod.build_state_message(session, ondemand=False)["df_meta"]["stats"] == pending + + @pytest.mark.parametrize("stats_tier, delivery, limits, reason", [ + ("scalar", "deferred", {}, "host"), ("schema", "deferred", {}, "host"), ("schema", "inline", {}, "host"), + ("scalar", "inline", SCALAR_OVER_CEILING, "ceiling")]) + def test_a_tier_the_host_named_reaches_every_client(self, sp, stats_tier, delivery, limits, reason): + """The host chose it, so there is nothing for an older client to complete, + and it gets the status without the policy fields.""" + session = _session(sp, stats_tier, delivery, limits) + assert session_mod.stats_meta(session, ondemand=False) == {"status": "not_computed", "tier": "schema", + "gen": 1, "reason": reason} + + @pytest.mark.parametrize("stats_tier, limits", + [("auto", {}), ("auto", SCALAR_BY_SIZE), ("full", FULL_OVER_CEILING)]) + def test_a_completed_session_reports_no_policy(self, sp, stats_tier, limits): + session = _session(sp, stats_tier, "deferred", limits) + session.stats_status, session.stats_reason = "complete", None + assert session_mod.stats_meta(session) == {"status": "complete", "tier": "full", "gen": 1} + + def test_the_gen_follows_the_generation_and_the_status_restarts_from_the_policy(self, sp): + session = _session(sp, "auto", "deferred", SCALAR_BY_SIZE) + session.stats_status, session.stats_reason = "complete", None + session_mod.begin_stats_generation(session) + assert session_mod.stats_meta(session) == {"status": "not_computed", "tier": "schema", "gen": 2, + "reason": "size", "tier_target": "scalar", "estimate": ESTIMATE} + + def test_omitted_approx_and_demand_keys_ride_along_when_there_are_some(self, sp): + session = _session(sp, "auto", "deferred", SCALAR_BY_SIZE) + assert not {"omitted_keys", "approx_keys", "demand_columns"} & set(session_mod.stats_meta(session)) + session.stats_policy.update(omitted_keys=["value_counts"], approx_keys=["distinct_count"], + demand_columns=["a", "b"]) + stats = session_mod.stats_meta(session) + assert (stats["omitted_keys"], stats["approx_keys"], stats["demand_columns"]) == ( + ["value_counts"], ["distinct_count"], ["a", "b"]) + # A client reads what is left out as empty. + session.stats_policy.update(omitted_keys=[], approx_keys=[], demand_columns=[]) + assert not {"omitted_keys", "approx_keys", "demand_columns"} & set(session_mod.stats_meta(session)) + + def test_every_reason_the_server_sends_is_a_documented_one(self): + for *_, expected in POLICY_FRAMES: + assert expected.get("reason", "size") in session_mod.STATS_REASONS + + def test_the_policy_is_what_the_pure_function_resolved(self, sp): + session = _session(sp, "auto", "deferred", SCHEMA_BY_SIZE) + assert session.stats_policy == sp.resolve_stats_policy("xorq", "xorq_build", 5, 3, + limits=sp.StatsLimits(**SCHEMA_BY_SIZE)) + + +class TestPolicyForTheSession: + """``resolve_session_policy``: what a load handler resolves, and when.""" + + @pytest.mark.parametrize("stats_tier", ["auto", "full", "scalar", "schema"]) + def test_a_schema_dataflow_resolves_against_the_count_it_has(self, stats_tier): + out = stats_wire.resolve_session_policy(stats_tier, "schema", 5, 3) + assert out["estimate"] == ESTIMATE + assert out["tier_target"] == {"auto": "full", "full": "full", "scalar": "scalar", "schema": "schema"}[stats_tier] + + def test_a_dataflow_that_ran_its_stats_has_nothing_to_resolve(self): + assert stats_wire.resolve_session_policy("full", "full", 5, 3) is None + assert stats_wire.resolve_session_policy("auto", "full", 5_000_000_000, 40) is None + + def test_the_ceiling_clips_a_host_that_asks_for_full(self, sp): + out = stats_wire.resolve_session_policy("full", "schema", 78_000_000, 44) + assert (out["tier_target"], out["reason"]) == ("scalar", "ceiling") + + def test_the_limits_are_read_on_each_call(self, monkeypatch): + assert stats_wire.resolve_session_policy("auto", "schema", 5, 3)["tier_target"] == "full" + monkeypatch.setenv("BUCKAROO_STATS_FULL_AUTO_ROWS", "3") + assert stats_wire.resolve_session_policy("auto", "schema", 5, 3)["tier_target"] == "scalar" + + +class TestCapabilities: + def test_the_second_bit_is_named(self): + assert stats_wire.STATS_UPDATE_CAP == "stats_update" + assert stats_wire.STATS_ONDEMAND_CAP == "stats_ondemand" + assert stats_wire.parse_caps("stats_update,stats_ondemand") == ONDEMAND_CAPS + + @pytest.mark.parametrize("caps, expected", [ + (ONDEMAND_CAPS, True), (UPDATE_CAPS, False), (LEGACY_CAPS, False), + # The bit means the client also merges stats_update; alone it is not enough. + (frozenset({"stats_ondemand"}), False)]) + def test_only_stats_update_with_stats_ondemand_counts_as_ondemand(self, caps, expected): + assert stats_wire.client_has_ondemand(_client(caps)) is expected + + def test_a_client_with_no_caps_attribute_is_not_ondemand(self): + assert stats_wire.client_has_ondemand(SimpleNamespace()) is False + + @pytest.mark.parametrize("caps, expected", [(ONDEMAND_CAPS, False), (UPDATE_CAPS, True), (LEGACY_CAPS, True)]) + def test_a_policy_target_below_full_is_missing_stats_to_every_client_but_an_ondemand_one(self, sp, caps, expected): + session = _session(sp, "auto", "deferred", SCHEMA_BY_SIZE) + assert stats_wire.stats_to_pull(session, _client(caps)) is expected + + @pytest.mark.parametrize("caps", [ONDEMAND_CAPS, UPDATE_CAPS, LEGACY_CAPS]) + def test_a_pending_session_has_stats_to_pull_for_everyone(self, sp, caps): + assert stats_wire.stats_to_pull(_session(sp, "auto", "deferred"), _client(caps)) is True + + @pytest.mark.parametrize("caps", [ONDEMAND_CAPS, UPDATE_CAPS, LEGACY_CAPS]) + def test_a_tier_the_host_named_has_nothing_to_pull(self, sp, caps): + assert stats_wire.stats_to_pull(_session(sp, "schema", "deferred"), _client(caps)) is False + assert stats_wire.stats_to_pull(_session(sp, "scalar", "deferred"), _client(caps)) is False From 769b5c73d458eef16783783ba854adc352c3f9da Mon Sep 17 00:00:00 2001 From: Paddy Mullen Date: Sun, 4 Oct 2026 05:31:05 -0400 Subject: [PATCH 26/39] test(server): failing tests for unit tiers, the final display config and the compare reset under a stats policy (rows-first p33) Cases found untested after the first tests commit, each one a regression the first set did not catch: a stats_update client pulling units from a policy session gets full-tier updates and the final reply carries the rebuilt df_display_args, a scalar tier named with inline delivery builds a schema dataflow, and /load_compare clears the policy a session held. Co-Authored-By: Claude Sonnet 5.5 --- tests/unit/server/test_load_compare.py | 27 ++++++++++++++++++++++++++ tests/unit/server/test_load_expr.py | 14 +++++++++++++ 2 files changed, 41 insertions(+) diff --git a/tests/unit/server/test_load_compare.py b/tests/unit/server/test_load_compare.py index ba9be9198..aaea8b081 100644 --- a/tests/unit/server/test_load_compare.py +++ b/tests/unit/server/test_load_compare.py @@ -8,6 +8,7 @@ import tornado.testing from buckaroo.server.app import make_app as _make_app +from buckaroo.server.session import build_state_message pytestmark = pytest.mark.skipif(sys.platform == "win32", reason="Temp file locking prevents cleanup on Windows") @@ -114,6 +115,32 @@ def test_eqs_reports_diffs(self): os.unlink(f1.name) os.unlink(f2.name) + def test_load_compare_resolves_to_full_and_clears_a_stats_policy(self): + """A compare session is a viewer with no dataflow, so no stats are + deferred or refused: a policy that an earlier /load_expr left on the + session is dropped (rows-first p33).""" + session = self._app.settings["sessions"].create("cmp-policy", "") + session.stats_tier, session.stats_delivery = "auto", "deferred" + session.stats_policy = {"tier_target": "schema", "reason": "size"} + df1 = pd.DataFrame({"id": [1, 2], "score": [100, 200]}) + df2 = pd.DataFrame({"id": [1, 2], "score": [100, 999]}) + + with tempfile.NamedTemporaryFile(suffix=".csv", delete=False) as f1, \ + tempfile.NamedTemporaryFile(suffix=".csv", delete=False) as f2: + _write_df(df1, f1.name) + _write_df(df2, f2.name) + try: + resp = self._post_compare({"session": "cmp-policy", "path1": f1.name, "path2": f2.name, + "join_columns": ["id"]}) + self.assertEqual(resp.code, 200) + finally: + os.unlink(f1.name) + os.unlink(f2.name) + self.assertEqual((session.stats_tier, session.stats_delivery), ("full", "inline")) + self.assertIsNone(session.stats_policy) + self.assertEqual(session.stats_status, "complete") + self.assertNotIn("stats", build_state_message(session)["df_meta"]) + def test_missing_fields(self): resp = self._post_compare({"session": "x", "path1": "/a"}) self.assertEqual(resp.code, 400) diff --git a/tests/unit/server/test_load_expr.py b/tests/unit/server/test_load_expr.py index faa71435e..78411d6ca 100644 --- a/tests/unit/server/test_load_expr.py +++ b/tests/unit/server/test_load_expr.py @@ -2645,6 +2645,16 @@ async def test_every_stats_tier_value_is_accepted_and_stored(self): {"session": "pw-bad", "build_dir": self.build_path, "stats_tier": "none"}) self.assertEqual((resp.code, json.loads(resp.body)["error_code"]), (400, "invalid_stats_tier")) + @tornado.testing.gen_test + async def test_a_tier_named_below_full_builds_a_schema_dataflow_with_inline_delivery(self): + """No scalar-tier units exist yet, so a scalar target is built and held at + the schema tier.""" + await self._load("pw-scalar-inline", stats_tier="scalar") + session = self._session("pw-scalar-inline") + self.assertEqual((session.xorq_dataflow.stats_tier, session.stats_status, session.stats_reason), + ("schema", "not_computed", "host")) + self.assertEqual(session.stats_policy["tier_target"], "scalar") + @tornado.testing.gen_test async def test_the_default_tier_stays_full(self): await self._load("pw-default", stats_delivery="deferred") @@ -2824,6 +2834,9 @@ async def test_stats_update_only_is_served_as_a_session_headed_for_full(self): self.assertEqual((update["type"], update["final"], update["tier"]), ("stats_update", True, "full")) self.assertIn("histogram_bins", _rows_by_stat(update["payload"])) self.assertEqual(self._session("pw-skew-update").stats_status, "complete") + # The pending frame it was sent is what the final reply is compared with: + # the float column's minWidth changes with the stats, so the config follows. + self.assertIn("df_display_args", update) @tornado.testing.gen_test async def test_stats_update_only_pulls_units_for_a_policy_session(self): @@ -2841,6 +2854,7 @@ async def test_stats_update_only_pulls_units_for_a_policy_session(self): break 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["tier"] for r in replies}, {"full"}, "the units are full-tier whatever the host asked") self.assertEqual(self._session("pw-skew-units").stats_status, "complete") @tornado.testing.gen_test From c691825f5fdf045b1211bb46abb58d898265fdd8 Mon Sep 17 00:00:00 2001 From: Paddy Mullen Date: Sun, 4 Oct 2026 05:26:40 -0400 Subject: [PATCH 27/39] feat(server): stats policy wire contract and handler plumbing (rows-first p33) /load_expr and /reload_expr accept stats_tier auto, full, scalar or schema, stored with the pair and kept out of has_config, so the warm short-circuit holds. When the dataflow is built at the schema tier the handler resolves the policy (resolve_stats_policy) from the count it already has, stores it on the session and starts the stats generation from it: a target below full is not_computed with the policy's reason. df_meta.stats reports the policy as tier_target and estimate, plus auto_request, requestable, omitted_keys, approx_keys and demand_columns where they differ from their documented defaults. It is applied at WebSocket open only for a client that sends ?caps=stats_update,stats_ondemand. A client with stats_update only is told the session is pending and pulls the stats, and a client with no caps gets them at connect, as for any deferred session. A tier the host named (scalar, schema) reaches every client as before. A session on an explicit stats_tier full within the ceiling sends the message it always has. /load keeps resolving to full: it does not read the field and clears a policy left on the session. Co-Authored-By: Claude Sonnet 5.5 --- buckaroo/server/handlers.py | 44 +++++--- buckaroo/server/session.py | 136 +++++++++++++++++++++---- buckaroo/server/stats_wire.py | 79 +++++++++++--- tests/unit/server/test_load_expr.py | 3 + tests/unit/server/test_stats_policy.py | 16 ++- 5 files changed, 221 insertions(+), 57 deletions(-) diff --git a/buckaroo/server/handlers.py b/buckaroo/server/handlers.py index 48100161a..a3c292d1d 100644 --- a/buckaroo/server/handlers.py +++ b/buckaroo/server/handlers.py @@ -14,10 +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.dataflow.dataflow import STATS_TIERS from buckaroo.server.session import ( - STATS_DELIVERIES, begin_stats_generation, dataflow_stats_tier) -from buckaroo.server.stats_wire import broadcast_state, refresh_session_snapshot + STATS_DELIVERIES, STATS_TIER_REQUESTS, begin_stats_generation, dataflow_stats_tier) +from buckaroo.server.stats_wire import broadcast_state, refresh_session_snapshot, resolve_session_policy from buckaroo.server import telemetry from buckaroo.pluggable_analysis_framework import perf_log @@ -367,9 +366,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" + # /load serves the eager backends, whose stats cost does not grow with + # the file, so they resolve to full: it reads no stats_tier and 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, session.stats_policy = "full", "inline", None begin_stats_generation(session) # Notify connected clients and open browser @@ -382,7 +383,8 @@ async def post(self): def _stats_policy_from_body(body: dict, current_tier: str, current_delivery: str): - """Read ``stats_tier`` and ``stats_delivery`` from a request body. + """Read ``stats_tier`` (``auto``, ``full``, ``scalar`` or ``schema``) 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 @@ -391,9 +393,9 @@ def _stats_policy_from_body(body: dict, current_tier: str, current_delivery: str stats_tier = body.get("stats_tier") if stats_tier is None: stats_tier = current_tier - if stats_tier not in STATS_TIERS: + if stats_tier not in STATS_TIER_REQUESTS: return None, None, {"error_code": "invalid_stats_tier", - "message": f"stats_tier must be one of {list(STATS_TIERS)}, got {stats_tier!r}"} + "message": f"stats_tier must be one of {list(STATS_TIER_REQUESTS)}, got {stats_tier!r}"} stats_delivery = body.get("stats_delivery") if stats_delivery is None: stats_delivery = current_delivery @@ -562,6 +564,7 @@ async def post(self): "skip_stat_columns": body.get("skip_stat_columns")} project_root = body.get("project_root") + dataflow_tier = dataflow_stats_tier(stats_tier, stats_delivery) try: # session= correlates these spans across concurrent loads — the @@ -590,8 +593,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) + stats_tier=dataflow_tier, **dataflow_kwargs) # Spanning metadata too leaves only the small klass-load step # unmeasured inside the outer firstpull.load_expr total. with perf_log.perf_span("firstpull.metadata", session=session_id): @@ -607,6 +609,12 @@ async def post(self): self.write(resp) return + # The dataflow and the count exist: this is where the size policy is + # resolved. A dataflow built at the schema tier has run no stat query, so + # the policy can still keep it there; one that ran its stats in the + # constructor has nothing left to resolve (None). + stats_policy = resolve_session_policy(stats_tier, dataflow_tier, metadata["rows"], len(metadata["columns"])) + sessions = self.application.settings["sessions"] session = sessions.get_or_create(session_id, build_dir) session.mode = "buckaroo" @@ -618,6 +626,7 @@ async def post(self): session.dataflow_kwargs = dataflow_kwargs session.stats_tier = stats_tier session.stats_delivery = stats_delivery + session.stats_policy = stats_policy session.tele_sink = tele_sink session.xorq_dataflow = xorq_dataflow # Clear pandas-side state left by a prior /load on the same @@ -803,7 +812,7 @@ async def post(self): 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" + session.stats_tier, session.stats_delivery, session.stats_policy = "full", "inline", None begin_stats_generation(session) # Push to WebSocket clients. Reset per-client live search (#851). @@ -839,7 +848,7 @@ class ReloadExprHandler(tornado.web.RequestHandler): 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. + when the reload succeeds. The stats policy is resolved again for the pair. Returns 404 when the session does not exist, 400 when it is not a xorq session, has no project_root recorded or carries an invalid stats policy, @@ -896,10 +905,14 @@ async def post(self, session_id): xorq_loading.load_project_stat_klasses(session.project_root) + xorq_loading.load_project_post_processing_klasses(session.project_root) + xorq_loading.load_project_display_klasses(session.project_root)) + dataflow_tier = dataflow_stats_tier(stats_tier, stats_delivery) 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) + stats_tier=dataflow_tier, **session.dataflow_kwargs) + # Resolved again, with the count the load took: the expression is + # reused, so there is no new one to take. + stats_policy = resolve_session_policy(stats_tier, dataflow_tier, session.metadata["rows"], + len(session.metadata["columns"])) except Exception: tb = traceback.format_exc() log.error("reload_expr error session=%s: %s", session_id, tb) @@ -927,6 +940,7 @@ async def post(self, session_id): session.xorq_dataflow = xorq_dataflow session.stats_tier = stats_tier session.stats_delivery = stats_delivery + session.stats_policy = stats_policy refresh_session_snapshot(session, xorq_dataflow) begin_stats_generation(session) diff --git a/buckaroo/server/session.py b/buckaroo/server/session.py index 754ddc1ce..f36ee8d53 100644 --- a/buckaroo/server/session.py +++ b/buckaroo/server/session.py @@ -1,3 +1,4 @@ +import copy import logging import time from dataclasses import dataclass, field @@ -16,11 +17,35 @@ # first and leaves the rest to later requests. STATS_DELIVERIES = ("inline", "deferred") +# What a host may send as ``stats_tier`` on /load_expr and /reload_expr. ``auto`` +# lets the server decide from the size of the entry (``stats_policy.py``); the +# others name a tier, which the server can lower (the ceiling) but never raise. +# Not the dataflow's own tiers (``dataflow.STATS_TIERS``: ``full`` and ``schema``), +# which are what a dataflow is built at. +STATS_TIER_REQUESTS = ("auto", "full", "scalar", "schema") + +# Why a session is ``not_computed``, in ``df_meta.stats.reason``: the server's size +# rule, the host naming a tier, the cost guard pausing a run (phase 6a), or the +# ceiling lowering a request. (``error`` has its own reason, ``stats_failed``.) +STATS_REASONS = ("size", "host", "cost", "ceiling") + +# What a client assumes for a ``df_meta.stats`` field the server leaves out. The +# server omits a field that equals its default, so a message with no policy to +# report is the one it always was. +STATS_FIELD_DEFAULTS: Dict[str, Any] = {"auto_request": True, "requestable": ["full"], "omitted_keys": [], + "approx_keys": [], "demand_columns": []} + 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 + """The tier a session's dataflow is constructed at. Only a session whose + stats run inline and whose host named no lower tier builds at the full tier; + a deferred one starts at the schema tier whatever tier it is headed for, and + so does one headed for ``scalar`` (no scalar-tier units exist yet) or + ``schema``. ``auto`` is resolved against a schema-tier dataflow, which an + inline session never has, so inline ``auto`` is ``full``.""" + if stats_delivery == "deferred" or stats_tier in ("scalar", "schema"): + return "schema" + return "full" # What ``df_meta.stats.status`` says about a session's stats for its current @@ -32,10 +57,15 @@ def dataflow_stats_tier(stats_tier: str, stats_delivery: str) -> str: STATS_STATUSES = ("complete", "pending", "not_computed", "error") -def initial_stats_status(stats_tier: str, stats_delivery: str) -> tuple[str, Optional[str]]: +def initial_stats_status(stats_tier: str, stats_delivery: str, policy: Optional[dict] = None) -> 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": + started, from the policy pair and the policy resolved at load. A target below + ``full`` is ``not_computed`` for the reason the policy gives; with no policy + (a dataflow built with its stats inline) a tier named below ``full`` is the + host's.""" + if policy is not None and policy["tier_target"] != "full": + return "not_computed", policy["reason"] + if policy is None and stats_tier in ("scalar", "schema"): return "not_computed", "host" if stats_delivery == "deferred": return "pending", None @@ -73,13 +103,20 @@ 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). + # The /load_expr stats policy. ``stats_tier`` is what the host asked for (see + # STATS_TIER_REQUESTS) and ``stats_delivery`` how the stats get 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" + # What ``stats_policy.resolve_stats_policy`` made of the pair and the size of + # the entry (``tier_target``, ``auto_request``, ``requestable``, ``reason``, + # ``estimate``), resolved by /load_expr and /reload_expr once the schema-tier + # dataflow and the count exist. ``None`` for a dataflow built with its stats + # inline, and for anything /load serves. Later phases add ``omitted_keys``, + # ``approx_keys`` and ``demand_columns``, which ``stats_meta`` reports. + stats_policy: Optional[dict] = None # 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). @@ -139,30 +176,83 @@ def begin_stats_generation(session: "SessionState") -> None: dropped.""" session.stats_gen += 1 session.stats_status, session.stats_reason = initial_stats_status( - session.stats_tier, session.stats_delivery) + session.stats_tier, session.stats_delivery, session.stats_policy) session.stat_runs.clear() -def stats_meta(session: "SessionState") -> Optional[dict]: +def stats_deferred_by_policy(session: "SessionState") -> bool: + """Whether the server, not the host, put the session below ``full``: the host + asked for ``auto`` or ``full`` and the policy resolved lower (a size rule, or + the ceiling). A tier the host named (``scalar``, ``schema``) is its own + choice. A client that cannot take the ``not_computed`` state is served such a + session as one headed for ``full`` (see ``stats_wire.stats_to_pull``).""" + return session.stats_status == "not_computed" and session.stats_tier in ("auto", "full") + + +def _policy_fields(policy: dict) -> dict: + """The part of a resolved policy ``df_meta.stats`` reports: ``tier_target`` and + ``estimate`` always, and the other fields only where they differ from + ``STATS_FIELD_DEFAULTS``.""" + fields: dict = {"tier_target": policy["tier_target"], "estimate": dict(policy["estimate"])} + for name in ("auto_request", "requestable"): + if policy[name] != STATS_FIELD_DEFAULTS[name]: + fields[name] = list(policy[name]) if isinstance(policy[name], list) else policy[name] + for name in ("omitted_keys", "approx_keys", "demand_columns"): + if policy.get(name): + fields[name] = list(policy[name]) + return fields + + +def stats_meta(session: "SessionState", ondemand: bool = True) -> 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. + ``full`` only once the session is complete. + + ``ondemand`` says whether the recipient advertised ``stats_ondemand``, so it + can take a ``not_computed`` state the server chose and the policy fields. Any + other client is told a session the server put below ``full`` is ``pending``, + headed for ``full``, with no policy fields. For an ondemand client a session + that is ``pending`` or ``not_computed`` adds the policy it resolved + (``_policy_fields``), unless that is an explicit ``full`` within the ceiling, + which sends the message it always has. ``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.""" + ``stats`` as complete (``stats_with_defaults``).""" 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 + status, reason = session.stats_status, session.stats_reason + if not ondemand and stats_deferred_by_policy(session): + status, reason = "pending", None + stats: dict = {"status": status, "tier": "full" if status == "complete" else "schema", "gen": session.stats_gen} + if reason: + stats["reason"] = reason + policy = session.stats_policy + reported = policy is not None and not (session.stats_tier == "full" and policy["tier_target"] == "full") + if ondemand and reported and status in ("pending", "not_computed"): + stats.update(_policy_fields(policy)) return stats +def stats_with_defaults(df_meta: Optional[dict]) -> dict: + """``df_meta.stats`` as a client reads it, with every field filled in: what a + message leaves out takes its documented default. No ``stats`` at all is what + an old server sends and means ``complete``. ``tier_target`` defaults to the + tier reached when the status is ``complete`` or ``not_computed`` and to + ``full`` otherwise (a ``pending`` or ``error`` session is headed for + ``full``). ``gen``, ``reason`` and ``estimate`` are ``None`` when absent.""" + stats = (df_meta or {}).get("stats") or {"status": "complete", "tier": "full"} + status, tier = stats["status"], stats["tier"] + out: dict = {"status": status, "tier": tier, + "tier_target": stats.get("tier_target", tier if status in ("complete", "not_computed") else "full"), + "gen": stats.get("gen"), "reason": stats.get("reason"), "estimate": stats.get("estimate")} + for name, default in STATS_FIELD_DEFAULTS.items(): + out[name] = copy.deepcopy(stats[name] if name in stats else default) + return out + + def build_state_message(session: "SessionState", metadata: dict | None = None, - search_string: str = "") -> dict: + search_string: str = "", ondemand: bool = True) -> dict: """Build the full ``initial_state`` WebSocket payload from a session. Args: @@ -175,6 +265,8 @@ def build_state_message(session: "SessionState", metadata: dict | None = None, search box. The session itself never owns this value; callers must pass the right value per recipient (typically ``handler.search_string``). + ondemand: Whether the recipient advertised ``stats_ondemand`` (see + ``stats_meta``); ``stats_wire.build_state_message_for`` passes it. Returns: A dict ready to be JSON-serialised and sent to WebSocket clients. @@ -182,7 +274,7 @@ def build_state_message(session: "SessionState", metadata: dict | None = None, # 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) + stats = stats_meta(session, ondemand=ondemand) if stats is not None: df_meta = {**df_meta, "stats": stats} msg: dict = {"type": "initial_state", "protocol_version": PROTOCOL_VERSION, diff --git a/buckaroo/server/stats_wire.py b/buckaroo/server/stats_wire.py index 950daf29a..47ddfb4ee 100644 --- a/buckaroo/server/stats_wire.py +++ b/buckaroo/server/stats_wire.py @@ -13,6 +13,17 @@ missing stats synchronously before it builds a message for one, which is today's cost, paid on the loop. +A session also carries the policy resolved at load (``stats_policy.py``). A +target the server chose below ``full`` (a size rule, or the ceiling on a host +that asked for ``auto`` or ``full``) is applied only for a client that +advertised ``?caps=stats_update,stats_ondemand``: it gets the schema tier, +``df_meta.stats.status == "not_computed"`` and the policy fields. A client with +``stats_update`` only is told the session is ``pending`` and pulls the stats as +above, and one with no caps gets them synchronously, as for any deferred +session. A tier the host named below ``full`` (``scalar``, ``schema``) is its +choice and reaches every client as ``not_computed``. ``stats_to_pull`` decides, +per client, whether a session has stats left to compute. + 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. @@ -41,8 +52,9 @@ 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.session import SessionState, build_state_message, stats_deferred_by_policy from buckaroo.server.stat_run import StatCursor, StatRun +from buckaroo.server.stats_policy import resolve_stats_policy log = logging.getLogger("buckaroo.server.stats_wire") @@ -52,6 +64,11 @@ # first message before the client has said anything. STATS_UPDATE_CAP = "stats_update" +# The second capability bit: the client also renders ``not_computed``, honours +# ``auto_request`` and sends ``tier`` and ``force`` on ``stats_request``. It counts +# only together with ``stats_update``, which it builds on. +STATS_ONDEMAND_CAP = "stats_ondemand" + # 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 @@ -77,6 +94,32 @@ def client_has_cap(client: Any, cap: str) -> bool: return cap in getattr(client, "caps", ()) +def client_has_ondemand(client: Any) -> bool: + """Whether a WebSocket handler advertised ``stats_update`` and + ``stats_ondemand``, so a session the server put below ``full`` is applied + for it as such.""" + return client_has_cap(client, STATS_UPDATE_CAP) and client_has_cap(client, STATS_ONDEMAND_CAP) + + +def stats_to_pull(session: SessionState, client: Any) -> bool: + """Whether the session has stats left to compute from this client's side: it + is ``pending``, or the server put it below ``full`` and the client cannot + take that (``stats_deferred_by_policy``). A client with ``stats_update`` + pulls them with ``stats_request``, one with no caps gets them at connect.""" + return session.stats_status == "pending" or (stats_deferred_by_policy(session) and not client_has_ondemand(client)) + + +def resolve_session_policy(stats_tier: str, dataflow_tier: str, rows: int, cols: int) -> Optional[dict]: + """The policy for a session the handler has just built: ``stats_tier`` is the + host's request, ``dataflow_tier`` the tier its dataflow was built at, and + ``rows`` and ``cols`` the size load already took (the cached count, so no + query runs). ``None`` when the dataflow ran its stats in the constructor + (``full``): there is nothing left to defer or refuse.""" + if dataflow_tier != "schema": + return None + return resolve_stats_policy("xorq", "xorq_build", rows, cols, host_tier=stats_tier) + + def session_dataflow(session: Optional[SessionState]) -> Any: """The dataflow behind a buckaroo-mode session, or ``None`` (viewer and lazy sessions have none).""" @@ -257,7 +300,8 @@ def _fail_stats(session: SessionState) -> None: def complete_stats(session: SessionState) -> bool: - """Run the stats a pending session is missing and publish them: the final + """Run the stats a pending session is missing (or one the server put below + ``full``, for a client that cannot take that) 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. @@ -274,7 +318,7 @@ def complete_stats(session: SessionState) -> bool: if session.stats_status == "complete": return True dataflow = session_dataflow(session) - if session.stats_status != "pending" or dataflow is None: + if not (session.stats_status == "pending" or stats_deferred_by_policy(session)) or dataflow is None: return False run_key = (session.stats_gen, "raw") with ( @@ -303,22 +347,26 @@ def complete_stats(session: SessionState) -> bool: 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). + A client without the ``stats_update`` capability on a deferred session that + has stats to compute for it (``stats_to_pull``) gets them run first, so its + message is complete; a capable client gets the session snapshot as it is + (stats-free while the stats are to be pulled) and pulls the rest. A client + that also advertised ``stats_ondemand`` is told the session as the server + resolved it (``not_computed`` with the policy fields); the others are told + ``pending`` where the server chose a target below ``full``. The search term + is the recipient's 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" + if (session.stats_delivery == "deferred" and stats_to_pull(session, client) 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", "")) + if stats_to_pull(session, client) else None) + return build_state_message(session, metadata=metadata, search_string=getattr(client, "search_string", ""), + ondemand=client_has_ondemand(client)) def broadcast_state(session: SessionState, metadata: Optional[dict] = None, reset_search: bool = False) -> None: @@ -396,7 +444,7 @@ def _serve_units(session: SessionState, client: Any, prefer: Sequence[str], stat fragments = cursor.take(run) if run.status != "pending": return None - return {"type": "stats_update", "stats_gen": stats_gen, "scope": scope, "tier": session.stats_tier, + return {"type": "stats_update", "stats_gen": stats_gen, "scope": scope, "tier": "full", "final": False, "remaining": run.remaining, "payload": partial_payload(dataflow, run, fragments), "elapsed_ms": _elapsed_ms(started)} @@ -432,13 +480,14 @@ def _answer_stats_request(session: Optional[SessionState], msg: dict, client: An 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": + to_pull = stats_to_pull(session, client) + if session.stats_status == "not_computed" and not to_pull: return _aborted(stats_gen, scope, "not_requestable", session) - if session.stats_status == "pending" and incremental: + if to_pull 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": + if to_pull: complete_stats(session) if session.stats_status != "complete": return _aborted(stats_gen, scope, "error", session) diff --git a/tests/unit/server/test_load_expr.py b/tests/unit/server/test_load_expr.py index 78411d6ca..17ebe564c 100644 --- a/tests/unit/server/test_load_expr.py +++ b/tests/unit/server/test_load_expr.py @@ -2937,6 +2937,9 @@ async def test_an_explicit_full_changes_no_frame(self): self.assertEqual(_as_json(before["df_meta"]), _as_json(after["df_meta"]), (delivery, caps)) self.assertEqual(_comparable(before), _comparable(after), (delivery, caps)) self.assertEqual(list(before), list(after)) + if delivery == "deferred" and caps: + stats = self._stats(after) + self.assertEqual(stats, {"status": "pending", "tier": "schema", "gen": stats["gen"]}, caps) self.assertEqual(self._session("pw-same-deferred-b").stats_policy["tier_target"], "full") self.assertIsNone(self._session("pw-same-inline-b").stats_policy) diff --git a/tests/unit/server/test_stats_policy.py b/tests/unit/server/test_stats_policy.py index a16f0ca6e..71be53ce7 100644 --- a/tests/unit/server/test_stats_policy.py +++ b/tests/unit/server/test_stats_policy.py @@ -765,6 +765,17 @@ def test_a_tier_the_host_named_reaches_every_client(self, sp, stats_tier, delive assert session_mod.stats_meta(session, ondemand=False) == {"status": "not_computed", "tier": "schema", "gen": 1, "reason": reason} + @pytest.mark.parametrize("ondemand", [True, False]) + def test_an_explicit_full_within_the_ceiling_reports_no_policy(self, sp, ondemand): + """The message is the one sent before the policy existed, though the + policy is resolved and recorded.""" + session = _session(sp, "full", "deferred") + assert session.stats_policy["tier_target"] == "full" + assert session_mod.stats_meta(session, ondemand=ondemand) == {"status": "pending", "tier": "schema", "gen": 1} + inline = _session(sp, "full", "inline", resolve=False) + assert session_mod.stats_meta(inline, ondemand=ondemand) is None + assert session_mod.build_state_message(inline, ondemand=ondemand)["df_meta"] == {"total_rows": 5} + @pytest.mark.parametrize("stats_tier, limits", [("auto", {}), ("auto", SCALAR_BY_SIZE), ("full", FULL_OVER_CEILING)]) def test_a_completed_session_reports_no_policy(self, sp, stats_tier, limits): @@ -795,11 +806,6 @@ def test_every_reason_the_server_sends_is_a_documented_one(self): for *_, expected in POLICY_FRAMES: assert expected.get("reason", "size") in session_mod.STATS_REASONS - def test_the_policy_is_what_the_pure_function_resolved(self, sp): - session = _session(sp, "auto", "deferred", SCHEMA_BY_SIZE) - assert session.stats_policy == sp.resolve_stats_policy("xorq", "xorq_build", 5, 3, - limits=sp.StatsLimits(**SCHEMA_BY_SIZE)) - class TestPolicyForTheSession: """``resolve_session_policy``: what a load handler resolves, and when.""" From 3bc2c09be929f9fb82db507adc5a56d788112c36 Mon Sep 17 00:00:00 2001 From: Paddy Mullen Date: Sun, 4 Oct 2026 06:19:28 -0400 Subject: [PATCH 28/39] test(stats): failing tests for the xorq scalar tier units (rows-first p34) A xorq stat run at the scalar tier is the scalar class of the batch without approx_median and distinct_count, and no histogram query. Its histogram_bins come from the batch's min and max. The run behaves like a filtered one: fragments go to the clients and nothing is assigned to the dataflow, the session snapshot or summary_stats_cache, so scalar stats are never served as the complete ones. Tests cover the plan, the queries a scalar run issues, the fragment union against the full stats, the bins, the chunked batch, the StatRun key and assignment rule, which clients are served the tier, and the stats_update messages of an ondemand client on a scalar target. Co-Authored-By: Claude Sonnet 5.5 --- tests/unit/server/test_data_loading_polars.py | 28 ++ tests/unit/server/test_load_expr.py | 312 +++++++++++++++++- tests/unit/server/test_stats_policy.py | 36 ++ tests/unit/test_paf_v2.py | 8 + tests/unit/test_xorq_stats_v2.py | 114 ++++++- 5 files changed, 496 insertions(+), 2 deletions(-) diff --git a/tests/unit/server/test_data_loading_polars.py b/tests/unit/server/test_data_loading_polars.py index ed4bcfd16..45fbd3078 100644 --- a/tests/unit/server/test_data_loading_polars.py +++ b/tests/unit/server/test_data_loading_polars.py @@ -40,6 +40,7 @@ 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 stat_run as stat_run_mod 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) @@ -682,6 +683,33 @@ def run(self, unit, acc): assert run.run_next() is None, "a failed run is not retried" assert run.fragments == [{"x": {"v": 1}}] + def test_the_full_run_over_every_column_is_the_run_that_assigns(self, backend): + run, _dataflow = _stat_run(backend) + assert (run.tier, run.assigns) == ("full", True) + + def test_a_run_over_a_column_group_has_its_own_key_and_assigns_nothing(self, backend): + dataflow = _build_dataflow(backend, _frame(backend, _UNTIED_DATA)) + stats = dataflow.build_stats(dataflow.processed_df, run=False) + run = StatRun(4, "raw", stats, state=dataclasses.replace(stats.state, columns=("price",))) + assert run.key == (4, "raw", "full", ("price",)) + assert (run.tier, run.assigns) == ("full", False) + assert [u.columns for u in run.units] == [("price",)] + while run.run_next() is not None: + pass + assert list(run.acc.sd()) == ["price"] + + def test_the_key_of_a_run_names_its_tier_and_group_unless_it_is_the_full_run(self): + key = stat_run_mod.stat_run_key + assert key(4, "raw") == key(4, "raw", "full", None) == (4, "raw") + assert key(4, "raw", "scalar") == (4, "raw", "scalar", None) + assert key(4, "raw", "full", ["b", "a"]) == (4, "raw", "full", ("b", "a")) + + def test_a_pandas_or_polars_run_has_only_the_full_tier(self, backend): + dataflow = _build_dataflow(backend, _frame(backend, _UNTIED_DATA)) + stats = dataflow.build_stats(dataflow.processed_df, run=False) + with pytest.raises(ValueError, match="full tier"): + StatRun(4, "raw", stats, state=dataclasses.replace(stats.state, tier="scalar")) + class _FakeClock: """A clock a test moves by hand, so a budget can be spent without waiting.""" diff --git a/tests/unit/server/test_load_expr.py b/tests/unit/server/test_load_expr.py index 17ebe564c..5ee7ef4ae 100644 --- a/tests/unit/server/test_load_expr.py +++ b/tests/unit/server/test_load_expr.py @@ -30,6 +30,7 @@ from xorq.common.utils.graph_utils import replace_nodes, walk_nodes # noqa: E402 from xorq.common.utils.provenance_utils import read_parquet_provenance # noqa: E402 from xorq.expr.relations import CachedNode # noqa: E402 +from xorq.vendor.ibis.expr import operations as ops # noqa: E402 from buckaroo.dataflow.dataflow import assemble_merged_sd # noqa: E402 from buckaroo.dataflow.sd_cache import split_chain_by_scope # noqa: E402 @@ -41,7 +42,7 @@ from buckaroo.server import session as session_mod # 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.stat_run import StatCursor, 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 @@ -3060,6 +3061,315 @@ def test_the_split_is_refused_for_a_join(self, tmp_path): assert _as_json(split.merged_sd) == _as_json(whole.merged_sd) +# The stat rows the wire carries for the five-by-three table: the scalar tier is +# the full tier's rows without the three it leaves out. +_SCALAR_WIRE_ONLY = {"dtype", "min", "max", "null_count", "mean", "std", "non_null_count", "histogram_bins"} +_FULL_ONLY_WIRE = {"distinct_count", "median", "histogram"} + +# The three ways an entry reaches a scalar target, as (host tier, limits). +_SCALAR_TARGETS = {"by size": ("auto", _SCALAR_BY_SIZE), "named by the host": ("scalar", {}), + "full over the ceiling": ("full", {"ceiling_full_rows": 3})} + + +def _scalar_session(dataflow): + """A buckaroo-mode session over ``dataflow`` whose policy target is + ``scalar`` and has nothing computed, as ``/load_expr`` leaves one.""" + session = session_mod.SessionState(session_id="s", path="p") + session.mode, session.backend, session.xorq_dataflow = "buckaroo", "xorq", dataflow + session.stats_tier, session.stats_delivery = "auto", "deferred" + session.stats_policy = {"tier_target": "scalar", "auto_request": True, "requestable": ["full"], + "reason": "size", "estimate": {"rows": 12, "cols": 6}} + session_mod.begin_stats_generation(session) + return session + + +def _ondemand_client(): + return SimpleNamespace(caps=frozenset({"stats_update", "stats_ondemand"}), search_string="", + stats_cursor=StatCursor(), display_args_hash=None) + + +class TestScalarTierWire(tornado.testing.AsyncHTTPTestCase): + """The scalar tier on a ``/load_expr`` session whose policy target is + ``scalar`` (rows-first p34). A client that advertised ``stats_ondemand`` + pulls it with ``stats_request`` and gets ``stats_update`` messages whose + tier is ``scalar``. The run behaves like a filtered one: fragments go to + the clients and nothing is assigned to the dataflow, the session snapshot + or ``summary_stats_cache``, so the scalar stats can never be served as the + complete ones.""" + + 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) + 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, limits=None, **body): + with _stats_limits(**(limits or {})): + 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 _scalar_target(self, sid): + """A deferred ``auto`` session that resolved to a scalar target by size.""" + await self._load(sid, limits=_SCALAR_BY_SIZE, stats_tier="auto", stats_delivery="deferred") + + async def _connect(self, sid, caps=None): + 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): + await self._load(sid) + return (await self._connect(sid))[1] + + async def _pull(self, ws, gen, **fields): + """The reply to an incremental ``stats_request``.""" + ws.write_message(_stats_request(gen, incremental=True, **fields)) + return await _read_json(ws) + + async def _pull_update(self, ws, gen, **fields): + """The reply to an incremental ``stats_request``, which must be a ``stats_update``.""" + reply = await self._pull(ws, gen, **fields) + self.assertEqual(reply["type"], "stats_update", reply) + return reply + + @tornado.testing.gen_test + async def test_an_ondemand_client_pulls_the_scalar_tier_of_a_scalar_target(self): + for index, (name, (tier, limits)) in enumerate(_SCALAR_TARGETS.items()): + for incremental in (True, False): + sid = f"st-pull-{index}-{incremental}" + await self._load(sid, limits=limits, stats_tier=tier, stats_delivery="deferred") + ws, first = await self._connect(sid, caps=_ONDEMAND) + stats = self._stats(first) + self.assertEqual((stats["status"], stats["tier_target"]), ("not_computed", "scalar"), name) + + with _count_stat_queries() as queries: + ws.write_message(_stats_request(stats["gen"], **({"incremental": True} if incremental else {}))) + update = await _read_json(ws) + + label = f"{name}, incremental={incremental}" + self.assertEqual(update["type"], "stats_update", f"{label}: {update}") + self.assertEqual((update["tier"], update["final"], update["remaining"]), ("scalar", True, 0), label) + self.assertEqual((update["stats_gen"], update["scope"]), (stats["gen"], "raw"), label) + self.assertIsInstance(update["elapsed_ms"], (int, float)) + self.assertEqual(set(_rows_by_stat(update["payload"])), _SCALAR_WIRE_ONLY, label) + self.assertEqual([type(q.op()).__name__ for q in queries], ["Aggregate"], label) + self.assertFalse(list(queries[0].op().find((ops.ApproxMedian, ops.CountDistinct))), label) + + @tornado.testing.gen_test + async def test_the_scalar_fragments_equal_the_full_stats_for_the_keys_they_cover(self): + full = _rows_by_stat((await self._inline_frame("st-eq-inline"))["df_data_dict"]["all_stats"]) + await self._scalar_target("st-eq") + ws, first = await self._connect("st-eq", caps=_ONDEMAND) + scalar = _rows_by_stat((await self._pull_update(ws, self._stats(first)["gen"]))["payload"]) + + self.assertEqual(set(scalar), set(full) - _FULL_ONLY_WIRE) + for stat, row in scalar.items(): + for column, value in row.items(): + if column in ("index", "level_0") or (stat, column) == ("histogram_bins", "b"): + continue + self.assertEqual(_as_json(value), _as_json(full[stat][column]), (stat, column)) + # qty has three distinct values: unknown at this tier, so it gets bins that the full tier leaves out. + self.assertEqual((len(scalar["histogram_bins"]["b"]), full["histogram_bins"]["b"]), (11, [])) + + @tornado.testing.gen_test + async def test_a_scalar_run_assigns_nothing(self): + sid = "st-nothing" + await self._scalar_target(sid) + session = self._session(sid) + dataflow = session.xorq_dataflow + ws, first = await self._connect(sid, caps=_ONDEMAND) + gen = self._stats(first)["gen"] + held = (dataflow.summary_sd, session.df_data_dict, session.df_display_args, session.df_meta) + cached = set(dataflow.summary_stats_cache) + + await self._pull_update(ws, gen) + + self.assertTrue(all(a is b for a, b in zip(held, + (dataflow.summary_sd, session.df_data_dict, session.df_display_args, session.df_meta)))) + self.assertEqual(set(dataflow.summary_stats_cache), cached, "nothing is written to summary_stats_cache") + self.assertEqual((session.stats_status, session.stats_reason, dataflow.stats_tier), + ("not_computed", "size", "schema")) + self.assertEqual(list(session.stat_runs), [(gen, "raw", "scalar", None)], "the run is kept for later clients") + _, later = await self._connect(sid, caps=_ONDEMAND) + self.assertEqual((self._stats(later)["status"], list(_rows_by_stat(later["df_data_dict"]["all_stats"]))), + ("not_computed", ["dtype"]), "a client that connects later is told the same as the first one") + + @tornado.testing.gen_test + async def test_a_client_that_asks_later_reads_the_fragments_the_run_holds(self): + sid = "st-later" + await self._scalar_target(sid) + a, first = await self._connect(sid, caps=_ONDEMAND) + gen = self._stats(first)["gen"] + rows = _rows_by_stat((await self._pull_update(a, gen))["payload"]) + b, _ = await self._connect(sid, caps=_ONDEMAND) + + with _count_stat_queries() as queries: + caught_up = await self._pull_update(b, gen) + + self.assertEqual(queries, [], "the run's units ran once") + self.assertEqual((caught_up["tier"], caught_up["final"]), ("scalar", True)) + self.assertEqual(_as_json(_rows_by_stat(caught_up["payload"])), _as_json(rows)) + + @tornado.testing.gen_test + async def test_a_scalar_run_is_never_served_as_the_complete_stats(self): + inline = await self._inline_frame("st-never-inline") + await self._scalar_target("st-never") + a, first = await self._connect("st-never", caps=_ONDEMAND) + gen = self._stats(first)["gen"] + self.assertEqual((await self._pull_update(a, gen))["tier"], "scalar") + + # The policy put the session below full for a client that cannot take that: it is owed the full stats. + _, legacy = await self._connect("st-never") + stats = self._stats(legacy) + self.assertEqual((stats["status"], stats["tier"]), ("complete", "full")) + self.assertEqual(_comparable(legacy), _comparable(inline)) + self.assertEqual(set(_rows_by_stat(legacy["df_data_dict"]["all_stats"])), + set(_rows_by_stat(inline["df_data_dict"]["all_stats"]))) + self.assertTrue(_FULL_ONLY_WIRE <= set(_rows_by_stat(legacy["df_data_dict"]["all_stats"]))) + # Once the full stats exist a request for the session is answered with them. + update = await self._pull_update(a, gen) + self.assertEqual((update["tier"], update["final"]), ("full", True)) + self.assertTrue(_FULL_ONLY_WIRE <= set(_rows_by_stat(update["payload"]))) + + @tornado.testing.gen_test + async def test_a_scalar_run_and_a_column_scoped_run_are_stored_apart_from_the_full_run(self): + await self._load("st-keys", stats_delivery="deferred") + session = self._session("st-keys") + runs = {"full": start_stat_run(session), "scalar": start_stat_run(session, tier="scalar"), + "scalar group": start_stat_run(session, tier="scalar", columns=("qty",)), + "full group": start_stat_run(session, tier="full", columns=("qty",))} + gen = session.stats_gen + self.assertEqual({name: run.key for name, run in runs.items()}, { + "full": (gen, "raw"), "scalar": (gen, "raw", "scalar", None), + "scalar group": (gen, "raw", "scalar", ("qty",)), "full group": (gen, "raw", "full", ("qty",))}) + self.assertEqual(session.stat_runs, {run.key: run for run in runs.values()}) + self.assertIs(start_stat_run(session, tier="scalar"), runs["scalar"], "one run for a tier and group") + self.assertEqual({name: run.assigns for name, run in runs.items()}, + {"full": True, "scalar": False, "scalar group": False, "full group": False}) + self.assertEqual([u.id for u in runs["scalar"].units], ["batch"]) + self.assertEqual([u.id for u in runs["scalar group"].units], ["batch"]) + self.assertEqual([u.columns for u in runs["full group"].units], [("qty",), ("qty",)]) + + @tornado.testing.gen_test + async def test_runs_that_assign_nothing_leave_the_dataflow_the_snapshot_and_the_cache_alone(self): + await self._load("st-scoped", stats_delivery="deferred") + session = self._session("st-scoped") + dataflow = session.xorq_dataflow + held = (dataflow.summary_sd, session.df_data_dict) + cached = set(dataflow.summary_stats_cache) + key = dataflow._scope_cache_key(split_chain_by_scope(dataflow.operations)["filt"], tier="full") + + for tier, columns in (("scalar", None), ("scalar", ("qty",)), ("full", ("qty",))): + run = start_stat_run(session, tier=tier, columns=columns) + stats_wire.run_units(run, None) + self.assertEqual(run.status, "complete") + rows = _rows_by_stat(stats_wire.partial_payload(dataflow, run, run.fragments)) + columns_sent = {c for row in rows.values() for c in row} - {"index", "level_0"} + self.assertEqual(columns_sent, {"a", "b", "c"} if columns is None else {"b"}, (tier, columns)) + self.assertTrue(all(a is b for a, b in zip(held, (dataflow.summary_sd, session.df_data_dict)))) + self.assertEqual(set(dataflow.summary_stats_cache), cached) + self.assertEqual((session.stats_status, dataflow.stats_tier), ("pending", "schema")) + + # The run that is assigned is a whole full run, whatever else ran. + self.assertTrue(stats_wire.complete_stats(session)) + self.assertEqual(set(dataflow.summary_stats_cache) - cached, {key}) + self.assertTrue(_FULL_ONLY_WIRE <= set(_rows_by_stat(session.df_data_dict["all_stats"]))) + + @tornado.testing.gen_test + async def test_a_scalar_run_is_not_assigned_even_if_it_is_stored_under_the_full_key(self): + await self._load("st-guard", stats_delivery="deferred") + session = self._session("st-guard") + dataflow = session.xorq_dataflow + cached = set(dataflow.summary_stats_cache) + run = start_stat_run(session, tier="scalar") + stats_wire.run_units(run, None) + session.stat_runs[(session.stats_gen, "raw")] = run + + with self.assertLogs("buckaroo.server.stats_wire", level="ERROR"): + self.assertFalse(stats_wire.complete_stats(session)) + + self.assertEqual((session.stats_status, session.stats_reason), ("error", "stats_failed")) + self.assertEqual(set(dataflow.summary_stats_cache), cached) + + +class TestScalarTierRequests: + """An incremental scalar request on a session whose batch is cut into + chunks (rows-first p34): one chunk per request, each reply reporting the + ``scalar`` tier, and the last one ``final``.""" + + @staticmethod + def _session(tmp_path, **kwargs): + expr = _wide_parquet_expr(tmp_path) + dataflow = xorq_loading.XorqServerDataflow(expr, skip_main_serial=True, stats_tier="schema", **kwargs) + return _scalar_session(dataflow) + + @staticmethod + def _request(session, client, **fields): + return stats_wire.handle_stats_request( + session, {"type": "stats_request", "stats_gen": session.stats_gen, "scope": "raw", **fields}, client) + + def test_each_request_runs_one_chunk_and_reports_the_scalar_tier(self, tmp_path, monkeypatch): + monkeypatch.setattr(stats_wire, "STATS_BUDGET_S", 0) + # 24 cells at 12 rows is two columns per chunk: three batch units. + session, client = self._session(tmp_path, stat_chunk_cells=24), _ondemand_client() + replies = [self._request(session, client, incremental=True) for _ in range(3)] + assert [r["type"] for r in replies] == ["stats_update"] * 3, replies + assert [(r["tier"], r["final"], r["remaining"]) for r in replies] == [ + ("scalar", False, 2), ("scalar", False, 1), ("scalar", True, 0)] + assert [{c for row in _rows_by_stat(r["payload"]).values() for c in row} - {"index", "level_0"} + for r in replies] == [{"a", "b"}, {"c", "d"}, {"e", "f"}] + whole = self._request(self._session(tmp_path), _ondemand_client()) + assert _as_json(_merge_stat_payloads(r["payload"] for r in replies)) == _as_json( + _merge_stat_payloads([whole["payload"]])) + + def test_a_request_after_the_last_chunk_replays_nothing_and_runs_nothing(self, tmp_path, monkeypatch): + monkeypatch.setattr(stats_wire, "STATS_BUDGET_S", 0) + session, client = self._session(tmp_path, stat_chunk_cells=24), _ondemand_client() + for _ in range(3): + self._request(session, client, incremental=True) + with _count_stat_queries() as queries: + again = self._request(session, client, incremental=True) + assert queries == [] + assert again["type"] == "stats_update", again + assert (again["tier"], again["final"], again["remaining"]) == ("scalar", True, 0) + + def test_a_unit_that_fails_aborts_the_request_and_not_the_session(self, tmp_path, caplog): + session, client = self._session(tmp_path), _ondemand_client() + dataflow = session.xorq_dataflow + stats = dataflow.build_stats(dataflow.processed_df, run=False) + run = StatRun(session.stats_gen, "raw", _SlowStats(stats, fail_on="batch"), + state=dataclasses.replace(stats.state, tier="scalar")) + session.stat_runs[run.key] = run + for _ in range(2): + reply = self._request(session, client, incremental=True) + assert (reply["type"], reply["reason"]) == ("stats_aborted", "error") + assert run.status == "error", "a failed run is not retried" + assert (session.stats_status, session.stats_reason) == ("not_computed", "size"), "the full tier is still open" + assert any("stat unit failed" in record.getMessage() for record in caplog.records) + + class TestReloadExpr(tornado.testing.AsyncHTTPTestCase): def get_app(self): return make_app() diff --git a/tests/unit/server/test_stats_policy.py b/tests/unit/server/test_stats_policy.py index 71be53ce7..f3427a483 100644 --- a/tests/unit/server/test_stats_policy.py +++ b/tests/unit/server/test_stats_policy.py @@ -859,3 +859,39 @@ def test_a_pending_session_has_stats_to_pull_for_everyone(self, sp, caps): def test_a_tier_the_host_named_has_nothing_to_pull(self, sp, caps): assert stats_wire.stats_to_pull(_session(sp, "schema", "deferred"), _client(caps)) is False assert stats_wire.stats_to_pull(_session(sp, "scalar", "deferred"), _client(caps)) is False + + +class TestScalarTierServing: + """``serves_scalar_tier`` (rows-first p34): which client a ``stats_request`` + is answered with the scalar tier for. A session whose policy target is + ``scalar`` and has nothing computed serves it to a client that advertised + ``stats_ondemand``, which takes tiers. Every other client is served as + before.""" + + @pytest.mark.parametrize("tier, limits, expected", [ + ("auto", SCALAR_BY_SIZE, True), + ("scalar", {}, True), + ("full", FULL_OVER_CEILING, True), + ("auto", SCHEMA_BY_SIZE, False), + ("schema", {}, False), + ("scalar", SCALAR_OVER_CEILING, False), + ("auto", {}, False), + ("full", {}, False)]) + def test_an_ondemand_client_is_served_a_session_whose_target_is_scalar(self, sp, tier, limits, expected): + session = _session(sp, tier, "deferred", limits) + assert stats_wire.serves_scalar_tier(session, _client(ONDEMAND_CAPS)) is expected + + @pytest.mark.parametrize("caps", [UPDATE_CAPS, LEGACY_CAPS, frozenset({"stats_ondemand"})]) + @pytest.mark.parametrize("tier, limits", [("auto", SCALAR_BY_SIZE), ("scalar", {})]) + def test_a_client_without_both_bits_is_not(self, sp, caps, tier, limits): + assert stats_wire.serves_scalar_tier(_session(sp, tier, "deferred", limits), _client(caps)) is False + + def test_a_session_whose_stats_are_complete_serves_no_tier(self, sp): + session = _session(sp, "auto", "deferred", SCALAR_BY_SIZE) + session.stats_status, session.stats_reason = "complete", None + assert stats_wire.serves_scalar_tier(session, _client(ONDEMAND_CAPS)) is False + + def test_a_session_with_no_policy_has_no_target_to_serve(self, sp): + session = _session(sp, "scalar", "inline", resolve=False) + assert (session.stats_status, session.stats_policy) == ("not_computed", None) + assert stats_wire.serves_scalar_tier(session, _client(ONDEMAND_CAPS)) is False diff --git a/tests/unit/test_paf_v2.py b/tests/unit/test_paf_v2.py index e60b4e892..8f07c4938 100644 --- a/tests/unit/test_paf_v2.py +++ b/tests/unit/test_paf_v2.py @@ -969,3 +969,11 @@ def test_priority_reads_each_name_once_in_its_namespace(self, namespace, priorit 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')) + + @pytest.mark.parametrize('backend', ['pandas', 'polars']) + def test_the_scalar_tier_has_no_pandas_or_polars_units(self, backend): + """The scalar tier is xorq's. A state that asks for it must not be + answered with full-tier units.""" + pipeline = StatPipeline(_stat_lists()[backend], unit_test=False) + with pytest.raises(ValueError, match='full tier'): + pipeline.plan(StatState(_mixed_frames()[backend], tier='scalar')) diff --git a/tests/unit/test_xorq_stats_v2.py b/tests/unit/test_xorq_stats_v2.py index 3e420cec9..76a720d7f 100644 --- a/tests/unit/test_xorq_stats_v2.py +++ b/tests/unit/test_xorq_stats_v2.py @@ -14,12 +14,14 @@ xo = pytest.importorskip("xorq.api") +from xorq.vendor.ibis.expr import operations as ops # noqa: E402 + from buckaroo.pluggable_analysis_framework import perf_log # noqa: E402 from buckaroo.pluggable_analysis_framework.xorq_stat_pipeline import ( # noqa: E402 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.pluggable_analysis_framework.stat_units import StatState, StatUnit, merge_fragments # noqa: E402 from buckaroo.customizations.xorq_stats_v2 import ( # noqa: E402 XORQ_STATS_V2) @@ -1270,3 +1272,113 @@ def test_the_split_is_refused_for_anything_but_a_plain_parquet_scan(self, kind, 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) + + +# ============================================================ +# The scalar tier (rows-first p34) +# ============================================================ + +# What the full tier reports and the scalar tier does not. +_FULL_ONLY = {"median", "distinct_count", "distinct_per", "histogram"} + + +def _tier_run(table, tier, **pipeline_kwargs): + """Every unit of a run of ``tier``, run in order: ``(acc, [(unit, fragment)])``.""" + return _run_units(XorqStatPipeline(XORQ_STATS_V2, unit_test=False, **pipeline_kwargs), StatState(table, tier=tier)) + + +def _union(run): + return merge_fragments(fragment for _unit, fragment in run[1]) + + +class TestScalarTier: + """The scalar tier of a xorq run: the scalar class of the batch without + ``approx_median`` and ``distinct_count``, and no histogram query. Its + ``histogram_bins`` come from the batch's ``min`` and ``max``, so a + ``color_map`` column needs no further query.""" + + _ALL = ("ints", "floats", "strs", "bools") + + def test_an_unknown_tier_is_refused(self): + with pytest.raises(ValueError, match="tier"): + StatState(_make_table(), tier="median") + + def test_the_scalar_plan_is_the_batch_alone(self): + units = XorqStatPipeline(XORQ_STATS_V2, unit_test=False).plan(StatState(_make_table(), tier="scalar")) + assert [(u.id, u.phase, u.cost, u.after) for u in units] == [("batch", "batch", "scan", ())] + assert units[0].columns == self._ALL + + def test_a_scalar_run_issues_the_batch_query_and_no_histogram_query(self, monkeypatch): + spy = _QuerySpy(monkeypatch) + _tier_run(_make_table(), "scalar") + assert spy.kinds() == ["batch"] + + def test_the_scalar_batch_has_no_approx_median_and_no_distinct_count(self, monkeypatch): + spy = _QuerySpy(monkeypatch) + _tier_run(_make_table(), "scalar") + (batch,) = spy.queries + # ApproxCountDistinct is a CountDistinct, so this finds both distinct aggregates. + assert not list(batch.op().find((ops.ApproxMedian, ops.CountDistinct))) + # The same search finds them in the full batch, so it is looking at the right thing. + spy.queries.clear() + _tier_run(_make_table(), "full") + assert list(spy.queries[0].op().find(ops.ApproxMedian)) and list(spy.queries[0].op().find(ops.CountDistinct)) + + def test_the_scalar_stats_are_the_full_stats_without_the_keys_the_tier_leaves_out(self): + table = _make_table() + scalar, full = _union(_tier_run(table, "scalar")), _union(_tier_run(table, "full")) + for col in self._ALL: + assert set(scalar[col]) == set(full[col]) - _FULL_ONLY, col + assert _stable(scalar[col]) == _stable({key: full[col][key] for key in scalar[col]}), col + assert {"length", "null_count", "min", "max", "mean", "std", "non_null_count", "nan_per", + "histogram_bins", "_type"} <= set(scalar["ints"]) + + def test_the_accumulator_reports_what_the_fragments_do(self): + acc, run = _tier_run(_make_table_with_nulls(), "scalar") + assert _stable(acc.sd()) == _stable(_union((acc, run))) + assert not any(_FULL_ONLY & set(stats) for stats in acc.sd().values()) + assert acc.errors == [] + + def test_histogram_bins_come_with_the_batch_and_equal_the_full_tier_s(self, monkeypatch): + table = _make_table() + spy = _QuerySpy(monkeypatch) + scalar = _union(_tier_run(table, "scalar")) + assert spy.kinds() == ["batch"] + full = _union(_tier_run(table, "full")) + for col in ("ints", "floats"): + assert len(scalar[col]["histogram_bins"]) == 11, col + assert scalar[col]["histogram_bins"] == full[col]["histogram_bins"], col + assert scalar["strs"]["histogram_bins"] == scalar["bools"]["histogram_bins"] == [] + + def test_a_low_cardinality_integer_gets_bins_the_full_tier_does_not(self): + """Without ``distinct_count`` the cardinality is unknown and the column + is treated as one that wants bins. The full tier counts, finds five or + fewer values, and sends none.""" + table = xo.memtable(pd.DataFrame({"few": [1, 2, 3, 1, 2, 3, 1], "many": [1, 2, 3, 4, 5, 6, 7]})) + scalar, full = _union(_tier_run(table, "scalar")), _union(_tier_run(table, "full")) + assert len(scalar["few"]["histogram_bins"]) == 11 + assert full["few"]["histogram_bins"] == [] + assert scalar["many"]["histogram_bins"] == full["many"]["histogram_bins"] + + def test_a_skipped_column_has_no_scalar_unit_and_keeps_its_structural_entry(self): + pipeline = XorqStatPipeline(XORQ_STATS_V2, unit_test=False) + state = StatState(_make_table(), frozenset({"strs"}), tier="scalar") + assert pipeline.plan(state)[0].columns == ("ints", "floats", "bools") + acc, _run = _run_units(pipeline, state) + assert set(acc.sd()["strs"]) == {"orig_col_name", "rewritten_col_name", "dtype", "length", "min", "max"} + assert acc.sd()["strs"]["length"] == 7 + + def test_a_chunked_scalar_batch_equals_the_single_batch(self, tmp_path): + scan = _six_column_scan(tmp_path) + single = _union(_tier_run(scan, "scalar")) + pipeline = _chunked(2 * _SIX_ROWS) + state = StatState(scan, rows=_SIX_ROWS, tier="scalar") + assert [u.id for u in pipeline.plan(state)] == ["batch:0", "batch:1", "batch:2"] + assert _stable(_union(_run_units(pipeline, state))) == _stable(single) + + def test_the_scalar_tier_runs_no_histogram_unit(self): + pipeline = XorqStatPipeline(XORQ_STATS_V2, unit_test=False) + acc = pipeline.new_accumulator(StatState(_make_table(), tier="scalar")) + histogram = StatUnit("histogram:ints", ("ints",), "histogram", after=("batch",), cost="query") + with pytest.raises(ValueError, match="scalar"): + pipeline.run(histogram, acc) From e214db0989911e74858ef341d3df25f2f61dc893 Mon Sep 17 00:00:00 2001 From: Paddy Mullen Date: Sun, 4 Oct 2026 06:34:09 -0400 Subject: [PATCH 29/39] test(stats): failing tests for the scalar tier's dependents, spans and generations (rows-first p34) A deliberate regression of the implementation showed four behaviours the first set of tests does not cover: a stat that reads a key the scalar tier leaves out is left out with it, a stat that reads distinct_count runs on the unknown, a scalar request emits the request and unit spans with the scalar tier and no completion span, and a dataflow-field change drops the scalar run and starts one over the new state. Co-Authored-By: Claude Sonnet 5.5 --- tests/unit/server/test_load_expr.py | 40 +++++++++++++++++++++++++++++ tests/unit/test_xorq_stats_v2.py | 23 +++++++++++++++++ 2 files changed, 63 insertions(+) diff --git a/tests/unit/server/test_load_expr.py b/tests/unit/server/test_load_expr.py index 5ee7ef4ae..4b179ca39 100644 --- a/tests/unit/server/test_load_expr.py +++ b/tests/unit/server/test_load_expr.py @@ -3253,6 +3253,46 @@ async def test_a_scalar_run_is_never_served_as_the_complete_stats(self): self.assertEqual((update["tier"], update["final"]), ("full", True)) self.assertTrue(_FULL_ONLY_WIRE <= set(_rows_by_stat(update["payload"]))) + @tornado.testing.gen_test + async def test_a_scalar_request_emits_the_request_and_unit_spans_with_the_scalar_tier(self): + captured: list = [] + with patch.object(telemetry, "make_http_sink", lambda url, **kw: captured.append): + await self._load("st-span", limits=_SCALAR_BY_SIZE, stats_tier="auto", stats_delivery="deferred", + telemetry_url="http://companion.invalid/internal/telemetry") + ws, first = await self._connect("st-span", caps=_ONDEMAND) + gen = self._stats(first)["gen"] + await self._pull_update(ws, gen) + + (request,) = [r for r in captured if r["name"] == "stats.request"] + (unit,) = [r for r in captured if r["name"] == "stats.unit"] + self.assertEqual((request["attrs"]["tier"], request["attrs"]["outcome"], request["attrs"]["final"]), + ("scalar", "update", True)) + self.assertEqual((request["attrs"]["stats_gen"], request["attrs"]["units"]), (gen, 1)) + self.assertEqual((unit["attrs"]["unit"], unit["attrs"]["phase"], unit["attrs"]["cost"]), + ("batch", "batch", "scan")) + self.assertNotIn("firstpull.stats_total", [r["name"] for r in captured], "nothing is completed") + + @tornado.testing.gen_test + async def test_a_dataflow_field_change_starts_a_new_scalar_run_over_the_new_state(self): + sid = "st-change" + await self._scalar_target(sid) + session = self._session(sid) + ws, first = await self._connect(sid, caps=_ONDEMAND) + gen = self._stats(first)["gen"] + await self._pull_update(ws, gen) + self.assertEqual(session.stat_runs[(gen, "raw", "scalar", None)].acc.sd()["qty"]["length"], 5) + + ws.write_message(_state_change(quick_command_args={"search": ["a"]})) + changed = await _read_json(ws) + self.assertEqual((self._stats(changed)["status"], self._stats(changed)["gen"]), ("not_computed", gen + 1)) + self.assertEqual(session.stat_runs, {}, "the run of the old generation is dropped") + stale = await self._pull(ws, gen) + self.assertEqual((stale["type"], stale["reason"]), ("stats_aborted", "stale")) + after = await self._pull_update(ws, gen + 1) + self.assertEqual((after["stats_gen"], after["tier"]), (gen + 1, "scalar")) + self.assertEqual(session.stat_runs[(gen + 1, "raw", "scalar", None)].acc.sd()["qty"]["length"], 2, + "the run is over the filtered state") + @tornado.testing.gen_test async def test_a_scalar_run_and_a_column_scoped_run_are_stored_apart_from_the_full_run(self): await self._load("st-keys", stats_delivery="deferred") diff --git a/tests/unit/test_xorq_stats_v2.py b/tests/unit/test_xorq_stats_v2.py index 76a720d7f..9f5c35882 100644 --- a/tests/unit/test_xorq_stats_v2.py +++ b/tests/unit/test_xorq_stats_v2.py @@ -1382,3 +1382,26 @@ def test_the_scalar_tier_runs_no_histogram_unit(self): histogram = StatUnit("histogram:ints", ("ints",), "histogram", after=("batch",), cost="query") with pytest.raises(ValueError, match="scalar"): pipeline.run(histogram, acc) + + def test_a_stat_that_reads_an_omitted_key_is_left_out_with_it(self): + @stat() + def double_median(median: float) -> float: + return median * 2 + + pipeline = XorqStatPipeline([*XORQ_STATS_V2, double_median], unit_test=False) + acc, run = _run_units(pipeline, StatState(_make_table(), tier="scalar")) + assert not any("double_median" in stats for _unit, fragment in run for stats in fragment.values()) + assert acc.errors == [], "a stat with an input the tier does not have is not run, so it is no error" + full = _union(_run_units(pipeline, StatState(_make_table()))) + assert full["ints"]["double_median"] == 2 * full["ints"]["median"] + + def test_a_stat_that_reads_distinct_count_runs_on_the_unknown(self): + @stat() + def cardinality_unknown(distinct_count: int) -> bool: + return distinct_count is None + + pipeline = XorqStatPipeline([*XORQ_STATS_V2, cardinality_unknown], unit_test=False) + scalar = _union(_run_units(pipeline, StatState(_make_table(), tier="scalar"))) + full = _union(_run_units(pipeline, StatState(_make_table()))) + assert (scalar["ints"]["cardinality_unknown"], full["ints"]["cardinality_unknown"]) == (True, False) + assert "distinct_count" not in scalar["ints"] From 60abb8f310f158ed94ee1bdc314594e506a32361 Mon Sep 17 00:00:00 2001 From: Paddy Mullen Date: Sun, 4 Oct 2026 06:29:02 -0400 Subject: [PATCH 30/39] feat(stats): xorq scalar tier units, served as fragments and never assigned (rows-first p34) StatState takes a tier. A xorq run at the scalar tier is the batch alone, in the same column chunks, without the aggregates behind median and distinct_count and with no histogram unit. Its histogram_bins come out of the batch's min and max, so color_map needs no histogram query. The accumulator withholds median, distinct_count, distinct_per and histogram, so no fragment and no assembled sd carries a key the tier did not compute. The pandas and polars pipelines refuse a non-full state. StatRun takes a state, and reports its tier and whether it assigns (only the full tier over every column does). The full run keeps its (stats_gen, scope) key; every other run has a key of its own, so it never takes the full run's place. start_stat_run takes a tier and a column group. _run_summary refuses a run that does not assign, so a scalar or column-scoped sd is never written to summary_stats_cache as the complete one. A stats_request from a client that advertised stats_update and stats_ondemand, on a session whose policy target is scalar and has nothing computed, runs the scalar units and answers stats_update with tier scalar, final and remaining. The run behaves like a filtered one: the fragments go to the clients that ask (each through its own cursor) and nothing is assigned to the dataflow, the session snapshot, summary_stats_cache or the status. Other clients are served as before. Co-Authored-By: Claude Sonnet 5.5 --- .../stat_pipeline.py | 5 +- .../stat_units.py | 11 +- .../xorq_stat_pipeline.py | 83 ++++++++++---- buckaroo/server/session.py | 14 ++- buckaroo/server/stat_run.py | 45 ++++++-- buckaroo/server/stats_wire.py | 103 +++++++++++++++--- tests/unit/server/test_load_expr.py | 29 ++++- 7 files changed, 233 insertions(+), 57 deletions(-) diff --git a/buckaroo/pluggable_analysis_framework/stat_pipeline.py b/buckaroo/pluggable_analysis_framework/stat_pipeline.py index 812553ae4..22197f870 100644 --- a/buckaroo/pluggable_analysis_framework/stat_pipeline.py +++ b/buckaroo/pluggable_analysis_framework/stat_pipeline.py @@ -288,7 +288,10 @@ def _skipped(state: StatState, pairs: Sequence[Tuple[Any, str]]) -> set: 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.""" + column, so no unit waits for another. Nothing is computed. These + pipelines have only the full tier.""" + if state.tier != "full": + raise ValueError(f"the pandas and polars pipelines have only the full tier, not {state.tier!r}") if len(state.data) == 0: return [] pairs = columns_in_scope(state) diff --git a/buckaroo/pluggable_analysis_framework/stat_units.py b/buckaroo/pluggable_analysis_framework/stat_units.py index c5a02c1c4..611bb16ce 100644 --- a/buckaroo/pluggable_analysis_framework/stat_units.py +++ b/buckaroo/pluggable_analysis_framework/stat_units.py @@ -42,6 +42,11 @@ class (``column``, ``scan`` or ``query``) a scheduler can read. ``after`` cost: str = "column" +# The tiers a stat run is planned at: ``full`` is every stat, and ``scalar`` the +# scalar class of them (the xorq pipeline says which, ``SCALAR_OMITTED_KEYS``). +# The ``schema`` tier is the dataflow's own and runs no units. +UNIT_TIERS = ("scalar", "full") + # 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``. @@ -87,7 +92,8 @@ class StatState: 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. + needs and which is never queried for. ``tier`` is the tier the run is at + (see ``UNIT_TIERS``); a pipeline that has only the full tier refuses another. """ data: Any skip_columns: frozenset = frozenset() @@ -95,9 +101,12 @@ class StatState: priority: Tuple[Any, ...] = () rows: Optional[int] = None namespace: str = "any" + tier: str = "full" def __post_init__(self) -> None: check_namespace(self.namespace) + if self.tier not in UNIT_TIERS: + raise ValueError(f"tier must be one of {UNIT_TIERS}, not {self.tier!r}") def columns_in_scope(state: StatState) -> List[Tuple[Any, str]]: diff --git a/buckaroo/pluggable_analysis_framework/xorq_stat_pipeline.py b/buckaroo/pluggable_analysis_framework/xorq_stat_pipeline.py index 9601e2e03..cec1a672a 100644 --- a/buckaroo/pluggable_analysis_framework/xorq_stat_pipeline.py +++ b/buckaroo/pluggable_analysis_framework/xorq_stat_pipeline.py @@ -13,6 +13,10 @@ typed-DAG executor with results written into the per-column accumulator. One ``histogram:`` unit per column. +At the scalar tier (``StatState.tier``) a run is the batch alone, without the +aggregates that cost more than a pass over a column (``approx_median``, and +the ``approx_nunique`` behind ``distinct_count``) and with no histogram query. + Errors are captured into ``StatError`` via the standard Ok/Err mechanism; nothing is silently swallowed. Construction validates the DAG up front and raises ``DAGConfigError`` on bad configurations. @@ -26,7 +30,7 @@ import os import time from contextlib import nullcontext -from typing import Any, Dict, List, Optional, Sequence, Tuple +from typing import Any, Dict, Iterable, List, Optional, Sequence, Tuple import pandas as pd @@ -111,14 +115,17 @@ 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.""" + it reported once; ``phases`` caches how each tier and dtype's stats split + (see ``XorqStatPipeline._column_phases``). ``withheld`` are keys the tier + does not report: they are in ``working`` (as the unknown ``None`` that + ``histogram_bins`` reads) but in no fragment and not in ``sd()``. ``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] = ()): + skipped: Sequence[Any] = (), withheld: Iterable[str] = ()): super().__init__(state, columns, rewritten) self.skipped: Tuple[Any, ...] = tuple(skipped) + self.withheld: frozenset = frozenset(withheld) self.working: Dict[Any, Dict[str, StatResult]] = {} self.emitted: Dict[Any, set] = {} self.phases: Dict[Any, Any] = {} @@ -126,7 +133,7 @@ def __init__(self, state: StatState, columns: Sequence[Any] = (), rewritten: Opt 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} + for key, result in self.working[col].items() if key not in self.withheld} for col in self.order} class XorqStatPipeline(UnitPipeline): @@ -146,6 +153,9 @@ class XorqStatPipeline(UnitPipeline): 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``). + + ``StatState.tier`` is ``full`` (every stat) or ``scalar``: the batch alone, + in the same chunks, without the stats in ``SCALAR_OMITTED_KEYS``. """ # Keys that the pipeline pre-populates per column. Listed as external @@ -157,6 +167,15 @@ class XorqStatPipeline(UnitPipeline): {"orig_col_name", "rewritten_col_name", "dtype", "length", "min", "max", "distinct_count"}) + # What the scalar tier does not report. ``median`` (approx_median) and + # ``distinct_count`` (approx_nunique) are the batch aggregates it leaves + # out, ``distinct_per`` is the ratio of the second and ``histogram`` is the + # per-column query. ``histogram_bins`` is still produced, from ``min`` and + # ``max`` with ``distinct_count`` unknown, so a low-cardinality integer + # column gets bins that the full tier then replaces with none. A stat that + # reads an omitted key other than ``distinct_count`` is left out with it. + SCALAR_OMITTED_KEYS = frozenset({"median", "distinct_count", "distinct_per", "histogram"}) + def __init__(self, stat_funcs: list, backend: Any = None, unit_test: bool = True, cache_storage=None, chunk_cells: Optional[int] = None): if not HAS_XORQ: @@ -426,13 +445,24 @@ def _batch_chunks(self, state: StatState, columns: List[Any]) -> List[List[Any]] 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]]: + def _scalar_omits(self, sf: StatFunc) -> bool: + """Whether the scalar tier leaves ``sf`` out: it provides a key the + tier does not report.""" + return any(sk.name in self.SCALAR_OMITTED_KEYS for sk in sf.provides) + + def _column_phases(self, cache: Dict[Any, Any], dtype, + tier: str = "full") -> 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) + stat key -> func map errors are reported against. At the scalar tier the + funcs that provide a key in ``SCALAR_OMITTED_KEYS`` are left out, and + so are those that read a key only they provide (``distinct_count`` is + external, so what reads it stays).""" + if (tier, dtype) not in cache: + funcs = self.all_stat_funcs if tier == "full" else [ + sf for sf in self.all_stat_funcs if not self._scalar_omits(sf)] + col_funcs = build_column_dag(funcs, dtype, external_keys=self.EXTERNAL_KEYS) pure: List[StatFunc] = [] query: List[StatFunc] = [] query_keys: set = set() @@ -443,15 +473,16 @@ def _column_phases(self, cache: Dict[Any, Any], dtype) -> Tuple[List[StatFunc], 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] + cache[(tier, dtype)] = (pure, query, key_to_func) + return cache[(tier, 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.""" + every histogram, since a histogram reads what its chunk computes. The + scalar tier has the batch units alone. Skipped columns get no unit. + Nothing is computed: no query is sent.""" table = state.data pairs = columns_in_scope(state) if not pairs: @@ -468,7 +499,7 @@ def plan(self, state: StatState) -> List[StatUnit]: 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: + for col in active if state.tier == "full" else (): 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")) @@ -479,7 +510,9 @@ def new_accumulator(self, state: StatState) -> "XorqAccumulator": 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`` + keys (the scalar tier withholds the ones it does not report, so + ``distinct_count`` stays the unknown ``None`` and is never returned). + ``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 @@ -492,7 +525,8 @@ def new_accumulator(self, state: StatState) -> "XorqAccumulator": 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]) + skipped=[orig for orig, _rewritten in pairs if orig in state.skip_columns], + withheld=self.SCALAR_OMITTED_KEYS if state.tier == "scalar" else ()) 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)} @@ -507,6 +541,8 @@ def run(self, unit: StatUnit, acc: StatAccumulator) -> Fragment: assert isinstance(acc, XorqAccumulator) if unit.phase == "batch": return self._run_batch(unit, acc) + if acc.state.tier == "scalar": + raise ValueError(f"the scalar tier runs no {unit.phase} unit ({unit.id})") return self._run_histogram(unit, acc) def _run_funcs(self, funcs: List[StatFunc], col: Any, acc: "XorqAccumulator") -> None: @@ -525,10 +561,12 @@ def _run_funcs(self, funcs: List[StatFunc], col: Any, acc: "XorqAccumulator") -> 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.""" + """The stats of ``col`` that no earlier unit has returned and the tier + reports, 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]} + fresh = {key: result for key, result in working.items() + if key not in acc.emitted[col] and key not in acc.withheld} plain, errors = resolve_accumulator(fresh, col, key_to_func) acc.emitted[col].update(fresh) acc.record({col: plain}, errors) @@ -537,6 +575,7 @@ def _emit(self, acc: "XorqAccumulator", col: Any, key_to_func: Dict[str, StatFun def _run_batch(self, unit: StatUnit, acc: "XorqAccumulator") -> Fragment: table = acc.state.data schema = acc.schema + tier = acc.state.tier 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 @@ -544,7 +583,7 @@ def _run_batch(self, unit: StatUnit, acc: "XorqAccumulator") -> Fragment: TOTAL_LENGTH_KEY = "__total_length__" batch_items: List[Tuple[str, StatFunc, Any]] = [] for sf in self.ordered_stat_funcs: - if not _is_batch_func(sf): + if not _is_batch_func(sf) or (tier == "scalar" and self._scalar_omits(sf)): continue xorq_col_param = next(r.name for r in sf.requires if r.type is XorqColumn) for col in columns: @@ -603,7 +642,7 @@ def _run_batch(self, unit: StatUnit, acc: "XorqAccumulator") -> Fragment: # histogram_bins (and with it color_map) needs no histogram query. fragment: Fragment = {} for col in columns: - pure, _query, key_to_func = self._column_phases(acc.phases, schema[col]) + pure, _query, key_to_func = self._column_phases(acc.phases, schema[col], tier) self._run_funcs(pure, col, acc) fragment[col] = self._emit(acc, col, key_to_func) return fragment diff --git a/buckaroo/server/session.py b/buckaroo/server/session.py index f36ee8d53..ec72c758a 100644 --- a/buckaroo/server/session.py +++ b/buckaroo/server/session.py @@ -40,9 +40,10 @@ def dataflow_stats_tier(stats_tier: str, stats_delivery: str) -> str: """The tier a session's dataflow is constructed at. Only a session whose stats run inline and whose host named no lower tier builds at the full tier; a deferred one starts at the schema tier whatever tier it is headed for, and - so does one headed for ``scalar`` (no scalar-tier units exist yet) or - ``schema``. ``auto`` is resolved against a schema-tier dataflow, which an - inline session never has, so inline ``auto`` is ``full``.""" + so does one headed for ``scalar`` (its stats are fragments of a scalar + ``StatRun``, never assigned to the dataflow) or ``schema``. ``auto`` is + resolved against a schema-tier dataflow, which an inline session never has, + so inline ``auto`` is ``full``.""" if stats_delivery == "deferred" or stats_tier in ("scalar", "schema"): return "schema" return "full" @@ -51,9 +52,10 @@ def dataflow_stats_tier(stats_tier: str, stats_delivery: str) -> str: # 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). +# them yet), ``not_computed`` (the snapshot has no stats, and none will be +# assigned to it: the session is headed for the schema tier, or for ``scalar``, +# whose stats reach a client that asks as ``stats_update`` fragments) and +# ``error`` (the run failed; sticky until the next generation). STATS_STATUSES = ("complete", "pending", "not_computed", "error") diff --git a/buckaroo/server/stat_run.py b/buckaroo/server/stat_run.py index 82a13fba7..ff5682ce9 100644 --- a/buckaroo/server/stat_run.py +++ b/buckaroo/server/stat_run.py @@ -9,6 +9,12 @@ 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. + +Only the full run over every column is the session's stats: its results are the +final assignment. A run at another tier (``scalar``) or over a column group +behaves like a filtered run: its fragments go to the clients that ask and +nothing is assigned. ``stat_run_key`` keeps the full run's key as it was and +gives every other run a key of its own, so the two never meet. """ import itertools import time @@ -16,14 +22,27 @@ 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 +from buckaroo.pluggable_analysis_framework.stat_units import Fragment, StatState, StatUnit, resolve_names, rewrite_sd _run_ids = itertools.count(1) +def stat_run_key(stats_gen: int, scope: str, tier: str = "full", columns: Optional[Sequence[Any]] = None) -> tuple: + """The key a session keeps a run under. The full run over every column is + ``(stats_gen, scope)``, as it has always been. A run at another tier or over + a column group is ``(stats_gen, scope, tier, columns)`` (``columns`` a tuple, + or ``None`` for every column), so it never takes the full run's place.""" + group = None if columns is None else tuple(columns) + if tier == "full" and group is None: + return (stats_gen, scope) + return (stats_gen, scope, tier, group) + + 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. + ``run=False``: the run plans it once and runs its units through it. ``state`` + is what the run describes, the stats class's own unless a tier or a column + group says otherwise (see ``StatState``). ``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 @@ -32,13 +51,14 @@ class StatRun: every request that ran one. """ - def __init__(self, stats_gen: int, scope: str, stats: Any): + def __init__(self, stats_gen: int, scope: str, stats: Any, state: Optional[StatState] = None): 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.state: StatState = stats.state if state is None else state + self.units: List[StatUnit] = stats.plan(self.state) + self.acc = stats.new_accumulator(self.state) self.fragments: List[Fragment] = [] self.ran: List[str] = [] self.status = "pending" if self.units else "complete" @@ -48,8 +68,19 @@ def __init__(self, stats_gen: int, scope: str, stats: Any): self._frame_columns: Optional[List[Tuple[Any, str]]] = None @property - def key(self) -> Tuple[int, str]: - return (self.stats_gen, self.scope) + def tier(self) -> str: + return self.state.tier + + @property + def assigns(self) -> bool: + """Whether the run's results are the session's stats: the full tier over + every column. Any other run is a fragment source for clients and is + never assigned or cached as complete.""" + return self.state.tier == "full" and self.state.columns is None + + @property + def key(self) -> tuple: + return stat_run_key(self.stats_gen, self.scope, self.state.tier, self.state.columns) @property def remaining(self) -> int: diff --git a/buckaroo/server/stats_wire.py b/buckaroo/server/stats_wire.py index 47ddfb4ee..8dd6442ef 100644 --- a/buckaroo/server/stats_wire.py +++ b/buckaroo/server/stats_wire.py @@ -24,6 +24,15 @@ choice and reaches every client as ``not_computed``. ``stats_to_pull`` decides, per client, whether a session has stats left to compute. +A target of ``scalar`` (``serves_scalar_tier``) can be pulled by the client that +advertised both bits: its ``stats_request`` runs the units of a scalar ``StatRun`` +and is answered with ``stats_update`` messages whose tier is ``scalar``. Nothing +is assigned. The run's fragments go to the clients that ask, as a filtered run's +would, and the dataflow, the session snapshot, ``summary_stats_cache`` and the +status stay as they were, so the scalar stats are never served as the complete +ones. Only the full run over every column reaches the final assignment +(``StatRun.assigns``). + 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. @@ -40,6 +49,7 @@ cursor, so work is done once whoever asks. """ import copy +import dataclasses import hashlib import json import logging @@ -53,7 +63,7 @@ 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, stats_deferred_by_policy -from buckaroo.server.stat_run import StatCursor, StatRun +from buckaroo.server.stat_run import StatCursor, StatRun, stat_run_key from buckaroo.server.stats_policy import resolve_stats_policy log = logging.getLogger("buckaroo.server.stats_wire") @@ -109,6 +119,17 @@ def stats_to_pull(session: SessionState, client: Any) -> bool: return session.stats_status == "pending" or (stats_deferred_by_policy(session) and not client_has_ondemand(client)) +def serves_scalar_tier(session: SessionState, client: Any) -> bool: + """Whether a ``stats_request`` from ``client`` is answered with the scalar + tier: the client advertised both bits, so it takes tiers, and the session has + nothing computed (``not_computed``) and a policy target of ``scalar``, + whether the size rule, the host or the ceiling put it there. Any other client + of such a session is served as ``stats_to_pull`` says.""" + policy = session.stats_policy + return (client_has_ondemand(client) and session.stats_status == "not_computed" and policy is not None + and policy["tier_target"] == "scalar") + + def resolve_session_policy(stats_tier: str, dataflow_tier: str, rows: int, cols: int) -> Optional[dict]: """The policy for a session the handler has just built: ``stats_tier`` is the host's request, ``dataflow_tier`` the tier its dataflow was built at, and @@ -145,20 +166,25 @@ 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.""" +def start_stat_run(session: SessionState, scope: str = "raw", tier: str = "full", + columns: Optional[Sequence[Any]] = None) -> Optional[StatRun]: + """The ``StatRun`` of the session's current generation for ``scope``, ``tier`` + and column group: the one it holds, or a new one planned from the dataflow + and stored on the session under ``stat_run_key``. Planning sends no query. + The run analyzes the frame ``assign_full_stats`` does, so running every unit + of the full run gives the stats that call computes whole. A run at another + tier, or over a column group (``columns``, original names), is never + assigned (``StatRun.assigns``). ``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) + key = stat_run_key(session.stats_gen, scope, tier, columns) run = session.stat_runs.get(key) if run is None: - run = StatRun(session.stats_gen, scope, dataflow.build_stats(dataflow.processed_df, run=False)) + stats = dataflow.build_stats(dataflow.processed_df, run=False) + state = dataclasses.replace(stats.state, tier=tier, columns=None if columns is None else tuple(columns)) + run = StatRun(session.stats_gen, scope, stats, state) session.stat_runs[key] = run return run @@ -270,7 +296,12 @@ def assign_full_stats(dataflow: Any, computed: Optional[tuple] = None) -> None: 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).""" + (it raises on a failed stat in debug mode, and so does this). Only a run that + assigns has them: a scalar or column-scoped run is a part of the stats, and + caching it as the whole would serve it as complete.""" + if not run.assigns: + raise ValueError(f"a {run.tier} run over {'every column' if run.state.columns is None else 'some columns'} " + "is not the session's stats and cannot be assigned") errs = run.errs() if errs and dataflow.debug: raise Exception("Error executing analysis") @@ -401,13 +432,14 @@ 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.""" +def _plan_run(session: SessionState, tier: str = "full") -> Optional[StatRun]: + """The session's ``StatRun`` at ``tier`` 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). For the full tier the whole-run path then decides what the + stats are.""" try: - return start_stat_run(session) + return start_stat_run(session, tier=tier) except Exception: log.warning("stat units not planned session=%s stats_gen=%s: %s", session.session_id, session.stats_gen, traceback.format_exc()) @@ -449,6 +481,37 @@ def _serve_units(session: SessionState, client: Any, prefer: Sequence[str], stat "elapsed_ms": _elapsed_ms(started)} +def _serve_scalar(session: SessionState, client: Any, prefer: Sequence[str], incremental: bool, stats_gen: Any, + scope: Any, started: float, info: dict) -> dict: + """The reply to a ``stats_request`` on a session whose target is ``scalar``: + a ``stats_update`` with ``tier: "scalar"`` holding the fragments this client + has not seen, ``final`` once the run is done. An incremental request runs + units for the time budget, as the full run's does; any other runs every unit + still to run. + + Nothing is assigned. The run stays on the session for the generation, so a + client that asks later is caught up from the fragments with no unit run, and + the status, the snapshot and the caches are as they were. A unit that raises + aborts the request and the run (it is not retried) but not the session, whose + full tier is still open.""" + run = _plan_run(session, "scalar") + if run is None or run.status == "error": + return _aborted(stats_gen, scope, "error", session) + cursor = getattr(client, "stats_cursor", None) or StatCursor() + if cursor.caught_up(run): + try: + info["units"] = run_units(run, STATS_BUDGET_S if incremental else None, prefer, + session_id=session.session_id) + except Exception: + log.error("scalar stat unit failed session=%s stats_gen=%s: %s", session.session_id, session.stats_gen, + traceback.format_exc()) + return _aborted(stats_gen, scope, "error", session) + fragments = cursor.take(run) + return {"type": "stats_update", "stats_gen": stats_gen, "scope": scope, "tier": run.tier, + "final": run.status != "pending", "remaining": run.remaining, + "payload": partial_payload(session_dataflow(session), 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``) @@ -480,6 +543,8 @@ def _answer_stats_request(session: Optional[SessionState], msg: dict, client: An 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 serves_scalar_tier(session, client): + return _serve_scalar(session, client, prefer, bool(incremental), stats_gen, scope, started, info) to_pull = stats_to_pull(session, client) if session.stats_status == "not_computed" and not to_pull: return _aborted(stats_gen, scope, "not_requestable", session) @@ -519,7 +584,9 @@ def handle_stats_request(session: Optional[SessionState], msg: dict, client: Any 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. + ``bad_request``, and is not run as a whole run. On a session whose target is + ``scalar`` an ondemand client's request is answered as ``_serve_scalar`` + says: ``stats_update`` messages with ``tier: "scalar"`` and nothing assigned. The ``stats.request`` span records the request and its ``outcome`` (``update`` or the abort reason), which is where updates sent, requests diff --git a/tests/unit/server/test_load_expr.py b/tests/unit/server/test_load_expr.py index 4b179ca39..a4bf7aeac 100644 --- a/tests/unit/server/test_load_expr.py +++ b/tests/unit/server/test_load_expr.py @@ -3293,6 +3293,26 @@ async def test_a_dataflow_field_change_starts_a_new_scalar_run_over_the_new_stat self.assertEqual(session.stat_runs[(gen + 1, "raw", "scalar", None)].acc.sd()["qty"]["length"], 2, "the run is over the filtered state") + @tornado.testing.gen_test + async def test_the_clients_that_cannot_take_a_scalar_target_are_served_as_before(self): + await self._load("st-host", stats_tier="scalar", stats_delivery="deferred") + ws, update = await self._connect("st-host", caps="stats_update") + ws.write_message(_stats_request(self._stats(update)["gen"], incremental=True)) + self.assertEqual((await _read_json(ws))["reason"], "not_requestable") + + await self._scalar_target("st-skew") + ws, pending = await self._connect("st-skew", caps="stats_update") + self.assertEqual(self._stats(pending)["status"], "pending") + reply = await self._pull(ws, self._stats(pending)["gen"]) + self.assertEqual((reply["type"], reply["tier"]), ("stats_update", "full")) + + @tornado.testing.gen_test + async def test_a_schema_target_still_refuses_the_request(self): + await self._load("st-schema", limits=_SCHEMA_BY_SIZE, stats_tier="auto", stats_delivery="deferred") + ws, first = await self._connect("st-schema", caps=_ONDEMAND) + aborted = await self._pull(ws, self._stats(first)["gen"]) + self.assertEqual((aborted["type"], aborted["reason"]), ("stats_aborted", "not_requestable")) + @tornado.testing.gen_test async def test_a_scalar_run_and_a_column_scoped_run_are_stored_apart_from_the_full_run(self): await self._load("st-keys", stats_delivery="deferred") @@ -3381,8 +3401,13 @@ def test_each_request_runs_one_chunk_and_reports_the_scalar_tier(self, tmp_path, assert [{c for row in _rows_by_stat(r["payload"]).values() for c in row} - {"index", "level_0"} for r in replies] == [{"a", "b"}, {"c", "d"}, {"e", "f"}] whole = self._request(self._session(tmp_path), _ondemand_client()) - assert _as_json(_merge_stat_payloads(r["payload"] for r in replies)) == _as_json( - _merge_stat_payloads([whole["payload"]])) + + def held(payloads): + """What a client holds, less the empty cells: a chunk with no column that has a stat sends no row for it.""" + return {stat: {col: value for col, value in cells.items() if value is not None} + for stat, cells in _merge_stat_payloads(payloads).items()} + + assert _as_json(held(r["payload"] for r in replies)) == _as_json(held([whole["payload"]])) def test_a_request_after_the_last_chunk_replays_nothing_and_runs_nothing(self, tmp_path, monkeypatch): monkeypatch.setattr(stats_wire, "STATS_BUDGET_S", 0) From d4477c469044eae62e827264981f49c492ca42b8 Mon Sep 17 00:00:00 2001 From: Paddy Mullen Date: Sun, 4 Oct 2026 08:47:07 -0400 Subject: [PATCH 31/39] test(server): failing tests for limits, override, cost guard and the demand scan (rows-first p36a) A stats_request from a client that advertised both capability bits can now say tier, force and columns. The tests pin what the server does with them: the ceiling refuses a tier for the cells asked for and runs nothing, requestable and the target decide which tier a whole-table request may name, force raises the target through a session override that survives a dataflow-field change for the unfiltered scope only, an automatic unit over the budget pauses the session (not_computed for cost) until a force request, and the pause survives a generation and /reload_expr. The demand scan reads the color_map rules of the built config (overrides and rules a klass adds at style time) with no query, and a scoped request computes those columns only. Every test here fails on the earlier code. The guard tests that pass there (a client without both bits is served as before, a session with no policy ignores the fields) go in with the implementation. Co-Authored-By: Claude Sonnet 5.5 --- tests/unit/server/test_load_expr.py | 712 +++++++++++++++++++++++++ tests/unit/server/test_stats_policy.py | 382 ++++++++++++- 2 files changed, 1091 insertions(+), 3 deletions(-) diff --git a/tests/unit/server/test_load_expr.py b/tests/unit/server/test_load_expr.py index a4bf7aeac..b1080cadd 100644 --- a/tests/unit/server/test_load_expr.py +++ b/tests/unit/server/test_load_expr.py @@ -3435,6 +3435,718 @@ def test_a_unit_that_fails_aborts_the_request_and_not_the_session(self, tmp_path assert any("stat unit failed" in record.getMessage() for record in caplog.records) +# --------------------------------------------------------------------------- +# Limits, override, cost guard and the demand scan (rows-first p36a) +# --------------------------------------------------------------------------- + + +@contextmanager +def _slow_stat_queries(delay): + """Every query ``XorqStatPipeline`` sends while the block runs takes + ``delay`` seconds longer, as a slow source would make it.""" + original = XorqStatPipeline._execute + + def slow(pipeline, query): + time.sleep(delay) + return original(pipeline, query) + + with patch.object(XorqStatPipeline, "_execute", slow): + yield + + +class _LimitsWire(tornado.testing.AsyncHTTPTestCase): + """What the p36a wire tests share: a ``/load_expr`` session over the + five-row table (deferred unless a test says otherwise) and clients that + advertise both capability bits. A request is judged against the thresholds + again, so a test that sets them holds them in the environment for the whole + block (``_stats_limits``), not only around the load.""" + + DISPLAY_FILES: dict = {} + + 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") + if self.DISPLAY_FILES: + display_dir = os.path.join(self.project_root, "display") + os.makedirs(display_dir) + for name, source in self.DISPLAY_FILES.items(): + with open(os.path.join(display_dir, name), "w") as f: + f.write(source) + 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) + + def _patch(self, target, name, value): + patcher = patch.object(target, name, value, create=True) + patcher.start() + self.addCleanup(patcher.stop) + + def _one_unit_per_request(self): + self._patch(stats_wire, "STATS_BUDGET_S", 0) + + def _cost_budget(self, seconds): + self._patch(stats_wire, "STATS_COST_BUDGET_S", seconds) + + 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, "stats_delivery": "deferred", **body}) + self.assertEqual(resp.code, 200, resp.body) + + async def _reload(self, sid, **body): + resp = await _post(self.get_http_port(), f"/reload_expr/{sid}", body) + self.assertEqual(resp.code, 200, resp.body) + + async def _connect(self, sid, caps=_ONDEMAND): + 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): + """The complete message an inline session of the same build sends.""" + await self._load(sid, stats_delivery="inline") + return (await self._connect(sid, caps=None))[1] + + async def _state(self, ws, **changes): + """A state change from ``ws``, and the frame it broadcasts.""" + ws.write_message(_state_change(**changes)) + return await _read_json(ws) + + async def _ask(self, ws, gen, **fields): + """The reply to an incremental ``stats_request``.""" + ws.write_message(_stats_request(gen, incremental=True, **fields)) + return await _read_json(ws) + + async def _ask_to_the_end(self, ws, gen, limit=12, **fields): + """Incremental requests, the same each time, until a reply is final.""" + replies = [] + for _ in range(limit): + replies.append(await self._ask(ws, gen, **fields)) + if replies[-1]["type"] != "stats_update" or replies[-1]["final"]: + break + return replies + + +# (host tier, limits, request fields): a request the ceiling refuses, whoever +# sends it. The table is 15 cells; a ceiling_full_rows of 3 refuses full for it +# and a ceiling_scalar_cells of 10 refuses scalar for the whole table. +_OVER_THE_CEILING = { + "full over a ceiling on full": ("auto", {**_SCHEMA_BY_SIZE, "ceiling_full_rows": 3}, + {"tier": "full", "force": True}), + "full, not forced": ("auto", {**_SCHEMA_BY_SIZE, "ceiling_full_rows": 3}, {"tier": "full"}), + "full, the tier the host asked for": ("full", {"ceiling_full_rows": 3}, {"tier": "full", "force": True}), + "scalar over the ceiling on scalar": ("schema", {"ceiling_scalar_cells": 10}, {"tier": "scalar", "force": True}), + "full over the ceiling on scalar": ("schema", {"ceiling_scalar_cells": 10}, {"tier": "full", "force": True}), + "every column of the table": ("schema", {"ceiling_scalar_cells": 10}, + {"tier": "scalar", "columns": ["a", "b", "c"]})} + + +class TestForcedRequestsWire(_LimitsWire): + """``stats_request`` with ``tier`` and ``force`` (rows-first p36a): the + ceiling, ``requestable`` and the session override, from a client that + advertised both bits.""" + + @tornado.testing.gen_test + async def test_a_request_over_the_ceiling_runs_no_unit(self): + for index, (name, (tier, limits, fields)) in enumerate(_OVER_THE_CEILING.items()): + with _stats_limits(**limits): + sid = f"fo-over-{index}" + await self._load(sid, stats_tier=tier) + session = self._session(sid) + ws, first = await self._connect(sid) + gen = self._stats(first)["gen"] + held = (session.stats_status, session.stats_reason) + with _count_stat_queries() as queries: + reply = await self._ask(ws, gen, **fields) + self.assertEqual(queries, [], name) + self.assertEqual({key: reply.get(key) for key in ("type", "final", "status", "reason", "remaining")}, + {"type": "stats_update", "final": True, "status": "not_computed", "reason": "ceiling", + "remaining": 0}, name) + self.assertEqual((reply["stats_gen"], reply["scope"]), (gen, "raw"), name) + self.assertNotIn("payload", reply, name) + self.assertEqual((session.stats_status, session.stats_reason, session.stat_runs, session.stats_override), + (*held, {}, None), name) + + @tornado.testing.gen_test + async def test_a_forced_scalar_request_under_the_ceiling_runs_and_raises_the_target(self): + with _stats_limits(**_SCHEMA_BY_SIZE): + await self._load("fo-scalar", stats_tier="auto") + session = self._session("fo-scalar") + ws, first = await self._connect("fo-scalar") + gen = self._stats(first)["gen"] + self.assertEqual((self._stats(first)["tier_target"], self._stats(first)["requestable"]), ( + "schema", ["scalar", "full"])) + with _count_stat_queries() as queries: + update = await self._ask(ws, gen, tier="scalar", force=True) + self.assertEqual((update["type"], update["tier"], update["final"], update["remaining"]), + ("stats_update", "scalar", True, 0), update) + self.assertEqual(set(_rows_by_stat(update["payload"])), _SCALAR_WIRE_ONLY) + self.assertEqual([type(q.op()).__name__ for q in queries], ["Aggregate"]) + self.assertEqual((update["status"], update["reason"]), ("not_computed", "size"), + "the final reply says the session still has no complete stats") + self.assertIsInstance(update["elapsed_ms"], (int, float)) + self.assertEqual(session.stats_override, "scalar") + self.assertEqual((session.stats_status, session.stats_reason, session.stats_policy["tier_target"]), + ("not_computed", "size", "schema")) + # A client that connects later is told the target the override set. + _, later = await self._connect("fo-scalar") + self.assertEqual(self._stats(later), {"status": "not_computed", "tier": "schema", "gen": gen, + "reason": "size", "tier_target": "scalar", "estimate": _WIRE_ESTIMATE}) + + @tornado.testing.gen_test + async def test_a_forced_full_request_runs_every_unit_and_completes_the_session(self): + self._one_unit_per_request() + inline = await self._inline_frame("fo-full-inline") + with _stats_limits(**_SCHEMA_BY_SIZE): + await self._load("fo-full", stats_tier="auto") + session = self._session("fo-full") + ws, first = await self._connect("fo-full") + gen = self._stats(first)["gen"] + self.assertEqual(self._stats(first)["status"], "not_computed") + + opening = await self._ask(ws, gen, tier="full", force=True) + # Run for real: the session is now headed for full, and says so. + _, mid = await self._connect("fo-full") + self.assertEqual({key: self._stats(mid).get(key) for key in ("status", "tier_target", "requestable")}, + {"status": "pending", "tier_target": "full", "requestable": []}) + rest = await self._ask_to_the_end(ws, gen, tier="full", force=True) + replies = [opening, *rest] + self.assertEqual([r["type"] for r in replies], ["stats_update"] * 4) + self.assertEqual([(r["tier"], r["final"], r["remaining"]) for r in replies], + [("full", False, 3), ("full", False, 2), ("full", False, 1), ("full", True, 0)]) + final = replies[-1] + self.assertNotIn("status", final, "a final reply that completes the session says nothing more") + self.assertTrue(_FULL_ONLY_WIRE <= set(_rows_by_stat(final["payload"]))) + self.assertIn("df_display_args", final, "the config the complete stats give is not the one the client holds") + self.assertEqual((session.stats_status, session.stats_override), ("complete", "full")) + _, legacy = await self._connect("fo-full", caps=None) + self.assertEqual(_comparable(legacy), _comparable(inline)) + + @tornado.testing.gen_test + async def test_a_request_above_the_target_without_force_is_not_requestable(self): + with _stats_limits(**_SCHEMA_BY_SIZE): + await self._load("fo-noforce", stats_tier="auto") + session = self._session("fo-noforce") + ws, first = await self._connect("fo-noforce") + gen = self._stats(first)["gen"] + for fields in ({"tier": "scalar"}, {"tier": "full"}, {"tier": "full", "force": False}): + with _count_stat_queries() as queries: + aborted = await self._ask(ws, gen, **fields) + self.assertEqual((aborted["type"], aborted["reason"]), ("stats_aborted", "not_requestable"), fields) + self.assertEqual(queries, [], fields) + self.assertEqual((session.stats_override, session.stat_runs), (None, {})) + + @tornado.testing.gen_test + async def test_a_request_below_the_target_is_not_requestable(self): + await self._load("fo-below", stats_tier="auto") + ws, first = await self._connect("fo-below") + self.assertEqual(self._stats(first)["tier_target"], "full") + aborted = await self._ask(ws, self._stats(first)["gen"], tier="scalar", force=True) + self.assertEqual((aborted["type"], aborted["reason"]), ("stats_aborted", "not_requestable")) + + @tornado.testing.gen_test + async def test_the_override_survives_a_dataflow_field_change_for_the_unfiltered_scope_only(self): + with _stats_limits(**_SCHEMA_BY_SIZE): + await self._load("fo-survive", stats_tier="auto") + ws, first = await self._connect("fo-survive") + gen = self._stats(first)["gen"] + # No force yet: a field change leaves the session where the policy put it. + plain = self._stats(await self._state(ws, post_processing="first_three")) + self.assertEqual((plain["status"], plain["reason"], plain["tier_target"], plain["gen"]), + ("not_computed", "size", "schema", gen + 1)) + gen += 1 + replies = await self._ask_to_the_end(ws, gen, tier="full", force=True) + self.assertEqual((replies[-1]["type"], replies[-1].get("final")), ("stats_update", True)) + + with _count_stat_queries() as queries: + unfiltered = self._stats(await self._state(ws, post_processing="")) + self.assertEqual((unfiltered["status"], unfiltered["tier_target"], unfiltered["gen"]), + ("pending", "full", gen + 1)) + self.assertNotIn("reason", unfiltered) + self.assertEqual(queries, [], "the new state is the schema tier again, and runs no stat query") + gen += 1 + + filtered = self._stats(await self._state(ws, quick_command_args={"search": ["a"]})) + self.assertEqual((filtered["status"], filtered["reason"], filtered["tier_target"], filtered["gen"]), + ("not_computed", "size", "schema", gen + 1), "a search is the filtered scope: the policy alone") + gen += 1 + + cleared = self._stats(await self._state(ws)) + self.assertEqual((cleared["status"], cleared["tier_target"], cleared["gen"]), ("pending", "full", gen + 1)) + + @tornado.testing.gen_test + async def test_a_force_in_a_filtered_state_is_served_for_that_state_and_waits_for_an_unfiltered_one(self): + with _stats_limits(**_SCHEMA_BY_SIZE): + await self._load("fo-filtered", stats_tier="auto") + session = self._session("fo-filtered") + ws, _first = await self._connect("fo-filtered") + searched = self._stats(await self._state(ws, quick_command_args={"search": ["a"]})) + self.assertEqual(searched["status"], "not_computed") + + replies = await self._ask_to_the_end(ws, searched["gen"], tier="full", force=True) + self.assertEqual((replies[-1]["type"], replies[-1].get("final")), ("stats_update", True), replies[-1]) + self.assertEqual(session.stats_status, "complete") + + cleared = self._stats(await self._state(ws)) + self.assertEqual((cleared["status"], cleared["tier_target"]), ("pending", "full")) + + @tornado.testing.gen_test + async def test_a_reload_keeps_the_override_unless_it_names_a_tier(self): + with _stats_limits(**_SCHEMA_BY_SIZE): + await self._load("fo-reload", stats_tier="auto") + session = self._session("fo-reload") + ws, first = await self._connect("fo-reload") + await self._ask_to_the_end(ws, self._stats(first)["gen"], tier="full", force=True) + self.assertEqual(session.stats_override, "full") + + await self._reload("fo-reload") + kept = self._stats(await _read_json(ws)) + self.assertEqual((kept["status"], kept["tier_target"]), ("pending", "full")) + self.assertEqual(session.stats_override, "full") + + await self._reload("fo-reload", stats_tier="auto") + reset = self._stats(await _read_json(ws)) + self.assertEqual((reset["status"], reset["reason"], reset["tier_target"]), + ("not_computed", "size", "schema")) + self.assertIsNone(session.stats_override) + + @tornado.testing.gen_test + async def test_a_new_expression_forgets_the_override_and_a_rebuild_of_the_same_one_keeps_it(self): + with _stats_limits(**_SCHEMA_BY_SIZE): + await self._load("fo-new", stats_tier="auto") + session = self._session("fo-new") + ws, first = await self._connect("fo-new") + await self._ask_to_the_end(ws, self._stats(first)["gen"], tier="full", force=True) + self.assertEqual(session.stats_override, "full") + + resp = await _post(self.get_http_port(), "/load_expr", {"session": "fo-new", "build_dir": self.build_path, + "stats_delivery": "deferred", "force_reload": True}) + self.assertEqual(resp.code, 200, resp.body) + self.assertEqual(session.stats_override, "full", "the same expression, rebuilt, is the same dataset") + + other_root = tempfile.mkdtemp() + self.addCleanup(shutil.rmtree, other_root, ignore_errors=True) + resp = await _post(self.get_http_port(), "/load_expr", {"session": "fo-new", + "build_dir": _build_stats_wire_dir(other_root), "stats_delivery": "deferred"}) + self.assertEqual(resp.code, 200, resp.body) + self.assertIsNone(session.stats_override) + self.assertEqual((session.stats_status, session.stats_reason), ("not_computed", "size")) + + @tornado.testing.gen_test + async def test_load_and_load_compare_forget_the_override_and_the_pause(self): + """They swap the session to a dataset with no policy, so what was set under the old one must not wait + for the next /load_expr.""" + paths = [] + try: + for index in range(3): + fd, path = tempfile.mkstemp(suffix=".csv") + os.close(fd) + pd.DataFrame({"id": [1, 2], "v": [10 * (index + 1), 20]}).to_csv(path, index=False) + paths.append(path) + for sid, route, body in ( + ("fo-load", "/load", {"path": paths[0], "mode": "buckaroo"}), + ("fo-compare", "/load_compare", {"path1": paths[1], "path2": paths[2], "join_columns": ["id"]})): + await self._load(sid, stats_tier="auto") + session = self._session(sid) + session.stats_override, session.stats_override_gen, session.cost_paused = "full", session.stats_gen, True + resp = await _post(self.get_http_port(), route, {"session": sid, **body}) + self.assertEqual(resp.code, 200, resp.body) + self.assertEqual((session.stats_override, session.stats_override_gen, session.cost_paused), + (None, 0, False), route) + self.assertEqual((session.stats_policy, session.stats_status), (None, "complete"), route) + finally: + for path in paths: + os.unlink(path) + + @tornado.testing.gen_test + async def test_every_stats_update_carries_elapsed_ms_and_it_is_the_time_the_request_took(self): + self._one_unit_per_request() + self._cost_budget(3600) + with _stats_limits(**{**_SCHEMA_BY_SIZE, "ceiling_scalar_cells": 10}): + await self._load("fo-elapsed", stats_tier="auto") + ws, first = await self._connect("fo-elapsed") + gen = self._stats(first)["gen"] + with _slow_stat_queries(0.05): + scoped = await self._ask(ws, gen, tier="scalar", columns=["a"]) + refused = await self._ask(ws, gen, tier="full", force=True) + for name, reply in (("scoped", scoped), ("refused", refused)): + self.assertEqual(reply["type"], "stats_update", name) + self.assertIsInstance(reply["elapsed_ms"], (int, float), name) + self.assertGreaterEqual(scoped["elapsed_ms"], 50, "the request ran a query that took 50 ms") + self.assertLess(refused["elapsed_ms"], 50, "a refusal runs nothing") + + +class TestCostGuardWire(_LimitsWire): + """The cost guard (rows-first p36a): an automatic unit over the budget + pauses the session, held on the session, until a ``force`` request.""" + + async def _paused_session(self, sid): + """A deferred session whose first unit ran slow, one unit per request.""" + self._one_unit_per_request() + self._cost_budget(0.02) + await self._load(sid) + ws, first = await self._connect(sid) + gen = self._stats(first)["gen"] + self.assertEqual(self._stats(first)["status"], "pending") + with _slow_stat_queries(0.06): + update = await self._ask(ws, gen) + self.assertEqual((update["type"], update["final"], update["remaining"]), ("stats_update", False, 3), update) + return ws, gen + + @tornado.testing.gen_test + async def test_an_automatic_unit_over_the_budget_pauses_the_session(self): + ws, gen = await self._paused_session("cg-pause") + session = self._session("cg-pause") + self.assertTrue(session.cost_paused) + self.assertEqual((session.stats_status, session.stats_reason), ("not_computed", "cost")) + with _count_stat_queries() as queries: + refused = await self._ask(ws, gen) + self.assertEqual(queries, [], "the next request runs nothing") + self.assertEqual({key: refused.get(key) for key in ("type", "final", "status", "reason", "tier", "remaining")}, + {"type": "stats_update", "final": True, "status": "not_computed", "reason": "cost", "tier": "full", + "remaining": 0}) + self.assertNotIn("payload", refused) + self.assertIsInstance(refused["elapsed_ms"], (int, float)) + self.assertEqual(len(session.stat_runs[(gen, "raw")].ran), 1, "the unit that ran stays run") + + @tornado.testing.gen_test + async def test_a_unit_under_the_budget_does_not_pause_the_session(self): + self._one_unit_per_request() + self._cost_budget(3600) + await self._load("cg-fast") + ws, first = await self._connect("cg-fast") + with _slow_stat_queries(0.03): + replies = await self._ask_to_the_end(ws, self._stats(first)["gen"]) + self.assertEqual([r["final"] for r in replies], [False, False, False, True]) + session = self._session("cg-fast") + self.assertEqual((session.cost_paused, session.stats_status), (False, "complete")) + + @tornado.testing.gen_test + async def test_a_forced_unit_over_the_budget_does_not_pause_the_session(self): + self._one_unit_per_request() + self._cost_budget(0.02) + with _stats_limits(**_SCHEMA_BY_SIZE): + await self._load("cg-forced", stats_tier="auto") + ws, first = await self._connect("cg-forced") + gen = self._stats(first)["gen"] + with _slow_stat_queries(0.06): + opening = await self._ask(ws, gen, tier="full", force=True) + session = self._session("cg-forced") + self.assertEqual((opening["final"], opening["remaining"]), (False, 3)) + # Looked at here: a later forced request would clear a pause this one had set. + self.assertEqual((session.cost_paused, session.stats_status), (False, "pending")) + with _slow_stat_queries(0.06): + rest = await self._ask_to_the_end(ws, gen, tier="full", force=True) + self.assertEqual([r["final"] for r in rest], [False, False, True]) + self.assertFalse(session.cost_paused) + + @tornado.testing.gen_test + async def test_a_request_that_finishes_the_run_does_not_pause_the_session(self): + """Nothing is left to hold back, and the session is complete.""" + self._patch(stats_wire, "STATS_BUDGET_S", 3600) + self._cost_budget(0.02) + await self._load("cg-done") + ws, first = await self._connect("cg-done") + with _slow_stat_queries(0.03): + reply = await self._ask(ws, self._stats(first)["gen"]) + self.assertEqual((reply["type"], reply["final"]), ("stats_update", True)) + session = self._session("cg-done") + self.assertEqual((session.cost_paused, session.stats_status), (False, "complete")) + + @tornado.testing.gen_test + async def test_the_pause_survives_a_generation_bump_a_reconnect_and_a_reload(self): + ws, gen = await self._paused_session("cg-survive") + session = self._session("cg-survive") + + changed = self._stats(await self._state(ws, post_processing="first_three")) + self.assertEqual((changed["status"], changed["reason"], changed["gen"], changed["tier_target"]), + ("not_computed", "cost", gen + 1, "full")) + self.assertTrue(session.cost_paused) + + _, again = await self._connect("cg-survive") + held = session_mod.stats_with_defaults(again["df_meta"]) + self.assertEqual((held["status"], held["reason"], held["tier_target"]), ("not_computed", "cost", "full"), + "an explicit full reports no policy, but a paused one says which tier a forced request continues") + + await self._reload("cg-survive") + reloaded = self._stats(await _read_json(ws)) + self.assertEqual((reloaded["status"], reloaded["reason"], reloaded["gen"]), ("not_computed", "cost", gen + 2)) + refused = await self._ask(ws, gen + 2) + self.assertEqual((refused["type"], refused["reason"]), ("stats_update", "cost")) + + await self._reload("cg-survive", stats_tier="auto") + reset = self._stats(await _read_json(ws)) + self.assertEqual((reset["status"], reset["gen"]), ("pending", gen + 3)) + self.assertFalse(session.cost_paused) + + @tornado.testing.gen_test + async def test_a_new_expression_forgets_the_pause(self): + _ws, _gen = await self._paused_session("cg-new") + session = self._session("cg-new") + other_root = tempfile.mkdtemp() + self.addCleanup(shutil.rmtree, other_root, ignore_errors=True) + resp = await _post(self.get_http_port(), "/load_expr", {"session": "cg-new", + "build_dir": _build_stats_wire_dir(other_root), "stats_delivery": "deferred"}) + self.assertEqual(resp.code, 200, resp.body) + self.assertEqual((session.cost_paused, session.stats_status), (False, "pending")) + + @tornado.testing.gen_test + async def test_force_clears_the_pause_and_the_run_goes_on_where_it_stopped(self): + ws, gen = await self._paused_session("cg-resume") + session = self._session("cg-resume") + run = session.stat_runs[(gen, "raw")] + first_unit = list(run.ran) + + reply = await self._ask(ws, gen, force=True) + self.assertEqual((reply["type"], reply["tier"], reply["final"]), ("stats_update", "full", False), reply) + self.assertFalse(session.cost_paused) + self.assertEqual((session.stats_status, session.stats_reason), ("pending", None)) + self.assertEqual((run.ran[:1], len(run.ran)), (first_unit, 2), "the unit that ran is not run again") + + # Back to automatic requests, with a budget no unit exceeds. + self._cost_budget(3600) + rest = await self._ask_to_the_end(ws, gen) + self.assertEqual((rest[-1]["type"], rest[-1]["final"]), ("stats_update", True)) + self.assertEqual((session.stats_status, session.cost_paused), ("complete", False)) + + @tornado.testing.gen_test + async def test_a_request_with_no_policy_is_never_paused(self): + """An inline session has no policy, so its ondemand client's requests are answered as before.""" + self._cost_budget(0.0) + await self._load("cg-inline", stats_delivery="inline") + ws, frame = await self._connect("cg-inline") + self.assertNotIn("stats", frame["df_meta"]) + with _slow_stat_queries(0.03): + reply = await self._ask(ws, self._session("cg-inline").stats_gen, tier="full", force=True) + self.assertEqual((reply["type"], reply["final"], reply["remaining"]), ("stats_update", True, 0)) + self.assertTrue(_FULL_ONLY_WIRE <= set(_rows_by_stat(reply["payload"]))) + self.assertFalse(self._session("cg-inline").cost_paused) + + +class TestCostGuardUnits: + """The guard on the scalar tier: a chunk of the batch is a unit, and the + first one over the budget pauses the run's session.""" + + def test_a_scalar_chunk_over_the_budget_pauses_the_session_and_force_resumes_it(self, tmp_path, monkeypatch): + monkeypatch.setattr(stats_wire, "STATS_BUDGET_S", 0) + monkeypatch.setattr(stats_wire, "STATS_COST_BUDGET_S", 0.02, raising=False) + # 24 cells at 12 rows is two columns per chunk: three batch units. + session, client = TestScalarTierRequests._session(tmp_path, stat_chunk_cells=24), _ondemand_client() + request = TestScalarTierRequests._request + with _slow_stat_queries(0.05): + opening = request(session, client, incremental=True) + assert (opening["type"], opening["tier"], opening["final"], opening["remaining"]) == ( + "stats_update", "scalar", False, 2), opening + assert session.cost_paused is True + assert (session.stats_status, session.stats_reason) == ("not_computed", "cost") + + with _count_stat_queries() as queries: + refused = request(session, client, incremental=True) + assert queries == [] + assert (refused["type"], refused["status"], refused["reason"], refused["tier"]) == ( + "stats_update", "not_computed", "cost", "scalar") + + forced = request(session, client, incremental=True, tier="scalar", force=True) + assert (forced["type"], forced["tier"], forced["final"], forced["remaining"]) == ("stats_update", "scalar", + False, 1) + assert session.cost_paused is False + assert (session.stats_status, session.stats_reason) == ("not_computed", "size") + last = request(session, client, incremental=True) + assert (last["final"], last["remaining"]) == (True, 0) + assert (last["status"], last["reason"]) == ("not_computed", "size") + + +# A diff-styling klass as the live diff route has one: the numeric colour rule +# is not in the overrides (the route strips it), and the klass adds it at style +# time with the short alias of the delta column as its val_column. +_DELTA_STYLING = ''' +class DeltaStyling(DefaultMainStyling): + df_display_name = "main" + + @classmethod + def style_column(cls, col, column_metadata): + base_config = super().style_column(col, column_metadata) + if column_metadata.get("orig_col_name") == "price": + base_config["color_map_config"] = {"color_rule": "color_map", "map_name": "DIVERGING_RED_WHITE_BLUE", + "val_column": str(col)} + return base_config +''' + +# What remains in the overrides once the live route has stripped the numeric rules. +_CATEGORICAL_OVERRIDE = {"category": {"color_map_config": {"color_rule": "color_categorical", + "map_name": "BLUE_TO_YELLOW"}}} + +# A promoted entry keeps its numeric rule in the overrides, with the delta column +# named as the original column name the host knows. +_NUMERIC_OVERRIDE = {"qty": {"color_map_config": {"color_rule": "color_map", "map_name": "BLUE_TO_YELLOW", + "val_column": "price"}}} + + +class TestDemandScanWire(_LimitsWire): + """The demand scan and the scoped requests it becomes (rows-first p36a): + the columns a built config's ``color_map`` rules need, found with no data + query at load and computed on request, for those columns only.""" + + DISPLAY_FILES = {"delta_styling.py": _DELTA_STYLING} + + @tornado.testing.gen_test + async def test_the_scan_finds_the_delta_column_a_klass_adds_at_style_time(self): + with _count_stat_queries() as queries: + await self._load("dm-klass", stats_tier="schema", column_config_overrides=_CATEGORICAL_OVERRIDE) + self.assertEqual(queries, [], "the scan reads the built config and sends no data query") + session = self._session("dm-klass") + self.assertEqual(session.stats_policy["demand_columns"], ["a"]) + ws, ondemand = await self._connect("dm-klass") + self.assertEqual(self._stats(ondemand)["demand_columns"], ["a"]) + self.assertEqual(self._stats(ondemand)["auto_request"], False) + _, update_only = await self._connect("dm-klass", caps="stats_update") + self.assertNotIn("demand_columns", self._stats(update_only)) + + @tornado.testing.gen_test + async def test_a_numeric_rule_in_the_overrides_is_scanned_too(self): + await self._load("dm-override", stats_tier="schema", project_root=None, + column_config_overrides=_NUMERIC_OVERRIDE) + _, frame = await self._connect("dm-override") + self.assertEqual(self._stats(frame)["demand_columns"], ["a"], + "the host's name for the column is the stats' key") + + @tornado.testing.gen_test + async def test_the_demand_follows_a_state_change_and_a_reload(self): + await self._load("dm-follow", stats_tier="schema", column_config_overrides=_CATEGORICAL_OVERRIDE) + ws, first = await self._connect("dm-follow") + changed = self._stats(await self._state(ws, post_processing="first_three")) + self.assertEqual(changed["demand_columns"], ["a"]) + await self._reload("dm-follow") + self.assertEqual(self._stats(await _read_json(ws))["demand_columns"], ["a"]) + + @tornado.testing.gen_test + async def test_a_demand_request_runs_the_scalar_units_of_the_named_columns_only(self): + await self._load("dm-run", stats_tier="schema", column_config_overrides=_CATEGORICAL_OVERRIDE) + session = self._session("dm-run") + ws, first = await self._connect("dm-run") + stats = self._stats(first) + held = (session.df_data_dict, session.df_display_args, session.df_meta) + + with _count_stat_queries() as queries: + reply = await self._ask(ws, stats["gen"], tier="scalar", columns=stats["demand_columns"]) + + self.assertEqual((reply["type"], reply["tier"], reply["final"], reply["remaining"]), ("stats_update", "scalar", + True, 0), reply) + self.assertEqual((reply["status"], reply["reason"]), ("not_computed", "host")) + rows = _rows_by_stat(reply["payload"]) + self.assertEqual(set(rows), _SCALAR_WIRE_ONLY) + self.assertEqual({c for row in rows.values() for c in row} - {"index", "level_0"}, {"a"}) + self.assertEqual(len(rows["histogram_bins"]["a"]), 11, "what the colour map reads") + self.assertEqual([type(q.op()).__name__ for q in queries], ["Aggregate"]) + self.assertEqual({name.split("|")[0] for name in queries[0].op().metrics} - {"__total_length__"}, {"price"}) + self.assertEqual(list(session.stat_runs), [(stats["gen"], "raw", "scalar", ("price",))]) + self.assertTrue(all(a is b for a, b in zip(held, (session.df_data_dict, session.df_display_args, + session.df_meta))), "nothing is assigned") + self.assertEqual((session.stats_status, session.stats_reason), ("not_computed", "host")) + + # Another client asking for the same columns reads the run's fragments. + other, other_first = await self._connect("dm-run") + with _count_stat_queries() as queries: + replay = await self._ask(other, self._stats(other_first)["gen"], tier="scalar", columns=["a"]) + self.assertEqual(queries, []) + self.assertEqual(_as_json(_rows_by_stat(replay["payload"])), _as_json(rows)) + + @tornado.testing.gen_test + async def test_a_named_columns_request_is_not_forced_and_needs_no_requestable_tier(self): + """A demand request is automatic. The whole table's tier above a schema target needs a force.""" + with _stats_limits(**{**_SCHEMA_BY_SIZE, "ceiling_full_rows": 3}): + await self._load("dm-auto", stats_tier="auto", column_config_overrides=_CATEGORICAL_OVERRIDE) + ws, first = await self._connect("dm-auto") + gen = self._stats(first)["gen"] + scoped = await self._ask(ws, gen, tier="scalar", columns=["b"]) + self.assertEqual((scoped["type"], scoped["tier"], scoped["final"]), ("stats_update", "scalar", True), + scoped) + whole = await self._ask(ws, gen, tier="scalar") + self.assertEqual((whole["type"], whole["reason"]), ("stats_aborted", "not_requestable")) + + @tornado.testing.gen_test + async def test_the_ceiling_judges_a_named_columns_request_by_its_own_cells(self): + """The table is 15 cells and the ceiling on scalar is 10: two columns fit, three do not.""" + with _stats_limits(ceiling_scalar_cells=10): + await self._load("dm-cells", stats_tier="schema") + ws, first = await self._connect("dm-cells") + gen = self._stats(first)["gen"] + self.assertEqual(self._stats(first)["requestable"], []) + fits = await self._ask(ws, gen, tier="scalar", columns=["a", "b"]) + self.assertEqual((fits["type"], fits["final"]), ("stats_update", True), fits) + self.assertEqual({c for row in _rows_by_stat(fits["payload"]).values() for c in row} - {"index", "level_0"}, + {"a", "b"}) + over = await self._ask(ws, gen, tier="scalar", columns=["a", "b", "c"]) + self.assertEqual((over["type"], over["reason"]), ("stats_update", "ceiling")) + + @tornado.testing.gen_test + async def test_a_named_columns_request_at_the_full_tier_runs_those_columns_units_and_assigns_nothing(self): + self._one_unit_per_request() + await self._load("dm-full", stats_tier="schema") + session = self._session("dm-full") + ws, first = await self._connect("dm-full") + gen = self._stats(first)["gen"] + replies = await self._ask_to_the_end(ws, gen, tier="full", columns=["b"]) + self.assertEqual([(r["type"], r["tier"], r["final"], r["remaining"]) for r in replies], + [("stats_update", "full", False, 1), ("stats_update", "full", True, 0)]) + self.assertEqual((replies[-1]["status"], replies[-1]["reason"]), ("not_computed", "host")) + merged = _merge_stat_payloads(r["payload"] for r in replies) + self.assertTrue(_FULL_ONLY_WIRE <= set(merged)) + self.assertEqual({c for row in merged.values() for c in row}, {"b"}) + self.assertEqual(list(session.stat_runs), [(gen, "raw", "full", ("qty",))]) + self.assertEqual((session.stats_status, session.xorq_dataflow.stats_tier), ("not_computed", "schema")) + + @tornado.testing.gen_test + async def test_two_requests_for_the_same_columns_in_any_order_share_a_run(self): + await self._load("dm-share", stats_tier="schema") + session = self._session("dm-share") + ws, first = await self._connect("dm-share") + gen = self._stats(first)["gen"] + await self._ask(ws, gen, tier="scalar", columns=["b", "a", "a"]) + with _count_stat_queries() as queries: + await self._ask(ws, gen, tier="scalar", columns=["a", "b"]) + self.assertEqual(queries, []) + self.assertEqual(list(session.stat_runs), [(gen, "raw", "scalar", ("price", "qty"))]) + +class TestTierFieldsAndOlderClients(_LimitsWire): + """``tier``, ``force`` and a ``columns`` that scopes the run belong to a + client that advertised ``stats_ondemand``; every other client is served as + before.""" + + @tornado.testing.gen_test + async def test_force_from_a_client_without_both_bits_records_no_override(self): + with _stats_limits(**_SCHEMA_BY_SIZE): + await self._load("sk-no-override", stats_tier="auto") + ws, first = await self._connect("sk-no-override", caps="stats_update") + ws.write_message(_stats_request(self._stats(first)["gen"], tier="full", force=True)) + reply = await _read_json(ws) + self.assertEqual((reply["type"], reply["final"]), ("stats_update", True)) + session = self._session("sk-no-override") + self.assertEqual((session.stats_override, session.cost_paused), (None, False)) + class TestReloadExpr(tornado.testing.AsyncHTTPTestCase): def get_app(self): return make_app() diff --git a/tests/unit/server/test_stats_policy.py b/tests/unit/server/test_stats_policy.py index f3427a483..bc0367e6b 100644 --- a/tests/unit/server/test_stats_policy.py +++ b/tests/unit/server/test_stats_policy.py @@ -2,9 +2,11 @@ Everything here is pure Python: the policy is a function of numbers, the probes read a schema or a parquet footer, and nothing needs a server. The last -sections (rows-first p33) test how the session carries the result and how -``df_meta.stats`` reports it, still without a server; the HTTP and WebSocket -side is in test_load_expr.py. +sections (rows-first p33 and p36a) test how the session carries the result and how +``df_meta.stats`` reports it, and how a ``stats_request`` is judged against it (the +ceiling, ``requestable``, ``force``, the cost pause, the policy a forced tier +raises it to, and the demand scan), still without a server; the HTTP and +WebSocket side is in test_load_expr.py. The thresholds are the provisional values proposed by the phase-0 measurements (p31b). The boundary tables run on the module defaults, so a @@ -27,6 +29,8 @@ import pyarrow.parquet as pq import pytest +from buckaroo.jlisp.lisp_utils import s as lisp_sym +from buckaroo.jlisp.lisp_utils import sQ from buckaroo.server import session as session_mod from buckaroo.server import stats_wire @@ -895,3 +899,375 @@ def test_a_session_with_no_policy_has_no_target_to_serve(self, sp): session = _session(sp, "scalar", "inline", resolve=False) assert (session.stats_status, session.stats_policy) == ("not_computed", None) assert stats_wire.serves_scalar_tier(session, _client(ONDEMAND_CAPS)) is False + + +# --------------------------------------------------------------------------- +# Limits, override, cost guard and the demand scan (rows-first p36a) +# --------------------------------------------------------------------------- + +# The frame a request resolves a client's column names against: price, qty and +# category are the columns a client calls a, b and c. +FRAME = pd.DataFrame({"price": [12.5, 18.9, 7.4, 22.1, 1e9], "qty": [1, 2, 1, 3, 2], "category": list("abaca")}) +PAIRS = [("price", "a"), ("qty", "b"), ("category", "c")] +SEARCH_OP = [sQ("search"), lisp_sym("df"), "col", "a"] +CLEANING_OP = [lisp_sym("dropcol"), lisp_sym("df"), "col"] + + +def _with_frame(session, operations=()): + """Give ``session`` the dataflow a request is answered against, as far as a + refusal or a state test needs it: a frame to read column names from and the + operations that say whether it is filtered. No unit can run on it.""" + session.backend = "xorq" + session.xorq_dataflow = SimpleNamespace(operations=list(operations), processed_df=FRAME) + return session + + +def _limited(sp, monkeypatch, stats_tier="auto", limits=None, operations=()): + """A deferred session that resolved ``limits``, with the same thresholds in + the environment, where a request reads them again for its own check.""" + for name, value in (limits or {}).items(): + monkeypatch.setenv(f"BUCKAROO_STATS_{name.upper()}", str(value)) + return _with_frame(_session(sp, stats_tier, "deferred", limits), operations) + + +def _request(session, client, **fields): + return stats_wire.handle_stats_request( + session, {"type": "stats_request", "stats_gen": session.stats_gen, "scope": "raw", **fields}, client) + + +def _refusal(session, tier, reason): + return {"stats_gen": session.stats_gen, "scope": "raw", "tier": tier, "final": True, "remaining": 0, + "status": "not_computed", "reason": reason} + + +def _refusal_fields(reply): + return {key: reply[key] for key in ("stats_gen", "scope", "tier", "final", "remaining", "status", "reason")} + + +# (host tier, limits, the request fields, the tier the reply says was asked for). +# The five-by-three table is 15 cells: a ceiling_full_rows of 3 refuses full for +# it, and a ceiling_scalar_cells of 10 refuses scalar for the whole table. +OVER_THE_CEILING = [ + ("auto", {**SCHEMA_BY_SIZE, "ceiling_full_rows": 3}, {"tier": "full", "force": True}, "full"), + ("auto", {**SCHEMA_BY_SIZE, "ceiling_full_rows": 3}, {"tier": "full"}, "full"), + ("full", FULL_OVER_CEILING, {"tier": "full", "force": True}, "full"), + ("schema", SCALAR_OVER_CEILING, {"tier": "scalar", "force": True}, "scalar"), + ("schema", SCALAR_OVER_CEILING, {"tier": "full", "force": True}, "full"), + # Named columns are checked at their own cell count: all three are 15 cells. + ("scalar", SCALAR_OVER_CEILING, {"tier": "scalar", "columns": ["a", "b", "c"]}, "scalar")] + + +class TestRequestsOverTheCeiling: + """A request for a tier the ceiling refuses is answered ``stats_update + {final: true, status: "not_computed", reason: "ceiling"}`` and runs nothing, + whoever asks and whether or not it is forced.""" + + @pytest.mark.parametrize("stats_tier, limits, fields, tier", OVER_THE_CEILING) + def test_a_request_over_the_ceiling_is_refused_and_runs_nothing(self, sp, monkeypatch, stats_tier, limits, fields, + tier): + session = _limited(sp, monkeypatch, stats_tier, limits) + held = (session.stats_status, session.stats_reason, dict(session.stat_runs)) + reply = _request(session, _client(ONDEMAND_CAPS), incremental=True, **fields) + assert reply["type"] == "stats_update", reply + assert _refusal_fields(reply) == _refusal(session, tier, "ceiling") + assert "payload" not in reply and "df_display_args" not in reply + assert isinstance(reply["elapsed_ms"], (int, float)) and reply["elapsed_ms"] >= 0 + assert (session.stats_status, session.stats_reason, session.stat_runs) == held + assert session.stats_override is None + + def test_the_refusal_does_not_depend_on_the_request_being_incremental(self, sp, monkeypatch): + session = _limited(sp, monkeypatch, "full", FULL_OVER_CEILING) + reply = _request(session, _client(ONDEMAND_CAPS), tier="full", force=True) + assert _refusal_fields(reply) == _refusal(session, "full", "ceiling") + + +class TestMalformedTierRequests: + @pytest.mark.parametrize("fields", [ + {"tier": "schema"}, {"tier": "medium"}, {"tier": 3}, {"tier": ["full"]}, {"force": "yes"}, {"force": 1}, + {"tier": "scalar", "columns": "a"}, {"tier": "scalar", "columns": ["zz"]}, + {"tier": "scalar", "columns": [1, None]}, {"tier": "scalar", "columns": []}, + {"tier": "scalar", "columns": {"a": 1}}]) + def test_a_field_the_server_cannot_read_is_a_bad_request(self, sp, monkeypatch, fields): + session = _limited(sp, monkeypatch, "auto", SCHEMA_BY_SIZE) + reply = _request(session, _client(ONDEMAND_CAPS), **fields) + assert (reply["type"], reply["reason"], reply["current_gen"]) == ("stats_aborted", "bad_request", + session.stats_gen) + assert session.stat_runs == {} and session.stats_status == "not_computed" + + +class TestRequestableTiers: + """``requestable`` and the target are enforced by the server: a tier outside + them is not run, whatever the client believes.""" + + @pytest.mark.parametrize("stats_tier, limits, fields", [ + ("auto", {}, {"tier": "scalar"}), + ("auto", {}, {"tier": "scalar", "force": True}), + ("auto", SCALAR_BY_SIZE, {"tier": "full"}), + ("auto", SCHEMA_BY_SIZE, {"tier": "scalar"}), + ("auto", SCHEMA_BY_SIZE, {"tier": "full"}), + ("auto", SCHEMA_BY_SIZE, {}), + ("auto", SCHEMA_BY_SIZE, {"force": False})]) + def test_a_tier_that_is_not_the_target_or_above_it_with_force_is_not_requestable(self, sp, monkeypatch, stats_tier, + limits, fields): + """Below the target, or above it without ``force``, or a schema target that + nothing is requested for on its own.""" + session = _limited(sp, monkeypatch, stats_tier, limits) + reply = _request(session, _client(ONDEMAND_CAPS), incremental=True, **fields) + assert (reply["type"], reply["reason"]) == ("stats_aborted", "not_requestable"), reply + assert session.stat_runs == {} and session.stats_override is None + + +class TestCostPausedStatus: + @pytest.mark.parametrize("stats_tier, delivery, policy, expected", [ + ("auto", "deferred", {"tier_target": "full", "reason": None}, ("not_computed", "cost")), + ("auto", "deferred", {"tier_target": "scalar", "reason": "size"}, ("not_computed", "cost")), + ("auto", "deferred", {"tier_target": "schema", "reason": "size"}, ("not_computed", "cost")), + ("schema", "deferred", {"tier_target": "schema", "reason": "host"}, ("not_computed", "cost")), + # No policy, no pause: a dataflow built with its stats inline has no run to hold. + ("full", "inline", None, ("complete", None)), + ("schema", "deferred", None, ("not_computed", "host"))]) + def test_a_paused_session_starts_its_generation_not_computed_for_cost(self, stats_tier, delivery, policy, expected): + assert session_mod.initial_stats_status(stats_tier, delivery, policy, cost_paused=True) == expected + + def test_an_unpaused_session_is_as_before(self): + assert session_mod.initial_stats_status("auto", "deferred", {"tier_target": "full", "reason": None}, + cost_paused=False) == ("pending", None) + assert session_mod.initial_stats_status("auto", "deferred", {"tier_target": "scalar", "reason": "size"}) == ("not_computed", "size") + + def test_a_policy_target_of_full_is_pending_whatever_the_delivery(self): + """A policy exists only for a dataflow built at the schema tier, so a full + target has its stats still to run, as a forced request makes it for a + session the host named below full.""" + assert session_mod.initial_stats_status("schema", "inline", {"tier_target": "full", "reason": None}) == ("pending", None) + + @pytest.mark.parametrize("stats_tier, limits, expected", [ + ("auto", {}, {"status": "not_computed", "tier": "schema", "gen": 2, "reason": "cost", "tier_target": "full", + "requestable": [], "estimate": ESTIMATE}), + ("auto", SCALAR_BY_SIZE, {"status": "not_computed", "tier": "schema", "gen": 2, "reason": "cost", + "tier_target": "scalar", "estimate": ESTIMATE})]) + def test_the_pause_survives_a_new_generation_and_is_reported(self, sp, stats_tier, limits, expected): + session = _session(sp, stats_tier, "deferred", limits) + session.cost_paused = True + session_mod.begin_stats_generation(session) + assert session_mod.stats_meta(session) == expected + assert (session.stats_status, session.stats_reason) == ("not_computed", "cost") + + def test_clearing_the_pause_restores_the_status_the_policy_gives(self, sp): + session = _session(sp, "auto", "deferred", SCALAR_BY_SIZE) + session.cost_paused = True + session_mod.begin_stats_generation(session) + assert (session.stats_status, session.stats_reason) == ("not_computed", "cost") + session.cost_paused = False + session_mod.begin_stats_generation(session) + assert (session.stats_status, session.stats_reason) == ("not_computed", "size") + + @pytest.mark.parametrize("stats_tier, limits, tier", [("auto", {}, "full"), ("auto", SCALAR_BY_SIZE, "scalar"), + ("scalar", {}, "scalar")]) + def test_a_paused_session_refuses_a_request_that_is_not_forced(self, sp, monkeypatch, stats_tier, limits, tier): + session = _limited(sp, monkeypatch, stats_tier, limits) + session.cost_paused = True + session_mod.begin_stats_generation(session) + reply = _request(session, _client(ONDEMAND_CAPS), incremental=True) + assert reply["type"] == "stats_update", reply + assert _refusal_fields(reply) == _refusal(session, tier, "cost") + assert "payload" not in reply and isinstance(reply["elapsed_ms"], (int, float)) + assert session.stat_runs == {} + + def test_a_paused_session_refuses_a_named_columns_request_too(self, sp, monkeypatch): + """A demand request is automatic, so it is held by the pause like any other.""" + session = _limited(sp, monkeypatch, "auto", SCHEMA_BY_SIZE) + session.cost_paused = True + session_mod.begin_stats_generation(session) + reply = _request(session, _client(ONDEMAND_CAPS), tier="scalar", columns=["a"]) + assert _refusal_fields(reply) == _refusal(session, "scalar", "cost") + assert session.stat_runs == {} + + def test_a_paused_schema_target_still_has_nothing_to_request(self, sp, monkeypatch): + """Not requestable comes before the pause: a Continue would ask for nothing, and a force with no tier + must not clear a pause that holds back the demand requests.""" + session = _limited(sp, monkeypatch, "auto", SCHEMA_BY_SIZE) + session.cost_paused = True + session_mod.begin_stats_generation(session) + for fields in ({}, {"force": True}): + reply = _request(session, _client(ONDEMAND_CAPS), **fields) + assert (reply["type"], reply["reason"]) == ("stats_aborted", "not_requestable"), fields + assert session.cost_paused is True + assert (session.stats_status, session.stats_reason) == ("not_computed", "cost") + + def test_the_ceiling_is_answered_before_the_pause(self, sp, monkeypatch): + """Continuing would not help a request the ceiling refuses.""" + session = _limited(sp, monkeypatch, "full", FULL_OVER_CEILING) + session.cost_paused = True + session_mod.begin_stats_generation(session) + reply = _request(session, _client(ONDEMAND_CAPS), tier="full", force=True) + assert _refusal_fields(reply)["reason"] == "ceiling" + reply = _request(session, _client(ONDEMAND_CAPS), tier="full") + assert _refusal_fields(reply)["reason"] == "ceiling" + + +class TestEffectiveStatsPolicy: + """The policy in force for the session's state: the one resolved at load, + raised to the tier a forced request asked for while the state is unfiltered + (or the force was made in this generation).""" + + def test_with_no_override_it_is_the_policy_resolved_at_load(self, sp): + session = _session(sp, "auto", "deferred", SCHEMA_BY_SIZE) + assert session_mod.effective_stats_policy(session) is session.stats_policy + + def test_a_session_with_no_policy_has_none(self, sp): + assert session_mod.effective_stats_policy(_session(sp, "full", "inline", resolve=False)) is None + + @pytest.mark.parametrize("override, expected", [ + ("full", {"tier_target": "full", "auto_request": True, "requestable": [], "reason": None}), + ("scalar", {"tier_target": "scalar", "auto_request": True, "requestable": ["full"], "reason": "size"})]) + def test_a_forced_tier_raises_the_target_and_works_out_the_rest(self, sp, override, expected): + session = _with_frame(_session(sp, "auto", "deferred", SCHEMA_BY_SIZE)) + session.stats_override, session.stats_override_gen = override, session.stats_gen + policy = session_mod.effective_stats_policy(session) + assert {key: policy[key] for key in expected} == expected + assert policy["estimate"] == ESTIMATE + assert (session.stats_policy["tier_target"], session.stats_policy["auto_request"]) == ("schema", False), ( + "the policy resolved at load is not changed") + + @pytest.mark.parametrize("override", ["scalar", "full"]) + def test_an_override_never_lowers_the_target(self, sp, override): + session = _with_frame(_session(sp, "auto", "deferred")) + assert session.stats_policy["tier_target"] == "full" + session.stats_override, session.stats_override_gen = override, session.stats_gen + assert session_mod.effective_stats_policy(session)["tier_target"] == "full" + + def test_the_ceiling_still_holds_for_the_override(self, sp, monkeypatch): + session = _limited(sp, monkeypatch, "auto", {**SCHEMA_BY_SIZE, "ceiling_full_rows": 3}) + session.stats_override, session.stats_override_gen = "full", session.stats_gen + policy = session_mod.effective_stats_policy(session) + assert (policy["tier_target"], policy["requestable"]) == ("scalar", []) + + def test_the_override_does_not_reach_a_later_filtered_generation(self, sp): + session = _with_frame(_session(sp, "auto", "deferred", SCHEMA_BY_SIZE)) + session.stats_override, session.stats_override_gen = "full", session.stats_gen + assert session_mod.effective_stats_policy(session)["tier_target"] == "full" + _with_frame(session, [SEARCH_OP]) + session_mod.begin_stats_generation(session) + assert session_mod.effective_stats_policy(session) is session.stats_policy + assert (session.stats_status, session.stats_reason) == ("not_computed", "size") + + def test_the_override_reaches_a_later_unfiltered_generation(self, sp): + session = _with_frame(_session(sp, "auto", "deferred", SCHEMA_BY_SIZE)) + session.stats_override, session.stats_override_gen = "full", session.stats_gen + _with_frame(session, [CLEANING_OP]) + session_mod.begin_stats_generation(session) + assert session_mod.effective_stats_policy(session)["tier_target"] == "full" + assert (session.stats_status, session.stats_reason) == ("pending", None) + + def test_a_force_made_in_a_filtered_generation_applies_to_that_generation(self, sp): + session = _with_frame(_session(sp, "auto", "deferred", SCHEMA_BY_SIZE), [SEARCH_OP]) + session.stats_override, session.stats_override_gen = "full", session.stats_gen + assert session_mod.effective_stats_policy(session)["tier_target"] == "full" + session_mod.begin_stats_generation(session) + assert session_mod.effective_stats_policy(session)["tier_target"] == "schema" + # Back to an unfiltered state, the override applies again. + _with_frame(session, []) + session_mod.begin_stats_generation(session) + assert session_mod.effective_stats_policy(session)["tier_target"] == "full" + + def test_a_host_named_schema_session_headed_for_full_is_pending_even_with_inline_delivery(self, sp): + session = _with_frame(_session(sp, "schema", "inline")) + assert (session.stats_status, session.stats_reason) == ("not_computed", "host") + session.stats_override, session.stats_override_gen = "full", session.stats_gen + session_mod.begin_stats_generation(session) + assert (session.stats_status, session.stats_reason) == ("pending", None) + + def test_df_meta_stats_reports_the_effective_policy(self, sp): + session = _with_frame(_session(sp, "auto", "deferred", SCHEMA_BY_SIZE)) + session.stats_override, session.stats_override_gen = "full", session.stats_gen + session_mod.begin_stats_generation(session) + assert session_mod.stats_meta(session) == {"status": "pending", "tier": "schema", "gen": 2, + "tier_target": "full", "requestable": [], "estimate": ESTIMATE} + scalar = _with_frame(_session(sp, "auto", "deferred", SCHEMA_BY_SIZE)) + scalar.stats_override, scalar.stats_override_gen = "scalar", scalar.stats_gen + assert session_mod.stats_meta(scalar) == {"status": "not_computed", "tier": "schema", "gen": 1, + "reason": "size", "tier_target": "scalar", "estimate": ESTIMATE} + + def test_a_scalar_override_makes_scalar_the_tier_a_request_with_no_tier_is_served(self, sp): + session = _with_frame(_session(sp, "auto", "deferred", SCHEMA_BY_SIZE)) + assert stats_wire.serves_scalar_tier(session, _client(ONDEMAND_CAPS)) is False + session.stats_override, session.stats_override_gen = "scalar", session.stats_gen + assert stats_wire.serves_scalar_tier(session, _client(ONDEMAND_CAPS)) is True + + def test_resetting_the_controls_clears_the_override_and_the_pause(self, sp): + session = _with_frame(_session(sp, "auto", "deferred", SCHEMA_BY_SIZE)) + session.stats_override, session.stats_override_gen, session.cost_paused = "full", 4, True + session_mod.reset_stats_controls(session) + assert (session.stats_override, session.stats_override_gen, session.cost_paused) == (None, 0, False) + + +class TestDemandColumns: + """``demand_columns``: the columns a built display config needs min and max + for now, read from the config alone (no data query).""" + + @staticmethod + def _display(*column_configs, summary=()): + return {"main": {"df_viewer_config": {"column_config": list(column_configs)}}, + "summary": {"df_viewer_config": {"column_config": list(summary)}}} + + @staticmethod + def _rule(val_column="a", color_rule="color_map", **extra): + rule = {"color_rule": color_rule, "map_name": "DIVERGING_RED_WHITE_BLUE", **extra} + if val_column is not None: + rule["val_column"] = val_column + return {"col_name": "x", "color_map_config": rule} + + def test_a_color_map_rule_names_the_column_it_reads(self, sp): + assert sp.demand_columns(self._display(self._rule("a")), PAIRS) == ["a"] + + def test_a_val_column_that_is_an_original_name_is_mapped_to_the_key_the_stats_use(self, sp): + """A klass writes the short alias or a name it read from the sd; the stats are keyed by the rewritten name.""" + assert sp.demand_columns(self._display(self._rule("price"), self._rule("category")), PAIRS) == ["a", "c"] + + def test_a_rewritten_name_wins_over_an_original_name_that_looks_like_one(self, sp): + pairs = [("b", "a"), ("a", "b")] + assert sp.demand_columns(self._display(self._rule("a")), pairs) == ["a"] + + @pytest.mark.parametrize("rule", [ + {"color_rule": "color_categorical", "map_name": "BLUE_TO_YELLOW", "val_column": "a"}, + {"color_rule": "color_categorical", "map_name": "BLUE_TO_YELLOW"}, + {"color_rule": "color_from_column", "val_column": "a"}, + {"color_rule": "color_not_null", "conditional_color": "red", "exist_column": "a"}, + {"color_rule": "color_static", "color": "red"}, + # The client reads nothing for a color_map with no val_column. + {"color_rule": "color_map", "map_name": "BLUE_TO_YELLOW"}]) + def test_only_a_numeric_color_map_rule_creates_demand(self, sp, rule): + assert sp.demand_columns(self._display({"col_name": "x", "color_map_config": rule}), PAIRS) == [] + + def test_a_tooltip_val_column_creates_no_demand(self, sp): + config = {"col_name": "x", "tooltip_config": {"tooltip_type": "simple", "val_column": "a"}} + assert sp.demand_columns(self._display(config), PAIRS) == [] + + def test_a_val_column_that_is_not_a_column_of_the_table_is_dropped(self, sp): + assert sp.demand_columns(self._display(self._rule("zz"), self._rule("a")), PAIRS) == ["a"] + + def test_the_columns_come_once_each_in_the_frame_s_order(self, sp): + display = self._display(self._rule("c"), self._rule("a"), self._rule("c"), self._rule("price")) + assert sp.demand_columns(display, PAIRS) == ["a", "c"] + + def test_every_display_is_scanned(self, sp): + display = self._display(self._rule("a"), summary=[self._rule("b")]) + assert sp.demand_columns(display, PAIRS) == ["a", "b"] + + @pytest.mark.parametrize("display", [None, {}, {"main": None}, {"main": {}}, {"main": {"df_viewer_config": None}}, + {"main": {"df_viewer_config": {"column_config": None}}}, + {"main": {"df_viewer_config": {"column_config": ["x", None, {"color_map_config": None}, + {"color_map_config": "a"}, {"color_map_config": {"color_rule": "color_map", "val_column": 3}}, + {"color_map_config": {"color_rule": "color_map", "val_column": ["a"]}}, + {"color_map_config": {"color_rule": "color_map", "val_column": {"a": 1}}}]}}}]) + def test_a_config_that_is_not_the_usual_shape_yields_nothing_and_does_not_raise(self, sp, display): + assert sp.demand_columns(display, PAIRS) == [] + + def test_the_config_is_not_changed(self, sp): + display = self._display(self._rule("price")) + before = json.dumps(display, sort_keys=True) + sp.demand_columns(display, PAIRS) + assert json.dumps(display, sort_keys=True) == before + + def test_no_columns_no_demand(self, sp): + assert sp.demand_columns(self._display(self._rule("a")), []) == [] From 9af4042c653ef7951e1b854d6d41f26a7368c69a Mon Sep 17 00:00:00 2001 From: Paddy Mullen Date: Sun, 4 Oct 2026 09:42:52 -0400 Subject: [PATCH 32/39] feat(server): stats request limits, session override, cost guard and the demand scan (rows-first p36a) A stats_request from a client that advertised both capability bits reads tier, force and, with a tier, columns. The server judges it: the ceiling for the cells asked for (resolve_stats_policy with the tier as the host's) answers a refusal and runs nothing, a whole-table tier must be the target or above it with force, and a paused session refuses a request that is not forced. A force above the target stores stats_override on the session, and effective_stats_policy raises the target for the unfiltered scope and for the generation it was forced in. begin_stats_generation, stats_meta and serves_scalar_tier read the effective policy. An automatic request that ran a unit over STATS_COST_BUDGET_S and left units to run sets cost_paused, which begin_stats_generation reads, so the pause survives a stats_gen bump and /reload_expr until a force request. The override and the pause are forgotten by /load, /load_compare, another expression and a body that names a tier. stats_policy.demand_columns scans the built display config for color_map rules and names the columns they read; the snapshot refresh writes them to stats_policy["demand_columns"] for a schema target, and a scoped request computes those columns only, with a run that is never assigned. The final reply of a run that does not assign now carries the session's status. Guard tests that pass on the earlier code go in with this commit, and the elapsed_ms test has wider margins on both sides. Co-Authored-By: Claude Sonnet 5.5 --- buckaroo/server/handlers.py | 25 ++- buckaroo/server/session.py | 123 +++++++++++-- buckaroo/server/stats_policy.py | 39 +++- buckaroo/server/stats_wire.py | 264 ++++++++++++++++++++++++---- tests/unit/server/test_load_expr.py | 108 +++++++++++- 5 files changed, 497 insertions(+), 62 deletions(-) diff --git a/buckaroo/server/handlers.py b/buckaroo/server/handlers.py index a3c292d1d..4789ccdad 100644 --- a/buckaroo/server/handlers.py +++ b/buckaroo/server/handlers.py @@ -15,8 +15,9 @@ from buckaroo.df_util import old_col_new_col from buckaroo.server.focus import find_or_create_session_window from buckaroo.server.session import ( - STATS_DELIVERIES, STATS_TIER_REQUESTS, begin_stats_generation, dataflow_stats_tier) -from buckaroo.server.stats_wire import broadcast_state, refresh_session_snapshot, resolve_session_policy + STATS_DELIVERIES, STATS_TIER_REQUESTS, begin_stats_generation, dataflow_stats_tier, reset_stats_controls) +from buckaroo.server.stats_wire import ( + broadcast_state, refresh_demand_columns, refresh_session_snapshot, resolve_session_policy) from buckaroo.server import telemetry from buckaroo.pluggable_analysis_framework import perf_log @@ -371,6 +372,7 @@ async def post(self): # 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, session.stats_policy = "full", "inline", None + reset_stats_controls(session) begin_stats_generation(session) # Notify connected clients and open browser @@ -506,6 +508,11 @@ async def post(self): self.set_status(400) self.write(policy_error) return + # What a client forced or the cost guard paused (SessionState.stats_override, + # cost_paused) belongs to one expression: a rebuild of the same one keeps it, + # and so does a re-POST that names no stats_tier. + same_expression = (existing is not None and existing.backend == "xorq" and existing.build_dir == build_dir + and existing.cache_dir == cache_dir) # /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 @@ -627,6 +634,8 @@ async def post(self): session.stats_tier = stats_tier session.stats_delivery = stats_delivery session.stats_policy = stats_policy + if not same_expression or body.get("stats_tier") is not None: + reset_stats_controls(session) session.tele_sink = tele_sink session.xorq_dataflow = xorq_dataflow # Clear pandas-side state left by a prior /load on the same @@ -668,6 +677,9 @@ async def post(self): **dvc.get("component_config", {}), **component_config} + # The demand scan reads the config just stored; no query is sent. + refresh_demand_columns(session, xorq_dataflow) + # A new expression is a new stats generation; a deferred session starts # it pending. begin_stats_generation(session) @@ -813,6 +825,7 @@ async def post(self): session.mode = "viewer" # A viewer session has no dataflow and so no deferred stats. session.stats_tier, session.stats_delivery, session.stats_policy = "full", "inline", None + reset_stats_controls(session) begin_stats_generation(session) # Push to WebSocket clients. Reset per-client live search (#851). @@ -849,6 +862,9 @@ class ReloadExprHandler(tornado.web.RequestHandler): 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. The stats policy is resolved again for the pair. + A forced tier and a cost pause (``SessionState.stats_override``, + ``cost_paused``) are kept, unless the body names a ``stats_tier``, which + starts the session's stats decisions over. Returns 404 when the session does not exist, 400 when it is not a xorq session, has no project_root recorded or carries an invalid stats policy, @@ -884,8 +900,9 @@ async def post(self, session_id): "message": "Session has no project_root — pass project_root to /load_expr first"}) return + body = self._optional_body() stats_tier, stats_delivery, policy_error = _stats_policy_from_body( - self._optional_body(), session.stats_tier, session.stats_delivery) + body, session.stats_tier, session.stats_delivery) if policy_error is not None: self.set_status(400) self.write(policy_error) @@ -941,6 +958,8 @@ async def post(self, session_id): session.stats_tier = stats_tier session.stats_delivery = stats_delivery session.stats_policy = stats_policy + if body.get("stats_tier") is not None: + reset_stats_controls(session) refresh_session_snapshot(session, xorq_dataflow) begin_stats_generation(session) diff --git a/buckaroo/server/session.py b/buckaroo/server/session.py index ec72c758a..685507253 100644 --- a/buckaroo/server/session.py +++ b/buckaroo/server/session.py @@ -5,6 +5,10 @@ from typing import Any, Callable, Dict, Optional import pandas as pd + +from buckaroo.dataflow.sd_cache import split_chain_by_scope +from buckaroo.server.stats_policy import TIERS, resolve_stats_policy + # polars is optional — only used in lazy mode log = logging.getLogger("buckaroo.server.session") @@ -25,8 +29,8 @@ STATS_TIER_REQUESTS = ("auto", "full", "scalar", "schema") # Why a session is ``not_computed``, in ``df_meta.stats.reason``: the server's size -# rule, the host naming a tier, the cost guard pausing a run (phase 6a), or the -# ceiling lowering a request. (``error`` has its own reason, ``stats_failed``.) +# rule, the host naming a tier, the cost guard pausing a run, or the ceiling +# lowering a request. (``error`` has its own reason, ``stats_failed``.) STATS_REASONS = ("size", "host", "cost", "ceiling") # What a client assumes for a ``df_meta.stats`` field the server leaves out. The @@ -59,15 +63,24 @@ def dataflow_stats_tier(stats_tier: str, stats_delivery: str) -> str: STATS_STATUSES = ("complete", "pending", "not_computed", "error") -def initial_stats_status(stats_tier: str, stats_delivery: str, policy: Optional[dict] = None) -> tuple[str, Optional[str]]: +def initial_stats_status(stats_tier: str, stats_delivery: str, policy: Optional[dict] = None, + cost_paused: bool = False) -> tuple[str, Optional[str]]: """The status, and its reason, of a session whose stats generation has just - started, from the policy pair and the policy resolved at load. A target below - ``full`` is ``not_computed`` for the reason the policy gives; with no policy - (a dataflow built with its stats inline) a tier named below ``full`` is the - host's.""" - if policy is not None and policy["tier_target"] != "full": - return "not_computed", policy["reason"] - if policy is None and stats_tier in ("scalar", "schema"): + started, from the policy pair, the policy in force and the cost guard. A + session with a policy is ``not_computed`` for ``cost`` while its automatic + stats are paused (``cost_paused``), else for the reason the policy gives when + its target is below ``full``, else ``pending``: a policy exists only for a + dataflow built at the schema tier, so a ``full`` target has its stats still + to run. With no policy (a dataflow built with its stats inline) a tier named + below ``full`` is the host's, and the stats are complete unless delivery is + deferred.""" + if policy is not None: + if cost_paused: + return "not_computed", "cost" + if policy["tier_target"] != "full": + return "not_computed", policy["reason"] + return "pending", None + if stats_tier in ("scalar", "schema"): return "not_computed", "host" if stats_delivery == "deferred": return "pending", None @@ -116,8 +129,10 @@ class SessionState: # the entry (``tier_target``, ``auto_request``, ``requestable``, ``reason``, # ``estimate``), resolved by /load_expr and /reload_expr once the schema-tier # dataflow and the count exist. ``None`` for a dataflow built with its stats - # inline, and for anything /load serves. Later phases add ``omitted_keys``, - # ``approx_keys`` and ``demand_columns``, which ``stats_meta`` reports. + # inline, and for anything /load serves. ``demand_columns`` is written to it + # whenever the display config changes (``stats_wire.refresh_demand_columns``); + # later phases add ``omitted_keys`` and ``approx_keys``. ``stats_meta`` reports + # them, for the policy in force (``effective_stats_policy``). stats_policy: Optional[dict] = None # The stats generation: a counter the server owns, bumped whenever the # dataflow state the stats describe changes (/load, /load_expr, /load_compare, @@ -129,6 +144,18 @@ class SessionState: stats_gen: int = 0 stats_status: str = "complete" stats_reason: Optional[str] = None + # What a client's forced ``stats_request`` set on the session (rows-first + # p36a). ``stats_override`` is the tier it forced above the policy's target + # (``scalar`` or ``full``) and ``stats_override_gen`` the generation it was + # made in: it raises the target of a later generation as long as the state + # is not filtered (``effective_stats_policy``). ``cost_paused`` is set when an + # automatic unit ran over the budget, and holds back further automatic + # requests until a forced one. Both survive a generation and /reload_expr, + # which is why they live here and not on a client; ``reset_stats_controls`` + # clears them. + stats_override: Optional[str] = None + stats_override_gen: int = 0 + cost_paused: bool = False # 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 @@ -171,14 +198,73 @@ def touch(self) -> None: this field is the runtime escape hatch.""" +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 _filtered(session: "SessionState") -> bool: + """Whether the session's current state has a quick command (a search) on it, + which makes its stats a filtered scope's: the dataflow's ``filt`` chain is + the ``clean`` chain plus the quick commands (``sd_cache.split_chain_by_scope``).""" + scopes = split_chain_by_scope(getattr(session_dataflow(session), "operations", None) or []) + return len(scopes["filt"]) != len(scopes["clean"]) + + +def effective_stats_policy(session: "SessionState") -> Optional[dict]: + """The policy in force for the session's current state: the one resolved at + load (``session.stats_policy``, itself, when nothing changes it), or, when a + forced request raised the tier (``stats_override``), the policy that tier + resolves to for the same entry. The override never lowers the target and the + ceiling still holds for it. It applies to an unfiltered state, and to the + generation it was forced in (a filtered state's forced stats are that + state's); another filtered generation has the policy of the load. + + The thresholds are read again, as every resolution does. The reason stays the + load's while the target is below ``full``, and the demand columns are dropped: + they are the columns a ``schema`` target still wants, and a raised target + computes them.""" + policy = session.stats_policy + override = session.stats_override + if policy is None or override is None or TIERS.index(override) <= TIERS.index(policy["tier_target"]): + return policy + if session.stats_override_gen != session.stats_gen and _filtered(session): + return policy + estimate = policy["estimate"] + raised = resolve_stats_policy("xorq", "xorq_build", estimate["rows"], estimate["cols"], host_tier=override) + if TIERS.index(raised["tier_target"]) <= TIERS.index(policy["tier_target"]): + return policy + return {**raised, "reason": None if raised["tier_target"] == "full" else policy["reason"], + "estimate": dict(estimate)} + + +def reset_stats_controls(session: "SessionState") -> None: + """Forget what clients set with ``force`` and the cost guard's pause: the + session now holds another dataset, or its host named a tier again.""" + session.stats_override, session.stats_override_gen, session.cost_paused = None, 0, False + + +def restore_stats_status(session: "SessionState") -> None: + """Take a ``pending`` or ``not_computed`` session's status back to what its + policy in force and the cost guard give, without starting a generation: a + forced request has raised the target or cleared the pause, and the stats of + this generation are the same state's.""" + if session.stats_policy is not None and session.stats_status in ("pending", "not_computed"): + session.stats_status, session.stats_reason = initial_stats_status( + session.stats_tier, session.stats_delivery, effective_stats_policy(session), session.cost_paused) + + 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, and the stat runs of the old generation are - dropped.""" + from the session's policy pair, the policy in force and the cost guard, 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.stats_policy) + session.stats_tier, session.stats_delivery, effective_stats_policy(session), session.cost_paused) session.stat_runs.clear() @@ -229,8 +315,11 @@ def stats_meta(session: "SessionState", ondemand: bool = True) -> Optional[dict] stats: dict = {"status": status, "tier": "full" if status == "complete" else "schema", "gen": session.stats_gen} if reason: stats["reason"] = reason - policy = session.stats_policy - reported = policy is not None and not (session.stats_tier == "full" and policy["tier_target"] == "full") + policy = effective_stats_policy(session) + # An explicit ``full`` within the ceiling reports no policy, unless the cost + # guard has paused it: the client then needs the target to continue. + reported = policy is not None and (session.cost_paused or not (session.stats_tier == "full" + and policy["tier_target"] == "full")) if ondemand and reported and status in ("pending", "not_computed"): stats.update(_policy_fields(policy)) return stats diff --git a/buckaroo/server/stats_policy.py b/buckaroo/server/stats_policy.py index 04d171bc2..9567f936c 100644 --- a/buckaroo/server/stats_policy.py +++ b/buckaroo/server/stats_policy.py @@ -1,8 +1,9 @@ """Which stats tier an entry should reach, decided from its size. -Pure functions. Nothing in the handlers, the session or the dataflow imports -this module yet; it exists so the policy can be tested without a backend -before it is wired in. +Pure functions, so the policy can be tested without a backend. The load +handlers resolve it (``stats_wire.resolve_session_policy``), the session applies +a forced tier on top of it (``session.effective_stats_policy``) and a +``stats_request`` is judged against it. Tiers, lowest to highest: @@ -23,6 +24,9 @@ :func:`route_polars_entry` is the separate size rule for polars: an entry above R rows belongs on xorq's ``/load_expr``, not eager ``/load``. +:func:`demand_columns` is the demand scan: which columns a built display config +needs ``min`` and ``max`` for even when the policy computes no stats. + The thresholds are PROVISIONAL. They are the values the phase-0 measurements proposed (one Apple M4 Pro, 115 tallyman telemetry loads, xorq stats on parquet and CSV slices, eager polars RSS), and they rest on gaps: the @@ -263,6 +267,35 @@ def route_polars_entry(rows: int, cols: int, limits: StatsLimits | None = None) return "xorq" if rows > limits.polars_route_rows else "eager" +def demand_columns(display_args: Any, pairs: Any) -> list[str]: + """The columns a display config reads ``histogram_bins`` of: the + ``val_column`` of every ``color_map_config`` whose ``color_rule`` is + ``color_map``, in any display of ``display_args`` (``df_display_args``). The + scan reads the config alone, so it sends no query. It covers the rules a host + put in the column overrides and the ones a klass adds at style time, since + both are in the config once it is built. + + ``pairs`` are the ``(original, rewritten)`` names of the table's columns. A + ``val_column`` may be either, and a rewritten name wins over an original + that looks like one, as the client reads it; one that is neither (or not a + string) names no column and is dropped. A ``color_map`` rule with no + ``val_column`` reads nothing, and neither does a categorical rule or a + tooltip. The result is the rewritten names, once each, in the table's order; + ``display_args`` is not changed and may have any shape.""" + wanted = set() + for display in display_args.values() if isinstance(display_args, dict) else (): + config = display.get("df_viewer_config") if isinstance(display, dict) else None + columns = config.get("column_config") if isinstance(config, dict) else None + for column in columns if isinstance(columns, list) else (): + rule = column.get("color_map_config") if isinstance(column, dict) else None + val_column = rule.get("val_column") if isinstance(rule, dict) else None + if isinstance(val_column, str) and rule.get("color_rule") == "color_map": + wanted.add(val_column) + by_name = {**{orig: new for orig, new in pairs}, **{new: new for _orig, new in pairs}} + named = {by_name[name] for name in wanted if name in by_name} + return [new for _orig, new in pairs if new in named] + + def probe_dtypes(obj: Any) -> dict[str, str]: """Column name to dtype string, from the schema alone. diff --git a/buckaroo/server/stats_wire.py b/buckaroo/server/stats_wire.py index 8dd6442ef..04a76e2c0 100644 --- a/buckaroo/server/stats_wire.py +++ b/buckaroo/server/stats_wire.py @@ -33,6 +33,24 @@ ones. Only the full run over every column reaches the final assignment (``StatRun.assigns``). +That client can also say what it wants (rows-first p36a). ``stats_request`` takes +``tier`` (``scalar`` or ``full``), ``force`` and, with a ``tier``, ``columns`` that +scope the run to those columns (a run that is never assigned, like the scalar one). +The server judges every such request itself, in this order: a tier the ceiling +refuses for the cells asked for is answered ``stats_update {final: true, status: +"not_computed", reason: "ceiling"}`` and runs nothing; a whole-table tier below the +target, or above it without ``force``, is ``not_requestable``; while the cost guard +has paused the session, a request without ``force`` is answered with reason +``cost``. A ``force`` above the target sets the session's ``stats_override`` +(``session.effective_stats_policy``) and clears the pause. An automatic request, +one without ``force``, that ran a unit over ``STATS_COST_BUDGET_S`` and left units +to run pauses the session (``cost_paused``): the status is ``not_computed`` for +``cost`` until a ``force`` request. The fields mean nothing to any other client, +or on a session with no policy. The demand scan (``stats_policy.demand_columns``) +names the columns the display config's ``color_map`` rules read, and +``df_meta.stats.demand_columns`` carries them to a client, which asks for them as +a scoped scalar request. + 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. @@ -55,6 +73,7 @@ import logging import time import traceback +from dataclasses import dataclass from typing import Any, Callable, Dict, Optional, Sequence from buckaroo.dataflow.sd_cache import split_chain_by_scope @@ -62,9 +81,11 @@ 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, stats_deferred_by_policy +from buckaroo.server.session import ( + SessionState, build_state_message, effective_stats_policy, restore_stats_status, session_dataflow, + stats_deferred_by_policy) from buckaroo.server.stat_run import StatCursor, StatRun, stat_run_key -from buckaroo.server.stats_policy import resolve_stats_policy +from buckaroo.server.stats_policy import TIERS, demand_columns, resolve_stats_policy log = logging.getLogger("buckaroo.server.stats_wire") @@ -92,6 +113,19 @@ # longest unit is what a waiting ``infinite_request`` stalls behind. STATS_BUDGET_S = 0.075 +# The longest an automatic unit may take before the cost guard pauses the +# session, in seconds. PROVISIONAL: a unit cannot be cut short, so this bounds +# how long the loop is held by one unit that nobody asked for. The number is the +# tallyman client's base timeout (10 s), not a measurement of units. For scale: +# the largest unit on parking_2017 (10.8M rows x 43 columns), the scalar batch, +# takes 0.9 to 1.0 s, and the slowest in-limit telemetry load (an 11.8M-row CSV +# union) took 22 s for all of its units together. +STATS_COST_BUDGET_S = 10.0 + +# The tiers a ``stats_request`` can name. ``schema`` is what a session has until +# something is computed, so there is nothing to request at it. +REQUEST_TIERS = ("scalar", "full") + def parse_caps(raw: str) -> frozenset: """The capabilities in a ``?caps=a,b`` value.""" @@ -123,9 +157,10 @@ def serves_scalar_tier(session: SessionState, client: Any) -> bool: """Whether a ``stats_request`` from ``client`` is answered with the scalar tier: the client advertised both bits, so it takes tiers, and the session has nothing computed (``not_computed``) and a policy target of ``scalar``, - whether the size rule, the host or the ceiling put it there. Any other client - of such a session is served as ``stats_to_pull`` says.""" - policy = session.stats_policy + whether the size rule, the host or the ceiling put it there, or a forced + request raised it (``effective_stats_policy``). Any other client of such a + session is served as ``stats_to_pull`` says.""" + policy = effective_stats_policy(session) return (client_has_ondemand(client) and session.stats_status == "not_computed" and policy is not None and policy["tier_target"] == "scalar") @@ -141,18 +176,11 @@ def resolve_session_policy(stats_tier: str, dataflow_tier: str, rows: int, cols: return resolve_stats_policy("xorq", "xorq_build", rows, cols, host_tier=stats_tier) -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.""" + settings survive. The dataflow reaches no client until this runs. The demand + columns follow the config the snapshot now holds.""" refreshed = get_buckaroo_display_state(dataflow) session.df_display_args = refreshed["df_display_args"] session.df_data_dict = refreshed["df_data_dict"] @@ -164,6 +192,26 @@ def refresh_session_snapshot(session: SessionState, dataflow: Any) -> None: 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} + refresh_demand_columns(session, dataflow) + + +def refresh_demand_columns(session: SessionState, dataflow: Any) -> None: + """Put the columns the session's display config wants min and max for in + ``session.stats_policy["demand_columns"]``, where ``df_meta.stats`` reports + them: the demand scan (``stats_policy.demand_columns``) over the snapshot's + config, read from the config alone with no query. Only a ``schema`` target + has demand, since a higher one computes those columns anyway. The key is + absent when there are none.""" + policy = session.stats_policy + if policy is None: + return + columns: list = [] + if policy["tier_target"] == "schema" and dataflow.processed_df is not None: + columns = demand_columns(session.df_display_args, old_col_new_col(dataflow.processed_df)) + if columns: + policy["demand_columns"] = columns + else: + policy.pop("demand_columns", None) def start_stat_run(session: SessionState, scope: str = "raw", tier: str = "full", @@ -216,7 +264,8 @@ def highlighted_display_args(display_args: dict, term: str) -> dict: def run_units(run: StatRun, budget_s: Optional[float], prefer: Sequence[str] = (), - clock: Callable[[], float] = time.perf_counter, session_id: Optional[str] = None) -> int: + clock: Callable[[], float] = time.perf_counter, session_id: Optional[str] = None, + timings: Optional[list] = 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 @@ -225,7 +274,8 @@ def run_units(run: StatRun, budget_s: Optional[float], prefer: Sequence[str] = ( 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.""" + propagates. Each unit is a ``stats.unit`` span. ``timings``, when given, + gets the seconds each unit took (the run's own timer), for the cost guard.""" started = clock() ran = 0 while run.status == "pending": @@ -235,9 +285,12 @@ def run_units(run: StatRun, budget_s: Optional[float], prefer: Sequence[str] = ( if unit is None: run.run_next() # no unit is left: this marks the run complete break + spent = run.elapsed_s 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") + if timings is not None: + timings.append(run.elapsed_s - spent) ran += 1 return ran @@ -432,14 +485,14 @@ def _elapsed_ms(started: float) -> float: return round((time.perf_counter() - started) * 1000, 1) -def _plan_run(session: SessionState, tier: str = "full") -> Optional[StatRun]: - """The session's ``StatRun`` at ``tier`` 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). For the full tier the whole-run path then decides what the - stats are.""" +def _plan_run(session: SessionState, tier: str = "full", columns: Optional[Sequence[Any]] = None) -> Optional[StatRun]: + """The session's ``StatRun`` at ``tier`` (over ``columns``, original names, + when a request scoped it) 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). For + the full tier the whole-run path then decides what the stats are.""" try: - return start_stat_run(session, tier=tier) + return start_stat_run(session, tier=tier, columns=columns) except Exception: log.warning("stat units not planned session=%s stats_gen=%s: %s", session.session_id, session.stats_gen, traceback.format_exc()) @@ -467,7 +520,8 @@ def _serve_units(session: SessionState, client: Any, prefer: Sequence[str], stat 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) + info["units"] = run_units(run, STATS_BUDGET_S, prefer, session_id=session.session_id, + timings=info.setdefault("timings", [])) except Exception: log.error("stat unit failed session=%s stats_gen=%s: %s", session.session_id, session.stats_gen, traceback.format_exc()) @@ -481,11 +535,15 @@ def _serve_units(session: SessionState, client: Any, prefer: Sequence[str], stat "elapsed_ms": _elapsed_ms(started)} -def _serve_scalar(session: SessionState, client: Any, prefer: Sequence[str], incremental: bool, stats_gen: Any, - scope: Any, started: float, info: dict) -> dict: - """The reply to a ``stats_request`` on a session whose target is ``scalar``: - a ``stats_update`` with ``tier: "scalar"`` holding the fragments this client - has not seen, ``final`` once the run is done. An incremental request runs +def _serve_run(session: SessionState, client: Any, prefer: Sequence[str], incremental: bool, stats_gen: Any, + scope: Any, started: float, info: dict, tier: str = "scalar", + columns: Optional[Sequence[Any]] = None) -> dict: + """The reply to a ``stats_request`` served from a run that is not the + session's stats: the ``scalar`` run over the whole table, or a run over named + ``columns`` (original names) at ``tier``. A ``stats_update`` with that tier + holding the fragments this client has not seen, ``final`` once the run is + done; the final one also says what the session still is (``status``, and the + ``reason``), since this run did not change it. An incremental request runs units for the time budget, as the full run's does; any other runs every unit still to run. @@ -494,22 +552,133 @@ def _serve_scalar(session: SessionState, client: Any, prefer: Sequence[str], inc the status, the snapshot and the caches are as they were. A unit that raises aborts the request and the run (it is not retried) but not the session, whose full tier is still open.""" - run = _plan_run(session, "scalar") + run = _plan_run(session, tier, columns) if run is None or run.status == "error": return _aborted(stats_gen, scope, "error", session) cursor = getattr(client, "stats_cursor", None) or StatCursor() if cursor.caught_up(run): try: info["units"] = run_units(run, STATS_BUDGET_S if incremental else None, prefer, - session_id=session.session_id) + session_id=session.session_id, timings=info.setdefault("timings", [])) except Exception: - log.error("scalar stat unit failed session=%s stats_gen=%s: %s", session.session_id, session.stats_gen, - traceback.format_exc()) + log.error("%s stat unit failed session=%s stats_gen=%s: %s", run.tier, session.session_id, + session.stats_gen, traceback.format_exc()) return _aborted(stats_gen, scope, "error", session) fragments = cursor.take(run) - return {"type": "stats_update", "stats_gen": stats_gen, "scope": scope, "tier": run.tier, + reply = {"type": "stats_update", "stats_gen": stats_gen, "scope": scope, "tier": run.tier, "final": run.status != "pending", "remaining": run.remaining, "payload": partial_payload(session_dataflow(session), run, fragments), "elapsed_ms": _elapsed_ms(started)} + if reply["final"]: + reply["status"] = session.stats_status + if session.stats_reason: + reply["reason"] = session.stats_reason + return reply + + +@dataclass(frozen=True) +class TierRequest: + """The part of a ``stats_request`` an ondemand client's ``tier``, ``force`` + and ``columns`` say. ``tier`` is ``None`` for the session's own target; + ``columns`` (original names, in frame order) is set only for a request that + scopes the run, which needs a ``tier``.""" + tier: Optional[str] + force: bool + columns: Optional[tuple] + + +def tier_fields_apply(session: SessionState, client: Any) -> bool: + """Whether ``tier``, ``force`` and a scoping ``columns`` mean anything for + this request: the client advertised both bits, and the session has a policy + to enforce them against. For any other client they are ignored.""" + return client_has_ondemand(client) and session.stats_policy is not None + + +def _tier_request(msg: dict, dataflow: Any) -> Optional[TierRequest]: + """The request's ``TierRequest``, or ``None`` when a field is not one the + server can read: a ``tier`` other than ``scalar`` or ``full``, a ``force`` + that is not a boolean, or ``columns`` (with a ``tier``) that is not a + non-empty list of the grid's column names. Without a ``tier``, ``columns`` is + only the hint it always was.""" + tier, force, columns = msg.get("tier"), msg.get("force"), msg.get("columns") + force = False if force is None else force + if (tier is not None and tier not in REQUEST_TIERS) or not isinstance(force, bool): + return None + if tier is None or columns is None: + return TierRequest(tier, force, None) + pairs = old_col_new_col(dataflow.processed_df) + if not (isinstance(columns, list) and columns and all(isinstance(c, str) for c in columns) + and set(columns) <= {rewritten for _orig, rewritten in pairs}): + return None + wanted = set(columns) + return TierRequest(tier, force, tuple(orig for orig, rewritten in pairs if rewritten in wanted)) + + +def _refused(session: SessionState, stats_gen: Any, scope: Any, tier: str, reason: str, started: float, + info: dict) -> dict: + """The reply to a request the server will not run (reason ``ceiling`` or + ``cost``): a final ``stats_update`` with no payload that says the session has + no more stats than it had, so a client can show why and stop asking.""" + info["refused"] = reason + return {"type": "stats_update", "stats_gen": stats_gen, "scope": scope, "tier": tier, "final": True, + "remaining": 0, "status": "not_computed", "reason": reason, "elapsed_ms": _elapsed_ms(started)} + + +def _gate_request(session: SessionState, asked: TierRequest, stats_gen: Any, scope: Any, started: float, + info: dict) -> Optional[dict]: + """Judge a tier request against the policy in force, in the order the module + docstring gives: the ceiling for the cells asked for (``resolve_stats_policy`` + with the tier as the host's, the same function load uses), then, for a + whole-table request, whether the tier is the target or above it with + ``force``, then the cost pause. Returns the reply that refuses it, or + ``None`` when it may run. The count is the one load took.""" + policy = effective_stats_policy(session) + estimate, target = policy["estimate"], policy["tier_target"] + scoped = asked.columns is not None + tier = asked.tier or target + if not scoped and tier == "schema": + return _aborted(stats_gen, scope, "not_requestable", session) + cols = len(asked.columns) if scoped else estimate["cols"] + if resolve_stats_policy("xorq", "xorq_build", estimate["rows"], cols, host_tier=tier)["tier_target"] != tier: + return _refused(session, stats_gen, scope, tier, "ceiling", started, info) + if not scoped: + rank, target_rank = TIERS.index(tier), TIERS.index(target) + if rank < target_rank or (rank > target_rank and not asked.force): + return _aborted(stats_gen, scope, "not_requestable", session) + if session.cost_paused and not asked.force: + return _refused(session, stats_gen, scope, tier, "cost", started, info) + return None + + +def _apply_force(session: SessionState, client: Any, asked: TierRequest) -> None: + """What a ``force`` request does to the session before it runs: clear the + cost pause, and for a whole-table tier above the target, record the override + so the target (and the status) follow it for this and later unfiltered + generations. The client that forced a ``full`` run holds the schema-tier + display config, so its digest is kept for the final reply to compare.""" + if not asked.force: + return + session.cost_paused = False + target = effective_stats_policy(session)["tier_target"] + if asked.columns is None and asked.tier is not None and TIERS.index(asked.tier) > TIERS.index(target): + session.stats_override, session.stats_override_gen = asked.tier, session.stats_gen + if asked.tier == "full" and getattr(client, "display_args_hash", None) is None: + client.display_args_hash = display_args_hash(session.df_display_args) + restore_stats_status(session) + + +def _pause_if_slow(session: SessionState, reply: dict, info: dict) -> None: + """The cost guard. An automatic request (one with no ``force``, from a client + the fields apply to) whose reply leaves units to run, and which ran a unit + over ``STATS_COST_BUDGET_S``, pauses the session: the status is + ``not_computed`` for ``cost``, and the next automatic request is refused until + a ``force`` one clears it. A request that finished the run holds nothing back, + and a forced one is the user's own choice.""" + if not info.get("automatic") or reply["type"] != "stats_update" or reply["final"]: + return + if max(info.get("timings") or [0.0]) > STATS_COST_BUDGET_S: + session.cost_paused = True + session.stats_status, session.stats_reason = "not_computed", "cost" + info["paused"] = True def _rebuilt_display_args(session: SessionState, client: Any) -> Optional[dict]: @@ -543,8 +712,23 @@ def _answer_stats_request(session: Optional[SessionState], msg: dict, client: An 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 tier_fields_apply(session, client): + asked = _tier_request(msg, dataflow) + if asked is None: + return _aborted(stats_gen, scope, "bad_request", session) + # A complete session answers from the dataflow, and a failed one is an + # error below, whatever the request names. + if session.stats_status in ("pending", "not_computed"): + refusal = _gate_request(session, asked, stats_gen, scope, started, info) + if refusal is not None: + return refusal + info["automatic"] = not asked.force + _apply_force(session, client, asked) + if asked.columns is not None: + return _serve_run(session, client, prefer, bool(incremental), stats_gen, scope, started, info, + tier=asked.tier, columns=asked.columns) if serves_scalar_tier(session, client): - return _serve_scalar(session, client, prefer, bool(incremental), stats_gen, scope, started, info) + return _serve_run(session, client, prefer, bool(incremental), stats_gen, scope, started, info) to_pull = stats_to_pull(session, client) if session.stats_status == "not_computed" and not to_pull: return _aborted(stats_gen, scope, "not_requestable", session) @@ -585,8 +769,11 @@ def handle_stats_request(session: Optional[SessionState], msg: dict, client: Any 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. On a session whose target is - ``scalar`` an ondemand client's request is answered as ``_serve_scalar`` + ``scalar`` an ondemand client's request is answered as ``_serve_run`` says: ``stats_update`` messages with ``tier: "scalar"`` and nothing assigned. + ``tier``, ``force`` and ``columns`` are read for an ondemand client of a + session with a policy (see the module docstring), and ``elapsed_ms`` is in + every ``stats_update``. The ``stats.request`` span records the request and its ``outcome`` (``update`` or the abort reason), which is where updates sent, requests @@ -603,10 +790,15 @@ def handle_stats_request(session: Optional[SessionState], msg: dict, client: Any 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) + _pause_if_slow(session, reply, info) if reply["type"] == "stats_update": 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"]) + if "refused" in info: + span.set_attr(refused=info["refused"]) + if info.get("paused"): + span.set_attr(paused=True) return reply diff --git a/tests/unit/server/test_load_expr.py b/tests/unit/server/test_load_expr.py index b1080cadd..0299b4d23 100644 --- a/tests/unit/server/test_load_expr.py +++ b/tests/unit/server/test_load_expr.py @@ -3713,6 +3713,16 @@ async def test_a_force_in_a_filtered_state_is_served_for_that_state_and_waits_fo cleared = self._stats(await self._state(ws)) self.assertEqual((cleared["status"], cleared["tier_target"]), ("pending", "full")) + @tornado.testing.gen_test + async def test_a_session_nobody_forced_stays_not_computed_through_a_field_change(self): + with _stats_limits(**_SCHEMA_BY_SIZE): + await self._load("fo-plain", stats_tier="auto") + ws, _first = await self._connect("fo-plain") + for changes in ({"post_processing": "first_three"}, {"post_processing": ""}, + {"quick_command_args": {"search": ["a"]}}, {}): + stats = self._stats(await self._state(ws, **changes)) + self.assertEqual((stats["status"], stats["tier_target"]), ("not_computed", "schema"), changes) + @tornado.testing.gen_test async def test_a_reload_keeps_the_override_unless_it_names_a_tier(self): with _stats_limits(**_SCHEMA_BY_SIZE): @@ -3789,14 +3799,15 @@ async def test_every_stats_update_carries_elapsed_ms_and_it_is_the_time_the_requ await self._load("fo-elapsed", stats_tier="auto") ws, first = await self._connect("fo-elapsed") gen = self._stats(first)["gen"] - with _slow_stat_queries(0.05): + with _slow_stat_queries(0.3): scoped = await self._ask(ws, gen, tier="scalar", columns=["a"]) refused = await self._ask(ws, gen, tier="full", force=True) for name, reply in (("scoped", scoped), ("refused", refused)): self.assertEqual(reply["type"], "stats_update", name) self.assertIsInstance(reply["elapsed_ms"], (int, float), name) - self.assertGreaterEqual(scoped["elapsed_ms"], 50, "the request ran a query that took 50 ms") - self.assertLess(refused["elapsed_ms"], 50, "a refusal runs nothing") + # The margins are wide on both sides: a loaded machine stalls a refusal for tens of milliseconds. + self.assertGreaterEqual(scoped["elapsed_ms"], 250, "the request ran a query that took 300 ms") + self.assertLess(refused["elapsed_ms"], 250, "a refusal runs nothing") class TestCostGuardWire(_LimitsWire): @@ -3932,6 +3943,18 @@ async def test_force_clears_the_pause_and_the_run_goes_on_where_it_stopped(self) self.assertEqual((rest[-1]["type"], rest[-1]["final"]), ("stats_update", True)) self.assertEqual((session.stats_status, session.cost_paused), ("complete", False)) + @tornado.testing.gen_test + async def test_a_client_without_both_bits_is_not_held_by_the_pause(self): + _ws, gen = await self._paused_session("cg-legacy") + session = self._session("cg-legacy") + ws, frame = await self._connect("cg-legacy", caps="stats_update") + self.assertEqual(self._stats(frame)["status"], "pending", + "a client that cannot show the pause is told it is on its way") + ws.write_message(_stats_request(gen)) + reply = await _read_json(ws) + self.assertEqual((reply["type"], reply["tier"], reply["final"]), ("stats_update", "full", True), reply) + self.assertEqual(session.stats_status, "complete") + @tornado.testing.gen_test async def test_a_request_with_no_policy_is_never_paused(self): """An inline session has no policy, so its ondemand client's requests are answered as before.""" @@ -4025,6 +4048,14 @@ async def test_the_scan_finds_the_delta_column_a_klass_adds_at_style_time(self): _, update_only = await self._connect("dm-klass", caps="stats_update") self.assertNotIn("demand_columns", self._stats(update_only)) + @tornado.testing.gen_test + async def test_categorical_rules_alone_yield_no_demand(self): + await self._load("dm-categorical", stats_tier="schema", project_root=None, + column_config_overrides=_CATEGORICAL_OVERRIDE) + self.assertEqual(self._session("dm-categorical").stats_policy.get("demand_columns", []), []) + _, frame = await self._connect("dm-categorical") + self.assertNotIn("demand_columns", self._stats(frame)) + @tornado.testing.gen_test async def test_a_numeric_rule_in_the_overrides_is_scanned_too(self): await self._load("dm-override", stats_tier="schema", project_root=None, @@ -4033,6 +4064,15 @@ async def test_a_numeric_rule_in_the_overrides_is_scanned_too(self): self.assertEqual(self._stats(frame)["demand_columns"], ["a"], "the host's name for the column is the stats' key") + @tornado.testing.gen_test + async def test_a_session_that_will_compute_the_columns_anyway_has_no_demand(self): + for index, (tier, limits) in enumerate((("scalar", {}), ("auto", {}), ("auto", _SCALAR_BY_SIZE))): + with _stats_limits(**limits): + sid = f"dm-covered-{index}" + await self._load(sid, stats_tier=tier) + _, frame = await self._connect(sid) + self.assertNotIn("demand_columns", self._stats(frame), (tier, limits)) + @tornado.testing.gen_test async def test_the_demand_follows_a_state_change_and_a_reload(self): await self._load("dm-follow", stats_tier="schema", column_config_overrides=_CATEGORICAL_OVERRIDE) @@ -4131,11 +4171,38 @@ async def test_two_requests_for_the_same_columns_in_any_order_share_a_run(self): self.assertEqual(queries, []) self.assertEqual(list(session.stat_runs), [(gen, "raw", "scalar", ("price", "qty"))]) + @tornado.testing.gen_test + async def test_a_named_columns_request_on_a_complete_session_is_answered_from_the_dataflow(self): + await self._load("dm-complete", stats_tier="auto") + ws, first = await self._connect("dm-complete") + gen = self._stats(first)["gen"] + done = await self._ask_to_the_end(ws, gen, tier="full", force=True) + self.assertTrue(done[-1]["final"]) + self.assertEqual(self._session("dm-complete").stats_status, "complete") + with _count_stat_queries() as queries: + again = await self._ask(ws, gen, tier="scalar", columns=["a"]) + self.assertEqual((again["type"], again["final"], again["tier"]), ("stats_update", True, "full")) + self.assertEqual(queries, []) + + class TestTierFieldsAndOlderClients(_LimitsWire): """``tier``, ``force`` and a ``columns`` that scopes the run belong to a client that advertised ``stats_ondemand``; every other client is served as before.""" + @tornado.testing.gen_test + async def test_a_client_with_stats_update_only_has_the_fields_ignored(self): + with _stats_limits(**_SCHEMA_BY_SIZE): + await self._load("sk-update", stats_tier="auto") + ws, first = await self._connect("sk-update", caps="stats_update") + gen = self._stats(first)["gen"] + self.assertEqual(self._stats(first)["status"], "pending") + ws.write_message(_stats_request(gen, tier="scalar", force=True, columns=["a"])) + reply = await _read_json(ws) + self.assertEqual((reply["type"], reply["tier"], reply["final"]), ("stats_update", "full", True), reply) + self.assertNotIn("status", reply) + self.assertTrue(_FULL_ONLY_WIRE <= set(_rows_by_stat(reply["payload"]))) + @tornado.testing.gen_test async def test_force_from_a_client_without_both_bits_records_no_override(self): with _stats_limits(**_SCHEMA_BY_SIZE): @@ -4147,6 +4214,41 @@ async def test_force_from_a_client_without_both_bits_records_no_override(self): session = self._session("sk-no-override") self.assertEqual((session.stats_override, session.cost_paused), (None, False)) + @tornado.testing.gen_test + async def test_a_host_named_schema_session_refuses_a_stats_update_only_client_whatever_it_sends(self): + await self._load("sk-host", stats_tier="schema") + ws, first = await self._connect("sk-host", caps="stats_update") + gen = self._stats(first)["gen"] + for fields in ({}, {"tier": "full", "force": True}): + ws.write_message(_stats_request(gen, incremental=True, **fields)) + aborted = await _read_json(ws) + self.assertEqual((aborted["type"], aborted["reason"]), ("stats_aborted", "not_requestable"), fields) + + @tornado.testing.gen_test + async def test_a_session_with_no_policy_ignores_the_fields_of_an_ondemand_client(self): + await self._load("sk-inline", stats_delivery="inline") + ws, frame = await self._connect("sk-inline") + self.assertNotIn("stats", frame["df_meta"]) + gen = self._session("sk-inline").stats_gen + for fields in ({}, {"tier": "scalar"}, {"tier": "full", "force": True}, {"tier": "scalar", "columns": ["a"]}, + {"tier": "medium"}, {"force": "yes"}, {"tier": "scalar", "columns": ["zz"]}): + reply = await self._ask(ws, gen, **fields) + self.assertEqual((reply["type"], reply["final"], reply["tier"]), ("stats_update", True, "full"), fields) + self.assertTrue(_FULL_ONLY_WIRE <= set(_rows_by_stat(reply["payload"])), fields) + self.assertEqual(self._session("sk-inline").stat_runs, {}) + + @tornado.testing.gen_test + async def test_an_explicit_full_session_serves_an_ondemand_client_its_plain_requests_as_before(self): + self._one_unit_per_request() + await self._load("sk-full", stats_tier="full") + ws, first = await self._connect("sk-full") + self.assertEqual(self._stats(first), {"status": "pending", "tier": "schema", "gen": self._stats(first)["gen"]}) + replies = await self._ask_to_the_end(ws, self._stats(first)["gen"]) + self.assertEqual([(r["tier"], r["final"], r["remaining"]) for r in replies], + [("full", False, 3), ("full", False, 2), ("full", False, 1), ("full", True, 0)]) + self.assertTrue(all("status" not in r for r in replies)) + + class TestReloadExpr(tornado.testing.AsyncHTTPTestCase): def get_app(self): return make_app() From 6e85bd0f9216fafc7841a0807ad7da7b282d723d Mon Sep 17 00:00:00 2001 From: Paddy Mullen Date: Sun, 4 Oct 2026 18:11:54 -0400 Subject: [PATCH 33/39] test(server): failing tests for the display config a client holds when it continues after a cost pause (rows-first p36a) A client that continues a cost-paused session (force with no tier, or a tier equal to the target) never has its display_args_hash recorded, so the final stats_update carries no df_display_args and the client keeps the config its not_computed frame carried. Two clients are in that state: one that connected while the session was paused, and one that changed state (a search) while it was. Both complete the session with force and expect the rebuilt config in the final reply. The client that connected pending before the pause is covered by the existing tests. Both tests fail on the earlier code: the final reply has no df_display_args. Co-Authored-By: Claude Sonnet 5.5 --- tests/unit/server/test_load_expr.py | 40 +++++++++++++++++++++++++++++ 1 file changed, 40 insertions(+) diff --git a/tests/unit/server/test_load_expr.py b/tests/unit/server/test_load_expr.py index 0299b4d23..063261d4f 100644 --- a/tests/unit/server/test_load_expr.py +++ b/tests/unit/server/test_load_expr.py @@ -3943,6 +3943,46 @@ async def test_force_clears_the_pause_and_the_run_goes_on_where_it_stopped(self) self.assertEqual((rest[-1]["type"], rest[-1]["final"]), ("stats_update", True)) self.assertEqual((session.stats_status, session.cost_paused), ("complete", False)) + @tornado.testing.gen_test + async def test_a_client_that_connected_while_paused_is_sent_the_rebuilt_config_when_it_continues(self): + """Its frame was ``not_computed`` and carried the schema tier's config, so the + complete stats have to say it changed, whether the request names the + tier the session is headed for or leaves it out.""" + for index, fields in enumerate(({}, {"tier": "full"})): + sid = f"cg-late-{index}" + _ws, gen = await self._paused_session(sid) + session = self._session(sid) + late, frame = await self._connect(sid) + self.assertEqual({key: self._stats(frame).get(key) for key in ("status", "reason", "tier_target")}, + {"status": "not_computed", "reason": "cost", "tier_target": "full"}, fields) + replies = await self._ask_to_the_end(late, gen, force=True, **fields) + final = replies[-1] + self.assertEqual((final["type"], final["final"], session.stats_status), ("stats_update", True, "complete"), + fields) + for reply in replies[:-1]: + self.assertNotIn("df_display_args", reply, fields) + self.assertIn("df_display_args", final, fields) + self.assertEqual(_as_json(final["df_display_args"]), _as_json(session.df_display_args), fields) + self.assertNotEqual(_as_json(frame["df_display_args"]), _as_json(session.df_display_args), + "the complete stats change the config the paused frame carried") + + @tornado.testing.gen_test + async def test_a_client_that_changed_state_while_paused_is_sent_the_rebuilt_config_when_it_continues(self): + """A state change while paused sends the client a new ``not_computed`` frame, with + the schema tier's config and its own search highlight.""" + ws, _gen = await self._paused_session("cg-search") + session = self._session("cg-search") + frame = await self._state(ws, quick_command_args={"search": ["b"]}, search_string="b") + stats = self._stats(frame) + self.assertEqual((stats["status"], stats["reason"]), ("not_computed", "cost")) + replies = await self._ask_to_the_end(ws, stats["gen"], force=True) + final = replies[-1] + self.assertEqual((final["type"], final["final"], session.stats_status), ("stats_update", True, "complete")) + self.assertIn("df_display_args", final) + self.assertEqual(_as_json(final["df_display_args"]), + _as_json(stats_wire.highlighted_display_args(session.df_display_args, "b"))) + self.assertNotEqual(_as_json(frame["df_display_args"]), _as_json(final["df_display_args"])) + @tornado.testing.gen_test async def test_a_client_without_both_bits_is_not_held_by_the_pause(self): _ws, gen = await self._paused_session("cg-legacy") From 909ed59cb65fded2b9a1ccaccda1337ce332b8af Mon Sep 17 00:00:00 2001 From: Paddy Mullen Date: Sun, 4 Oct 2026 18:28:54 -0400 Subject: [PATCH 34/39] fix(server): keep the display digest of a client that continues a paused run (rows-first p36a) _apply_force recorded the client's display_args_hash only when a force raised the target above the policy's. A continue after a cost pause names no tier, or the target, so the digest was never recorded: build_state_message_for had already cleared it for the paused not_computed frame, and the final stats_update left out df_display_args. A client that connected while the session was paused, or changed state while it was, kept the schema-tier config its frame carried after the session completed. The digest is now recorded for any whole-table force that runs the full tier, by raising the target or by continuing to it. A force for the scalar tier or for named columns still records none, since neither run ends in an assignment. A client that already holds a digest keeps it. Co-Authored-By: Claude Sonnet 5.5 --- buckaroo/server/stats_wire.py | 17 +++++++++++------ 1 file changed, 11 insertions(+), 6 deletions(-) diff --git a/buckaroo/server/stats_wire.py b/buckaroo/server/stats_wire.py index 04a76e2c0..70db2dafb 100644 --- a/buckaroo/server/stats_wire.py +++ b/buckaroo/server/stats_wire.py @@ -653,16 +653,21 @@ def _apply_force(session: SessionState, client: Any, asked: TierRequest) -> None """What a ``force`` request does to the session before it runs: clear the cost pause, and for a whole-table tier above the target, record the override so the target (and the status) follow it for this and later unfiltered - generations. The client that forced a ``full`` run holds the schema-tier - display config, so its digest is kept for the final reply to compare.""" + generations. The client that forced a whole-table ``full`` run, whether it + raised the target or continued a paused run to it, holds the schema-tier + display config its ``not_computed`` frame carried (a frame that is not + ``pending`` records no digest), so its digest is kept for the final reply to + compare.""" if not asked.force: return session.cost_paused = False target = effective_stats_policy(session)["tier_target"] - if asked.columns is None and asked.tier is not None and TIERS.index(asked.tier) > TIERS.index(target): - session.stats_override, session.stats_override_gen = asked.tier, session.stats_gen - if asked.tier == "full" and getattr(client, "display_args_hash", None) is None: - client.display_args_hash = display_args_hash(session.df_display_args) + tier = asked.tier or target + whole_table = asked.columns is None + if whole_table and TIERS.index(tier) > TIERS.index(target): + session.stats_override, session.stats_override_gen = tier, session.stats_gen + if whole_table and tier == "full" and getattr(client, "display_args_hash", None) is None: + client.display_args_hash = display_args_hash(session.df_display_args) restore_stats_status(session) From c7fb1ff97beb80db15c7e3501c4bca26dee10b8d Mon Sep 17 00:00:00 2001 From: Paddy Mullen Date: Sun, 4 Oct 2026 19:02:46 -0400 Subject: [PATCH 35/39] test(server): failing tests for the sort and search guards and the searched count memo (rows-first p37) Above the sort threshold every column of the grid served by infinite_request has ag_grid_specs sortable false, including a column a klass or a host override styles, and a sorted infinite_request gets error_code sort_disabled. df_meta carries sort and search flags. A repeated searched xorq window issues one count. The thresholds come from the stats policy module. Co-Authored-By: Claude Sonnet 5.5 --- tests/unit/server/test_load_expr.py | 355 +++++++++++++++++++++++++ tests/unit/server/test_stats_policy.py | 347 +++++++++++++++++++++++- 2 files changed, 701 insertions(+), 1 deletion(-) diff --git a/tests/unit/server/test_load_expr.py b/tests/unit/server/test_load_expr.py index 063261d4f..ae1900d48 100644 --- a/tests/unit/server/test_load_expr.py +++ b/tests/unit/server/test_load_expr.py @@ -31,6 +31,7 @@ from xorq.common.utils.provenance_utils import read_parquet_provenance # noqa: E402 from xorq.expr.relations import CachedNode # noqa: E402 from xorq.vendor.ibis.expr import operations as ops # noqa: E402 +from xorq.vendor.ibis.expr.types.core import Expr # noqa: E402 from buckaroo.dataflow.dataflow import assemble_merged_sd # noqa: E402 from buckaroo.dataflow.sd_cache import split_chain_by_scope # noqa: E402 @@ -4289,6 +4290,360 @@ async def test_an_explicit_full_session_serves_an_ondemand_client_its_plain_requ self.assertTrue(all("status" not in r for r in replies)) +# --------------------------------------------------------------------------- +# Guards for huge sources (rows-first p37) +# --------------------------------------------------------------------------- + +# A display klass that turns sorting back on for every column, as a klass that +# sets its own ag_grid_specs might. +_SORTABLE_STYLING = ''' +class SortableStyling(DefaultMainStyling): + df_display_name = "main" + + @classmethod + def style_column(cls, col, column_metadata): + base_config = super().style_column(col, column_metadata) + base_config["ag_grid_specs"] = {**base_config.get("ag_grid_specs", {}), "sortable": True, "minWidth": 90} + return base_config +''' + +# The same from the host's side: an override applied after the klasses have styled. +_SORTABLE_OVERRIDE = {"qty": {"ag_grid_specs": {"sortable": True}}} + + +@contextmanager +def _guard_limits(**limits): + """Sort and search thresholds for the block, as ``BUCKAROO_*`` environment + overrides, which are read each time a load resolves its guards. The table + ``_build_stats_wire_dir`` builds has five rows, so a ``sort_disable_rows`` of + 3 puts it above the sort threshold.""" + env = {f"BUCKAROO_{name.upper()}": str(value) for name, value in limits.items()} + with patch.dict(os.environ, env): + yield + + +@contextmanager +def _count_backend_queries(): + """Record the outermost op of every query sent to the backend (``execute`` + and ``to_pyarrow``) while the block runs: a count is ``CountStar``.""" + queries = [] + original_execute, original_to_pyarrow = Expr.execute, Expr.to_pyarrow + + def execute(self, *args, **kwargs): + queries.append(type(self.op()).__name__) + return original_execute(self, *args, **kwargs) + + def to_pyarrow(self, *args, **kwargs): + queries.append(type(self.op()).__name__) + return original_to_pyarrow(self, *args, **kwargs) + + with patch.object(Expr, "execute", execute), patch.object(Expr, "to_pyarrow", to_pyarrow): + yield queries + + +def _grid_columns(frame, display="main"): + """Every column config the grid of ``display`` is built from, the index + columns included.""" + config = frame["df_display_args"][display]["df_viewer_config"] + return config["column_config"] + config["left_col_configs"] + + +def _sortable(column): + return column.get("ag_grid_specs", {}).get("sortable") + + +class TestSortGuardDataflow: + """The sort guard on the dataflow (rows-first p37): ``set_sort_enabled`` turns + sorting off in the display config as a final pass, after the klasses and the + host's overrides, and the state of the dataflow survives a rebuild.""" + + @staticmethod + def _main(dataflow): + config = dataflow.df_display_args["main"]["df_viewer_config"] + return config["column_config"] + config["left_col_configs"] + + def test_a_dataflow_sorts_by_default(self): + dataflow = _build_dataflow() + assert dataflow.sort_enabled is True + assert all(_sortable(c) is not False for c in self._main(dataflow)) + + def test_disabling_sorting_covers_every_column_including_one_an_override_sets(self): + dataflow = _build_dataflow(column_config_overrides=_SORTABLE_OVERRIDE) + assert [_sortable(c) for c in self._main(dataflow) if c.get("header_name") == "qty"] == [True] + dataflow.set_sort_enabled(False) + columns = self._main(dataflow) + assert len(columns) == 4 + assert [_sortable(c) for c in columns] == [False] * 4 + + def test_the_summary_display_is_not_changed(self): + dataflow = _build_dataflow() + before = json.dumps(dataflow.df_display_args["summary"], sort_keys=True) + dataflow.set_sort_enabled(False) + assert json.dumps(dataflow.df_display_args["summary"], sort_keys=True) == before + + def test_enabling_sorting_again_restores_the_config(self): + dataflow = _build_dataflow() + before = json.dumps(dataflow.df_display_args, sort_keys=True) + dataflow.set_sort_enabled(False) + assert json.dumps(dataflow.df_display_args, sort_keys=True) != before + dataflow.set_sort_enabled(True) + assert json.dumps(dataflow.df_display_args, sort_keys=True) == before + + def test_the_guard_survives_a_state_change_and_new_stats(self): + dataflow = _build_dataflow(stats_tier="schema") + dataflow.set_sort_enabled(False) + dataflow.quick_command_args = {"search": ["a"]} + assert [_sortable(c) for c in self._main(dataflow)] == [False] * 4 + stats_wire.assign_full_stats(dataflow) + assert [_sortable(c) for c in self._main(dataflow)] == [False] * 4 + + +class TestSortGuardWire(_LimitsWire): + """A session a host opened with a stats policy, over the sort threshold + (rows-first p37): sorting is off in its config, a sorted window is refused + with ``sort_disabled``, and ``df_meta`` says so. The table has five rows, so a + ``sort_disable_rows`` of 3 is over.""" + + DISPLAY_FILES = {"sortable_styling.py": _SORTABLE_STYLING} + + async def _window(self, ws, **fields): + ws.write_message(json.dumps({"type": "infinite_request", "payload_args": {"start": 0, "end": 5, + "sourceName": "default", "origEnd": 5, **fields}})) + return await _read_json(ws) + + async def _served(self, ws, **fields): + """A window the server serves: the ``infinite_resp``, then its parquet frame.""" + resp = await self._window(ws, **fields) + self.assertEqual((resp["type"], resp.get("error_info")), ("infinite_resp", None), resp) + self.assertIsInstance(await ws.read_message(), bytes) + return resp + + async def _refused(self, ws, **fields): + """A window the server refuses, with nothing after it: the next frame + read is the answer to the next request, not a parquet frame.""" + resp = await self._refusal(ws, **fields) + again = await self._window(ws) + self.assertEqual((again["type"], again["length"]), ("infinite_resp", 5), again) + self.assertIsInstance(await ws.read_message(), bytes) + return resp + + async def _refusal(self, ws, **fields): + resp = await self._window(ws, **fields) + self.assertEqual((resp["type"], resp.get("error_code")), ("infinite_resp", "sort_disabled"), resp) + return resp + + @tornado.testing.gen_test + async def test_every_column_stops_sorting_whatever_a_klass_or_an_override_says(self): + with _guard_limits(sort_disable_rows=3): + await self._load("sg-columns", column_config_overrides=_SORTABLE_OVERRIDE) + _, frame = await self._connect("sg-columns") + columns = _grid_columns(frame) + self.assertEqual(len(columns), 4) + self.assertEqual([_sortable(c) for c in columns], [False] * 4, columns) + by_name = {c["header_name"]: c["ag_grid_specs"] for c in columns if "header_name" in c and c["col_name"] != "index"} + self.assertEqual(by_name["price"], {"minWidth": 90, "sortable": False}, "the klass's other specs are kept") + self.assertEqual(by_name["qty"], {"sortable": False}, + "the override replaced the specs, and the guard has the last word") + self.assertEqual(frame["df_display_args"], self._session("sg-columns").df_display_args) + + + @tornado.testing.gen_test + async def test_a_sorted_window_is_refused_and_no_query_runs(self): + with _guard_limits(sort_disable_rows=3): + await self._load("sg-refuse") + ws, _ = await self._connect("sg-refuse") + with _count_backend_queries() as queries: + refusal = await self._refusal(ws, sort="a", sort_direction="asc") + self.assertEqual(queries, []) + self.assertEqual(refusal["length"], 0) + self.assertEqual(refusal["key"]["sort"], "a") + self.assertTrue(isinstance(refusal["error_info"], str) and refusal["error_info"]) + again = await self._window(ws) + self.assertEqual((again["type"], again["length"]), ("infinite_resp", 5), again) + self.assertIsInstance(await ws.read_message(), bytes, "no frame followed the refusal") + + @tornado.testing.gen_test + async def test_the_second_window_of_a_request_is_judged_too(self): + with _guard_limits(sort_disable_rows=3): + await self._load("sg-second") + ws, _ = await self._connect("sg-second") + second = {"start": 5, "end": 10, "sourceName": "default", "origEnd": 10, "sort": "a", "sort_direction": "asc"} + first = await self._window(ws, second_request=second) + self.assertEqual((first["type"], first.get("error_info")), ("infinite_resp", None), first) + self.assertIsInstance(await ws.read_message(), bytes) + refusal = await _read_json(ws) + self.assertEqual((refusal["type"], refusal["error_code"], refusal["key"]), ("infinite_resp", "sort_disabled", + second)) + + + @tornado.testing.gen_test + async def test_df_meta_carries_the_sort_flag_for_every_client(self): + with _guard_limits(sort_disable_rows=3): + await self._load("sg-meta") + for caps in (_ONDEMAND, "stats_update", None): + _, frame = await self._connect("sg-meta", caps=caps) + self.assertEqual(frame["df_meta"].get("sort"), "disabled", caps) + self.assertNotIn("search", frame["df_meta"], caps) + + @tornado.testing.gen_test + async def test_df_meta_carries_the_search_flag_when_its_threshold_is_passed_and_search_is_still_served(self): + with _guard_limits(sort_disable_rows=100, search_disable_rows=3): + await self._load("sg-search") + ws, frame = await self._connect("sg-search") + self.assertEqual(frame["df_meta"].get("search"), "disabled") + self.assertNotIn("sort", frame["df_meta"]) + self.assertTrue(all(_sortable(c) is not False for c in _grid_columns(frame))) + await self._state(ws, search_string="a") + self.assertEqual((await self._served(ws))["length"], 2, "the server has no switch for search") + + + @tornado.testing.gen_test + async def test_the_guard_follows_a_state_change(self): + with _guard_limits(sort_disable_rows=3): + await self._load("sg-state") + ws, _ = await self._connect("sg-state") + frame = await self._state(ws, post_processing="first_three") + self.assertEqual(frame["df_meta"].get("sort"), "disabled") + self.assertEqual([_sortable(c) for c in _grid_columns(frame)], [False] * 4) + await self._refused(ws, sort="a", sort_direction="desc") + + @tornado.testing.gen_test + async def test_the_guard_holds_once_the_stats_are_complete(self): + with _guard_limits(sort_disable_rows=3): + await self._load("sg-complete") + ws, first = await self._connect("sg-complete") + replies = await self._ask_to_the_end(ws, self._stats(first)["gen"]) + self.assertTrue(replies[-1]["final"]) + session = self._session("sg-complete") + self.assertEqual(session.stats_status, "complete") + config = session.df_display_args["main"]["df_viewer_config"] + self.assertEqual([_sortable(c) for c in config["column_config"] + config["left_col_configs"]], [False] * 4) + if "df_display_args" in replies[-1]: + held = replies[-1]["df_display_args"]["main"]["df_viewer_config"] + self.assertEqual([_sortable(c) for c in held["column_config"] + held["left_col_configs"]], [False] * 4) + await self._refused(ws, sort="b", sort_direction="asc") + + @tornado.testing.gen_test + async def test_a_reload_resolves_the_guard_again(self): + with _guard_limits(sort_disable_rows=100): + await self._load("sg-reload") + ws, first = await self._connect("sg-reload") + self.assertNotIn("sort", first["df_meta"]) + with _guard_limits(sort_disable_rows=3): + await self._reload("sg-reload") + frame = await _read_json(ws) + self.assertEqual(frame["df_meta"].get("sort"), "disabled") + self.assertEqual([_sortable(c) for c in _grid_columns(frame)], [False] * 4) + await self._refused(ws, sort="a", sort_direction="asc") + with _guard_limits(sort_disable_rows=100): + await self._reload("sg-reload") + frame = await _read_json(ws) + self.assertNotIn("sort", frame["df_meta"]) + self.assertTrue(all(_sortable(c) is not False for c in _grid_columns(frame))) + await self._served(ws, sort="a", sort_direction="asc") + + @tornado.testing.gen_test + async def test_a_load_of_a_file_ends_the_guard(self): + sid = "sg-load" + with _guard_limits(sort_disable_rows=3): + await self._load(sid) + ws, _ = await self._connect(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) + frame = await _read_json(ws) + self.assertNotIn("sort", frame["df_meta"]) + self.assertIsNone(getattr(self._session(sid), "source_guards", "missing")) + self.assertEqual((await self._served(ws, sort="a", sort_direction="asc"))["length"], 3) + + +class TestSearchedCountMemo: + """A searched xorq window counts the filtered rows, and ``search_expr`` builds + a new expression each time, which ``_expr_count`` caches by object, so every + request counted again. The count is now held for (base expression, term) + (rows-first p37, plan 2 section 4.3).""" + + WINDOW = {"start": 0, "end": 5, "sourceName": "default", "origEnd": 5} + + @staticmethod + def _counts(queries): + return queries.count("CountStar") + + def test_a_repeated_searched_window_issues_one_count(self): + dataflow = _build_dataflow() + with _count_backend_queries() as queries: + first, _ = xorq_loading.handle_infinite_request_xorq(dataflow, self.WINDOW, search_string="a") + second, _ = xorq_loading.handle_infinite_request_xorq(dataflow, self.WINDOW, search_string="a") + assert (first["length"], second["length"]) == (2, 2) + assert self._counts(queries) == 1 + + def test_another_window_of_the_same_search_reuses_the_count(self): + dataflow = _build_dataflow() + with _count_backend_queries() as queries: + for start, end in ((0, 1), (1, 2), (0, 2)): + resp, parquet = xorq_loading.handle_infinite_request_xorq(dataflow, + {**self.WINDOW, "start": start, "end": end}, search_string="a") + assert (resp["length"], pq.read_table(io.BytesIO(parquet)).num_rows) == (2, end - start) + assert self._counts(queries) == 1 + + def test_each_term_is_counted_once(self): + dataflow = _build_dataflow() + with _count_backend_queries() as queries: + lengths = [xorq_loading.handle_infinite_request_xorq(dataflow, self.WINDOW, search_string=term)[0]["length"] + for term in ("a", "b", "a", "b", "c", "a")] + assert lengths == [2, 2, 2, 2, 1, 2] + assert self._counts(queries) == 3 + + def test_a_new_base_expression_is_counted_again(self): + dataflow = _build_dataflow() + before, _ = xorq_loading.handle_infinite_request_xorq(dataflow, self.WINDOW, search_string="a") + dataflow.quick_command_args = {"search": ["b"]} + with _count_backend_queries() as queries: + after, _ = xorq_loading.handle_infinite_request_xorq(dataflow, self.WINDOW, search_string="a") + xorq_loading.handle_infinite_request_xorq(dataflow, self.WINDOW, search_string="a") + assert (before["length"], after["length"]) == (2, 0), "the rows the quick search left hold no a" + assert self._counts(queries) == 1 + + + +class TestSearchedWindowCountsWire(_LimitsWire): + """The same through the WebSocket: the term a client types is per client, and + the count it costs is shared.""" + + async def _window(self, ws, **fields): + ws.write_message(json.dumps({"type": "infinite_request", "payload_args": {"start": 0, "end": 5, + "sourceName": "default", "origEnd": 5, **fields}})) + resp = await _read_json(ws) + self.assertIsInstance(await ws.read_message(), bytes) + return resp + + @tornado.testing.gen_test + async def test_a_repeated_searched_window_issues_one_count(self): + await self._load("sc-repeat", stats_delivery="inline") + ws, _ = await self._connect("sc-repeat", caps=None) + await self._state(ws, search_string="a") + with _count_backend_queries() as queries: + lengths = [(await self._window(ws, start=start, end=end))["length"] for start, end in ((0, 5), (0, 5), (1, 2))] + self.assertEqual(lengths, [2, 2, 2]) + self.assertEqual(queries.count("CountStar"), 1) + + @tornado.testing.gen_test + async def test_two_clients_searching_the_same_term_share_the_count(self): + await self._load("sc-share", stats_delivery="inline") + a, _ = await self._connect("sc-share", caps=None) + b, _ = await self._connect("sc-share", caps=None) + for ws in (a, b): + await self._state(ws, search_string="a") + with _count_backend_queries() as queries: + self.assertEqual([(await self._window(ws))["length"] for ws in (a, b, a, b)], [2, 2, 2, 2]) + self.assertEqual(queries.count("CountStar"), 1) + + class TestReloadExpr(tornado.testing.AsyncHTTPTestCase): def get_app(self): return make_app() diff --git a/tests/unit/server/test_stats_policy.py b/tests/unit/server/test_stats_policy.py index bc0367e6b..598a8fd2b 100644 --- a/tests/unit/server/test_stats_policy.py +++ b/tests/unit/server/test_stats_policy.py @@ -6,7 +6,10 @@ ``df_meta.stats`` reports it, and how a ``stats_request`` is judged against it (the ceiling, ``requestable``, ``force``, the cost pause, the policy a forced tier raises it to, and the demand scan), still without a server; the HTTP and -WebSocket side is in test_load_expr.py. +WebSocket side is in test_load_expr.py. The last section (rows-first p37) tests the +guards for huge sources: the sort and search thresholds, the flags a session +reports, the refusal of a sorted window, and the pass that turns sorting off in a +display config. The thresholds are the provisional values proposed by the phase-0 measurements (p31b). The boundary tables run on the module defaults, so a @@ -1271,3 +1274,345 @@ def test_the_config_is_not_changed(self, sp): def test_no_columns_no_demand(self, sp): assert sp.demand_columns(self._display(self._rule("a")), []) == [] + + +# --------------------------------------------------------------------------- +# Guards for huge sources (rows-first p37) +# --------------------------------------------------------------------------- + +GUARD_ENV_FIELDS = {"BUCKAROO_SORT_DISABLE_ROWS": "sort_disable_rows", + "BUCKAROO_SEARCH_DISABLE_ROWS": "search_disable_rows"} +# The provisional defaults, written out so a change to a module constant has to +# change a test too. Search has no default threshold: the server never refuses a +# search (plan 3 question 7), so the flag is only ever turned off by a host. +GUARD_DEFAULTS = {"sort_disable_rows": 25_000_000, "search_disable_rows": None} +ENABLED = {"sort": "enabled", "search": "enabled"} + + +@pytest.fixture(autouse=True) +def clean_guard_env(monkeypatch): + """A developer's own sort and search thresholds must not leak in.""" + for name in GUARD_ENV_FIELDS: + monkeypatch.delenv(name, raising=False) + + +def _guard_api(sp, name): + api = getattr(sp, name, None) + assert api is not None, f"buckaroo.server.stats_policy.{name} does not exist" + return api + + +def _guards(sp, backend, rows, cols=10, source_kind="parquet", **limits): + """``resolve_source_guards`` on the module defaults, or on the limits given.""" + resolve = _guard_api(sp, "resolve_source_guards") + if limits: + return resolve(backend, source_kind, rows, cols, limits=_guard_api(sp, "GuardLimits")(**limits)) + return resolve(backend, source_kind, rows, cols) + + +class TestGuardLimits: + def test_the_defaults_are_the_provisional_values(self, sp): + assert dataclasses.asdict(_guard_api(sp, "GuardLimits")()) == GUARD_DEFAULTS + assert getattr(sp, "DEFAULT_SORT_DISABLE_ROWS", None) == GUARD_DEFAULTS["sort_disable_rows"] + assert getattr(sp, "DEFAULT_SEARCH_DISABLE_ROWS", "missing") is None + + def test_from_env_with_a_clean_environment_is_the_default_set(self, sp): + assert dataclasses.asdict(_guard_api(sp, "GuardLimits").from_env()) == GUARD_DEFAULTS + + @pytest.mark.parametrize("name, field", GUARD_ENV_FIELDS.items()) + def test_each_variable_overrides_its_field(self, sp, monkeypatch, name, field): + monkeypatch.setenv(name, "12345") + assert getattr(_guard_api(sp, "GuardLimits").from_env(), field) == 12345 + + def test_underscored_integers_are_accepted(self, sp, monkeypatch): + monkeypatch.setenv("BUCKAROO_SORT_DISABLE_ROWS", "2_000_000") + assert _guard_api(sp, "GuardLimits").from_env().sort_disable_rows == 2_000_000 + + @pytest.mark.parametrize("name, field", GUARD_ENV_FIELDS.items()) + @pytest.mark.parametrize("bad", ["lots", "-5", "1.5"]) + def test_invalid_value_falls_back_to_the_default_and_warns(self, sp, monkeypatch, caplog, name, field, bad): + monkeypatch.setenv(name, bad) + with caplog.at_level(logging.WARNING): + limits = _guard_api(sp, "GuardLimits").from_env() + assert getattr(limits, field) == GUARD_DEFAULTS[field] + assert name in caplog.text + + def test_empty_value_means_unset(self, sp, monkeypatch): + monkeypatch.setenv("BUCKAROO_SORT_DISABLE_ROWS", "") + assert _guard_api(sp, "GuardLimits").from_env().sort_disable_rows == GUARD_DEFAULTS["sort_disable_rows"] + + +# (backend, rows, limits) -> (sort, search). 10.8M x 43 is parking_2017; the +# 42.3M, 54.1M and 78.0M entries are the three tallyman loads over 70 s. +SOURCE_GUARD_TABLE = [ + ("xorq", 0, {}, ("enabled", "enabled")), + ("xorq", 52_814, {}, ("enabled", "enabled")), + ("xorq", 10_800_000, {}, ("enabled", "enabled")), + ("xorq", 24_999_999, {}, ("enabled", "enabled")), + ("xorq", 25_000_000, {}, ("enabled", "enabled")), + ("xorq", 25_000_001, {}, ("disabled", "enabled")), + ("xorq", 42_300_000, {}, ("disabled", "enabled")), + ("xorq", 54_100_000, {}, ("disabled", "enabled")), + ("xorq", 78_000_000, {}, ("disabled", "enabled")), + ("xorq", 10**12, {}, ("disabled", "enabled")), + # Eager backends keep both: their windows come from a bounded frame. + ("pandas", 78_000_000, {}, ("enabled", "enabled")), + ("polars", 78_000_000, {}, ("enabled", "enabled")), + ("polars", 78_000_000, {"sort_disable_rows": 1, "search_disable_rows": 1}, ("enabled", "enabled")), + # The sort threshold: equal, one above, and none at all. + ("xorq", 1000, {"sort_disable_rows": 1000}, ("enabled", "enabled")), + ("xorq", 1001, {"sort_disable_rows": 1000}, ("disabled", "enabled")), + ("xorq", 10**12, {"sort_disable_rows": None}, ("enabled", "enabled")), + # The search threshold is its own number. + ("xorq", 1000, {"search_disable_rows": 1000}, ("enabled", "enabled")), + ("xorq", 1001, {"search_disable_rows": 1000}, ("enabled", "disabled")), + ("xorq", 1001, {"sort_disable_rows": 5000, "search_disable_rows": 1000}, ("enabled", "disabled")), + ("xorq", 5001, {"sort_disable_rows": 5000, "search_disable_rows": 1000}, ("disabled", "disabled")), +] + + +def _source_guard_id(case): + backend, rows, limits, _expected = case + return f"{backend}-{rows}-{'-'.join(f'{k}={v}' for k, v in limits.items()) or 'defaults'}" + + +class TestResolveSourceGuards: + @pytest.mark.parametrize("backend, rows, limits, expected", SOURCE_GUARD_TABLE, + ids=[_source_guard_id(case) for case in SOURCE_GUARD_TABLE]) + def test_the_table(self, sp, backend, rows, limits, expected): + assert _guards(sp, backend, rows, **limits) == dict(zip(("sort", "search"), expected)) + + def test_the_result_has_exactly_the_two_flags(self, sp): + assert set(_guards(sp, "xorq", 5)) == {"sort", "search"} + + def test_the_column_count_does_not_move_a_row_threshold(self, sp): + assert _guards(sp, "xorq", 25_000_000, cols=1) == _guards(sp, "xorq", 25_000_000, cols=500) == ENABLED + + def test_the_guards_ignore_the_stats_thresholds(self, sp, monkeypatch): + """A ceiling that refuses full stats for the entry leaves its sort alone.""" + monkeypatch.setenv("BUCKAROO_STATS_FULL_AUTO_ROWS", "1") + monkeypatch.setenv("BUCKAROO_STATS_CEILING_FULL_ROWS", "1") + monkeypatch.setenv("BUCKAROO_STATS_CEILING_SCALAR_CELLS", "1") + assert _guards(sp, "xorq", 5_000) == ENABLED + + def test_the_environment_is_read_on_each_call(self, sp): + resolve = _guard_api(sp, "resolve_source_guards") + assert resolve("xorq", "parquet", 5_000, 10) == ENABLED + with pytest.MonkeyPatch.context() as env: + env.setenv("BUCKAROO_SORT_DISABLE_ROWS", "1000") + assert resolve("xorq", "parquet", 5_000, 10)["sort"] == "disabled" + assert resolve("xorq", "parquet", 5_000, 10) == ENABLED + + def test_the_environment_sets_the_search_flag(self, sp, monkeypatch): + monkeypatch.setenv("BUCKAROO_SEARCH_DISABLE_ROWS", "1000") + assert _guard_api(sp, "resolve_source_guards")("xorq", "parquet", 5_000, 10) == {"sort": "enabled", + "search": "disabled"} + + @pytest.mark.parametrize("rows, cols, error", [(-1, 3, ValueError), (5, -1, ValueError), (1.5, 3, TypeError), + ("5", 3, TypeError), (5, None, TypeError)]) + def test_a_bad_count_is_refused_as_the_stats_policy_refuses_it(self, sp, rows, cols, error): + with pytest.raises(error): + _guard_api(sp, "resolve_source_guards")("xorq", "parquet", rows, cols) + + def test_a_source_kind_that_is_not_a_string_is_refused(self, sp): + with pytest.raises(TypeError): + _guard_api(sp, "resolve_source_guards")("xorq", None, 5, 3) + + +def _guarded(sort="enabled", search="enabled"): + """A buckaroo-mode session whose load resolved these guard flags.""" + session = session_mod.SessionState(session_id="s", path="p") + session.mode = "buckaroo" + session.df_meta = {"total_rows": 5} + session.source_guards = {"sort": sort, "search": search} + return session + + +class TestGuardFlagsInDfMeta: + """``df_meta.sort`` and ``df_meta.search``: sent when a session has turned + one off, left out when it is on, which is what a client assumes.""" + + def test_a_session_with_no_guards_by_default(self): + assert getattr(session_mod.SessionState(session_id="s", path="p"), "source_guards", "missing") is None + + def test_the_documented_defaults(self): + assert getattr(session_mod, "SOURCE_GUARD_DEFAULTS", None) == ENABLED + + def test_a_disabled_sort_is_reported(self): + assert session_mod.build_state_message(_guarded(sort="disabled"))["df_meta"] == {"total_rows": 5, + "sort": "disabled"} + + def test_a_disabled_search_is_reported(self): + assert session_mod.build_state_message(_guarded(search="disabled"))["df_meta"] == {"total_rows": 5, + "search": "disabled"} + + def test_both_flags_are_reported_when_both_are_off(self): + assert session_mod.build_state_message(_guarded("disabled", "disabled"))["df_meta"] == {"total_rows": 5, + "sort": "disabled", "search": "disabled"} + + @pytest.mark.parametrize("ondemand", [True, False]) + def test_every_client_is_told(self, ondemand): + """The server refuses the sort whoever asks, so no capability gates the flag.""" + message = session_mod.build_state_message(_guarded(sort="disabled"), ondemand=ondemand) + assert message["df_meta"]["sort"] == "disabled" + + def test_the_flags_ride_with_the_stats_status(self, sp): + session = _session(sp, "auto", "deferred", SCALAR_BY_SIZE) + session.source_guards = {"sort": "disabled", "search": "enabled"} + df_meta = session_mod.build_state_message(session)["df_meta"] + assert df_meta == {"total_rows": 5, "sort": "disabled", "stats": session_mod.stats_meta(session)} + + def test_a_client_reads_what_is_left_out_as_enabled(self): + read = getattr(session_mod, "guards_with_defaults", None) + assert read is not None, "session.guards_with_defaults does not exist" + assert read({"total_rows": 5}) == ENABLED + assert read({}) == ENABLED + assert read(None) == ENABLED + assert read({"sort": "disabled"}) == {"sort": "disabled", "search": "enabled"} + assert read({"search": "disabled", "sort": "enabled"}) == {"sort": "enabled", "search": "disabled"} + + +class TestSortRefusal: + """``sort_refusal``: the ``infinite_resp`` a sorted window gets when the + session's sort is off, and nothing for any other request.""" + + WINDOW = {"start": 0, "end": 5, "sourceName": "default", "origEnd": 5} + + @staticmethod + def _refusal(session, payload): + refuse = getattr(session_mod, "sort_refusal", None) + assert refuse is not None, "session.sort_refusal does not exist" + return refuse(session, payload) + + def test_a_sorted_window_is_refused_with_a_code(self): + payload = {**self.WINDOW, "sort": "a", "sort_direction": "asc"} + refusal = self._refusal(_guarded(sort="disabled"), payload) + assert refusal is not None + assert (refusal["type"], refusal["key"], refusal["length"]) == ("infinite_resp", payload, 0) + assert refusal["error_code"] == "sort_disabled" + assert isinstance(refusal["error_info"], str) and refusal["error_info"] + assert "payload" not in refusal, "no rows travel with a refusal" + + @pytest.mark.parametrize("payload", [{}, {"sort": None}, {"sort": ""}, {"sort_direction": "asc"}]) + def test_a_window_with_no_sort_is_served(self, payload): + assert self._refusal(_guarded(sort="disabled"), {**self.WINDOW, **payload}) is None + + def test_a_session_with_the_sort_on_serves_a_sorted_window(self): + assert self._refusal(_guarded(), {**self.WINDOW, "sort": "a"}) is None + assert self._refusal(_guarded(search="disabled"), {**self.WINDOW, "sort": "a"}) is None + + def test_a_session_that_resolved_no_guards_serves_a_sorted_window(self): + session = session_mod.SessionState(session_id="s", path="p") + assert self._refusal(session, {**self.WINDOW, "sort": "a"}) is None + + @pytest.mark.parametrize("payload", [None, "sort", ["sort"], 5]) + def test_a_payload_that_is_not_a_dict_is_not_refused(self, payload): + assert self._refusal(_guarded(sort="disabled"), payload) is None + + +class TestResolveSessionGuards: + """``stats_wire.resolve_session_guards``: what a load handler resolves, and + for which sessions.""" + + @staticmethod + def _resolve(*args): + resolve = getattr(stats_wire, "resolve_session_guards", None) + assert resolve is not None, "stats_wire.resolve_session_guards does not exist" + return resolve(*args) + + def test_a_schema_dataflow_resolves_against_the_count_it_has(self, monkeypatch): + monkeypatch.setenv("BUCKAROO_SORT_DISABLE_ROWS", "3") + assert self._resolve("schema", 5, 3) == {"sort": "disabled", "search": "enabled"} + assert self._resolve("schema", 3, 3) == ENABLED + + def test_a_dataflow_that_ran_its_stats_has_no_guards(self, monkeypatch): + """Only a session a host opened with a stats policy is guarded. One built + the way it always was, with its stats inline, stays as it was.""" + monkeypatch.setenv("BUCKAROO_SORT_DISABLE_ROWS", "3") + assert self._resolve("full", 5, 3) is None + assert self._resolve("full", 5_000_000_000, 40) is None + + def test_the_defaults_guard_a_session_of_the_size_the_stats_ceiling_refuses(self): + assert self._resolve("schema", 78_000_000, 44)["sort"] == "disabled" + assert self._resolve("schema", 10_800_000, 43) == ENABLED + + def test_the_thresholds_are_read_on_each_call(self, monkeypatch): + assert self._resolve("schema", 5, 3) == ENABLED + monkeypatch.setenv("BUCKAROO_SORT_DISABLE_ROWS", "3") + assert self._resolve("schema", 5, 3)["sort"] == "disabled" + + +def _grid_display(main_columns, left=(), summary_columns=()): + """A ``df_display_args`` as the dataflow builds it: ``main`` is the display + the infinite grid serves from the server, ``summary`` is a client-side grid.""" + def display(data_key, columns, left_columns): + return {"data_key": data_key, "summary_stats_key": "all_stats", + "df_viewer_config": {"pinned_rows": [], "column_config": list(columns), + "left_col_configs": list(left_columns), "extra_grid_config": {}, "component_config": {}}} + return {"main": display("main", main_columns, left), "summary": display("empty", summary_columns, ())} + + +class TestDisableSorting: + """``disable_sorting``: the final pass that turns sorting off for every column + of a display served by ``infinite_request``, after the klasses and the + overrides have had their say.""" + + @staticmethod + def _disable(sp, display): + return _guard_api(sp, "disable_sorting")(display) + + @staticmethod + def _columns(display, key="main"): + config = display[key]["df_viewer_config"] + return config["column_config"] + config["left_col_configs"] + + def test_every_column_of_the_infinite_display_stops_sorting(self, sp): + display = _grid_display( + [{"col_name": "a", "header_name": "price"}, + {"col_name": "b", "header_name": "qty", "ag_grid_specs": {"minWidth": 90}}, + {"col_name": "c", "header_name": "category", "ag_grid_specs": {"sortable": True}}, + {"col_path": ("x", "y"), "field": "d", "ag_grid_specs": {"sortable": True, "minWidth": 70}}], + left=[{"col_name": "index", "header_name": "index"}]) + out = self._disable(sp, display) + assert [c["ag_grid_specs"]["sortable"] for c in self._columns(out)] == [False] * 5 + + def test_the_other_column_specs_are_kept(self, sp): + out = self._disable(sp, _grid_display([{"col_name": "b", "ag_grid_specs": {"minWidth": 90, "sortable": True}}])) + assert out["main"]["df_viewer_config"]["column_config"][0]["ag_grid_specs"] == {"minWidth": 90, + "sortable": False} + + def test_a_display_the_client_sorts_itself_is_left_alone(self, sp): + summary = [{"col_name": "a", "header_name": "price", "ag_grid_specs": {"sortable": True}}, {"col_name": "b"}] + display = _grid_display([{"col_name": "a"}], summary_columns=summary) + out = self._disable(sp, display) + assert out["summary"] == display["summary"] + + def test_the_rest_of_the_display_is_unchanged(self, sp): + display = _grid_display([{"col_name": "a", "displayer_args": {"displayer": "obj"}}]) + out = self._disable(sp, display) + main, held = out["main"], display["main"] + assert (main["data_key"], main["summary_stats_key"]) == (held["data_key"], held["summary_stats_key"]) + for key in ("pinned_rows", "extra_grid_config", "component_config"): + assert main["df_viewer_config"][key] == held["df_viewer_config"][key] + assert main["df_viewer_config"]["column_config"][0]["displayer_args"] == {"displayer": "obj"} + + def test_the_display_given_is_not_changed(self, sp): + display = _grid_display([{"col_name": "a", "ag_grid_specs": {"sortable": True}}, {"col_name": "b"}], + left=[{"col_name": "index"}]) + before = json.dumps(display, sort_keys=True) + out = self._disable(sp, display) + assert json.dumps(display, sort_keys=True) == before + assert out is not display + + @pytest.mark.parametrize("display", [{}, {"main": None}, {"main": {}}, {"main": {"data_key": "main"}}, + {"main": {"data_key": "main", "df_viewer_config": None}}, + {"main": {"data_key": "main", "df_viewer_config": {"column_config": None}}}, + {"main": {"data_key": "main", "df_viewer_config": {"column_config": []}}}]) + def test_a_display_that_is_not_the_usual_shape_does_not_raise(self, sp, display): + assert isinstance(self._disable(sp, display), dict) + + def test_a_display_with_no_left_columns_is_handled(self, sp): + display = _grid_display([{"col_name": "a"}]) + del display["main"]["df_viewer_config"]["left_col_configs"] + out = self._disable(sp, display) + assert out["main"]["df_viewer_config"]["column_config"][0]["ag_grid_specs"]["sortable"] is False From d9be4e902443a0c7782d2abbff539c55ee1f3680 Mon Sep 17 00:00:00 2001 From: Paddy Mullen Date: Sun, 4 Oct 2026 20:04:29 -0400 Subject: [PATCH 36/39] test(server): more failing tests for the sort guard and the searched count memo (rows-first p37) A regression check on the first tests commit left cases that no test pinned. A searched expression a base holds is bounded and the least recently asked term goes first, a /load_compare ends the guard like /load does, and the pass that turns sorting off leaves a column entry that is not a dict as it is and replaces ag_grid_specs that is not a dict. The wire test of a state change reads the three rows first_three leaves, not five. Co-Authored-By: Claude Sonnet 5.5 --- tests/unit/server/test_load_expr.py | 63 ++++++++++++++++++++++++-- tests/unit/server/test_stats_policy.py | 10 ++++ 2 files changed, 69 insertions(+), 4 deletions(-) diff --git a/tests/unit/server/test_load_expr.py b/tests/unit/server/test_load_expr.py index ae1900d48..9ccf8f5c4 100644 --- a/tests/unit/server/test_load_expr.py +++ b/tests/unit/server/test_load_expr.py @@ -4418,12 +4418,13 @@ async def _served(self, ws, **fields): self.assertIsInstance(await ws.read_message(), bytes) return resp - async def _refused(self, ws, **fields): + async def _refused(self, ws, rows=5, **fields): """A window the server refuses, with nothing after it: the next frame - read is the answer to the next request, not a parquet frame.""" + read is the answer to the next request, not a parquet frame. ``rows`` is + the length of that next answer.""" resp = await self._refusal(ws, **fields) again = await self._window(ws) - self.assertEqual((again["type"], again["length"]), ("infinite_resp", 5), again) + self.assertEqual((again["type"], again["length"]), ("infinite_resp", rows), again) self.assertIsInstance(await ws.read_message(), bytes) return resp @@ -4505,7 +4506,7 @@ async def test_the_guard_follows_a_state_change(self): frame = await self._state(ws, post_processing="first_three") self.assertEqual(frame["df_meta"].get("sort"), "disabled") self.assertEqual([_sortable(c) for c in _grid_columns(frame)], [False] * 4) - await self._refused(ws, sort="a", sort_direction="desc") + await self._refused(ws, rows=3, sort="a", sort_direction="desc") @tornado.testing.gen_test async def test_the_guard_holds_once_the_stats_are_complete(self): @@ -4561,6 +4562,29 @@ async def test_a_load_of_a_file_ends_the_guard(self): self.assertIsNone(getattr(self._session(sid), "source_guards", "missing")) self.assertEqual((await self._served(ws, sort="a", sort_direction="asc"))["length"], 3) + @tornado.testing.gen_test + async def test_a_comparison_ends_the_guard(self): + sid = "sg-compare" + with _guard_limits(sort_disable_rows=3): + await self._load(sid) + ws, _ = await self._connect(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) + frame = await _read_json(ws) + self.assertNotIn("sort", frame["df_meta"]) + self.assertIsNone(getattr(self._session(sid), "source_guards", "missing")) + class TestSearchedCountMemo: """A searched xorq window counts the filtered rows, and ``search_expr`` builds @@ -4609,6 +4633,37 @@ def test_a_new_base_expression_is_counted_again(self): assert (before["length"], after["length"]) == (2, 0), "the rows the quick search left hold no a" assert self._counts(queries) == 1 + @staticmethod + def _terms_per_base(): + per_base = getattr(xorq_loading, "_SEARCHED_EXPRS_PER_BASE", None) + assert per_base is not None, "xorq_loading._SEARCHED_EXPRS_PER_BASE does not exist" + return per_base + + def test_a_base_keeps_a_bounded_number_of_terms(self): + dataflow = _build_dataflow() + terms = [f"t{i}" for i in range(self._terms_per_base() + 5)] + for term in terms: + xorq_loading.handle_infinite_request_xorq(dataflow, self.WINDOW, search_string=term) + with _count_backend_queries() as queries: + xorq_loading.handle_infinite_request_xorq(dataflow, self.WINDOW, search_string=terms[-1]) + assert self._counts(queries) == 0, "the latest term is held" + xorq_loading.handle_infinite_request_xorq(dataflow, self.WINDOW, search_string=terms[0]) + assert self._counts(queries) == 1, "the oldest was dropped to make room" + + def test_a_term_asked_again_is_the_last_to_go(self): + dataflow = _build_dataflow() + terms = [f"t{i}" for i in range(self._terms_per_base())] + for term in terms: + xorq_loading.handle_infinite_request_xorq(dataflow, self.WINDOW, search_string=term) + # t0 is the oldest of a full base; asking for it again makes it the newest, + # so the next new term takes t1's place. + xorq_loading.handle_infinite_request_xorq(dataflow, self.WINDOW, search_string=terms[0]) + xorq_loading.handle_infinite_request_xorq(dataflow, self.WINDOW, search_string="newcomer") + with _count_backend_queries() as queries: + xorq_loading.handle_infinite_request_xorq(dataflow, self.WINDOW, search_string=terms[0]) + assert self._counts(queries) == 0, "a term that was asked again is still held" + xorq_loading.handle_infinite_request_xorq(dataflow, self.WINDOW, search_string=terms[1]) + assert self._counts(queries) == 1, "the least recently asked term was dropped" class TestSearchedWindowCountsWire(_LimitsWire): diff --git a/tests/unit/server/test_stats_policy.py b/tests/unit/server/test_stats_policy.py index 598a8fd2b..6ae0b78e6 100644 --- a/tests/unit/server/test_stats_policy.py +++ b/tests/unit/server/test_stats_policy.py @@ -1581,6 +1581,16 @@ def test_the_other_column_specs_are_kept(self, sp): assert out["main"]["df_viewer_config"]["column_config"][0]["ag_grid_specs"] == {"minWidth": 90, "sortable": False} + def test_specs_that_are_not_a_dict_are_replaced(self, sp): + out = self._disable(sp, _grid_display([{"col_name": "a", "ag_grid_specs": None}])) + assert out["main"]["df_viewer_config"]["column_config"] == [{"col_name": "a", + "ag_grid_specs": {"sortable": False}}] + + def test_a_column_entry_that_is_not_a_dict_is_left_as_it_is(self, sp): + out = self._disable(sp, _grid_display([None, "a", {"col_name": "b"}])) + assert out["main"]["df_viewer_config"]["column_config"] == [None, "a", {"col_name": "b", + "ag_grid_specs": {"sortable": False}}] + def test_a_display_the_client_sorts_itself_is_left_alone(self, sp): summary = [{"col_name": "a", "header_name": "price", "ag_grid_specs": {"sortable": True}}, {"col_name": "b"}] display = _grid_display([{"col_name": "a"}], summary_columns=summary) From 0181386e1a4a54580c81bb5a7332873a4a1c3a3a Mon Sep 17 00:00:00 2001 From: Paddy Mullen Date: Sun, 4 Oct 2026 20:24:14 -0400 Subject: [PATCH 37/39] feat(server): sort and search guards for huge sources and a searched count memo (rows-first p37) stats_policy gains GuardLimits (BUCKAROO_SORT_DISABLE_ROWS, default 25M rows; BUCKAROO_SEARCH_DISABLE_ROWS, default none) and resolve_source_guards, apart from the stats thresholds. A session a host opened with a stats policy resolves them at /load_expr and /reload_expr from the count load took. Above the sort threshold the xorq dataflow finishes its display config with disable_sorting, so no klass or override leaves a sortable column in a display infinite_request serves, and a sorted infinite_request is refused with error_code sort_disabled before any query runs. df_meta carries sort and search when one is disabled. handle_infinite_request_xorq holds the searched expression per (base expression, term), so the second window of a search is a hit in _expr_count's cache and issues no count. Co-Authored-By: Claude Sonnet 5.5 --- buckaroo/server/handlers.py | 13 ++- buckaroo/server/session.py | 52 ++++++++++- buckaroo/server/stats_policy.py | 123 +++++++++++++++++++++++-- buckaroo/server/stats_wire.py | 29 +++++- buckaroo/server/websocket_handler.py | 7 +- buckaroo/server/xorq_loading.py | 62 ++++++++++++- tests/unit/server/test_load_expr.py | 89 ++++++++++++++++++ tests/unit/server/test_stats_policy.py | 21 +++++ 8 files changed, 381 insertions(+), 15 deletions(-) diff --git a/buckaroo/server/handlers.py b/buckaroo/server/handlers.py index 4789ccdad..161f4bea5 100644 --- a/buckaroo/server/handlers.py +++ b/buckaroo/server/handlers.py @@ -17,7 +17,8 @@ from buckaroo.server.session import ( STATS_DELIVERIES, STATS_TIER_REQUESTS, begin_stats_generation, dataflow_stats_tier, reset_stats_controls) from buckaroo.server.stats_wire import ( - broadcast_state, refresh_demand_columns, refresh_session_snapshot, resolve_session_policy) + apply_source_guards, broadcast_state, refresh_demand_columns, refresh_session_snapshot, resolve_session_guards, + resolve_session_policy) from buckaroo.server import telemetry from buckaroo.pluggable_analysis_framework import perf_log @@ -372,6 +373,7 @@ async def post(self): # 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, session.stats_policy = "full", "inline", None + session.source_guards = None reset_stats_controls(session) begin_stats_generation(session) @@ -621,6 +623,9 @@ async def post(self): # the policy can still keep it there; one that ran its stats in the # constructor has nothing left to resolve (None). stats_policy = resolve_session_policy(stats_tier, dataflow_tier, metadata["rows"], len(metadata["columns"])) + # The guards for huge sources ride with the policy: a session a host opened + # with one has them, any other has none. + source_guards = resolve_session_guards(dataflow_tier, metadata["rows"], len(metadata["columns"])) sessions = self.application.settings["sessions"] session = sessions.get_or_create(session_id, build_dir) @@ -636,6 +641,8 @@ async def post(self): session.stats_policy = stats_policy if not same_expression or body.get("stats_tier") is not None: reset_stats_controls(session) + # Before the snapshot below: it holds the config a client sorts by. + apply_source_guards(session, xorq_dataflow, source_guards) session.tele_sink = tele_sink session.xorq_dataflow = xorq_dataflow # Clear pandas-side state left by a prior /load on the same @@ -825,6 +832,7 @@ async def post(self): session.mode = "viewer" # A viewer session has no dataflow and so no deferred stats. session.stats_tier, session.stats_delivery, session.stats_policy = "full", "inline", None + session.source_guards = None reset_stats_controls(session) begin_stats_generation(session) @@ -930,6 +938,8 @@ async def post(self, session_id): # reused, so there is no new one to take. stats_policy = resolve_session_policy(stats_tier, dataflow_tier, session.metadata["rows"], len(session.metadata["columns"])) + source_guards = resolve_session_guards(dataflow_tier, session.metadata["rows"], + len(session.metadata["columns"])) except Exception: tb = traceback.format_exc() log.error("reload_expr error session=%s: %s", session_id, tb) @@ -960,6 +970,7 @@ async def post(self, session_id): session.stats_policy = stats_policy if body.get("stats_tier") is not None: reset_stats_controls(session) + apply_source_guards(session, xorq_dataflow, source_guards) refresh_session_snapshot(session, xorq_dataflow) begin_stats_generation(session) diff --git a/buckaroo/server/session.py b/buckaroo/server/session.py index 685507253..1cfa16210 100644 --- a/buckaroo/server/session.py +++ b/buckaroo/server/session.py @@ -39,6 +39,12 @@ STATS_FIELD_DEFAULTS: Dict[str, Any] = {"auto_request": True, "requestable": ["full"], "omitted_keys": [], "approx_keys": [], "demand_columns": []} +# What a client assumes for the ``df_meta.sort`` and ``df_meta.search`` flags +# (rows-first p37) a message leaves out. The server sends a flag only when a +# source is too large for that request, so a session under the thresholds sends +# the message it always has. +SOURCE_GUARD_DEFAULTS: Dict[str, str] = {"sort": "enabled", "search": "enabled"} + def dataflow_stats_tier(stats_tier: str, stats_delivery: str) -> str: """The tier a session's dataflow is constructed at. Only a session whose @@ -134,6 +140,14 @@ class SessionState: # later phases add ``omitted_keys`` and ``approx_keys``. ``stats_meta`` reports # them, for the policy in force (``effective_stats_policy``). stats_policy: Optional[dict] = None + # Whether sort and search stay on for the entry (``stats_policy.resolve_source_guards``: + # ``{"sort": "enabled" | "disabled", "search": ...}``), resolved by /load_expr and + # /reload_expr for a session a host opened with a stats policy. ``None`` for a + # dataflow built with its stats inline, and for anything /load serves, which are + # not guarded. ``build_state_message`` reports a disabled flag in ``df_meta``, + # and a sorted window of a session whose sort is disabled is refused + # (``sort_refusal``). + source_guards: Optional[dict] = None # 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). @@ -342,6 +356,35 @@ def stats_with_defaults(df_meta: Optional[dict]) -> dict: return out +def guards_with_defaults(df_meta: Optional[dict]) -> dict: + """The ``sort`` and ``search`` flags of a ``df_meta`` as a client reads them: + one the message leaves out is ``enabled``, which is also what an old server + sends.""" + meta = df_meta or {} + return {name: meta.get(name, default) for name, default in SOURCE_GUARD_DEFAULTS.items()} + + +def _guard_flags(session: "SessionState") -> dict: + """The flags ``df_meta`` reports for the session: the ones that are not at their default.""" + return {name: value for name, value in (session.source_guards or {}).items() + if name in SOURCE_GUARD_DEFAULTS and value != SOURCE_GUARD_DEFAULTS[name]} + + +def sort_refusal(session: "SessionState", payload_args: Any) -> Optional[dict]: + """The ``infinite_resp`` that refuses a window with a ``sort`` on a session + whose sort is disabled, or ``None`` for any other request. It carries + ``error_code: "sort_disabled"`` beside the ``error_info`` every error response + has, and no rows, so the grid sees an error for that window and no parquet frame + follows.""" + guards = session.source_guards + if not guards or guards.get("sort") != "disabled": + return None + if not isinstance(payload_args, dict) or not payload_args.get("sort"): + return None + return {"type": "infinite_resp", "key": payload_args, "length": 0, "error_code": "sort_disabled", + "error_info": "Sorting is turned off for this source: it has too many rows to sort in a request."} + + def build_state_message(session: "SessionState", metadata: dict | None = None, search_string: str = "", ondemand: bool = True) -> dict: """Build the full ``initial_state`` WebSocket payload from a session. @@ -362,12 +405,13 @@ 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. + # The dataflow rebuilds df_meta wholesale, so the stats status and the guard + # flags are injected here, into a copy, rather than stored in it. df_meta = session.df_meta stats = stats_meta(session, ondemand=ondemand) - if stats is not None: - df_meta = {**df_meta, "stats": stats} + flags = _guard_flags(session) + if stats is not None or flags: + df_meta = {**df_meta, **flags, **({"stats": stats} if stats is not None else {})} 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, diff --git a/buckaroo/server/stats_policy.py b/buckaroo/server/stats_policy.py index 9567f936c..25b44c4a1 100644 --- a/buckaroo/server/stats_policy.py +++ b/buckaroo/server/stats_policy.py @@ -27,6 +27,12 @@ :func:`demand_columns` is the demand scan: which columns a built display config needs ``min`` and ``max`` for even when the policy computes no stats. +:func:`resolve_source_guards` is the rule for the two requests that stay expensive +whether or not stats are computed, a sorted window and a searched one (rows-first +p37). Its thresholds are apart from the stats ones (:class:`GuardLimits`), because a +sort or a search scales with rows where the stats tiers scale with cells, and +:func:`disable_sorting` is the pass that turns sorting off in a display config. + The thresholds are PROVISIONAL. They are the values the phase-0 measurements proposed (one Apple M4 Pro, 115 tallyman telemetry loads, xorq stats on parquet and CSV slices, eager polars RSS), and they rest on gaps: the @@ -52,6 +58,15 @@ * ``BUCKAROO_POLARS_ROUTE_ROWS`` (8,000,000 rows): R for :func:`route_polars_entry`. A memory threshold, set apart from the stats ones, and sized for ``pre_limit`` False (see the constant below). + +The guard thresholds are PROVISIONAL too and read the same way (on every call, a +bad value ignored with a warning): + +* ``BUCKAROO_SORT_DISABLE_ROWS`` (25,000,000 rows): a source with more rows has + sorting turned off. +* ``BUCKAROO_SEARCH_DISABLE_ROWS`` (none): a source with more rows reports its + search as disabled. The flag is advice for a client; the server serves a search + either way. """ from __future__ import annotations @@ -95,12 +110,29 @@ # be about 14M rows. DEFAULT_POLARS_ROUTE_ROWS: int = 8_000_000 +# Sorting a window sorts the whole source on the loop thread. At 10.8M rows x 43 +# columns that took 261-287 ms warm on parquet, and at 3M rows of CSV 326-337 ms +# (phase-0 measurements, item 4). Extrapolated linearly (nothing above 12M rows was +# measured) a 25M-row source costs about 0.7 s on parquet and 2.8 s on CSV, and the +# 42.3M, 54.1M and 78.0M entries 1.1 to 2.1 s and 4.6 to 8.5 s. 25M is the ceiling +# on full stats (the geometric midpoint of the telemetry gap between 11.8M and 42.3M +# rows), so an entry whose full stats the server refuses is one whose sort it +# refuses. A guess in the same gap, not a measurement. +DEFAULT_SORT_DISABLE_ROWS: int | None = 25_000_000 +# The server never refuses a search (plan 3 question 7), so there is no number for +# this one: the flag stays "enabled" until a host sets a threshold. +DEFAULT_SEARCH_DISABLE_ROWS: int | None = None + # Environment variable -> StatsLimits field. _ENV_FIELDS = {"BUCKAROO_STATS_FULL_AUTO_ROWS": "full_auto_rows", "BUCKAROO_STATS_SCALAR_AUTO_CELLS": "scalar_auto_cells", "BUCKAROO_STATS_CEILING_FULL_ROWS": "ceiling_full_rows", "BUCKAROO_STATS_CEILING_SCALAR_CELLS": "ceiling_scalar_cells", "BUCKAROO_POLARS_ROUTE_ROWS": "polars_route_rows", "BUCKAROO_STATS_FULL_AUTO_CELLS": "full_auto_cells", "BUCKAROO_STATS_CEILING_FULL_CELLS": "ceiling_full_cells"} +# Environment variable -> GuardLimits field. +_GUARD_ENV_FIELDS = {"BUCKAROO_SORT_DISABLE_ROWS": "sort_disable_rows", + "BUCKAROO_SEARCH_DISABLE_ROWS": "search_disable_rows"} + # Lowest to highest. TIERS = ("schema", "scalar", "full") @@ -122,6 +154,17 @@ def _env_int(name: str) -> int | None: return value +def _env_overrides(env_fields: dict[str, str]) -> dict[str, int]: + """The ``{field: value}`` of every variable in ``env_fields`` that is set to a + non-negative integer.""" + overrides = {} + for name, field in env_fields.items(): + value = _env_int(name) + if value is not None: + overrides[field] = value + return overrides + + @dataclass(frozen=True) class StatsLimits: """The thresholds :func:`resolve_stats_policy` and @@ -139,12 +182,20 @@ class StatsLimits: @classmethod def from_env(cls) -> StatsLimits: """The defaults with any ``BUCKAROO_*`` overrides applied.""" - overrides = {} - for name, field in _ENV_FIELDS.items(): - value = _env_int(name) - if value is not None: - overrides[field] = value - return cls(**overrides) + return cls(**_env_overrides(_ENV_FIELDS)) + + +@dataclass(frozen=True) +class GuardLimits: + """The thresholds :func:`resolve_source_guards` reads, in rows. ``None`` means + the flag is never turned off.""" + sort_disable_rows: int | None = DEFAULT_SORT_DISABLE_ROWS + search_disable_rows: int | None = DEFAULT_SEARCH_DISABLE_ROWS + + @classmethod + def from_env(cls) -> GuardLimits: + """The defaults with any ``BUCKAROO_*`` overrides applied.""" + return cls(**_env_overrides(_GUARD_ENV_FIELDS)) def _count(name: str, value: Any) -> int: @@ -267,6 +318,66 @@ def route_polars_entry(rows: int, cols: int, limits: StatsLimits | None = None) return "xorq" if rows > limits.polars_route_rows else "eager" +def resolve_source_guards(backend: str, source_kind: str, rows: int, cols: int, + limits: GuardLimits | None = None) -> dict[str, str]: + """Whether sort and search stay on for a source of this size. + + Returns ``{"sort": ..., "search": ...}``, each ``"enabled"`` or ``"disabled"``: + a flag is ``"disabled"`` when the source has more rows than its threshold in + ``limits`` (default :meth:`GuardLimits.from_env`) and the threshold is set. + ``rows`` is the count load already took; nothing is counted here. Eager + backends always keep both, since their windows come from a bounded frame. + ``cols`` and ``source_kind`` are accepted so a rule can read them without a + signature change; only ``rows`` decides today. Inputs are checked as + :func:`resolve_stats_policy` checks them. + """ + if not isinstance(source_kind, str): + raise TypeError(f"source_kind must be a str, got {source_kind!r}") + rows = _count("rows", rows) + _count("cols", cols) + if limits is None: + limits = GuardLimits.from_env() + if backend in EAGER_BACKENDS: + return {"sort": "enabled", "search": "enabled"} + + def flag(threshold: int | None) -> str: + return "disabled" if threshold is not None and rows > threshold else "enabled" + + return {"sort": flag(limits.sort_disable_rows), "search": flag(limits.search_disable_rows)} + + +def _unsortable(column: Any) -> Any: + """A column config with ``ag_grid_specs.sortable`` false, its other specs kept.""" + if not isinstance(column, dict): + return column + specs = column.get("ag_grid_specs") + return {**column, "ag_grid_specs": {**(specs if isinstance(specs, dict) else {}), "sortable": False}} + + +def disable_sorting(display_args: Any) -> Any: + """A copy of ``df_display_args`` in which no column of a display served by + ``infinite_request`` can sort: every ``column_config`` and ``left_col_configs`` + entry of a display whose ``data_key`` is ``main`` gets ``ag_grid_specs.sortable`` + false, whatever a klass or an override set. A display with another + ``data_key`` (the summary grid) is a client-side grid and is left as it is. + + It is the last step of building the config (``XorqServerDataflow``), after the + klasses styled it and ``column_config_overrides`` were merged, so neither can + bypass it. ``display_args`` is not changed and may have any shape.""" + if not isinstance(display_args, dict): + return display_args + out = {} + for name, display in display_args.items(): + config = display.get("df_viewer_config") if isinstance(display, dict) else None + if not isinstance(config, dict) or display.get("data_key") != "main": + out[name] = display + continue + columns = {key: [_unsortable(column) for column in config[key]] + for key in ("column_config", "left_col_configs") if isinstance(config.get(key), list)} + out[name] = {**display, "df_viewer_config": {**config, **columns}} + return out + + def demand_columns(display_args: Any, pairs: Any) -> list[str]: """The columns a display config reads ``histogram_bins`` of: the ``val_column`` of every ``color_map_config`` whose ``color_rule`` is diff --git a/buckaroo/server/stats_wire.py b/buckaroo/server/stats_wire.py index 70db2dafb..c0a315650 100644 --- a/buckaroo/server/stats_wire.py +++ b/buckaroo/server/stats_wire.py @@ -51,6 +51,12 @@ ``df_meta.stats.demand_columns`` carries them to a client, which asks for them as a scoped scalar request. +A session opened with a stats policy also resolves the guards for huge sources +(``resolve_session_guards``, rows-first p37): above the sort threshold the dataflow +turns sorting off in its display config (``apply_source_guards``) and a sorted +window is refused (``session.sort_refusal``). Neither depends on the client or on +what stats the session has. + 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. @@ -85,7 +91,7 @@ SessionState, build_state_message, effective_stats_policy, restore_stats_status, session_dataflow, stats_deferred_by_policy) from buckaroo.server.stat_run import StatCursor, StatRun, stat_run_key -from buckaroo.server.stats_policy import TIERS, demand_columns, resolve_stats_policy +from buckaroo.server.stats_policy import TIERS, demand_columns, resolve_source_guards, resolve_stats_policy log = logging.getLogger("buckaroo.server.stats_wire") @@ -176,6 +182,27 @@ def resolve_session_policy(stats_tier: str, dataflow_tier: str, rows: int, cols: return resolve_stats_policy("xorq", "xorq_build", rows, cols, host_tier=stats_tier) +def resolve_session_guards(dataflow_tier: str, rows: int, cols: int) -> Optional[dict]: + """The guards for a session the handler has just built, from the size load took + (``rows`` and ``cols``, so no query runs): ``{"sort": ..., "search": ...}``, each + ``"enabled"`` or ``"disabled"`` (``stats_policy.resolve_source_guards``). ``None`` + unless the dataflow was built at the schema tier, which is to say unless the host + opened the session with a stats policy, as ``resolve_session_policy`` does: a + session built with its stats inline is not guarded.""" + if dataflow_tier != "schema": + return None + return resolve_source_guards("xorq", "xorq_build", rows, cols) + + +def apply_source_guards(session: SessionState, dataflow: Any, guards: Optional[dict]) -> None: + """Record ``guards`` on the session and set the dataflow's sort switch to match, + which rebuilds its display config when the switch moves. Called before the + session snapshot is taken from the dataflow, since the config in the snapshot is + the one a client sorts by.""" + session.source_guards = guards + dataflow.set_sort_enabled(guards is None or guards["sort"] != "disabled") + + 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 diff --git a/buckaroo/server/websocket_handler.py b/buckaroo/server/websocket_handler.py index ec394f599..39d9254b7 100644 --- a/buckaroo/server/websocket_handler.py +++ b/buckaroo/server/websocket_handler.py @@ -9,7 +9,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.session import begin_stats_generation, dataflow_stats_tier, sort_refusal from buckaroo.server.stat_run import StatCursor from buckaroo.server.stats_wire import (broadcast_state, build_state_message_for, handle_stats_request, highlighted_display_args, parse_caps, refresh_session_snapshot) @@ -220,6 +220,11 @@ def _dispatch(pa): # each backend can apply the filter in its native # expression layer. search = self.search_string or "" + # A source too large to sort in a request refuses the window before any + # query runs, the second window of a request included. + refusal = sort_refusal(session, pa) + if refusal is not None: + return refusal, b"" if session.mode == "lazy" and session.ldf is not None: return handle_infinite_request_lazy(session.ldf, session.orig_to_rw, session.rw_to_orig, session.metadata.get("rows", 0), pa) diff --git a/buckaroo/server/xorq_loading.py b/buckaroo/server/xorq_loading.py index f60c9cc0e..286dc7847 100644 --- a/buckaroo/server/xorq_loading.py +++ b/buckaroo/server/xorq_loading.py @@ -14,11 +14,15 @@ import os import traceback import uuid +import weakref +from collections import OrderedDict from pathlib import Path import pyarrow.parquet as pq +from buckaroo.customizations.xorq_commands import search_expr from buckaroo.server.git_state_guard import install_git_state_guard +from buckaroo.server.stats_policy import disable_sorting from buckaroo.server.window import clamp_window from buckaroo.serialization_utils import make_infinite_resp from buckaroo.xorq_buckaroo import ( @@ -69,8 +73,15 @@ class XorqServerDataflow(XorqDataflow): 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``). + + ``sort_enabled`` is whether the grid served by ``infinite_request`` may sort. + The server turns it off (``set_sort_enabled``) for a source with too many rows + to sort in a request (``stats_policy.resolve_source_guards``); the config is + then built with ``stats_policy.disable_sorting`` as its last step, so no + klass or override can turn sorting back on, and every rebuild keeps it. """ + sort_enabled = True sampling_klass = XorqInfiniteSampling autocleaning_klass = XorqAutocleaning autoclean_conf = (NoCleaningConfXorq,) @@ -87,6 +98,20 @@ def __init__(self, expr, *args, extra_klasses=None, cache_storage_path=None, sta self.stat_chunk_cells = stat_chunk_cells super().__init__(expr, *args, **kwargs) + def set_sort_enabled(self, enabled: bool) -> None: + """Turn sorting on or off in the display config, rebuilding it now when + the switch moves.""" + if enabled == self.sort_enabled: + return + self.sort_enabled = enabled + _unused, processed_df, merged_sd = self.widget_args_tuple + if processed_df is not None: + self.df_display_args = self._build_df_display_args(processed_df, merged_sd) + + def _build_df_display_args(self, processed_df, merged_sd): + display_args = super()._build_df_display_args(processed_df, merged_sd) + return display_args if self.sort_enabled else disable_sorting(display_args) + def load_expr_build_dir(build_dir: str, cache_dir=None): """Rehydrate an ibis expression from a xorq build directory. @@ -562,6 +587,40 @@ def _compile_project_display(path: Path, ColAnalysis, *extra_bases): return found +# The searched expression of each (base expression, term). ``search_expr`` builds +# a new expression on every call and ``_expr_count`` caches by expression object, so +# every request for a searched window counted its rows again (0.5 to 1 s on 10.8M +# rows). Holding the expression makes the second request a hit in ``_expr_count``'s +# cache, which still does not cache a count that failed. Keyed weakly by the base +# expression, the dataflow's ``processed_df``, which a state change replaces, so an +# entry lasts as long as its base; a base keeps at most this many terms, the least +# recently used going first. +_SEARCHED_EXPRS_PER_BASE = 16 +_searched_exprs: weakref.WeakKeyDictionary = weakref.WeakKeyDictionary() + + +def _searched_expr(base, term): + """``search_expr(base, term)``: the same object each time for the same pair.""" + if not term: + return base + terms = _searched_exprs.get(base) + expr = terms.get(term) if terms is not None else None + if expr is not None: + terms.move_to_end(term) + return expr + expr = search_expr(base, term) + if expr is base: + # No string column to search: there is nothing to hold, and the base held as + # its own value would never be freed. + return base + if terms is None: + terms = _searched_exprs[base] = OrderedDict() + terms[term] = expr + if len(terms) > _SEARCHED_EXPRS_PER_BASE: + terms.popitem(last=False) + return expr + + def handle_infinite_request_xorq(xorq_dataflow: XorqServerDataflow, payload_args: dict, search_string: str = "") -> tuple[dict, bytes]: """Drive one infinite_request window against a xorq expression. @@ -573,13 +632,12 @@ def handle_infinite_request_xorq(xorq_dataflow: XorqServerDataflow, delegates to ``window_to_parquet`` and returns the ``(json_msg, parquet_bytes)`` pair the WebSocket handler ships as a text + binary frame pair (matching the pandas/polars paths).""" - from buckaroo.customizations.xorq_commands import search_expr _unused, processed_df, merged_sd = xorq_dataflow.widget_args_tuple if processed_df is None: return ({"type": "infinite_resp", "key": payload_args, "length": 0}, b"") try: - filtered_expr = search_expr(processed_df, search_string) if search_string else processed_df + filtered_expr = _searched_expr(processed_df, search_string) sort = payload_args.get("sort") sort_col = None diff --git a/tests/unit/server/test_load_expr.py b/tests/unit/server/test_load_expr.py index 9ccf8f5c4..27a10b824 100644 --- a/tests/unit/server/test_load_expr.py +++ b/tests/unit/server/test_load_expr.py @@ -2,6 +2,7 @@ XorqBuckarooInfiniteWidget over a xorq/ibis expression.""" import dataclasses import datetime +import gc import io import json import os @@ -9,6 +10,7 @@ import sys import tempfile import time +import weakref from contextlib import contextmanager from pathlib import Path from types import SimpleNamespace @@ -4447,6 +4449,51 @@ async def test_every_column_stops_sorting_whatever_a_klass_or_an_override_says(s "the override replaced the specs, and the guard has the last word") self.assertEqual(frame["df_display_args"], self._session("sg-columns").df_display_args) + @tornado.testing.gen_test + async def test_a_session_built_the_old_way_is_not_guarded(self): + """No stats policy, no guard: inline stats is what every host gets today.""" + with _guard_limits(sort_disable_rows=3): + await self._load("sg-inline", stats_delivery="inline") + ws, frame = await self._connect("sg-inline") + self.assertNotIn("sort", frame["df_meta"]) + self.assertTrue(all(_sortable(c) is not False for c in _grid_columns(frame))) + await self._served(ws, sort="a", sort_direction="asc") + + @tornado.testing.gen_test + async def test_a_session_under_both_thresholds_sends_no_flags(self): + with _guard_limits(sort_disable_rows=100, search_disable_rows=100): + await self._load("sg-none") + _, frame = await self._connect("sg-none") + self.assertFalse({"sort", "search"} & set(frame["df_meta"]), frame["df_meta"]) + + @tornado.testing.gen_test + async def test_a_session_within_the_threshold_is_as_it_was(self): + with _guard_limits(sort_disable_rows=5): + await self._load("sg-within") + ws, frame = await self._connect("sg-within") + self.assertNotIn("sort", frame["df_meta"]) + self.assertTrue(all(_sortable(c) is not False for c in _grid_columns(frame))) + self.assertEqual([_sortable(c) for c in frame["df_display_args"]["main"]["df_viewer_config"]["column_config"]], + [True] * 3, "the klass's own setting stands") + await self._served(ws, sort="a", sort_direction="asc") + + @tornado.testing.gen_test + async def test_the_summary_display_still_sorts(self): + with _guard_limits(sort_disable_rows=3): + await self._load("sg-summary") + _, frame = await self._connect("sg-summary") + self.assertTrue(_grid_columns(frame, "summary")) + self.assertTrue(all(_sortable(c) is not False for c in _grid_columns(frame, "summary"))) + + @tornado.testing.gen_test + async def test_unsorted_and_searched_windows_are_served(self): + with _guard_limits(sort_disable_rows=3): + await self._load("sg-serve") + ws, _ = await self._connect("sg-serve") + self.assertEqual((await self._served(ws))["length"], 5) + await self._state(ws, search_string="a") + self.assertEqual((await self._served(ws))["length"], 2) + @tornado.testing.gen_test async def test_a_sorted_window_is_refused_and_no_query_runs(self): @@ -4666,6 +4713,48 @@ def test_a_term_asked_again_is_the_last_to_go(self): assert self._counts(queries) == 1, "the least recently asked term was dropped" + def test_a_window_with_no_term_counts_nothing_new(self): + dataflow = _build_dataflow() + with _count_backend_queries() as queries: + for _ in range(3): + resp, _ = xorq_loading.handle_infinite_request_xorq(dataflow, self.WINDOW, search_string="") + assert resp["length"] == 5 + assert self._counts(queries) == 0 + + def test_a_count_that_failed_is_not_remembered(self): + dataflow = _build_dataflow() + failures = [RuntimeError("the backend is down")] + original = Expr.execute + + def execute(expr, *args, **kwargs): + if type(expr.op()).__name__ == "CountStar" and failures: + raise failures.pop() + return original(expr, *args, **kwargs) + + with patch.object(Expr, "execute", execute): + first, _ = xorq_loading.handle_infinite_request_xorq(dataflow, self.WINDOW, search_string="a") + second, _ = xorq_loading.handle_infinite_request_xorq(dataflow, self.WINDOW, search_string="a") + assert (first["length"], second["length"]) == (0, 2), "the next request counts again and gets the real count" + @pytest.mark.parametrize("table", [{"x": [1, 2, 3], "name": ["a", "b", "a"]}, {"x": [1, 2, 3]}], + ids=["with a string column", "no string column"]) + def test_the_memo_does_not_keep_the_base_expression_alive(self, table): + dataflow = _build_dataflow(xo.memtable(table, name="t")) + resp, _ = xorq_loading.handle_infinite_request_xorq(dataflow, self.WINDOW, search_string="a") + assert resp["length"] == (2 if "name" in table else 3) + base = weakref.ref(dataflow.widget_args_tuple[1]) + del dataflow + gc.collect() + assert base() is None + + def test_a_source_with_no_string_column_is_not_searched_and_counts_nothing_new(self): + dataflow = _build_dataflow(xo.memtable({"x": [1, 2, 3]}, name="t")) + with _count_backend_queries() as queries: + lengths = [xorq_loading.handle_infinite_request_xorq(dataflow, self.WINDOW, search_string="a")[0]["length"] + for _ in range(2)] + assert lengths == [3, 3] + assert self._counts(queries) == 0 + + class TestSearchedWindowCountsWire(_LimitsWire): """The same through the WebSocket: the term a client types is per client, and the count it costs is shared.""" diff --git a/tests/unit/server/test_stats_policy.py b/tests/unit/server/test_stats_policy.py index 6ae0b78e6..8ef6c0192 100644 --- a/tests/unit/server/test_stats_policy.py +++ b/tests/unit/server/test_stats_policy.py @@ -1319,6 +1319,11 @@ def test_the_defaults_are_the_provisional_values(self, sp): def test_from_env_with_a_clean_environment_is_the_default_set(self, sp): assert dataclasses.asdict(_guard_api(sp, "GuardLimits").from_env()) == GUARD_DEFAULTS + def test_the_stats_limits_do_not_carry_the_guard_thresholds(self, sp): + """They are separate numbers (plan 3 section 3.5): a threshold on stats + says nothing about the cost of a sorted window.""" + assert not {"sort_disable_rows", "search_disable_rows"} & {f.name for f in dataclasses.fields(sp.StatsLimits)} + @pytest.mark.parametrize("name, field", GUARD_ENV_FIELDS.items()) def test_each_variable_overrides_its_field(self, sp, monkeypatch, name, field): monkeypatch.setenv(name, "12345") @@ -1395,6 +1400,11 @@ def test_the_guards_ignore_the_stats_thresholds(self, sp, monkeypatch): monkeypatch.setenv("BUCKAROO_STATS_CEILING_SCALAR_CELLS", "1") assert _guards(sp, "xorq", 5_000) == ENABLED + def test_the_stats_policy_ignores_the_guard_thresholds(self, sp, monkeypatch): + monkeypatch.setenv("BUCKAROO_SORT_DISABLE_ROWS", "1") + monkeypatch.setenv("BUCKAROO_SEARCH_DISABLE_ROWS", "1") + assert sp.resolve_stats_policy("xorq", "parquet", 5_000, 10)["tier_target"] == "full" + def test_the_environment_is_read_on_each_call(self, sp): resolve = _guard_api(sp, "resolve_source_guards") assert resolve("xorq", "parquet", 5_000, 10) == ENABLED @@ -1450,6 +1460,12 @@ def test_both_flags_are_reported_when_both_are_off(self): assert session_mod.build_state_message(_guarded("disabled", "disabled"))["df_meta"] == {"total_rows": 5, "sort": "disabled", "search": "disabled"} + @pytest.mark.parametrize("session", [_guarded(), session_mod.SessionState(session_id="s", path="p")], + ids=["enabled", "unresolved"]) + def test_enabled_flags_leave_the_message_as_it_was(self, session): + session.df_meta = {"total_rows": 5} + assert session_mod.build_state_message(session)["df_meta"] == {"total_rows": 5} + @pytest.mark.parametrize("ondemand", [True, False]) def test_every_client_is_told(self, ondemand): """The server refuses the sort whoever asks, so no capability gates the flag.""" @@ -1462,6 +1478,11 @@ def test_the_flags_ride_with_the_stats_status(self, sp): df_meta = session_mod.build_state_message(session)["df_meta"] assert df_meta == {"total_rows": 5, "sort": "disabled", "stats": session_mod.stats_meta(session)} + def test_the_session_s_df_meta_is_not_changed(self): + session = _guarded("disabled", "disabled") + session_mod.build_state_message(session) + assert session.df_meta == {"total_rows": 5} + def test_a_client_reads_what_is_left_out_as_enabled(self): read = getattr(session_mod, "guards_with_defaults", None) assert read is not None, "session.guards_with_defaults does not exist" From dbef9308d3f2f2eff13c58124971f278cdff1691 Mon Sep 17 00:00:00 2001 From: Paddy Mullen Date: Mon, 5 Oct 2026 07:32:59 -0400 Subject: [PATCH 38/39] test(server): failing tests for the stat_chunk_cells host opt-in on /load_expr (rows-first s6) /load_expr should accept stat_chunk_cells (a positive integer cell count, or null for off) and /reload_expr should replay it, so a host can turn on the xorq stats column-chunk split over HTTP. Tests cover the field reaching XorqServerDataflow.stat_chunk_cells, several batch units and unchanged stats on a parquet scan, one batch unit on a join, the warm short-circuit rules, 400 invalid_stat_chunk_cells, the reload replay, and /load and /load_compare dropping a stored value. Co-Authored-By: Claude Sonnet 5.5 --- tests/unit/server/test_load_compare.py | 21 ++ tests/unit/server/test_load_expr.py | 253 +++++++++++++++++++++++++ tests/unit/server/test_server.py | 16 ++ 3 files changed, 290 insertions(+) diff --git a/tests/unit/server/test_load_compare.py b/tests/unit/server/test_load_compare.py index aaea8b081..a0909feec 100644 --- a/tests/unit/server/test_load_compare.py +++ b/tests/unit/server/test_load_compare.py @@ -141,6 +141,27 @@ def test_load_compare_resolves_to_full_and_clears_a_stats_policy(self): self.assertEqual(session.stats_status, "complete") self.assertNotIn("stats", build_state_message(session)["df_meta"]) + def test_load_compare_drops_a_stat_chunk_cells_left_by_load_expr(self): + """A compare session runs no xorq stats batch, so the count a prior + /load_expr stored on the session is dropped (rows-first s6).""" + session = self._app.settings["sessions"].create("cmp-chunk-cells", "") + session.stat_chunk_cells = 60_000_000 + df1 = pd.DataFrame({"id": [1, 2], "score": [100, 200]}) + df2 = pd.DataFrame({"id": [1, 2], "score": [100, 999]}) + + with tempfile.NamedTemporaryFile(suffix=".csv", delete=False) as f1, \ + tempfile.NamedTemporaryFile(suffix=".csv", delete=False) as f2: + _write_df(df1, f1.name) + _write_df(df2, f2.name) + try: + resp = self._post_compare({"session": "cmp-chunk-cells", "path1": f1.name, "path2": f2.name, + "join_columns": ["id"]}) + self.assertEqual(resp.code, 200) + finally: + os.unlink(f1.name) + os.unlink(f2.name) + self.assertIsNone(session.stat_chunk_cells) + def test_missing_fields(self): resp = self._post_compare({"session": "x", "path1": "/a"}) self.assertEqual(resp.code, 400) diff --git a/tests/unit/server/test_load_expr.py b/tests/unit/server/test_load_expr.py index 27a10b824..ad99f7e0d 100644 --- a/tests/unit/server/test_load_expr.py +++ b/tests/unit/server/test_load_expr.py @@ -1385,6 +1385,259 @@ async def test_reload_expr_accepts_a_new_pair_and_stores_it(self): shutil.rmtree(project_root, ignore_errors=True) +def _build_scan_dir(data_root, builds_root, joined=False): + """Build the twelve-row, six-column parquet read of ``_wide_parquet_expr`` to + ``builds_root`` and return the build path. ``joined`` builds a join of it + with a second parquet read instead, an expression the batch split refuses.""" + expr = _wide_parquet_expr(data_root, "left" if joined else "wide") + if joined: + pd.DataFrame({"n0": [float(i) for i in range(12)], "extra": list(range(12))}).to_parquet( + data_root / "right.parquet") + expr = expr.join(xo.deferred_read_parquet(str(data_root / "right.parquet")), "n0") + return str(xo.build_expr(expr, builds_dir=builds_root)) + + +@contextmanager +def _spy_units(): + """Record the id of every unit ``XorqStatPipeline`` runs while the block runs.""" + ran = [] + original = XorqStatPipeline.run + + def spy(pipeline, unit, acc): + ran.append(unit.id) + return original(pipeline, unit, acc) + + with patch.object(XorqStatPipeline, "run", spy): + yield ran + + +def _batch_ids(ran): + return [unit_id for unit_id in ran if unit_id.startswith("batch")] + + +class TestLoadExprStatChunkCells(tornado.testing.AsyncHTTPTestCase): + """``stat_chunk_cells`` on POST /load_expr (rows-first s6): the host's opt-in + to cutting the xorq stats batch into column chunks, so that no single stat + unit holds the loop for a whole scan. It is stored beside the stats policy + fields, replayed by /reload_expr and kept out of the has_config tuple. The + pipeline's scan guard still decides which sources are cut. The fixture has + twelve rows and six columns, so 36 cells is three columns per chunk.""" + + def get_app(self): + return make_app() + + def setUp(self): + super().setUp() + data_root = Path(tempfile.mkdtemp()) + self.builds_root = tempfile.mkdtemp() + self.project_root = tempfile.mkdtemp() + for root in (data_root, self.builds_root, self.project_root): + self.addCleanup(shutil.rmtree, root, ignore_errors=True) + self.scan_build = _build_scan_dir(data_root, self.builds_root) + self.join_build = _build_scan_dir(data_root, self.builds_root, joined=True) + # Zero runs exactly one unit per incremental request. + budget = patch.object(stats_wire, "STATS_BUDGET_S", 0) + budget.start() + self.addCleanup(budget.stop) + 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 _post_load(self, sid, build_dir=None, **body): + return await _post(self.get_http_port(), "/load_expr", + {"session": sid, "build_dir": build_dir or self.scan_build, **body}) + + async def _load(self, sid, build_dir=None, **body): + resp = await self._post_load(sid, build_dir, **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) + + async def _pull(self, sid, build_dir=None, **body): + """Load ``sid`` with deferred stats and pull them one unit per request, + as a client with ``stats_update`` would. Returns the first + ``initial_state``, the ids of the units that ran and every reply.""" + await self._load(sid, build_dir, stats_delivery="deferred", **body) + ws, first = await self._connect(sid, caps="stats_update") + gen = first["df_meta"]["stats"]["gen"] + replies = [] + with _spy_units() as ran: + for _ in range(12): + ws.write_message(_stats_request(gen, incremental=True)) + replies.append(await _read_json(ws)) + if replies[-1]["type"] != "stats_update" or replies[-1]["final"]: + break + return first, ran, replies + + @tornado.testing.gen_test + async def test_the_field_reaches_the_dataflow_and_is_stored_beside_the_kwargs(self): + await self._load("cc-field", stat_chunk_cells=36) + session = self._session("cc-field") + self.assertEqual(session.xorq_dataflow.stat_chunk_cells, 36) + self.assertEqual(session.stat_chunk_cells, 36) + self.assertNotIn("stat_chunk_cells", session.dataflow_kwargs) + + @tornado.testing.gen_test + async def test_the_default_is_off_and_null_is_the_same_as_absent(self): + with _count_stat_queries() as queries: + await self._load("cc-default") + self.assertEqual(len(_batch_queries(queries)), 1) + await self._load("cc-null", stat_chunk_cells=None) + for sid in ("cc-default", "cc-null"): + self.assertIsNone(self._session(sid).stat_chunk_cells) + self.assertIsNone(self._session(sid).xorq_dataflow.stat_chunk_cells) + _, default = await self._connect("cc-default") + _, null = await self._connect("cc-null") + self.assertEqual(_comparable(default), _comparable(null)) + + @tornado.testing.gen_test + async def test_a_parquet_scan_runs_several_batch_units_and_assembles_the_unsplit_stats(self): + _, whole_ran, whole = await self._pull("cc-whole") + first, ran, replies = await self._pull("cc-split", stat_chunk_cells=36) + self.assertEqual(_batch_ids(whole_ran), ["batch"]) + self.assertEqual(_batch_ids(ran), ["batch:0", "batch:1"]) + self.assertEqual(len(ran), len(whole_ran) + 1) + self.assertEqual((first["df_meta"]["stats"]["status"], first["df_meta"]["stats"]["tier"]), + ("pending", "schema")) + self.assertEqual([r["type"] for r in replies], ["stats_update"] * len(ran)) + self.assertTrue(replies[-1]["final"]) + self.assertEqual(_as_json(_rows_by_stat(replies[-1]["payload"])), + _as_json(_rows_by_stat(whole[-1]["payload"]))) + self.assertEqual(_as_json(self._session("cc-split").xorq_dataflow.merged_sd), + _as_json(self._session("cc-whole").xorq_dataflow.merged_sd)) + + @tornado.testing.gen_test + async def test_inline_stats_are_cut_too_and_the_frame_is_unchanged(self): + await self._load("cc-inline-whole") + _, whole = await self._connect("cc-inline-whole") + with _count_stat_queries() as queries: + await self._load("cc-inline", stat_chunk_cells=36) + self.assertEqual(len(_batch_queries(queries)), 2) + _, frame = await self._connect("cc-inline") + self.assertEqual(_comparable(frame), _comparable(whole)) + + @tornado.testing.gen_test + async def test_a_join_keeps_the_single_batch_and_nothing_reports_a_problem(self): + _, _, whole = await self._pull("cc-join-whole", self.join_build) + first, ran, replies = await self._pull("cc-join", self.join_build, stat_chunk_cells=12) + self.assertEqual(self._session("cc-join").stat_chunk_cells, 12) + self.assertEqual(_batch_ids(ran), ["batch"]) + self.assertEqual(set(first["df_meta"]["stats"]), {"status", "tier", "gen"}) + self.assertEqual(first["df_meta"]["stats"]["status"], "pending") + for reply in replies: + self.assertEqual(reply["type"], "stats_update") + self.assertNotIn("reason", reply) + self.assertNotIn("status", reply) + self.assertTrue(replies[-1]["final"]) + self.assertEqual(_as_json(_rows_by_stat(replies[-1]["payload"])), + _as_json(_rows_by_stat(whole[-1]["payload"]))) + _, done = await self._connect("cc-join", caps="stats_update") + self.assertEqual(done["df_meta"]["stats"]["status"], "complete") + self.assertNotIn("reason", done["df_meta"]["stats"]) + + @tornado.testing.gen_test + async def test_invalid_values_are_a_400(self): + for value in (0, -36, 1.5, 36.0, True, False, "36", [36], {"cells": 36}): + resp = await self._post_load("cc-bad", stat_chunk_cells=value) + self.assertEqual(resp.code, 400, value) + body = json.loads(resp.body) + self.assertEqual(body["error_code"], "invalid_stat_chunk_cells", value) + self.assertIn("stat_chunk_cells", body["message"]) + self.assertIsNone(self._session("cc-bad")) + + @tornado.testing.gen_test + async def test_an_invalid_value_leaves_the_stored_one(self): + await self._load("cc-kept", stat_chunk_cells=36) + before = self._session("cc-kept").xorq_dataflow + resp = await self._post_load("cc-kept", stat_chunk_cells=0) + self.assertEqual(resp.code, 400) + session = self._session("cc-kept") + self.assertEqual(session.stat_chunk_cells, 36) + self.assertIs(session.xorq_dataflow, before) + + @tornado.testing.gen_test + async def test_warm_repost_with_an_unchanged_value_short_circuits(self): + body = {"stat_chunk_cells": 36} + await self._load("cc-warm", **body) + with patch.object(xorq_loading, "load_expr_build_dir", + side_effect=AssertionError("an unchanged value must take the warm exit")): + same = await self._post_load("cc-warm", **body) + omitted = await self._post_load("cc-warm") + self.assertEqual(same.code, 200) + self.assertEqual(omitted.code, 200, "omitting the field keeps the session's") + session = self._session("cc-warm") + self.assertEqual(session.stat_chunk_cells, 36) + self.assertEqual(session.xorq_dataflow.stat_chunk_cells, 36) + + @tornado.testing.gen_test + async def test_the_off_value_sent_on_every_post_does_not_defeat_the_warm_exit(self): + """has_config tests truthiness; the field is not in it, so a host that + always sends null (or a count) does not rebuild on every POST (#944).""" + await self._load("cc-warm-null", stat_chunk_cells=None) + with patch.object(xorq_loading, "load_expr_build_dir", + side_effect=AssertionError("the warm exit must not see the field as config")): + resp = await self._post_load("cc-warm-null", stat_chunk_cells=None) + self.assertEqual(resp.code, 200) + + @tornado.testing.gen_test + async def test_warm_repost_with_a_changed_value_rebuilds(self): + 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 self._load("cc-changed") + self.assertEqual(len(calls), 1) + await self._load("cc-changed", stat_chunk_cells=36) + self.assertEqual(len(calls), 2, "turning the split on must rebuild") + self.assertEqual(self._session("cc-changed").xorq_dataflow.stat_chunk_cells, 36) + await self._load("cc-changed", stat_chunk_cells=24) + self.assertEqual(len(calls), 3, "a different count must rebuild") + await self._load("cc-changed", stat_chunk_cells=None) + self.assertEqual(len(calls), 4, "an explicit null turns the split off, which must rebuild") + session = self._session("cc-changed") + self.assertIsNone(session.stat_chunk_cells) + self.assertIsNone(session.xorq_dataflow.stat_chunk_cells) + + @tornado.testing.gen_test + async def test_reload_expr_replays_the_stored_value(self): + await self._load("cc-reload", project_root=self.project_root, stat_chunk_cells=36) + before = self._session("cc-reload").xorq_dataflow + with _count_stat_queries() as queries: + resp = await _post(self.get_http_port(), "/reload_expr/cc-reload", {}) + self.assertEqual(resp.code, 200, resp.body) + session = self._session("cc-reload") + self.assertIsNot(session.xorq_dataflow, before) + self.assertEqual(session.xorq_dataflow.stat_chunk_cells, 36) + self.assertEqual(session.stat_chunk_cells, 36) + self.assertEqual(len(_batch_queries(queries)), 2, "the rebuilt dataflow must split its batch") + + @tornado.testing.gen_test + async def test_reload_expr_without_a_stored_value_builds_the_single_batch(self): + await self._load("cc-reload-off", project_root=self.project_root) + with _count_stat_queries() as queries: + resp = await _post(self.get_http_port(), "/reload_expr/cc-reload-off", {}) + self.assertEqual(resp.code, 200, resp.body) + self.assertIsNone(self._session("cc-reload-off").xorq_dataflow.stat_chunk_cells) + self.assertEqual(len(_batch_queries(queries)), 1) + + 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 diff --git a/tests/unit/server/test_server.py b/tests/unit/server/test_server.py index 5d4c1f87d..f8fc2d7f7 100644 --- a/tests/unit/server/test_server.py +++ b/tests/unit/server/test_server.py @@ -267,6 +267,22 @@ def test_load_does_not_take_stats_tier_and_resolves_to_full(self): finally: os.unlink(f.name) + def test_load_drops_a_stat_chunk_cells_left_by_load_expr(self): + """The count belongs to the xorq stats batch, which /load does not run: a + session that /load_expr gave one starts over (rows-first s6).""" + with tempfile.NamedTemporaryFile(suffix=".csv", delete=False) as f: + _write_test_csv(f.name) + try: + session = self._app.settings["sessions"].create("chunk-cells-1", "") + session.stat_chunk_cells = 60_000_000 + resp = self.fetch("/load", method="POST", + body=json.dumps({"session": "chunk-cells-1", "path": f.name, "mode": "buckaroo"}), + headers={"Content-Type": "application/json"}) + self.assertEqual(resp.code, 200) + self.assertIsNone(session.stat_chunk_cells) + finally: + os.unlink(f.name) + class TestSessionPage(tornado.testing.AsyncHTTPTestCase): def get_app(self): From 852713168e739b6f6f4e3f3d7aa84c30e4cb7334 Mon Sep 17 00:00:00 2001 From: Paddy Mullen Date: Mon, 5 Oct 2026 07:43:05 -0400 Subject: [PATCH 39/39] feat(server): let a host enable the xorq stats column-chunk split over /load_expr (rows-first s6) /load_expr accepts stat_chunk_cells, a positive integer cell count or null for off, and passes it to XorqServerDataflow, whose scan guard still keeps a join, aggregate, diff, CSV read or registered table on the single batch. It is stored on the session beside stats_tier and stats_delivery and outside has_config, so a warm re-POST with the same value short-circuits and a changed value (null included) rebuilds; omitting the field keeps the stored value. /reload_expr replays it, and /load and /load_compare reset it. A zero, negative, non-integer or bool value answers 400 invalid_stat_chunk_cells. Co-Authored-By: Claude Sonnet 5.5 --- buckaroo/server/handlers.py | 46 +++++++++++++++++++++++++++++++------ buckaroo/server/session.py | 6 +++++ 2 files changed, 45 insertions(+), 7 deletions(-) diff --git a/buckaroo/server/handlers.py b/buckaroo/server/handlers.py index 161f4bea5..60440c7bf 100644 --- a/buckaroo/server/handlers.py +++ b/buckaroo/server/handlers.py @@ -373,6 +373,7 @@ async def post(self): # 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, session.stats_policy = "full", "inline", None + session.stat_chunk_cells = None session.source_guards = None reset_stats_controls(session) begin_stats_generation(session) @@ -409,6 +410,26 @@ def _stats_policy_from_body(body: dict, current_tier: str, current_delivery: str return stats_tier, stats_delivery, None +def _stat_chunk_cells_from_body(body: dict, current: int | None): + """Read ``stat_chunk_cells`` from a request body: the cell count the xorq + stats batch is cut by, or ``None`` for off. + + A body that omits the field keeps ``current`` (the session's, or ``None`` + for a new session), as the other stats fields do. An explicit ``null`` turns + the split off. Returns ``(stat_chunk_cells, None)``, or ``(None, error)`` + with the 400 response body for anything that is not a positive integer (a + bool and an integral float included).""" + if "stat_chunk_cells" not in body: + return current, None + stat_chunk_cells = body["stat_chunk_cells"] + if stat_chunk_cells is None: + return None, None + if isinstance(stat_chunk_cells, bool) or not isinstance(stat_chunk_cells, int) or stat_chunk_cells <= 0: + return None, {"error_code": "invalid_stat_chunk_cells", + "message": f"stat_chunk_cells must be a positive integer cell count or null, got {stat_chunk_cells!r}"} + return stat_chunk_cells, 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. @@ -467,10 +488,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. 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. + # ignore the new config. stats_tier, stats_delivery and stat_chunk_cells + # are not in this tuple: it tests truthiness, and a host that sends + # them on every POST would never get the warm exit (#944). They are + # compared with the session's stored values below instead. has_config = any(body.get(k) for k in ( "component_config", "column_config_overrides", "extra_grid_config", "init_sd", "skip_stat_columns")) @@ -510,6 +531,12 @@ async def post(self): self.set_status(400) self.write(policy_error) return + stat_chunk_cells, chunk_error = _stat_chunk_cells_from_body( + body, existing.stat_chunk_cells if existing is not None else None) + if chunk_error is not None: + self.set_status(400) + self.write(chunk_error) + return # What a client forced or the cost guard paused (SessionState.stats_override, # cost_paused) belongs to one expression: a rebuild of the same one keeps it, # and so does a re-POST that names no stats_tier. @@ -520,7 +547,8 @@ async def post(self): 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.stats_tier, existing.stats_delivery) == (stats_tier, stats_delivery)): + and (existing.stats_tier, existing.stats_delivery) == (stats_tier, stats_delivery) + and existing.stat_chunk_cells == stat_chunk_cells): # 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 @@ -602,7 +630,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_tier, **dataflow_kwargs) + stats_tier=dataflow_tier, stat_chunk_cells=stat_chunk_cells, **dataflow_kwargs) # Spanning metadata too leaves only the small klass-load step # unmeasured inside the outer firstpull.load_expr total. with perf_log.perf_span("firstpull.metadata", session=session_id): @@ -638,6 +666,7 @@ async def post(self): session.dataflow_kwargs = dataflow_kwargs session.stats_tier = stats_tier session.stats_delivery = stats_delivery + session.stat_chunk_cells = stat_chunk_cells session.stats_policy = stats_policy if not same_expression or body.get("stats_tier") is not None: reset_stats_controls(session) @@ -832,6 +861,7 @@ async def post(self): session.mode = "viewer" # A viewer session has no dataflow and so no deferred stats. session.stats_tier, session.stats_delivery, session.stats_policy = "full", "inline", None + session.stat_chunk_cells = None session.source_guards = None reset_stats_controls(session) begin_stats_generation(session) @@ -870,6 +900,8 @@ class ReloadExprHandler(tornado.web.RequestHandler): 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. The stats policy is resolved again for the pair. + The stored ``stat_chunk_cells`` is replayed as it is: the body has no field + for it. A forced tier and a cost pause (``SessionState.stats_override``, ``cost_paused``) are kept, unless the body names a ``stats_tier``, which starts the session's stats decisions over. @@ -933,7 +965,7 @@ async def post(self, session_id): dataflow_tier = dataflow_stats_tier(stats_tier, stats_delivery) xorq_dataflow = xorq_loading.XorqServerDataflow( session.expr, skip_main_serial=True, extra_klasses=extra_klasses, - stats_tier=dataflow_tier, **session.dataflow_kwargs) + stats_tier=dataflow_tier, stat_chunk_cells=session.stat_chunk_cells, **session.dataflow_kwargs) # Resolved again, with the count the load took: the expression is # reused, so there is no new one to take. stats_policy = resolve_session_policy(stats_tier, dataflow_tier, session.metadata["rows"], diff --git a/buckaroo/server/session.py b/buckaroo/server/session.py index 1cfa16210..5aa7ad1cc 100644 --- a/buckaroo/server/session.py +++ b/buckaroo/server/session.py @@ -131,6 +131,12 @@ class SessionState: # dataflow is built at follows from the pair (dataflow_stats_tier). stats_tier: str = "full" stats_delivery: str = "inline" + # The /load_expr host opt-in to cutting the xorq stats batch into column + # chunks of about this many cells (``None`` is off), so that no single stat + # unit holds the loop for a whole scan. Stored with the pair above, not in + # dataflow_kwargs, and replayed by /reload_expr; the dataflow's scan guard + # still decides which sources are cut. + stat_chunk_cells: Optional[int] = None # What ``stats_policy.resolve_stats_policy`` made of the pair and the size of # the entry (``tier_target``, ``auto_request``, ``requestable``, ``reason``, # ``estimate``), resolved by /load_expr and /reload_expr once the schema-tier