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/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..ea1650d9d 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 (
@@ -27,6 +28,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.
@@ -82,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]):
@@ -161,6 +226,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 +381,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 +391,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 +400,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 +460,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 = {}
@@ -491,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).
@@ -520,43 +604,17 @@ 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
+ 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
- # 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
+ return assemble_merged_sd(
+ self.init_sd, self.cleaned_sd, raw_sd, self.processed_sd, self.processed_df,
+ 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.
@@ -598,14 +656,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 +675,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 +718,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.
@@ -708,14 +782,25 @@ def _build_error_dataframe(self, e) -> pd.DataFrame:
### start summary stats block
- #TAny closer to some error type
- @override
- def _get_summary_sd(self, processed_df: DataFrameT) -> Tuple[SDType, ErrDict]:
- stats = self.DFStatsClass(
+ 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))
+ 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.build_stats(processed_df)
sdf = stats.sdf
if stats.errs:
if self.debug:
@@ -728,6 +813,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 +840,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 +884,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 +910,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/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 67e8c0f77..22197f870 100644
--- a/buckaroo/pluggable_analysis_framework/stat_pipeline.py
+++ b/buckaroo/pluggable_analysis_framework/stat_pipeline.py
@@ -9,16 +9,18 @@
from __future__ import annotations
import time
-from typing import Any, Dict, List, Optional, Tuple
+from typing import Any, Callable, Dict, List, Optional, Sequence, 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
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
@@ -172,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:
@@ -276,9 +278,59 @@ 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. 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)
+ skip = self._skipped(state, pairs)
+ 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:
+ """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
@@ -292,35 +344,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."""
@@ -432,3 +464,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/pluggable_analysis_framework/stat_units.py b/buckaroo/pluggable_analysis_framework/stat_units.py
new file mode 100644
index 000000000..611bb16ce
--- /dev/null
+++ b/buckaroo/pluggable_analysis_framework/stat_units.py
@@ -0,0 +1,230 @@
+"""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"
+
+
+# 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``.
+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
+ 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. ``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()
+ columns: Optional[Tuple[Any, ...]] = None
+ 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]]:
+ """``(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(resolve_names(pairs, state.columns, state.namespace))
+ return [pair for pair in pairs if pair[0] in wanted]
+
+
+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, 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:
+ """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..cec1a672a 100644
--- a/buckaroo/pluggable_analysis_framework/xorq_stat_pipeline.py
+++ b/buckaroo/pluggable_analysis_framework/xorq_stat_pipeline.py
@@ -1,14 +1,21 @@
"""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.
+
+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
@@ -23,7 +30,7 @@
import os
import time
from contextlib import nullcontext
-from typing import Any, Dict, List, Tuple
+from typing import Any, Dict, Iterable, List, Optional, Sequence, Tuple
import pandas as pd
@@ -33,6 +40,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 +50,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 +111,32 @@ 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 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] = (), 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] = {}
+ 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() if key not in self.withheld} 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 +146,16 @@ 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``).
+
+ ``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
@@ -122,8 +167,17 @@ class XorqStatPipeline:
{"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):
+ cache_storage=None, chunk_cells: Optional[int] = None):
if not HAS_XORQ:
raise ImportError(
"xorq is required for XorqStatPipeline. "
@@ -138,12 +192,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 +379,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,39 +404,189 @@ 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 _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. 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()
+ 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[(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. 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:
+ 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(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] = {}
+ 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 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"))
+ 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 (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
+ 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],
+ 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)}
-
- # ---- 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)
+ 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:
+ """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 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] 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)
+ return plain
+
+ 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
# per-column expressions.
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:
- 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 +594,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 +604,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 +620,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], tier)
+ 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 +694,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 +709,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 +721,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 +756,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/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`.
diff --git a/buckaroo/server/handlers.py b/buckaroo/server/handlers.py
index c60d99062..60440c7bf 100644
--- a/buckaroo/server/handlers.py
+++ b/buckaroo/server/handlers.py
@@ -14,7 +14,11 @@
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.server.session import (
+ STATS_DELIVERIES, STATS_TIER_REQUESTS, begin_stats_generation, dataflow_stats_tier, reset_stats_controls)
+from buckaroo.server.stats_wire import (
+ 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
@@ -192,17 +196,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."""
@@ -371,6 +368,16 @@ async def post(self):
**component_config,
}
+ # /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
+ session.stat_chunk_cells = None
+ session.source_guards = None
+ reset_stats_controls(session)
+ 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)
@@ -380,6 +387,49 @@ 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`` (``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
+ ``(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_TIER_REQUESTS:
+ return None, None, {"error_code": "invalid_stats_tier",
+ "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
+ 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
+
+
+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.
@@ -438,7 +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.
+ # 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"))
@@ -468,11 +521,34 @@ 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
+ 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.
+ 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
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)
+ 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
@@ -525,6 +601,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
@@ -553,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,
- **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):
@@ -569,6 +646,15 @@ 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"]))
+ # 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)
session.mode = "buckaroo"
@@ -578,6 +664,14 @@ 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.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)
+ # 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
@@ -619,17 +713,16 @@ 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)
+ # 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)
+
+ # 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"
@@ -766,18 +859,15 @@ 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, session.stats_policy = "full", "inline", None
+ session.stat_chunk_cells = None
+ session.source_guards = None
+ reset_stats_controls(session)
+ 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):
@@ -807,8 +897,27 @@ 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. 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.
+
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 +940,14 @@ 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(
+ 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:
@@ -845,9 +962,16 @@ 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,
- **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"],
+ 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)
@@ -872,29 +996,17 @@ 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.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}
-
- 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)
+ 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)
+ apply_source_guards(session, xorq_dataflow, source_guards)
+ refresh_session_snapshot(session, xorq_dataflow)
+ begin_stats_generation(session)
+
+ 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 563ce28fc..5aa7ad1cc 100644
--- a/buckaroo/server/session.py
+++ b/buckaroo/server/session.py
@@ -1,9 +1,14 @@
+import copy
import logging
import time
from dataclasses import dataclass, field
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")
@@ -11,6 +16,82 @@
_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")
+
+# 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, 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": []}
+
+# 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
+ 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`` (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"
+
+
+# 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 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")
+
+
+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, 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
+ return "complete", None
+
@dataclass
class SessionState:
@@ -43,6 +124,63 @@ 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 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"
+ # 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
+ # dataflow and the count exist. ``None`` for a dataflow built with its stats
+ # 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
+ # 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).
+ # 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
+ # 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
+ # 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
@@ -80,8 +218,181 @@ 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, 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, effective_stats_policy(session), session.cost_paused)
+ session.stat_runs.clear()
+
+
+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
+ ``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_with_defaults``)."""
+ if session.stats_status == "complete" and session.stats_delivery != "deferred":
+ return None
+ 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 = 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
+
+
+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 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 = "") -> dict:
+ search_string: str = "", ondemand: bool = True) -> dict:
"""Build the full ``initial_state`` WebSocket payload from a session.
Args:
@@ -94,14 +405,23 @@ 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.
"""
+ # 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)
+ 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,
- "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/stat_run.py b/buckaroo/server/stat_run.py
new file mode 100644
index 000000000..ff5682ce9
--- /dev/null
+++ b/buckaroo/server/stat_run.py
@@ -0,0 +1,166 @@
+"""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.
+
+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
+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, 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. ``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
+ inside a unit is not that: it is an entry in ``acc.errors`` and the unit
+ 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, state: Optional[StatState] = None):
+ self.id = next(_run_ids)
+ self.stats_gen = stats_gen
+ self.scope = scope
+ self.stats = stats
+ 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"
+ 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
+ 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:
+ """The number of units not yet run."""
+ return len(self.units) - len(self.ran)
+
+ 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``, 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:
+ 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 for col in unit.columns):
+ return unit
+ return ready[0] if ready else None
+
+ 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. ``prefer`` and
+ ``namespace`` are as in ``next_unit``."""
+ if self.status != "pending":
+ return None
+ unit = self.next_unit(prefer, namespace)
+ 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:
+ 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)
+
+ 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
+ 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_policy.py b/buckaroo/server/stats_policy.py
new file mode 100644
index 000000000..25b44c4a1
--- /dev/null
+++ b/buckaroo/server/stats_policy.py
@@ -0,0 +1,435 @@
+"""Which stats tier an entry should reach, decided from its size.
+
+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:
+
+* ``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``.
+
+: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
+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, 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
+
+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 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
+
+# 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")
+
+# 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
+
+
+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
+ :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
+ # 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:
+ """The defaults with any ``BUCKAROO_*`` overrides applied."""
+ 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:
+ 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 and cells <= limits.full_auto_cells:
+ 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"
+ 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,
+ 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.
+ ``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:
+ 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.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``,
+ 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)
+ if limits is None:
+ limits = StatsLimits.from_env()
+ 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
+ ``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.
+
+ 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)
diff --git a/buckaroo/server/stats_wire.py b/buckaroo/server/stats_wire.py
new file mode 100644
index 000000000..c0a315650
--- /dev/null
+++ b/buckaroo/server/stats_wire.py
@@ -0,0 +1,836 @@
+"""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
+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.
+
+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.
+
+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``).
+
+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.
+
+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.
+
+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 dataclasses
+import hashlib
+import json
+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
+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, 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_source_guards, resolve_stats_policy
+
+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 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
+# 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
+
+# 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."""
+ 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 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 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, 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")
+
+
+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 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
+ 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"]
+ 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}
+ 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",
+ 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 = stat_run_key(session.stats_gen, scope, tier, columns)
+ run = session.stat_runs.get(key)
+ if run is None:
+ 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
+
+
+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,
+ 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
+ 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. ``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":
+ 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
+ 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
+
+
+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``.
+
+ 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
+ key = _full_stats_key(dataflow)
+ dataflow.stats_tier = "full"
+ try:
+ sd = dataflow.summary_stats_cache.get(key)
+ errs = {}
+ if sd is None:
+ 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
+ except Exception:
+ dataflow.stats_tier = tier
+ 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). 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")
+ 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 (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.
+
+ 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
+ generation starts clean."""
+ if session.stats_status == "complete":
+ return True
+ dataflow = session_dataflow(session)
+ 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 (
+ 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) as span,
+ ):
+ try:
+ 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())
+ _fail_stats(session)
+ return False
+ session.stat_runs.pop(run_key, None)
+ 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 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 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 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:
+ """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 _elapsed_ms(started: float) -> float:
+ return round((time.perf_counter() - started) * 1000, 1)
+
+
+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, 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())
+ 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,
+ 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())
+ _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": "full",
+ "final": False, "remaining": run.remaining, "payload": partial_payload(dataflow, run, fragments),
+ "elapsed_ms": _elapsed_ms(started)}
+
+
+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.
+
+ 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, 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, timings=info.setdefault("timings", []))
+ except Exception:
+ 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)
+ 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 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"]
+ 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)
+
+
+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]:
+ """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")
+ 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 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_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)
+ if to_pull and incremental:
+ reply = _serve_units(session, client, prefer, stats_gen, scope, started, info)
+ if reply is not None:
+ return reply
+ if to_pull:
+ 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.
+ 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, 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. On a session whose target is
+ ``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
+ 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()
+ 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, 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)
+ _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/buckaroo/server/websocket_handler.py b/buckaroo/server/websocket_handler.py
index 2232211c7..39d9254b7 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
@@ -9,8 +8,10 @@
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, 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)
def _handle_infinite_request_xorq(xorq_dataflow, payload_args, search_string=""):
@@ -38,6 +39,18 @@ 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", ""))
+ # 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)
@@ -45,7 +58,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 +72,24 @@ 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 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
+ 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)
+ self.write_message(json.dumps(reply))
def _handle_buckaroo_state_change(self, new_state):
sessions = self.application.settings["sessions"]
@@ -99,6 +130,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 +145,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,23 +186,16 @@ def _send_highlight_overlay(self, session):
if not session.df_display_args:
return
term = self.search_string
- 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)
-
+ # 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(session, search_string=self.search_string)
+ msg = build_state_message_for(session, self)
+ overlay = highlighted_display_args(session.df_display_args, term)
+
msg["df_display_args"] = overlay
try:
self.write_message(json.dumps(msg))
@@ -209,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 60a9226cf..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 (
@@ -64,23 +68,50 @@ 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``).
+
+ ``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,)
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)
+ 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.
@@ -556,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.
@@ -567,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/buckaroo/xorq_buckaroo.py b/buckaroo/xorq_buckaroo.py
index 6b9149c40..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
@@ -27,13 +27,14 @@
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 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 = {
@@ -173,15 +180,10 @@ def _get_summary_sd(self, processed_df: "XorqExpr | pd.DataFrame"):
'orig_col_name': orig_col,
'rewritten_col_name': rewritten_col}
return empty, {}
- 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)
+ if self.stats_tier == "schema":
+ return self._get_schema_sd(processed_df), {}
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
@@ -205,13 +207,44 @@ 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
+ 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. 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)
+ 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,
+ **typing, 'length': length}
+ return sd
_XORQ_ANALYSIS_KLASSES = list(XORQ_STATS_V2) + [DefaultSummaryStatsStyling, DefaultMainStyling]
diff --git a/tests/unit/basic_widget_test.py b/tests/unit/basic_widget_test.py
index a028751ed..8852b8cdd 100644
--- a/tests/unit/basic_widget_test.py
+++ b/tests/unit/basic_widget_test.py
@@ -288,6 +288,43 @@ 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']) 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 = []
+ 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/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..bf908c6a4 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,7 +25,10 @@
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.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})
@@ -249,3 +254,98 @@ 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
+
+
+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."""
+ 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 58f06f5bb..45fbd3078 100644
--- a/tests/unit/server/test_data_loading_polars.py
+++ b/tests/unit/server/test_data_loading_polars.py
@@ -9,15 +9,44 @@
* ``.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 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.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 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 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)
+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)
def _payload(start=0, end=100):
@@ -89,3 +118,826 @@ 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"}
+
+
+# ---------------------------------------------------------------------------
+# 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 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):
+ 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):
+ (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 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
+ ``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):
+ dataflow = _build_dataflow(backend, _frame(backend, _UNTIED_DATA))
+ 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
+ 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_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:
+ 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):
+ 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):
+ 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):
+ 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):
+
+ 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}}]
+
+ 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."""
+
+ 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 _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)}
+
+
+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"], 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()
+ 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 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):
+ """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(_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),
+ 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_compare.py b/tests/unit/server/test_load_compare.py
index ba9be9198..a0909feec 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,53 @@ 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_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 bdbb47db6..ad99f7e0d 100644
--- a/tests/unit/server/test_load_expr.py
+++ b/tests/unit/server/test_load_expr.py
@@ -1,17 +1,25 @@
"""End-to-end tests for POST /load_expr — server load path for
XorqBuckarooInfiniteWidget over a xorq/ibis expression."""
+import dataclasses
+import datetime
+import gc
import io
import json
import os
import shutil
import sys
import tempfile
+import time
+import weakref
+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
@@ -24,9 +32,24 @@
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 buckaroo.server import telemetry, xorq_loading # 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
+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 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 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
+ _OverridingPostProcessing, _scope_inputs, _scope_sds_by_units)
pytestmark = pytest.mark.skipif(
sys.platform == "win32",
@@ -910,6 +933,4114 @@ 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_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")
+ 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)
+
+
+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
+ 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
+
+
+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 in ("index", "level_0") 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
+ 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_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
+ 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"]))
+
+
+ @tornado.testing.gen_test
+ async def test_a_stat_run_is_started_on_the_session_for_the_current_generation(self):
+ 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):
+ 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):
+ 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):
+ 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):
+ 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 _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 not in ("index", "level_0")}
+ 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")
+ 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)
+ 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")
+
+ @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, {})
+
+
+# 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_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")
+ 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")
+ # 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):
+ 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({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
+ 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))
+ 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)
+
+
+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):
+ 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):
+ 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)
+
+
+# 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_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_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")
+ 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())
+
+ 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)
+ 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)
+
+
+# ---------------------------------------------------------------------------
+# 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_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):
+ 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.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)
+ # 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):
+ """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_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")
+ 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."""
+ 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_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,
+ 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_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)
+ 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"))])
+
+ @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):
+ 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))
+
+ @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))
+
+
+# ---------------------------------------------------------------------------
+# 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, 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. ``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", rows), 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_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):
+ 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, rows=3, 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)
+
+ @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
+ 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
+
+ @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"
+
+
+ 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."""
+
+ 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_server.py b/tests/unit/server/test_server.py
index b1b8f9032..f8fc2d7f7 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,46 @@ 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)
+
+ 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):
diff --git a/tests/unit/server/test_stats_policy.py b/tests/unit/server/test_stats_policy.py
new file mode 100644
index 000000000..8ef6c0192
--- /dev/null
+++ b/tests/unit/server/test_stats_policy.py
@@ -0,0 +1,1649 @@
+"""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 last
+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 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
+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
+import json
+import logging
+from types import SimpleNamespace
+
+import numpy as np
+import pandas as pd
+import polars as pl
+import pyarrow as pa
+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
+
+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}
+
+
+@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 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):
+ 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.
+# 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 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)),
+ ("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_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 on full: 25M rows ...
+ ("xorq", "parquet", 12_000_000, 43, "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")),
+]
+
+
+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(13_000_000), np.int64(43))
+ assert out["tier_target"] == "scalar"
+ 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):
+ 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")
+
+ 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 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."""
+
+ 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", 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"
+ 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"
+
+
+ @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"
+
+
+ @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()
+
+ @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 == 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")
+ 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_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_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"
+ 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"
+
+ @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"
+
+ 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)
+
+
+# ---------------------------------------------------------------------------
+# 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("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):
+ 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
+
+
+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
+
+
+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
+
+
+# ---------------------------------------------------------------------------
+# 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")), []) == []
+
+
+# ---------------------------------------------------------------------------
+# 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
+
+ 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")
+ 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_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
+ 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("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."""
+ 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_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"
+ 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_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)
+ 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
diff --git a/tests/unit/skip_columns_test.py b/tests/unit/skip_columns_test.py
index bf0121be5..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
@@ -61,6 +62,63 @@ 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():
+ 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
+ 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.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 6ec2c252c..8f07c4938 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
@@ -736,3 +737,243 @@ 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')]
+
+
+# ============================================================================
+# 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):
+ 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):
+ 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):
+ 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):
+ 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):
+
+ @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."""
+ 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):
+ 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):
+ 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):
+ 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 == []
+
+ # 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'))
+
+ @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_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
diff --git a/tests/unit/test_xorq_stats_v2.py b/tests/unit/test_xorq_stats_v2.py
index 3565ef385..9f5c35882 100644
--- a/tests/unit/test_xorq_stats_v2.py
+++ b/tests/unit/test_xorq_stats_v2.py
@@ -14,11 +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, StatUnit, merge_fragments # noqa: E402
from buckaroo.customizations.xorq_stats_v2 import ( # noqa: E402
XORQ_STATS_V2)
@@ -876,3 +879,529 @@ 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)
+
+
+# ============================================================
+# 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, 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):
+ 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
+ if isinstance(value, float):
+ return float(f"{value:.9g}")
+ 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):
+ 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."""
+ 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):
+ 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):
+ 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."""
+ _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):
+ _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):
+ 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):
+
+ @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):
+ 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")})
+
+ @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)
+# ============================================================
+
+_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):
+ 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):
+ 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):
+ 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):
+ 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):
+ 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."""
+ 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):
+ 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):
+ 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)
+
+
+# ============================================================
+# 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)
+
+ 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"]