Merge #808: @stat(cost=...) decorator + tag known-expensive stats

Author: paddymul

buckaroo/customizations/pd_stats_v2.py

@@ -176,7 +176,7 @@ def vc_nth(pos):
 HistogramSeriesResult = TypedDict('HistogramSeriesResult', {'histogram_args': dict, 'histogram_bins': list})
 
 
-@stat()
+@stat(cost="aggregate")
 def histogram_series(ser: RawSeries) -> HistogramSeriesResult:
     """Compute histogram args from raw series (numeric path)."""
     if not pd.api.types.is_numeric_dtype(ser):
@@ -210,7 +210,7 @@ def histogram_series(ser: RawSeries) -> HistogramSeriesResult:
     }
 
 
-@stat()
+@stat(cost="aggregate")
 def histogram(value_counts: pd.Series, nan_per: float, is_numeric: bool, length: int, min: Any, max: Any,
         histogram_args: dict) -> list:
     """Compute histogram from summary stats and histogram args."""

buckaroo/customizations/pl_stats_v2.py

@@ -113,7 +113,7 @@ def pl_numeric_stats(ser: RawSeries) -> NumericStatsResult:
 # Histogram Series (polars series API)
 # ============================================================
 
-@stat()
+@stat(cost="aggregate")
 def pl_histogram_series(ser: RawSeries) -> HistogramSeriesResult:
     """Compute histogram args from raw polars series (numeric path)."""
     if not ser.dtype.is_numeric():

buckaroo/customizations/xorq_stats_v2.py

@@ -255,7 +255,7 @@ def _categorical_histogram(execute: Callable[[Any], pd.DataFrame], expr: Any, co
     return out
 
 
-@stat(default=[])
+@stat(default=[], cost="aggregate")
 def histogram(expr: XorqExpr, execute: XorqExecute, orig_col_name: str, is_numeric: bool, is_bool: bool, length: int,
         distinct_count: int, min: float, max: float) -> list:
     """10-bucket numeric histogram or top-10 categorical histogram.

buckaroo/pluggable_analysis_framework/stat_func.py

@@ -132,6 +132,9 @@ def __repr__(self):
 # StatFunc — a registered stat computation
 # ---------------------------------------------------------------------------
 
+VALID_COSTS = ("scalar", "aggregate")
+
+
 @dataclass
 class StatFunc:
     """A registered stat computation.
@@ -149,6 +152,11 @@ class StatFunc:
             engines can compute this stat via aggregation push-down without
             in-process materialization. Empty default means pandas-only /
             requires materialization.
+        cost: cost class — ``"scalar"`` (cheap, ships in the initial
+            state_change response) or ``"aggregate"`` (slow path that the
+            JS orchestrator fetches via a separate round-trip after a
+            debounce). Histograms, value_counts and other per-column
+            queries belong in ``"aggregate"``.
     """
     name: str
     func: Callable
@@ -159,6 +167,7 @@ class StatFunc:
     quiet: bool = False
     default: Any = field(default_factory=lambda: MISSING)
     pushdown: Tuple[str, ...] = ()
+    cost: str = "scalar"
     spread_dict_result: bool = False  # v1 compat: spread all dict keys into accumulator
     v1_computed: bool = False  # v1 compat: pass full accumulator as single dict arg
 
@@ -255,7 +264,7 @@ def _get_requires_from_params(sig: inspect.Signature, hints: dict) -> tuple:
 # @stat decorator
 # ---------------------------------------------------------------------------
 
-def stat(column_filter=None, quiet=False, default=MISSING, pushdown=()):
+def stat(column_filter=None, quiet=False, default=MISSING, pushdown=(), cost="scalar"):
     """Decorator that converts a function into a StatFunc.
 
     The function signature IS the contract:
@@ -274,6 +283,12 @@ def stat(column_filter=None, quiet=False, default=MISSING, pushdown=()):
     identifiers: ``"xorq"``, ``"polars"``. Normalised to a tuple so callers
     may pass a list.
 
+    ``cost=`` declares the compute-cost class. ``"scalar"`` (default)
+    means cheap — ships in the initial state_change response. ``"aggregate"``
+    means slow (histograms, value_counts, per-column queries) — the JS
+    orchestrator fetches these via a separate round-trip after a debounce.
+    See plans/js-driven-stat-debounce.md.
+
     Usage::
 
         @stat()
@@ -292,6 +307,10 @@ def safe_ratio(a: int, b: int) -> float:
         def mean(col):
             return col.mean()
 
+        @stat(cost="aggregate")
+        def histogram(...) -> list:
+            ...
+
         class TypingResult(MultipleProvides):
             is_numeric: bool
             is_integer: bool
@@ -300,6 +319,10 @@ class TypingResult(MultipleProvides):
         def typing_stats(dtype: str) -> TypingResult:
             ...
     """
+    if cost not in VALID_COSTS:
+        raise ValueError(
+            f"@stat(cost={cost!r}): invalid cost class. "
+            f"Must be one of {VALID_COSTS}.")
     def decorator(func):
         sig = inspect.signature(func)
         try:
@@ -317,7 +340,7 @@ def decorator(func):
         pushdown_norm = (pushdown,) if isinstance(pushdown, str) else tuple(pushdown)
         stat_func = StatFunc(name=func.__name__, func=func, requires=requires, provides=provides_keys,
             needs_raw=needs_raw, column_filter=column_filter, quiet=quiet, default=default,
-            pushdown=pushdown_norm)
+            pushdown=pushdown_norm, cost=cost)
 
         # Attach metadata to the function so pipeline can find it
         func._stat_func = stat_func

tests/unit/test_paf_v2.py

@@ -204,6 +204,49 @@ def pushdown_count(ser: RawSeries) -> int:
 
         assert pushdown_count._stat_func.pushdown == ("xorq",)
 
+    def test_stat_default_cost_is_scalar(self):
+        # Default cost class is "scalar" — the bulk of stats are cheap.
+        # Only known-expensive ones opt in to "aggregate".
+        sf = length._stat_func
+        assert sf.cost == "scalar"
+
+    def test_stat_explicit_cost_aggregate(self):
+        @stat(cost="aggregate")
+        def big_compute(ser: RawSeries) -> float:
+            return float(ser.mean())
+
+        assert big_compute._stat_func.cost == "aggregate"
+
+    def test_stat_invalid_cost_rejected(self):
+        # Only "scalar" and "aggregate" are recognised. A typo should
+        # fail loud at decoration time — silently dropping an invalid
+        # cost would leak into the cost-class router as an unscheduled
+        # stat group.
+        import pytest as _pt
+        with _pt.raises(ValueError, match="cost"):
+            @stat(cost="bigly")
+            def bad(ser: RawSeries) -> float:
+                return float(ser.mean())
+
+    def test_known_expensive_stats_marked_aggregate(self):
+        """The expensive built-in stat funcs (histogram producers across
+        all three engines) are tagged ``cost="aggregate"`` so a
+        downstream router can schedule them on the slow path."""
+        from buckaroo.customizations.pd_stats_v2 import histogram as pd_histogram
+        from buckaroo.customizations.pd_stats_v2 import histogram_series as pd_hs
+        from buckaroo.customizations.pl_stats_v2 import pl_histogram_series
+
+        assert pd_histogram._stat_func.cost == "aggregate"
+        assert pd_hs._stat_func.cost == "aggregate"
+        assert pl_histogram_series._stat_func.cost == "aggregate"
+
+        # xorq histogram (optional dep — skip if not installed)
+        try:
+            from buckaroo.customizations.xorq_stats_v2 import histogram as xq_histogram
+        except ImportError:
+            return
+        assert xq_histogram._stat_func.cost == "aggregate"
+
 
 class _MultiSizeStats(MultipleProvides):
     row_count: int