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..812553ae4 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,56 @@ def process_column(self, column_name: str, column_dtype, raw_series=None, sample
return resolve_accumulator(accumulator, column_name, col_key_to_func)
+ @staticmethod
+ def _skipped(state: StatState, pairs: Sequence[Tuple[Any, str]]) -> set:
+ """The columns of ``pairs`` named in ``state.skip_columns``, by original
+ or rewritten name."""
+ return {orig for orig, rewritten in pairs
+ if orig in state.skip_columns or rewritten in state.skip_columns}
+
+ def plan(self, state: StatState) -> List[StatUnit]:
+ """One unit per column that is not skipped, in column order (or the
+ order ``state.priority`` asks for). A column's stats read only its own
+ column, so no unit waits for another. Nothing is computed."""
+ if len(state.data) == 0:
+ return []
+ pairs = columns_in_scope(state)
+ skip = self._skipped(state, pairs)
+ active = prioritized(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 +341,15 @@ def process_df(self, df: pd.DataFrame, debug: bool = False,
if len(df) == 0:
return {}, []
- if self.record_timings:
- self.timings = []
-
- skip = set(skip_columns or ())
- summary: SDType = {}
- all_errors: List[StatError] = []
-
- for orig_col_name, rewritten_col_name in old_col_new_col(df):
- if orig_col_name in skip or rewritten_col_name in skip:
- # Provided externally — keep the column, skip computation.
- summary[rewritten_col_name] = {
- 'orig_col_name': orig_col_name, 'rewritten_col_name': rewritten_col_name}
- continue
- ser = df[orig_col_name]
- col_dtype = ser.dtype
-
- col_result, col_errors = self.process_column(column_name=rewritten_col_name, column_dtype=col_dtype,
- raw_series=ser, sampled_series=ser, raw_dataframe=df,
- initial_stats={'orig_col_name': orig_col_name, 'rewritten_col_name': rewritten_col_name})
-
- summary[rewritten_col_name] = col_result
- all_errors.extend(col_errors)
+ acc = self.run_all(StatState(df, frozenset(skip_columns or ())))
+ summary: SDType = {acc.rewritten[orig]: stats for orig, stats in acc.sd().items()}
if self.record_timings and perf_log.enabled() and not self._suppress_perf_summary:
rec = perf_log.PerfRecorder(label=f"stats rows={len(df)} cols={len(summary)}")
rec.extend_timings("pandas/polars", self.timings)
rec.summary()
- return summary, all_errors
+ return summary, acc.errors
def unit_test(self) -> Tuple[bool, List[StatError]]:
"""Test the pipeline against PERVERSE_DF."""
@@ -432,3 +461,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..c5a02c1c4
--- /dev/null
+++ b/buckaroo/pluggable_analysis_framework/stat_units.py
@@ -0,0 +1,221 @@
+"""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"
+
+
+# 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.
+ """
+ data: Any
+ skip_columns: frozenset = frozenset()
+ columns: Optional[Tuple[Any, ...]] = None
+ priority: Tuple[Any, ...] = ()
+ rows: Optional[int] = None
+ namespace: str = "any"
+
+ def __post_init__(self) -> None:
+ check_namespace(self.namespace)
+
+
+def columns_in_scope(state: StatState) -> List[Tuple[Any, str]]:
+ """``(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..9601e2e03 100644
--- a/buckaroo/pluggable_analysis_framework/xorq_stat_pipeline.py
+++ b/buckaroo/pluggable_analysis_framework/xorq_stat_pipeline.py
@@ -1,14 +1,17 @@
"""Xorq-backed stat pipeline for the v2 framework.
-Two-phase execution:
+Two-phase execution, each phase a kind of resumable unit (see ``stat_units``):
1. Batch aggregate — every @stat with an XorqColumn parameter contributes
one ibis scalar expression. All such expressions across all columns
are folded into a single ``table.aggregate(...)`` query and executed
- once.
- 2. Per-column post-batch — computed stats (deps only on other stats) and
- XorqExpr-param stats (e.g. histograms that need their own query)
- run through the standard typed-DAG executor with results written into
- the per-column accumulator.
+ once. This is the ``batch`` unit, and it also computes the stats that
+ are pure functions of its results (``histogram_bins``, ``nan_per``...).
+ Behind a host's opt-in the batch can be cut into one aggregate per chunk
+ of columns (``chunk_cells``), for a plain parquet scan only.
+ 2. Per-column post-batch — XorqExpr-param stats (e.g. histograms that need
+ their own query), and any stat reading one, run through the standard
+ typed-DAG executor with results written into the per-column
+ accumulator. One ``histogram:`` unit per column.
Errors are captured into ``StatError`` via the standard Ok/Err mechanism;
nothing is silently swallowed. Construction validates the DAG up front and
@@ -23,7 +26,7 @@
import os
import time
from contextlib import nullcontext
-from typing import Any, Dict, List, Tuple
+from typing import Any, Dict, List, Optional, Sequence, Tuple
import pandas as pd
@@ -33,6 +36,8 @@
from .stat_func import XorqColumn, XorqExpr, XorqExecute, RAW_MARKER_TYPES, StatFunc
from .stat_pipeline import _execute_stat_func, _normalize_inputs, errors_to_errdict
from .stat_result import Err, Ok, StatError, StatResult, resolve_accumulator
+from .stat_units import (Fragment, StatAccumulator, StatState, StatUnit, UnitPipeline, UnitStats, columns_in_scope,
+ prioritized)
from .typed_dag import build_column_dag, build_typed_dag
from .utils import PERVERSE_DF
@@ -41,10 +46,12 @@
try:
import xorq.api as xo
+ from xorq.expr.relations import Read
+ from xorq.vendor.ibis.expr import operations as ops
HAS_XORQ = True
except ImportError:
- xo = None
+ xo = Read = ops = None
HAS_XORQ = False
log = logging.getLogger(__name__)
@@ -100,7 +107,29 @@ def _is_batch_func(sf: StatFunc) -> bool:
return True
-class XorqStatPipeline:
+class XorqAccumulator(StatAccumulator):
+ """``StatAccumulator`` for a xorq run. ``working`` holds each column's
+ Ok/Err accumulator, which the units read and write; ``emitted`` the keys of
+ it already returned in a fragment, so each stat is returned and any error in
+ it reported once; ``phases`` caches how each dtype's stats split (see
+ ``XorqStatPipeline._column_phases``). ``sd()`` is every column's stats so
+ far in accumulator order, a failed stat as None."""
+
+ def __init__(self, state: StatState, columns: Sequence[Any] = (), rewritten: Optional[Dict[Any, str]] = None,
+ skipped: Sequence[Any] = ()):
+ super().__init__(state, columns, rewritten)
+ self.skipped: Tuple[Any, ...] = tuple(skipped)
+ self.working: Dict[Any, Dict[str, StatResult]] = {}
+ self.emitted: Dict[Any, set] = {}
+ self.phases: Dict[Any, Any] = {}
+ self.schema: Any = None
+
+ def sd(self) -> Dict[Any, Dict[str, Any]]:
+ return {col: {key: (result.value if isinstance(result, Ok) else None if isinstance(result, Err) else result)
+ for key, result in self.working[col].items()} for col in self.order}
+
+
+class XorqStatPipeline(UnitPipeline):
"""v2 stat pipeline for ``ibis.Table`` inputs.
Accepts the same kinds of inputs as ``StatPipeline``:
@@ -110,7 +139,13 @@ class XorqStatPipeline:
- ``ColAnalysis`` subclasses (via v1 adapter)
Use ``process_table(table)`` to run the pipeline; returns
- ``(SDType, List[StatError])``.
+ ``(SDType, List[StatError])``. ``plan(state)``, ``new_accumulator(state)``
+ and ``run(unit, acc)`` are the same run one unit at a time.
+
+ ``chunk_cells`` (off when ``None``) cuts the batch aggregate into one
+ aggregate per chunk of columns holding about that many cells, when the
+ caller gives the row count and the source is a plain parquet scan
+ (``chunk_refusal``).
"""
# Keys that the pipeline pre-populates per column. Listed as external
@@ -123,7 +158,7 @@ class XorqStatPipeline:
"distinct_count"})
def __init__(self, stat_funcs: list, backend: Any = None, unit_test: bool = True,
- cache_storage=None):
+ cache_storage=None, chunk_cells: Optional[int] = None):
if not HAS_XORQ:
raise ImportError(
"xorq is required for XorqStatPipeline. "
@@ -138,12 +173,13 @@ def __init__(self, stat_funcs: list, backend: Any = None, unit_test: bool = True
self._original_inputs = list(stat_funcs)
self.backend = backend
self.cache_storage = cache_storage
+ self.chunk_cells = chunk_cells
- # Per-run snapshot-cache counters, (re)initialised in process_table.
+ # Per-run snapshot-cache counters, (re)initialised in new_accumulator.
# Set here so the attribute always exists (e.g. for the unit_test()
# run kicked off below, which disables the cache).
self._cache_stats = _new_cache_run_stats()
- # Per-run perf recorder, (re)initialised in process_table when the
+ # Per-run perf recorder, (re)initialised in new_accumulator when the
# BUCKAROO_PERF toggle is on; None otherwise.
self._perf = None
# Set during the unit_test() DAG self-check so its PERVERSE_DF run
@@ -324,19 +360,21 @@ def _span(self, label, **fields):
return nullcontext()
return perf_log.perf_span(label, **fields)
- def process_table(self, table, skip_columns=None) -> Tuple[SDType, List[StatError]]:
+ def process_table(self, table, skip_columns=None, rows=None) -> Tuple[SDType, List[StatError]]:
# Each per-column query runs directly against ``table`` (the source
# expression). When a snapshot cache is set, queries are keyed against
# that source so the content-addressed key is stable across processes;
# ``_execute_cached`` serves a hit from the snapshot parquet or executes
- # and writes one on a miss.
- self._cache_stats = _new_cache_run_stats()
- self._perf = (perf_log.PerfRecorder()
- if perf_log.enabled() and not self._suppress_perf_summary else None)
+ # and writes one on a miss. ``rows`` is the row count when the caller
+ # has it; only the column-chunk split reads it.
+ state = StatState(table, frozenset(skip_columns or ()), rows=rows)
+ acc = self.new_accumulator(state)
_t0 = time.perf_counter()
with self._span("stat.xorq.total") as span:
try:
- return self._process_table_impl(table, skip_columns=skip_columns)
+ for _unit, _fragment in self.iter_units(state, acc):
+ pass
+ return acc.sd(), acc.errors
finally:
self._cache_stats["secs"] = round(time.perf_counter() - _t0, 4)
self._log_cache_stats(span)
@@ -347,27 +385,159 @@ def process_table(self, table, skip_columns=None) -> Tuple[SDType, List[StatErro
f"misses={self._cache_stats['misses']}")
self._perf.summary()
- def _process_table_impl(self, table, skip_columns=None) -> Tuple[SDType, List[StatError]]:
+ @staticmethod
+ def chunk_refusal(table) -> Optional[str]:
+ """Why ``table`` may not have its batch split by column chunk, or
+ ``None`` when it may.
+
+ Each chunk is its own aggregate over the source. On a parquet scan a
+ chunk reads only its columns, so the chunks cost about what one batch
+ does (0.90-1.16x at 10-12M rows) and each needs less memory. On
+ anything else a chunk runs the plan again: a join or an aggregate
+ executes in full per chunk, and a CSV read reparses the file (2.1-2.2x).
+ So the split is allowed for a parquet ``Read``, or a projection of
+ plain columns of one, and refused for everything else, including an
+ expression this does not recognise. A refusal can only keep the single
+ batch; it never turns the split on.
+ """
+ op = table.op()
+ if isinstance(op, ops.Project):
+ if not all(isinstance(value, ops.Field) for value in op.values.values()):
+ return "it computes columns instead of reading them"
+ op = op.parent
+ if not isinstance(op, Read):
+ return f"not a parquet scan ({type(op).__name__})"
+ if op.method_name != "read_parquet":
+ return f"not a parquet scan ({op.method_name})"
+ return None
+
+ def _batch_chunks(self, state: StatState, columns: List[Any]) -> List[List[Any]]:
+ """``columns`` cut into the chunks the batch runs in: one, unless the
+ split is on, the row count is known and the source allows it."""
+ if not self.chunk_cells:
+ return [columns]
+ if state.rows is None:
+ log.debug("xorq stat column chunking needs the row count; running one batch")
+ return [columns]
+ refusal = self.chunk_refusal(state.data)
+ if refusal is not None:
+ log.info("xorq stat column chunking refused: %s", refusal)
+ return [columns]
+ per_chunk = max(1, self.chunk_cells // max(state.rows, 1))
+ return [columns[i:i + per_chunk] for i in range(0, len(columns), per_chunk)] or [columns]
+
+ def _column_phases(self, cache: Dict[Any, Any], dtype) -> Tuple[List[StatFunc], List[StatFunc], Dict[str, StatFunc]]:
+ """A column's stat funcs split into those that need no further query
+ once the batch has run (``pure``) and those that run one or read the
+ result of one that does (``query``), each in dependency order, plus the
+ stat key -> func map errors are reported against."""
+ if dtype not in cache:
+ col_funcs = build_column_dag(self.all_stat_funcs, dtype, external_keys=self.EXTERNAL_KEYS)
+ pure: List[StatFunc] = []
+ query: List[StatFunc] = []
+ query_keys: set = set()
+ for sf in col_funcs:
+ if any(r.type in (XorqExpr, XorqExecute) or r.name in query_keys for r in sf.requires):
+ query.append(sf)
+ query_keys.update(sk.name for sk in sf.provides)
+ else:
+ pure.append(sf)
+ key_to_func = {sk.name: sf for sf in col_funcs for sk in sf.provides}
+ cache[dtype] = (pure, query, key_to_func)
+ return cache[dtype]
+
+ def plan(self, state: StatState) -> List[StatUnit]:
+ """The scalar batch first, then one histogram GROUP BY per column that
+ has one, columns named in ``state.priority`` first. The batch is cut
+ into chunks only as ``_batch_chunks`` allows, and every chunk precedes
+ every histogram, since a histogram reads what its chunk computes.
+ Skipped columns get no unit. Nothing is computed: no query is sent."""
+ table = state.data
+ pairs = columns_in_scope(state)
+ if not pairs:
+ return []
schema = table.schema()
- columns = list(table.columns)
- # Columns whose stats are supplied externally (via init_sd) keep their
- # structural metadata (name/dtype/length) but get no stat expressions
- # built — so the column's data is never scanned.
- skip = set(skip_columns or ())
-
- # Pre-populate every column accumulator with the externally-provided
- # keys. ``length`` is filled in by the batch query below. ``min`` /
- # ``max`` start as None so dependents (histogram) don't cascade-
- # exclude on non-numeric columns; ``min`` / ``max`` overwrite for
- # numeric cols. ``distinct_count`` likewise starts as None so float
- # columns (where the stat is column_filtered out) keep their
- # dependents (histogram, histogram_bins, distinct_per) runnable.
- accumulators: Dict[str, Dict[str, StatResult]] = {}
- for col in columns:
- accumulators[col] = {"orig_col_name": Ok(col), "rewritten_col_name": Ok(col), "dtype": Ok(str(schema[col])),
+ # xorq has always matched skip_columns on the column's own name only.
+ skip = {orig for orig, _rewritten in pairs if orig in state.skip_columns}
+ active = [orig for orig, _rewritten in prioritized(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 self._column_phases(phases, schema[col])[1]:
+ units.append(StatUnit(id=f"histogram:{col}", columns=(col,), phase="histogram",
+ after=(batch_of[col],), cost="query"))
+ return units
+
+ def new_accumulator(self, state: StatState) -> "XorqAccumulator":
+ """The accumulator for a run of ``state``, and the start of the run's
+ counters: the snapshot-cache hit/miss counts and the perf recorder.
+
+ Pre-populate every column accumulator with the externally-provided
+ keys. ``length`` is filled in by the batch query. ``min`` / ``max``
+ start as None so dependents (histogram) don't cascade-exclude on
+ non-numeric columns; ``min`` / ``max`` overwrite for numeric cols.
+ ``distinct_count`` likewise starts as None so float columns (where the
+ stat is column_filtered out) keep their dependents (histogram,
+ histogram_bins, distinct_per) runnable. A skipped column keeps these
+ (its stats are supplied via init_sd and its data is never scanned)."""
+ self._cache_stats = _new_cache_run_stats()
+ self._perf = (perf_log.PerfRecorder()
+ if perf_log.enabled() and not self._suppress_perf_summary else None)
+ schema = state.data.schema()
+ pairs = columns_in_scope(state)
+ acc = XorqAccumulator(state, columns=[orig for orig, _rewritten in pairs], rewritten=dict(pairs),
+ skipped=[orig for orig, _rewritten in pairs if orig in state.skip_columns])
+ for col in acc.order:
+ acc.working[col] = {"orig_col_name": Ok(col), "rewritten_col_name": Ok(col), "dtype": Ok(str(schema[col])),
"length": Ok(0), "min": Ok(None), "max": Ok(None), "distinct_count": Ok(None)}
-
- # ---- Phase 1: batch aggregate ----------------------------------
+ acc.emitted[col] = set()
+ acc.schema = schema
+ return acc
+
+ def run(self, unit: StatUnit, acc: StatAccumulator) -> Fragment:
+ """Run one unit and return the stats it produced, ``{col: {stat: value}}``
+ for its columns. A stat that fails is recorded on ``acc.errors`` once
+ and the run goes on."""
+ assert isinstance(acc, XorqAccumulator)
+ if unit.phase == "batch":
+ return self._run_batch(unit, acc)
+ return self._run_histogram(unit, acc)
+
+ def _run_funcs(self, funcs: List[StatFunc], col: Any, acc: "XorqAccumulator") -> None:
+ """Run ``funcs`` (in dependency order) for one column against its
+ accumulator, skipping stats whose results are already there (typically
+ the batch-phase stats)."""
+ col_accum = acc.working[col]
+ for sf in funcs:
+ if sf.provides and all(sk.name in col_accum for sk in sf.provides):
+ continue
+ if self._perf is not None:
+ t0 = time.perf_counter()
+ _execute_stat_func(sf, col_accum, col, raw_series=None, sampled_series=None, raw_dataframe=None,
+ xorq_expr=acc.state.data, xorq_execute=self._execute)
+ if self._perf is not None:
+ self._perf.record("xorq/per-column", col, sf.name, time.perf_counter() - t0)
+
+ def _emit(self, acc: "XorqAccumulator", col: Any, key_to_func: Dict[str, StatFunc]) -> Dict[str, Any]:
+ """The stats of ``col`` that no earlier unit has returned, as plain
+ values (a failed stat as None), recorded on ``acc`` with their errors."""
+ working = acc.working[col]
+ fresh = {key: result for key, result in working.items() if key not in acc.emitted[col]}
+ plain, errors = resolve_accumulator(fresh, col, key_to_func)
+ acc.emitted[col].update(fresh)
+ acc.record({col: plain}, errors)
+ return plain
+
+ def _run_batch(self, unit: StatUnit, acc: "XorqAccumulator") -> Fragment:
+ table = acc.state.data
+ schema = acc.schema
+ columns = list(unit.columns)
# ``length`` is a table-level scalar (same value for every column),
# so it goes in once as ``__total_length__`` rather than as N
# per-column expressions.
@@ -378,8 +548,6 @@ def _process_table_impl(self, table, skip_columns=None) -> Tuple[SDType, List[St
continue
xorq_col_param = next(r.name for r in sf.requires if r.type is XorqColumn)
for col in columns:
- if col in skip:
- continue
col_dtype = schema[col]
if sf.column_filter is not None and not sf.column_filter(col_dtype):
continue
@@ -387,7 +555,7 @@ def _process_table_impl(self, table, skip_columns=None) -> Tuple[SDType, List[St
expr = sf.func(**{xorq_col_param: table[col]})
except Exception as e:
for sk in sf.provides:
- accumulators[col][sk.name] = Err(error=e, stat_func_name=sf.name, column_name=col,
+ acc.working[col][sk.name] = Err(error=e, stat_func_name=sf.name, column_name=col,
inputs={"col": col})
continue
if expr is None:
@@ -397,7 +565,7 @@ def _process_table_impl(self, table, skip_columns=None) -> Tuple[SDType, List[St
expr = expr.name(f"{col}|{stat_name}")
except Exception as e:
for sk in sf.provides:
- accumulators[col][sk.name] = Err(error=e, stat_func_name=sf.name, column_name=col,
+ acc.working[col][sk.name] = Err(error=e, stat_func_name=sf.name, column_name=col,
inputs={"col": col})
continue
batch_items.append((col, sf, expr))
@@ -413,54 +581,40 @@ def _process_table_impl(self, table, skip_columns=None) -> Tuple[SDType, List[St
# length stays at the prepopulated 0 so consumers still see something.
for col, sf, _ in batch_items:
for sk in sf.provides:
- accumulators[col][sk.name] = Err(error=e, stat_func_name=sf.name, column_name=col, inputs={})
+ acc.working[col][sk.name] = Err(error=e, stat_func_name=sf.name, column_name=col, inputs={})
else:
total_length = _to_python_scalar(result_df[TOTAL_LENGTH_KEY].iloc[0])
if total_length is None:
total_length = 0
- for col in columns:
- accumulators[col]["length"] = Ok(total_length)
+ # A skipped column has no unit but has always been given the length.
+ for col in [*columns, *acc.skipped]:
+ acc.working[col]["length"] = Ok(total_length)
for col, sf, _ in batch_items:
stat_name = sf.provides[0].name
col_stat = f"{col}|{stat_name}"
if col_stat in result_df.columns:
raw_val = result_df[col_stat].iloc[0]
- accumulators[col][stat_name] = Ok(_to_python_scalar(raw_val))
+ acc.working[col][stat_name] = Ok(_to_python_scalar(raw_val))
else:
- accumulators[col][stat_name] = Err(error=KeyError(
+ acc.working[col][stat_name] = Err(error=KeyError(
f"missing aggregate column {col_stat!r} in result"), stat_func_name=sf.name, column_name=col, inputs={})
- # ---- Phase 2: per-column post-batch ----------------------------
- all_errors: List[StatError] = []
- summary: SDType = {}
-
+ # The stats that are pure functions of the ones just computed, so that
+ # histogram_bins (and with it color_map) needs no histogram query.
+ fragment: Fragment = {}
for col in columns:
- col_accum = accumulators[col]
- col_dtype = schema[col]
- col_funcs = build_column_dag(self.all_stat_funcs, col_dtype, external_keys=self.EXTERNAL_KEYS)
-
- for sf in col_funcs if col not in skip else []:
- # Skip stats whose results are already in the accumulator
- # (typically the batch-phase stats).
- if sf.provides and all(sk.name in col_accum for sk in sf.provides):
- continue
- if self._perf is not None:
- t0 = time.perf_counter()
- _execute_stat_func(sf, col_accum, col, raw_series=None, sampled_series=None, raw_dataframe=None,
- xorq_expr=table, xorq_execute=self._execute)
- if self._perf is not None:
- self._perf.record("xorq/per-column", col, sf.name, time.perf_counter() - t0)
-
- col_key_to_func: Dict[str, StatFunc] = {}
- for sf in col_funcs:
- for sk in sf.provides:
- col_key_to_func[sk.name] = sf
-
- plain, errors = resolve_accumulator(col_accum, col, col_key_to_func)
- summary[col] = plain
- all_errors.extend(errors)
-
- return summary, all_errors
+ pure, _query, key_to_func = self._column_phases(acc.phases, schema[col])
+ self._run_funcs(pure, col, acc)
+ fragment[col] = self._emit(acc, col, key_to_func)
+ return fragment
+
+ def _run_histogram(self, unit: StatUnit, acc: "XorqAccumulator") -> Fragment:
+ fragment: Fragment = {}
+ for col in unit.columns:
+ _pure, query, key_to_func = self._column_phases(acc.phases, acc.schema[col])
+ self._run_funcs(query, col, acc)
+ fragment[col] = self._emit(acc, col, key_to_func)
+ return fragment
def add_stat(self, stat_func_or_class) -> Tuple[bool, List[StatError]]:
"""Add a stat function or ColAnalysis class interactively.
@@ -501,7 +655,7 @@ def add_stat(self, stat_func_or_class) -> Tuple[bool, List[StatError]]:
return True, []
-class XorqDfStatsV2:
+class XorqDfStatsV2(UnitStats):
"""Stats wrapper for xorq table inputs.
Mirrors the ``DfStatsV2`` / ``PlDfStatsV2`` surface (``.sdf``, ``.errs``,
@@ -516,7 +670,9 @@ class XorqDfStatsV2:
Stats execute through ``XorqStatPipeline`` — a single batched
``table.aggregate(...)`` query plus per-column histogram queries —
pushing computation to the backend instead of materialising the
- entire table.
+ entire table. ``run=False`` builds the wrapper without running anything,
+ for a caller that drives ``plan`` / ``run`` itself. ``chunk_cells`` and
+ ``rows`` are the column-chunk split's size and the row count it needs.
"""
@classmethod
@@ -526,21 +682,26 @@ def verify_analysis_objects(cls, objs):
XorqStatPipeline(objs, unit_test=False)
def __init__(self, table, col_analysis_objs, operating_df_name=None, debug=False,
- cache_storage=None, skip_columns=None):
+ cache_storage=None, skip_columns=None, chunk_cells=None, rows=None, run=True):
self.table = table
# Skip the unit_test PERVERSE_DF run on each widget construction —
# it doubles the SQL query count (issue #709). The DAG-validation
# cost is already paid by verify_analysis_objects on first set up
# and by the test suite. Mirrors PlDfStatsV2.
self.ap = XorqStatPipeline(col_analysis_objs, unit_test=False,
- cache_storage=cache_storage)
+ cache_storage=cache_storage, chunk_cells=chunk_cells)
self.operating_df_name = operating_df_name
self.debug = debug
- self.sdf, errors = self.ap.process_table(self.table, skip_columns=skip_columns)
- self.errs = errors_to_errdict(errors)
+ self.rows = rows
+ self.state = StatState(table, frozenset(skip_columns or ()), rows=rows)
+ self.sdf: SDType = {}
+ self.errs = {}
self.stat_errors = []
- if self.errs:
- output_full_reproduce(self.errs, self.sdf, operating_df_name)
+ if run:
+ self.sdf, errors = self.ap.process_table(self.table, skip_columns=skip_columns, rows=rows)
+ self.errs = errors_to_errdict(errors)
+ if self.errs:
+ output_full_reproduce(self.errs, self.sdf, operating_df_name)
def cache_run_stats(self) -> dict:
"""The summary-stat cache outcome of the last pipeline run (#943) —
@@ -556,7 +717,7 @@ def add_analysis(self, a_obj):
so DataFlow.add_analysis works against a xorq-backed stats wrapper.
"""
passed, errors = self.ap.add_stat(a_obj)
- self.sdf, self.stat_errors = self.ap.process_table(self.table)
+ self.sdf, self.stat_errors = self.ap.process_table(self.table, rows=self.rows)
self.errs = errors_to_errdict(self.stat_errors)
if not passed:
print("DAG validation failed")
diff --git a/buckaroo/server/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..48100161a 100644
--- a/buckaroo/server/handlers.py
+++ b/buckaroo/server/handlers.py
@@ -14,7 +14,10 @@
from buckaroo.compare import col_join_dfs
from buckaroo.df_util import old_col_new_col
from buckaroo.server.focus import find_or_create_session_window
-from buckaroo.server.session import build_state_message
+from buckaroo.dataflow.dataflow import STATS_TIERS
+from buckaroo.server.session import (
+ STATS_DELIVERIES, begin_stats_generation, dataflow_stats_tier)
+from buckaroo.server.stats_wire import broadcast_state, refresh_session_snapshot
from buckaroo.server import telemetry
from buckaroo.pluggable_analysis_framework import perf_log
@@ -192,17 +195,10 @@ def _push_state_to_clients(self, session, metadata: dict):
if not session.ws_clients:
return
- for client in list(session.ws_clients):
- try:
- # Reset per-client live search first (dataset changed),
- # then build the msg so the injected ``buckaroo_state``
- # mirrors the post-reset value.
- client.search_string = ""
- msg = build_state_message(session, metadata=metadata,
- search_string=client.search_string)
- client.write_message(json.dumps(msg))
- except Exception:
- session.ws_clients.discard(client)
+ # Each client's search is reset first (dataset changed), then its
+ # message is built so the injected ``buckaroo_state`` mirrors the
+ # post-reset value.
+ broadcast_state(session, metadata=metadata, reset_search=True)
def _handle_browser_window(self, session_id: str) -> str:
"""Handle browser window management."""
@@ -371,6 +367,11 @@ async def post(self):
**component_config,
}
+ # /load builds no deferred stats, whatever policy a prior /load_expr on
+ # this session left behind, and the generation moves on with the data.
+ session.stats_tier, session.stats_delivery = "full", "inline"
+ begin_stats_generation(session)
+
# Notify connected clients and open browser
self._push_state_to_clients(session, metadata)
browser_action = "skipped" if no_browser else self._handle_browser_window(session_id)
@@ -380,6 +381,28 @@ async def post(self):
self.write({"session": session_id, "server_pid": os.getpid(), "browser_action": browser_action, **metadata})
+def _stats_policy_from_body(body: dict, current_tier: str, current_delivery: str):
+ """Read ``stats_tier`` and ``stats_delivery`` from a request body.
+
+ A field the body omits (or sends as null) keeps the current value (the
+ session's, or the default for a new session), as ``cache_dir`` does. Returns
+ ``(stats_tier, stats_delivery, None)``, or ``(None, None, error)`` with the
+ 400 response body for an unknown value."""
+ stats_tier = body.get("stats_tier")
+ if stats_tier is None:
+ stats_tier = current_tier
+ if stats_tier not in STATS_TIERS:
+ return None, None, {"error_code": "invalid_stats_tier",
+ "message": f"stats_tier must be one of {list(STATS_TIERS)}, got {stats_tier!r}"}
+ stats_delivery = body.get("stats_delivery")
+ if stats_delivery is None:
+ stats_delivery = current_delivery
+ if stats_delivery not in STATS_DELIVERIES:
+ return None, None, {"error_code": "invalid_stats_delivery",
+ "message": f"stats_delivery must be one of {list(STATS_DELIVERIES)}, got {stats_delivery!r}"}
+ return stats_tier, stats_delivery, None
+
+
class LoadExprHandler(tornado.web.RequestHandler):
"""POST /load_expr — load a xorq/ibis expression from a build dir
and serve it via the xorq-backed buckaroo dataflow.
@@ -438,7 +461,10 @@ async def post(self):
# Config-bearing fields that change how the result is computed or
# rendered. If the caller passes any of these on a warm POST we must
# re-run the pipeline — returning cached metadata would silently
- # ignore the new config.
+ # ignore the new config. stats_tier and stats_delivery are not in this
+ # tuple: it tests truthiness, and a host that sends the pair on every
+ # POST would never get the warm exit (#944). They are compared with the
+ # session's stored pair below instead.
has_config = any(body.get(k) for k in (
"component_config", "column_config_overrides", "extra_grid_config",
"init_sd", "skip_stat_columns"))
@@ -468,11 +494,22 @@ async def post(self):
# ~/.cache/xorq and recomputing the embedder's snapshots there.
if cache_dir is None and existing is not None:
cache_dir = existing.cache_dir
+ # The stats policy carries over the same way: a re-POST that omits it
+ # keeps the session's, and one that changes it rebuilds.
+ stats_tier, stats_delivery, policy_error = _stats_policy_from_body(
+ body,
+ existing.stats_tier if existing is not None else "full",
+ existing.stats_delivery if existing is not None else "inline")
+ if policy_error is not None:
+ self.set_status(400)
+ self.write(policy_error)
+ return
# /load swaps a session to pandas without clearing build_dir, so the
# backend is checked too — else its pandas metadata comes back here.
if (not force_reload and not has_config and existing
and existing.backend == "xorq" and existing.build_dir == build_dir
- and existing.cache_dir == cache_dir and existing.metadata):
+ and existing.cache_dir == cache_dir and existing.metadata
+ and (existing.stats_tier, existing.stats_delivery) == (stats_tier, stats_delivery)):
# The pipeline is skipped, but the refreshed page still opens a new
# WS and pulls a fresh time-to-first-rows. Re-arm first-pull telemetry
# on the existing session — rebind this request's sink and reset the
@@ -553,6 +590,7 @@ async def post(self):
with perf_log.perf_span("firstpull.dataflow_construct", session=session_id):
xorq_dataflow = xorq_loading.XorqServerDataflow(
expr, skip_main_serial=True, extra_klasses=extra_klasses,
+ stats_tier=dataflow_stats_tier(stats_tier, stats_delivery),
**dataflow_kwargs)
# Spanning metadata too leaves only the small klass-load step
# unmeasured inside the outer firstpull.load_expr total.
@@ -578,6 +616,8 @@ async def post(self):
session.cache_dir = cache_dir
session.project_root = project_root
session.dataflow_kwargs = dataflow_kwargs
+ session.stats_tier = stats_tier
+ session.stats_delivery = stats_delivery
session.tele_sink = tele_sink
session.xorq_dataflow = xorq_dataflow
# Clear pandas-side state left by a prior /load on the same
@@ -619,17 +659,13 @@ async def post(self):
**dvc.get("component_config", {}),
**component_config}
- if session.ws_clients:
- for client in list(session.ws_clients):
- try:
- # Reset per-client live search (#851): a term from
- # the prior expression would silently filter the new one.
- client.search_string = ""
- msg = build_state_message(session, metadata=metadata,
- search_string=client.search_string)
- client.write_message(json.dumps(msg))
- except Exception:
- session.ws_clients.discard(client)
+ # A new expression is a new stats generation; a deferred session starts
+ # it pending.
+ begin_stats_generation(session)
+
+ # Reset per-client live search (#851): a term from the prior
+ # expression would silently filter the new one.
+ broadcast_state(session, metadata=metadata, reset_search=True)
if no_browser or not self.application.settings.get("open_browser", False):
browser_action = "skipped"
@@ -766,18 +802,12 @@ async def post(self):
session.df_data_dict = display_state["df_data_dict"]
session.df_meta = display_state["df_meta"]
session.mode = "viewer"
+ # A viewer session has no dataflow and so no deferred stats.
+ session.stats_tier, session.stats_delivery = "full", "inline"
+ begin_stats_generation(session)
- # Push to WebSocket clients
- if session.ws_clients:
- for client in list(session.ws_clients):
- try:
- # Reset per-client live search (#851).
- client.search_string = ""
- msg = build_state_message(session,
- search_string=client.search_string)
- client.write_message(json.dumps(msg))
- except Exception:
- session.ws_clients.discard(client)
+ # Push to WebSocket clients. Reset per-client live search (#851).
+ broadcast_state(session, reset_search=True)
# Browser window
if no_browser or not self.application.settings.get("open_browser", False):
@@ -807,8 +837,22 @@ class ReloadExprHandler(tornado.web.RequestHandler):
``cache_storage_path`` store are cache hits, and column overrides,
extra grid config, init_sd and skip_stat_columns survive the reload.
+ The session's stored ``stats_tier`` / ``stats_delivery`` are replayed too.
+ The body is optional; a pair in it replaces the stored one (and is stored)
+ when the reload succeeds.
+
Returns 404 when the session does not exist, 400 when it is not a xorq
- session or has no project_root recorded, 501 when xorq is not installed."""
+ session, has no project_root recorded or carries an invalid stats policy,
+ 501 when xorq is not installed."""
+
+ def _optional_body(self) -> dict:
+ """The JSON object body, or ``{}`` when there is none or it is not an
+ object: the endpoint took no body before the stats policy fields."""
+ try:
+ body = json.loads(self.request.body)
+ except (json.JSONDecodeError, TypeError):
+ return {}
+ return body if isinstance(body, dict) else {}
async def post(self, session_id):
sessions = self.application.settings["sessions"]
@@ -831,6 +875,13 @@ async def post(self, session_id):
"message": "Session has no project_root — pass project_root to /load_expr first"})
return
+ stats_tier, stats_delivery, policy_error = _stats_policy_from_body(
+ self._optional_body(), session.stats_tier, session.stats_delivery)
+ if policy_error is not None:
+ self.set_status(400)
+ self.write(policy_error)
+ return
+
try:
from buckaroo.server import xorq_loading
except ImportError:
@@ -847,6 +898,7 @@ async def post(self, session_id):
+ xorq_loading.load_project_display_klasses(session.project_root))
xorq_dataflow = xorq_loading.XorqServerDataflow(
session.expr, skip_main_serial=True, extra_klasses=extra_klasses,
+ stats_tier=dataflow_stats_tier(stats_tier, stats_delivery),
**session.dataflow_kwargs)
except Exception:
tb = traceback.format_exc()
@@ -872,29 +924,13 @@ async def post(self, session_id):
if bs.get("quick_command_args"):
xorq_dataflow.quick_command_args = bs["quick_command_args"]
- refreshed = get_buckaroo_display_state(xorq_dataflow)
session.xorq_dataflow = xorq_dataflow
- session.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"]
+ session.stats_tier = stats_tier
+ session.stats_delivery = stats_delivery
+ refresh_session_snapshot(session, xorq_dataflow)
+ begin_stats_generation(session)
- if session.component_config and session.df_display_args:
- for key in session.df_display_args:
- dvc = session.df_display_args[key].get("df_viewer_config")
- if dvc is not None:
- dvc["component_config"] = {
- **dvc.get("component_config", {}),
- **session.component_config}
-
- for client in list(session.ws_clients):
- try:
- msg = build_state_message(session,
- search_string=getattr(client, "search_string", ""))
- client.write_message(json.dumps(msg))
- except Exception:
- session.ws_clients.discard(client)
+ broadcast_state(session)
klass_count = len(extra_klasses)
log.info("reload_expr session=%s project_root=%s klasses=%d",
diff --git a/buckaroo/server/session.py b/buckaroo/server/session.py
index 563ce28fc..754ddc1ce 100644
--- a/buckaroo/server/session.py
+++ b/buckaroo/server/session.py
@@ -11,6 +11,36 @@
_DEFAULT_SESSION_TTL_S = 3600.0 # 1 hour idle before eviction
_DEFAULT_EVICTION_INTERVAL_S = 300.0 # check every 5 minutes
+# How a session reaches its stats tier: ``inline`` computes them in the
+# dataflow constructor (today's behaviour); ``deferred`` builds the schema tier
+# first and leaves the rest to later requests.
+STATS_DELIVERIES = ("inline", "deferred")
+
+
+def dataflow_stats_tier(stats_tier: str, stats_delivery: str) -> str:
+ """The tier a session's dataflow is constructed at. A deferred session
+ starts at the schema tier whatever tier it is headed for."""
+ return "schema" if stats_delivery == "deferred" else stats_tier
+
+
+# What ``df_meta.stats.status`` says about a session's stats for its current
+# ``stats_gen``: ``complete`` (the stats of the tier the session is headed for
+# are in the snapshot), ``pending`` (a deferred session that has not produced
+# them yet), ``not_computed`` (the session is headed for the schema tier, so
+# none will arrive) and ``error`` (the run failed; sticky until the next
+# generation).
+STATS_STATUSES = ("complete", "pending", "not_computed", "error")
+
+
+def initial_stats_status(stats_tier: str, stats_delivery: str) -> tuple[str, Optional[str]]:
+ """The status, and its reason, of a session whose stats generation has just
+ started, from the policy pair alone."""
+ if stats_tier != "full":
+ return "not_computed", "host"
+ if stats_delivery == "deferred":
+ return "pending", None
+ return "complete", None
+
@dataclass
class SessionState:
@@ -43,6 +73,28 @@ class SessionState:
# extra_grid_config, init_sd, skip_stat_columns), replayed by /reload_expr so
# a reload keeps stat caching and column config (#957).
dataflow_kwargs: dict = field(default_factory=dict)
+ # The /load_expr stats policy. ``stats_tier`` is the tier the session is
+ # headed for (see dataflow.STATS_TIERS) and ``stats_delivery`` how it gets
+ # there (see STATS_DELIVERIES). Kept apart from dataflow_kwargs, which is
+ # splatted into the dataflow constructor, and replayed by /reload_expr. The
+ # tier the dataflow is built at follows from the pair (dataflow_stats_tier).
+ stats_tier: str = "full"
+ stats_delivery: str = "inline"
+ # 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
+ # 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,6 +132,35 @@ def touch(self) -> None:
this field is the runtime escape hatch."""
+def begin_stats_generation(session: "SessionState") -> None:
+ """Start a new stats generation: the session's dataflow state has changed,
+ so stats for the previous one no longer describe it. The status restarts
+ from the session's policy pair, and the stat runs of the old generation are
+ dropped."""
+ session.stats_gen += 1
+ session.stats_status, session.stats_reason = initial_stats_status(
+ session.stats_tier, session.stats_delivery)
+ session.stat_runs.clear()
+
+
+def stats_meta(session: "SessionState") -> Optional[dict]:
+ """The ``df_meta.stats`` value for a session: ``{status, tier, gen}`` plus
+ ``reason`` when there is one. ``tier`` is the tier reached so far, which is
+ the one the session is headed for only once it is complete.
+
+ ``None`` for a session on the default policy (inline delivery, full tier,
+ complete), which sends the message it always has; a client reads a missing
+ ``stats`` as complete."""
+ if session.stats_status == "complete" and session.stats_delivery != "deferred":
+ return None
+ stats: dict = {"status": session.stats_status,
+ "tier": session.stats_tier if session.stats_status == "complete" else "schema",
+ "gen": session.stats_gen}
+ if session.stats_reason:
+ stats["reason"] = session.stats_reason
+ return stats
+
+
def build_state_message(session: "SessionState", metadata: dict | None = None,
search_string: str = "") -> dict:
"""Build the full ``initial_state`` WebSocket payload from a session.
@@ -98,10 +179,16 @@ def build_state_message(session: "SessionState", metadata: dict | None = None,
Returns:
A dict ready to be JSON-serialised and sent to WebSocket clients.
"""
+ # The dataflow rebuilds df_meta wholesale, so the stats status is injected
+ # here, into a copy, rather than stored in it.
+ df_meta = session.df_meta
+ stats = stats_meta(session)
+ if stats is not None:
+ df_meta = {**df_meta, "stats": stats}
msg: dict = {"type": "initial_state", "protocol_version": PROTOCOL_VERSION,
"metadata": metadata if metadata is not None else session.metadata,
"prompt": session.prompt, "df_display_args": session.df_display_args, "df_data_dict": session.df_data_dict,
- "df_meta": session.df_meta, "mode": session.mode}
+ "df_meta": df_meta, "mode": session.mode}
if session.mode == "buckaroo":
# Per-client search_string overlay (#851 Codex P1): the snapshot
# on the session is search-agnostic; we re-inject the recipient's
diff --git a/buckaroo/server/stat_run.py b/buckaroo/server/stat_run.py
new file mode 100644
index 000000000..82a13fba7
--- /dev/null
+++ b/buckaroo/server/stat_run.py
@@ -0,0 +1,135 @@
+"""The stat run a session keeps for one stats generation (rows-first s4).
+
+A ``StatRun`` holds the planned units of one ``(stats_gen, scope)``, the
+fragments they produced in the order they finished, and the accumulator the
+units read. It runs a unit only when its caller asks: it has no thread, timer or
+callback, so there is nothing to cancel, and the session drops it when the
+generation changes (``begin_stats_generation``).
+
+The fragment list is append-only and shared. A ``StatCursor`` is one client's
+position in it, held by that client's WebSocket handler, so work is done once
+and every client reads all of it at its own pace.
+"""
+import itertools
+import time
+from typing import Any, List, Optional, Sequence, Tuple
+
+from buckaroo.df_util import old_col_new_col
+from buckaroo.pluggable_analysis_framework.stat_pipeline import errors_to_errdict
+from buckaroo.pluggable_analysis_framework.stat_units import Fragment, StatUnit, resolve_names, rewrite_sd
+
+_run_ids = itertools.count(1)
+
+
+class StatRun:
+ """``stats`` is a stats class (``DfStatsV2`` and its siblings) built with
+ ``run=False``: the run plans it once and runs its units through it.
+
+ ``status`` is ``pending`` until the last unit has run, then ``complete``;
+ ``error`` once a unit raised, and the run is not retried. A stat that fails
+ inside a unit is not that: it is an entry in ``acc.errors`` and the unit
+ still finishes. ``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):
+ self.id = next(_run_ids)
+ self.stats_gen = stats_gen
+ self.scope = scope
+ self.stats = stats
+ self.units: List[StatUnit] = stats.plan(stats.state)
+ self.acc = stats.new_accumulator(stats.state)
+ self.fragments: List[Fragment] = []
+ self.ran: List[str] = []
+ self.status = "pending" if self.units else "complete"
+ self.error: Optional[Exception] = None
+ self.created_at = time.time()
+ self.elapsed_s = 0.0
+ self._frame_columns: Optional[List[Tuple[Any, str]]] = None
+
+ @property
+ def key(self) -> Tuple[int, str]:
+ return (self.stats_gen, self.scope)
+
+ @property
+ def remaining(self) -> int:
+ """The number of units not yet run."""
+ return len(self.units) - len(self.ran)
+
+ def next_unit(self, prefer: Sequence[Any] = (), 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_wire.py b/buckaroo/server/stats_wire.py
new file mode 100644
index 000000000..950daf29a
--- /dev/null
+++ b/buckaroo/server/stats_wire.py
@@ -0,0 +1,496 @@
+"""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.
+
+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 hashlib
+import json
+import logging
+import time
+import traceback
+from typing import Any, Callable, Dict, Optional, Sequence
+
+from buckaroo.dataflow.sd_cache import split_chain_by_scope
+from buckaroo.df_util import old_col_new_col
+from buckaroo.pluggable_analysis_framework import perf_log
+from buckaroo.pluggable_analysis_framework.stat_units import Fragment
+from buckaroo.server.data_loading import get_buckaroo_display_state
+from buckaroo.server.session import SessionState, build_state_message
+from buckaroo.server.stat_run import StatCursor, StatRun
+
+log = logging.getLogger("buckaroo.server.stats_wire")
+
+# The capability a client advertises, as one of the comma-separated values of
+# ``?caps=`` on the WebSocket URL: it merges ``stats_update`` messages. The
+# connection is the only place it can be recorded, because ``open()`` sends the
+# first message before the client has said anything.
+STATS_UPDATE_CAP = "stats_update"
+
+# The scopes a ``stats_request`` can name. ``raw`` is the unfiltered table of
+# the session's current generation; filtered stats exist only through a
+# dataflow-field change (``quick_command_args.search``), which bumps the
+# generation.
+STATS_SCOPES = ("raw",)
+
+# The time an incremental ``stats_request`` may spend running units, in
+# seconds. A request always runs one unit and stops starting new ones once this
+# is spent, so a snapshot-cache hit (milliseconds) shares its request with the
+# others and a query is alone in its own. It cannot cut a unit short: the
+# longest unit is what a waiting ``infinite_request`` stalls behind.
+STATS_BUDGET_S = 0.075
+
+
+def parse_caps(raw: str) -> frozenset:
+ """The capabilities in a ``?caps=a,b`` value."""
+ return frozenset(cap for cap in (part.strip() for part in raw.split(",")) if cap)
+
+
+def client_has_cap(client: Any, cap: str) -> bool:
+ """Whether a WebSocket handler recorded ``cap`` at open. Anything else with
+ no ``caps`` (a test double, a future non-WebSocket client) has none."""
+ return cap in getattr(client, "caps", ())
+
+
+def session_dataflow(session: Optional[SessionState]) -> Any:
+ """The dataflow behind a buckaroo-mode session, or ``None`` (viewer and lazy
+ sessions have none)."""
+ if session is None or session.mode != "buckaroo":
+ return None
+ return session.xorq_dataflow if session.backend == "xorq" else session.dataflow
+
+
+def refresh_session_snapshot(session: SessionState, dataflow: Any) -> None:
+ """Copy the dataflow's display state onto the session snapshot that new
+ clients and every push read, and re-apply ``component_config`` so theme
+ settings survive. The dataflow reaches no client until this runs."""
+ refreshed = get_buckaroo_display_state(dataflow)
+ session.df_display_args = refreshed["df_display_args"]
+ session.df_data_dict = refreshed["df_data_dict"]
+ session.df_meta = refreshed["df_meta"]
+ session.buckaroo_options = refreshed["buckaroo_options"]
+ session.command_config = refreshed["command_config"]
+ if session.component_config and session.df_display_args:
+ for key in session.df_display_args:
+ dvc = session.df_display_args[key].get("df_viewer_config")
+ if dvc is not None:
+ dvc["component_config"] = {**dvc.get("component_config", {}), **session.component_config}
+
+
+def start_stat_run(session: SessionState, scope: str = "raw") -> Optional[StatRun]:
+ """The ``StatRun`` of the session's current generation for ``scope``: the one
+ it holds, or a new one planned from the dataflow and stored on the session.
+ Planning sends no query. The run analyzes the frame ``assign_full_stats``
+ does, so running every unit gives the stats that call computes whole.
+ ``None`` when there is no dataflow to plan from or the scope is not one a
+ request can name."""
+ dataflow = session_dataflow(session)
+ if dataflow is None or dataflow.processed_df is None or scope not in STATS_SCOPES:
+ return None
+ key = (session.stats_gen, scope)
+ run = session.stat_runs.get(key)
+ if run is None:
+ run = StatRun(session.stats_gen, scope, dataflow.build_stats(dataflow.processed_df, run=False))
+ session.stat_runs[key] = run
+ return run
+
+
+def display_args_hash(display_args: Any) -> str:
+ """A digest of a ``df_display_args`` value, equal for equal content whatever
+ the key order. JSON, because that is how the value reaches a client, and it
+ reads a NaN the same each time."""
+ return hashlib.sha1(json.dumps(display_args, sort_keys=True, default=str).encode()).hexdigest()
+
+
+def highlighted_display_args(display_args: dict, term: str) -> dict:
+ """A copy of ``df_display_args`` with ``term`` set as the highlight phrase
+ of every string column's displayer, or with the highlight removed when the
+ term is empty (#851). A copy, so the shared session snapshot is never
+ changed."""
+ overlay = copy.deepcopy(display_args)
+ for dva in overlay.values():
+ dvc = (dva or {}).get("df_viewer_config") or {}
+ for col in dvc.get("column_config", []) or []:
+ disp = col.get("displayer_args")
+ if not isinstance(disp, dict) or disp.get("displayer") != "string":
+ continue
+ if term:
+ disp["highlight_phrase"] = [term]
+ else:
+ disp.pop("highlight_phrase", None)
+ return overlay
+
+
+def run_units(run: StatRun, budget_s: Optional[float], prefer: Sequence[str] = (),
+ clock: Callable[[], float] = time.perf_counter, session_id: Optional[str] = None) -> int:
+ """Run units of ``run`` for one request, and return how many.
+
+ At least one runs while the run is pending. Another starts only while the
+ time spent is under ``budget_s``, so units that cost milliseconds (their
+ queries are snapshot-cache hits) run together and one that is a query ends
+ its request. ``budget_s=None`` runs every unit still to run. ``prefer`` are
+ column names as a client writes them (the rewritten ``a, b, c``), whose
+ units go first. A unit that raises fails the run and the exception
+ propagates. Each unit is a ``stats.unit`` span."""
+ started = clock()
+ ran = 0
+ while run.status == "pending":
+ if ran and budget_s is not None and clock() - started >= budget_s:
+ break
+ unit = run.next_unit(prefer, "rewritten")
+ if unit is None:
+ run.run_next() # no unit is left: this marks the run complete
+ break
+ with perf_log.perf_span("stats.unit", session=session_id, stats_gen=run.stats_gen, unit=unit.id,
+ phase=unit.phase, cost=unit.cost, columns=len(unit.columns)):
+ run.run_next(prefer, "rewritten")
+ ran += 1
+ return ran
+
+
+def partial_payload(dataflow: Any, run: StatRun, fragments: Sequence[Fragment]) -> Any:
+ """The ``all_stats`` payload (an inline wide ``DFEnvelope``) of the columns
+ ``fragments`` cover: their stats so far, assembled as ``merged_sd`` is, so
+ ``init_sd``, a processing step's sd and the ``cleaned_*`` and ``filtered_*``
+ layers apply as they will in the complete state. A client merges it into
+ its ``all_stats`` key by key, so a stat sent again with a later fragment
+ only replaces itself."""
+ rewritten = dict(old_col_new_col(run.acc.state.data))
+ touched = {rewritten[col] for fragment in fragments for col in fragment if col in rewritten}
+ merged = dataflow._assemble_merged_sd(run.raw_sd())
+ return dataflow._sd_to_jsondf({col: stats for col, stats in merged.items() if col in touched})
+
+
+def _full_stats_key(dataflow: Any) -> Any:
+ """The ``summary_stats_cache`` key of the current state's full-tier sd."""
+ return dataflow._scope_cache_key(split_chain_by_scope(dataflow.operations)["filt"], tier="full")
+
+
+def _cached_full_sd(dataflow: Any) -> Any:
+ """The full-tier sd of the current state if an earlier visit to it left one."""
+ return dataflow.summary_stats_cache.get(_full_stats_key(dataflow))
+
+
+def assign_full_stats(dataflow: Any, computed: Optional[tuple] = None) -> None:
+ """Take a schema-tier dataflow to the full tier: use the current state's
+ full stats (found in ``summary_stats_cache``, where an earlier visit to the
+ same state left them, or ``computed``, the ``(sd, errs)`` of a finished
+ ``StatRun``, or computed here whole), write them under the full-tier key,
+ then assign ``summary_sd``. Assigning runs the cascade, which fills the
+ other scopes' entries it still lacks and rebuilds ``merged_sd``,
+ ``df_data_dict`` and ``df_display_args``.
+
+ 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)."""
+ 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 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 session.stats_status != "pending" or dataflow is None:
+ return False
+ run_key = (session.stats_gen, "raw")
+ with (
+ perf_log.telemetry_context(session.session_id, session.tele_sink),
+ perf_log.perf_span("firstpull.stats_total", session=session.session_id, stats_gen=session.stats_gen) 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 pending deferred
+ session gets its missing stats run first, so its message is complete; a
+ capable client gets the session snapshot as it is (stats-free while the
+ session is pending) and pulls the rest. The search term is the recipient's
+ own (#851).
+
+ A capable client is also told apart from the display config it holds: the
+ digest of the one a pending frame carried is kept on it, so the final
+ ``stats_update`` can say whether the config it would send differs."""
+ if (session.stats_delivery == "deferred" and session.stats_status == "pending"
+ and not client_has_cap(client, STATS_UPDATE_CAP)):
+ complete_stats(session)
+ if client_has_cap(client, STATS_UPDATE_CAP):
+ client.display_args_hash = (display_args_hash(session.df_display_args)
+ if session.stats_status == "pending" else None)
+ return build_state_message(session, metadata=metadata, search_string=getattr(client, "search_string", ""))
+
+
+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) -> Optional[StatRun]:
+ """The session's ``StatRun`` for the current generation, or ``None`` when it
+ cannot be planned (there is no frame, or the dataflow's stats class cannot
+ plan the one it has: a post-processor that failed leaves an error frame).
+ The whole-run path then decides what the stats are."""
+ try:
+ return start_stat_run(session)
+ except Exception:
+ log.warning("stat units not planned session=%s stats_gen=%s: %s", session.session_id, session.stats_gen,
+ traceback.format_exc())
+ return None
+
+
+def _serve_units(session: SessionState, client: Any, prefer: Sequence[str], stats_gen: Any, scope: Any,
+ started: float, info: dict) -> Optional[dict]:
+ """The reply to an incremental request on a pending session, when the
+ answer is the fragments this client has not seen: a partial ``stats_update``
+ (``final`` false), or ``stats_aborted`` when a unit failed. ``None`` when
+ the request is not answered that way: the stats are cached or cannot be
+ planned (the whole-run path takes them), or the run is now finished and the
+ caller publishes it.
+
+ A client that is behind the run is caught up and no unit runs for it, so
+ work is done once and a client that asks late reads it all from its own
+ cursor."""
+ dataflow = session_dataflow(session)
+ if _cached_full_sd(dataflow) is not None:
+ return None
+ run = _plan_run(session)
+ if run is None:
+ return None
+ cursor = getattr(client, "stats_cursor", None) or StatCursor()
+ if cursor.caught_up(run):
+ try:
+ info["units"] = run_units(run, STATS_BUDGET_S, prefer, session_id=session.session_id)
+ except Exception:
+ log.error("stat unit failed session=%s stats_gen=%s: %s", session.session_id, session.stats_gen,
+ traceback.format_exc())
+ _fail_stats(session)
+ return _aborted(stats_gen, scope, "error", session)
+ fragments = cursor.take(run)
+ if run.status != "pending":
+ return None
+ return {"type": "stats_update", "stats_gen": stats_gen, "scope": scope, "tier": session.stats_tier,
+ "final": False, "remaining": run.remaining, "payload": partial_payload(dataflow, run, fragments),
+ "elapsed_ms": _elapsed_ms(started)}
+
+
+def _rebuilt_display_args(session: SessionState, client: Any) -> Optional[dict]:
+ """The session's display config, when it differs from the one this client
+ holds: the stats-derived parts of a config (a float column's ``minWidth``)
+ change when the stats arrive. The client's own highlight is kept on it,
+ since the config replaces the one that carries it. ``None`` when the client
+ holds the current config, or there is no record of one to compare with."""
+ held = getattr(client, "display_args_hash", None)
+ if held is None:
+ return None
+ client.display_args_hash = None
+ if display_args_hash(session.df_display_args) == held:
+ return None
+ term = getattr(client, "search_string", "")
+ return highlighted_display_args(session.df_display_args, term) if term else session.df_display_args
+
+
+def _answer_stats_request(session: Optional[SessionState], msg: dict, client: Any, started: float,
+ info: dict) -> dict:
+ stats_gen, scope = msg.get("stats_gen"), msg.get("scope", "raw")
+ columns = msg.get("columns")
+ prefer = [c for c in columns if isinstance(c, str)] if isinstance(columns, list) else []
+ incremental = msg.get("incremental")
+ dataflow = session_dataflow(session)
+ if session is None or dataflow is None:
+ return _aborted(stats_gen, scope, "no_data")
+ 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 session.stats_status == "not_computed":
+ return _aborted(stats_gen, scope, "not_requestable", session)
+ if session.stats_status == "pending" and incremental:
+ reply = _serve_units(session, client, prefer, stats_gen, scope, started, info)
+ if reply is not None:
+ return reply
+ if session.stats_status == "pending":
+ complete_stats(session)
+ if session.stats_status != "complete":
+ return _aborted(stats_gen, scope, "error", session)
+ # Complete: from here the answer is the dataflow's own all_stats, with no
+ # query, whether this request ran the stats or an earlier one did.
+ 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.
+
+ 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)
+ 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"])
+ return reply
diff --git a/buckaroo/server/websocket_handler.py b/buckaroo/server/websocket_handler.py
index 2232211c7..ec394f599 100644
--- a/buckaroo/server/websocket_handler.py
+++ b/buckaroo/server/websocket_handler.py
@@ -1,4 +1,3 @@
-import copy
import json
import logging
import os
@@ -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
+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))
diff --git a/buckaroo/server/xorq_loading.py b/buckaroo/server/xorq_loading.py
index 60a9226cf..f60c9cc0e 100644
--- a/buckaroo/server/xorq_loading.py
+++ b/buckaroo/server/xorq_loading.py
@@ -64,6 +64,11 @@ class XorqServerDataflow(XorqDataflow):
and project-authored post-processing classes discovered under
``/post_processing/*.py``; the built-in xorq stats are
kept first so collisions resolve to the built-in.
+
+ ``stat_chunk_cells`` is the host's request (off when ``None``) to run the
+ stats batch aggregate in chunks of about that many cells instead of one
+ query; the stats pipeline still refuses it for any source that is not a
+ plain parquet scan (``XorqStatPipeline.chunk_refusal``).
"""
sampling_klass = XorqInfiniteSampling
@@ -72,13 +77,14 @@ class XorqServerDataflow(XorqDataflow):
DFStatsClass = XorqDfStatsV2
analysis_klasses = _XORQ_ANALYSIS_KLASSES
- def __init__(self, expr, *args, extra_klasses=None, cache_storage_path=None, **kwargs):
+ def __init__(self, expr, *args, extra_klasses=None, cache_storage_path=None, stat_chunk_cells=None, **kwargs):
if extra_klasses:
# Per-instance override — class-level _XORQ_ANALYSIS_KLASSES is
# left untouched so other sessions / direct widget usage don't
# inherit one project's stats.
self.analysis_klasses = list(_XORQ_ANALYSIS_KLASSES) + list(extra_klasses)
self.cache_storage = _make_cache_storage(cache_storage_path)
+ self.stat_chunk_cells = stat_chunk_cells
super().__init__(expr, *args, **kwargs)
diff --git a/buckaroo/xorq_buckaroo.py b/buckaroo/xorq_buckaroo.py
index 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..ed4bcfd16 100644
--- a/tests/unit/server/test_data_loading_polars.py
+++ b/tests/unit/server/test_data_loading_polars.py
@@ -9,15 +9,43 @@
* ``.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 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 +117,799 @@ 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}}]
+
+
+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_expr.py b/tests/unit/server/test_load_expr.py
index bdbb47db6..813955234 100644
--- a/tests/unit/server/test_load_expr.py
+++ b/tests/unit/server/test_load_expr.py
@@ -1,17 +1,23 @@
"""End-to-end tests for POST /load_expr — server load path for
XorqBuckarooInfiniteWidget over a xorq/ibis expression."""
+import dataclasses
+import datetime
import io
import json
import os
import shutil
import sys
import tempfile
+import time
+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
@@ -25,8 +31,20 @@
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 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 stats_wire, telemetry, xorq_loading # noqa: E402
from buckaroo.server.app import make_app as _make_app # noqa: E402
+from buckaroo.server.stat_run import StatRun # noqa: E402
+from buckaroo.server.stats_wire import start_stat_run # noqa: E402
+from buckaroo.server.websocket_handler import DataStreamHandler # noqa: E402
+from tests.unit.dataflow.scoped_summary_stats_test import ( # noqa: E402
+ _OverridingPostProcessing, _scope_inputs, _scope_sds_by_units)
pytestmark = pytest.mark.skipif(
sys.platform == "win32",
@@ -910,6 +928,1719 @@ 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_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, {})
+
+
+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)
+
+
class TestReloadExpr(tornado.testing.AsyncHTTPTestCase):
def get_app(self):
return make_app()
diff --git a/tests/unit/skip_columns_test.py b/tests/unit/skip_columns_test.py
index bf0121be5..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..e60b4e892 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,235 @@ 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'))
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..3e420cec9 100644
--- a/tests/unit/test_xorq_stats_v2.py
+++ b/tests/unit/test_xorq_stats_v2.py
@@ -19,6 +19,7 @@
XorqStatPipeline,
XorqColumn)
from buckaroo.pluggable_analysis_framework.stat_func import stat # noqa: E402
+from buckaroo.pluggable_analysis_framework.stat_units import StatState, merge_fragments # noqa: E402
from buckaroo.customizations.xorq_stats_v2 import ( # noqa: E402
XORQ_STATS_V2)
@@ -876,3 +877,396 @@ 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)