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..68b7bd416 100644 --- a/buckaroo/dataflow/dataflow.py +++ b/buckaroo/dataflow/dataflow.py @@ -27,6 +27,63 @@ from .df_types import DataFrameT from .sd_cache import hash_chain, split_chain_by_scope +# Content levels a dataflow can be built at. ``full`` runs the summary-stats +# pipeline. ``schema`` publishes identity and typing keys derived from the +# schema alone and runs no data query (backends opt in, see ``_get_schema_sd``). +STATS_TIERS = ("full", "schema") + + +def assemble_merged_sd(init_sd: InitSD, cleaned_sd: SDType, raw_sd: SDType, processed_sd: SDType, + processed_df: TAny, chains: TDict[str, List[TAny]], + clean_sd: Optional[SDType] = None, filt_sd: Optional[SDType] = None) -> SDType: + """Build the wire ``merged_sd`` from the pieces ``_merged_sd`` reads. + + Bare keys come from the raw scope's sd (``raw_sd``, computed on + ``sampled_df``) merged under ``init_sd`` (keyed by original column name, + so rewritten here), ``cleaned_sd`` and ``processed_sd``. ``cleaned_*`` keys + are layered on top from ``clean_sd`` when cleaning is active, and + ``filtered_*`` keys from ``filt_sd`` when a search filter is active. + + Pure: the inputs are not mutated, so a caller holding an accumulating + ``raw_sd`` can assemble it repeatedly without the traitlets cascade. + ``chains`` is ``split_chain_by_scope(operations)``. + """ + if processed_df is None: + # on initial startup + return merge_sds(init_sd, cleaned_sd, raw_sd, processed_sd) + + # Gate each prefixed layer on a chain-shape diff between scopes. + # ``cleaned_*`` fires when the clean chain has ops the raw chain + # doesn't (cleaning is on); ``filtered_*`` fires when the filt + # chain has ops the clean chain doesn't (at least one + # quick-command op is present, i.e. a search filter is on). + # Keying off ``filt_sd_key != raw_sd_key`` would also fire + # ``filtered_*`` for cleaning-only states, mislabelling cleaned + # stats as filtered. + cleaning_active = chains['clean'] != chains['raw'] + filter_active = chains['filt'] != chains['clean'] + + #we do this to get rewrtten keys for init_sd + rewritten_init_sd = merge_sd_overrides({}, processed_df, init_sd) + intermediate_sd = merge_sds(rewritten_init_sd, cleaned_sd, raw_sd) + base = merge_sd_overrides(intermediate_sd, processed_df, processed_sd) + + # Layer ``cleaned_*`` keys on top when cleaning is active. + if cleaning_active and clean_sd: + for col, stats in clean_sd.items(): + col_dict = base.setdefault(col, {}) + for stat_name, val in stats.items(): + col_dict[f'cleaned_{stat_name}'] = val + + # Layer ``filtered_*`` keys on top when a filter is active. + if filter_active and filt_sd: + for col, stats in filt_sd.items(): + col_dict = base.setdefault(col, {}) + for stat_name, val in stats.items(): + col_dict[f'filtered_{stat_name}'] = val + + return base + class DfTrait(Any): """Any-trait for values that may contain a pandas/polars DataFrame. @@ -161,6 +218,10 @@ def __init__(self, raw_df): # are id()-based cache keys and the dead "foo"/"bar" test branches below, # all of which treat None and [] alike). analysis_klasses: List[Type[ColAnalysis]] = [] + # The content level the summary-stats cascade runs at (see STATS_TIERS). + # Instance-level on CustomizableDataflow; ``full`` here so the bare + # pipeline's dedupe key and cache keys can read it. + stats_tier: str = "full" summary_sd = Any() df_meta = Any() @@ -312,7 +373,7 @@ def _get_summary_sd(self, processed_df: DataFrameT) -> Tuple[SDType, TAny]: ret_summary[col] = {} return ret_summary, {} - _summary_sd_cache_key = (None, None) + _summary_sd_cache_key = (None, None, None) @observe('processed_result', 'analysis_klasses') @exception_protect('summary_sd-protector') @@ -322,7 +383,8 @@ def _summary_sd(self, change): # tuple wrapper, which makes this observer fire twice per widget # construction even when processed_df identity is unchanged. # Skip when neither the dataframe nor analysis_klasses has actually - # changed since the last run. See issue #709. + # changed since the last run. See issue #709. The tier is part of the + # key so a tier change reruns the summary even on the same frame. df = self.processed_df if df is None: # Nothing to summarize before the pipeline has produced a frame. @@ -330,9 +392,9 @@ def _summary_sd(self, change): # narrows Optional[DataFrameT] -> DataFrameT for _get_summary_sd. return klasses = self.analysis_klasses - if (id(df), id(klasses)) == self._summary_sd_cache_key: + if (id(df), id(klasses), self.stats_tier) == self._summary_sd_cache_key: return - self._summary_sd_cache_key = (id(df), id(klasses)) + self._summary_sd_cache_key = (id(df), id(klasses), self.stats_tier) result_summary_sd, errs = self._get_summary_sd(df) self.summary_sd = result_summary_sd self.errs = errs @@ -390,7 +452,11 @@ def __init__(self, orig_df, debug=False, pinned_rows:Union[Literal[None], PinnedRowConfig]=None, extra_grid_config=None, component_config:Union[Literal[None], ComponentConfig]=None, init_sd:Union[Literal[None], InitSD]=None, skip_main_serial=False, - skip_stat_columns=None): + skip_stat_columns=None, stats_tier:str="full"): + if stats_tier not in STATS_TIERS: + raise ValueError(f"stats_tier must be one of {STATS_TIERS}, got {stats_tier!r}") + # Set before super().__init__ — assigning raw_df runs the whole cascade. + self.stats_tier = stats_tier self.init_sd: InitSD if init_sd is None: self.init_sd = {} @@ -520,43 +586,9 @@ def _merged_sd(self, change): if filt_sd is None: filt_sd = self.summary_sd or {} - # Gate each prefixed layer on a chain-shape diff between scopes. - # ``cleaned_*`` fires when the clean chain has ops the raw chain - # doesn't (cleaning is on); ``filtered_*`` fires when the filt - # chain has ops the clean chain doesn't (at least one - # quick-command op is present, i.e. a search filter is on). - # Keying off ``filt_sd_key != raw_sd_key`` would also fire - # ``filtered_*`` for cleaning-only states, mislabelling cleaned - # stats as filtered. - chains = split_chain_by_scope(self.operations) - cleaning_active = chains['clean'] != chains['raw'] - filter_active = chains['filt'] != chains['clean'] - - if self.processed_df is None: - #on initial startup - self.merged_sd = merge_sds(self.init_sd, self.cleaned_sd, raw_sd, self.processed_sd) - return - - #we do this to get rewrtten keys for init_sd - rewritten_init_sd = merge_sd_overrides({}, self.processed_df, self.init_sd) - intermediate_sd = merge_sds(rewritten_init_sd, self.cleaned_sd, raw_sd) - base = merge_sd_overrides(intermediate_sd, self.processed_df, self.processed_sd) - - # Layer ``cleaned_*`` keys on top when cleaning is active. - if cleaning_active and clean_sd: - for col, stats in clean_sd.items(): - col_dict = base.setdefault(col, {}) - for stat_name, val in stats.items(): - col_dict[f'cleaned_{stat_name}'] = val - - # Layer ``filtered_*`` keys on top when a filter is active. - if filter_active and filt_sd: - for col, stats in filt_sd.items(): - col_dict = base.setdefault(col, {}) - for stat_name, val in stats.items(): - col_dict[f'filtered_{stat_name}'] = val - - self.merged_sd = base + self.merged_sd = assemble_merged_sd( + self.init_sd, self.cleaned_sd, raw_sd, self.processed_sd, self.processed_df, + split_chain_by_scope(self.operations), clean_sd=clean_sd, filt_sd=filt_sd) def _compute_scope_df(self, scope: str): """Return the df that scope's SD should be computed against. @@ -598,14 +630,14 @@ def _compute_scope_df(self, scope: str): return base return pp_result[0] if pp_result else base - def _scope_cache_key(self, chain): + def _scope_cache_key(self, chain, tier=None): """Hash that identifies a scope's SD-input identity. Includes the op chain *and* an identifier for the source dataframe (``id(sampled_df)``) *and* the post-processing method - *and* the analysis-klasses identity — all four are inputs to - the scope's SD, and a cache hit must mean "same SD-producing - inputs" not just "same chain". + *and* the analysis-klasses identity *and* the stats tier — all + five are inputs to the scope's SD, and a cache hit must mean + "same SD-producing inputs" not just "same chain". - sampled_df identity addresses codex P1 on #783: a ``raw_df`` swap with an unchanged chain must invalidate. @@ -617,11 +649,18 @@ def _scope_cache_key(self, chain): - analysis_klasses identity addresses codex P2 on #783: a klass-list swap with an unchanged chain must invalidate so new stat klasses surface in ``merged_sd``. + - the tier keeps a schema-tier sd from being read as a full one: + without it, assigning full stats after a schema-tier + construction hits the schema entries and never reaches + ``merged_sd``. ``tier`` defaults to the dataflow's own, and is + explicit so a caller can address the full-tier entry while the + dataflow is still at the schema tier. """ sampled_id = id(self.sampled_df) if self.sampled_df is not None else 0 pp = self.post_processing_method or '' klasses_id = id(self.analysis_klasses) - return hash_chain(chain, extra=f"{sampled_id}|{pp}|{klasses_id}") + tier = tier or self.stats_tier + return hash_chain(chain, extra=f"{sampled_id}|{pp}|{klasses_id}|{tier}") @observe('summary_sd', 'operations', 'analysis_klasses') @exception_protect('sd-cache-protector') @@ -653,9 +692,18 @@ def _populate_sd_cache(self, _change): new_cache = dict(self.summary_stats_cache) cache_grew = False - # filt scope reuses the SD that _summary_sd just produced. + # filt scope reuses the SD that _summary_sd just produced, but only if + # it was produced for this frame, klass list and tier. An sd left over + # from another (a tier changed since it ran, or an assignment ahead of + # the cascade) would be stored under this key and never repaired, since + # a present key is a hit. Compute the scope's SD instead. if keys['filt'] not in new_cache: - new_cache[keys['filt']] = dict(self.summary_sd or {}) + current = (id(self.processed_df), id(self.analysis_klasses), self.stats_tier) + if self._summary_sd_cache_key == current: + new_cache[keys['filt']] = dict(self.summary_sd or {}) + else: + filt_sd, _errs = self._get_summary_sd(self.processed_df) + new_cache[keys['filt']] = filt_sd cache_grew = True # raw + clean: fresh compute, but only on cache miss. @@ -711,6 +759,8 @@ def _build_error_dataframe(self, e) -> pd.DataFrame: #TAny closer to some error type @override def _get_summary_sd(self, processed_df: DataFrameT) -> Tuple[SDType, ErrDict]: + if self.stats_tier == "schema": + return self._get_schema_sd(processed_df), {} stats = self.DFStatsClass( processed_df, self.analysis_klasses, @@ -728,6 +778,16 @@ def _get_summary_sd(self, processed_df: DataFrameT) -> Tuple[SDType, ErrDict]: return sdf, {} + def _get_schema_sd(self, processed_df: DataFrameT) -> SDType: + """Identity and typing keys for every column, from the schema alone. + + ``orig_col_name``, ``rewritten_col_name``, ``dtype``, the ``is_*`` + flags, ``_type`` and ``length``: the keys column styling needs to pick + displayers, with no data query beyond the row count. A backend opts in + to ``stats_tier="schema"`` by overriding this. + """ + raise NotImplementedError(f"{type(self).__name__} has no schema stats tier") + # ### end summary stats block def _sd_to_jsondf(self, sd:SDType): @@ -745,35 +805,43 @@ def _sd_to_jsondf(self, sd:SDType): def _df_to_obj(self, df: DataFrameT) -> List[TDict[str, TAny]]: return pd_to_obj(self.sampling_klass.serialize_sample(df)) + def _analysis_klasses_with(self, analysis_klass:Type[ColAnalysis]) -> List[Type[ColAnalysis]]: + """The klass list the stats pipelines' ``add_stat`` builds: a class of the + same name is replaced, anything else is appended.""" + klasses = list(self.analysis_klasses) + if isinstance(analysis_klass, type): + klasses = [k for k in klasses + if not (isinstance(k, type) and k.__name__ == analysis_klass.__name__)] + klasses.append(analysis_klass) + return klasses + def add_analysis(self, analysis_klass:Type[ColAnalysis]) -> None: """ same as get_summary_sd, call whatever to set summary_sd and trigger further comps """ - stats = self.DFStatsClass( - self.processed_df, - self.analysis_klasses, - self.df_name, debug=self.debug) - stats.add_analysis(analysis_klass) - - self.analysis_klasses = stats.ap.ordered_a_objs + if self.stats_tier == "full": + stats = self.DFStatsClass( + self.processed_df, + self.analysis_klasses, + self.df_name, debug=self.debug) + stats.add_analysis(analysis_klass) + self.analysis_klasses = stats.ap.ordered_a_objs + else: + # Building DFStatsClass here would run the stats this tier exists + # to skip (twice, since the constructor and add_analysis each + # process the frame), and it does not go through _get_summary_sd. + self.analysis_klasses = self._analysis_klasses_with(analysis_klass) self.DFStatsClass.verify_analysis_objects(self.analysis_klasses) self.setup_options_from_analysis() #force recomputation self._handle_widget_change({}) - #final processing block - @observe('widget_args_tuple') - def _handle_widget_change(self, change): - """ - put together df_dict for consumption by the frontend - """ - # Tuple[TAny, pd.DataFrame, SDType] - _unused, processed_df, merged_sd = self.widget_args_tuple - if processed_df is None: - return - + def _build_df_data_dict(self, processed_df, merged_sd): + """The ``df_data_dict`` trait value: the main rows, the ``all_stats`` + payload and the empty placeholder. ``all_stats`` is the 5-11 ms part, + which is why it is built apart from ``_build_df_display_args``.""" # df_data_dict is still hardcoded for now # eventually processed_df will be able to add or alter values of df_data_dict # correlation would be added, filtered would probably be altered @@ -781,15 +849,15 @@ def _handle_widget_change(self, change): # to expedite processing maybe future provided dfs from # postprcoessing could default to empty until that is # selected, optionally - if self.skip_main_serial: - self.df_data_dict = {'main': [], - 'all_stats': self._sd_to_jsondf(merged_sd), - 'empty': []} - else: - self.df_data_dict = {'main': self._df_to_obj(processed_df), - 'all_stats': self._sd_to_jsondf(merged_sd), - 'empty': []} - + main = [] if self.skip_main_serial else self._df_to_obj(processed_df) + return {'main': main, + 'all_stats': self._sd_to_jsondf(merged_sd), + 'empty': []} + + def _build_df_display_args(self, processed_df, merged_sd): + """The ``df_display_args`` trait value: each display's column_config, + pinned_rows and keys, with the overrides applied. Reads ``merged_sd`` + but serializes none of it.""" temp_display_args = {} for display_name, A_Klass in self.df_display_klasses.items(): df_viewer_config = A_Klass.get_dfviewer_config(merged_sd, processed_df) @@ -807,8 +875,21 @@ def _handle_widget_change(self, change): temp_display_args['main']['df_viewer_config']['extra_grid_config'] = self.extra_grid_config if self.component_config: temp_display_args['main']['df_viewer_config']['component_config'] = self.component_config + return temp_display_args + + #final processing block + @observe('widget_args_tuple') + def _handle_widget_change(self, change): + """ + put together df_dict for consumption by the frontend + """ + # Tuple[TAny, pd.DataFrame, SDType] + _unused, processed_df, merged_sd = self.widget_args_tuple + if processed_df is None: + return - self.df_display_args = temp_display_args + self.df_data_dict = self._build_df_data_dict(processed_df, merged_sd) + self.df_display_args = self._build_df_display_args(processed_df, merged_sd) """ diff --git a/buckaroo/pluggable_analysis_framework/stat_pipeline.py b/buckaroo/pluggable_analysis_framework/stat_pipeline.py index 67e8c0f77..f121624e0 100644 --- a/buckaroo/pluggable_analysis_framework/stat_pipeline.py +++ b/buckaroo/pluggable_analysis_framework/stat_pipeline.py @@ -9,10 +9,11 @@ from __future__ import annotations import time -from typing import Any, Dict, List, Optional, Tuple +from typing import Any, Callable, Dict, List, Optional, Tuple import pandas as pd +from buckaroo.dataflow.df_types import DataFrameLike from buckaroo.df_util import old_col_new_col from . import perf_log @@ -432,3 +433,31 @@ def errors_to_errdict(errors: List[StatError]) -> ErrDict: err_key = (se.column, se.stat_func.name if se.stat_func else "unknown") errs[err_key] = (se.error, None) return errs + + +def schema_sd(df: DataFrameLike, column_typing: Callable[[Any], Dict[str, Any]], + skip_columns=None) -> SDType: + """The summary dict for ``df`` from its schema alone, with no stat run. + + Shaped like ``StatPipeline.process_df``'s output: one entry per column, + keyed by the rewritten name, holding ``orig_col_name``, + ``rewritten_col_name``, ``length`` and whatever ``column_typing`` returns + for the column. ``column_typing`` is handed the column's series and may read + its dtype only; it is what keeps this from touching a value. As in + ``process_df``, a frame with no rows gives an empty dict and a column in + ``skip_columns`` gets only its two names, so its typing comes from + ``init_sd``. + """ + length = len(df) + if length == 0: + return {} + skip = set(skip_columns or ()) + summary: SDType = {} + for orig_col_name, rewritten_col_name in old_col_new_col(df): + col_stats: Dict[str, Any] = { + 'orig_col_name': orig_col_name, 'rewritten_col_name': rewritten_col_name} + if orig_col_name not in skip and rewritten_col_name not in skip: + col_stats.update(column_typing(df[orig_col_name])) + col_stats['length'] = length + summary[rewritten_col_name] = col_stats + return summary diff --git a/buckaroo/server/data_loading.py b/buckaroo/server/data_loading.py index 1db3385cc..e162e6901 100644 --- a/buckaroo/server/data_loading.py +++ b/buckaroo/server/data_loading.py @@ -12,10 +12,11 @@ from buckaroo.dataflow.dataflow_extras import Sampling from buckaroo.dataflow.autocleaning import PandasAutocleaning from buckaroo.dataflow.styling_core import InitSD, StylingAnalysis -from buckaroo.customizations.pd_stats_v2 import PD_ANALYSIS_V2 +from buckaroo.customizations.pd_stats_v2 import PD_ANALYSIS_V2, schema_stats from buckaroo.customizations.styling import DefaultSummaryStatsStyling, DefaultMainStyling from buckaroo.customizations.pd_autoclean_conf import CleaningConf, NoCleaningConf from buckaroo.pluggable_analysis_framework.df_stats_v2 import DfStatsV2 +from buckaroo.pluggable_analysis_framework.stat_pipeline import schema_sd class ServerSampling(Sampling): @@ -49,6 +50,9 @@ def _df_to_obj(self, df): # No sampling — matches BuckarooInfiniteWidget._df_to_obj return pd_to_obj(df) + def _get_schema_sd(self, processed_df): + return schema_sd(processed_df, schema_stats, skip_columns=self.skip_stat_columns) + def create_dataflow(df: pd.DataFrame, column_config_overrides=None, extra_grid_config=None, init_sd: InitSD | None = None) -> ServerDataflow: diff --git a/buckaroo/server/data_loading_polars.py b/buckaroo/server/data_loading_polars.py index 64e211782..adadee42b 100644 --- a/buckaroo/server/data_loading_polars.py +++ b/buckaroo/server/data_loading_polars.py @@ -27,7 +27,9 @@ from buckaroo.dataflow.styling_core import InitSD from buckaroo.dataflow.autocleaning import PandasAutocleaning from buckaroo.customizations.pl_autocleaning_conf import NoCleaningConfPl +from buckaroo.customizations.pl_stats_v2 import pl_schema_stats from buckaroo.pluggable_analysis_framework.df_stats_v2 import PlDfStatsV2 +from buckaroo.pluggable_analysis_framework.stat_pipeline import schema_sd from buckaroo.polars_buckaroo import ( PLSampling, local_analysis_klasses, prepare_df_for_serialization) from buckaroo.serialization_utils import pd_to_obj, make_infinite_resp @@ -59,6 +61,9 @@ def _df_to_obj(self, df): return pd_to_obj(self.sampling_klass.serialize_sample(df)) return pd_to_obj(self.sampling_klass.serialize_sample(df.to_pandas())) + def _get_schema_sd(self, processed_df): + return schema_sd(processed_df, pl_schema_stats, skip_columns=self.skip_stat_columns) + def load_file_polars(path: str) -> pl.DataFrame: """Eager polars read. Extension dispatch mirrors :func:`load_file`. diff --git a/buckaroo/server/handlers.py b/buckaroo/server/handlers.py index 3e0802fa7..1d86d533f 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.server.session import ( + STATS_DELIVERIES, STATS_TIER_REQUESTS, begin_stats_generation, dataflow_stats_tier) +from buckaroo.server.stats_wire import ( + apply_source_guards, broadcast_state, refresh_session_snapshot, resolve_session_guards, resolve_session_policy) 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.""" @@ -387,6 +383,14 @@ async def post(self): **component_config, } + # /load serves the eager backends, whose stats cost does not grow with + # the file, so they resolve to full: it reads no stats_tier and builds + # no deferred stats, whatever policy a prior /load_expr on this session + # left behind, and the generation moves on with the data. + session.stats_tier, session.stats_delivery, session.stats_policy = "full", "inline", None + session.source_guards = None + 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) @@ -396,6 +400,29 @@ async def post(self): self.write({"session": session_id, "server_pid": os.getpid(), "browser_action": browser_action, **metadata}) +def _stats_policy_from_body(body: dict, current_tier: str, current_delivery: str): + """Read ``stats_tier`` (``auto``, ``full``, ``scalar`` or ``schema``) and + ``stats_delivery`` from a request body. + + A field the body omits (or sends as null) keeps the current value (the + session's, or the default for a new session), as ``cache_dir`` does. Returns + ``(stats_tier, stats_delivery, None)``, or ``(None, None, error)`` with the + 400 response body for an unknown value.""" + stats_tier = body.get("stats_tier") + if stats_tier is None: + stats_tier = current_tier + if stats_tier not in STATS_TIER_REQUESTS: + return None, None, {"error_code": "invalid_stats_tier", + "message": f"stats_tier must be one of {list(STATS_TIER_REQUESTS)}, got {stats_tier!r}"} + stats_delivery = body.get("stats_delivery") + if stats_delivery is None: + stats_delivery = current_delivery + if stats_delivery not in STATS_DELIVERIES: + return None, None, {"error_code": "invalid_stats_delivery", + "message": f"stats_delivery must be one of {list(STATS_DELIVERIES)}, got {stats_delivery!r}"} + return stats_tier, stats_delivery, None + + 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. @@ -454,7 +481,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")) @@ -479,11 +509,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 (#944). telemetry.arm_session(existing, tele_sink) @@ -532,6 +573,7 @@ async def post(self): "skip_stat_columns": body.get("skip_stat_columns")} project_root = body.get("project_root") + dataflow_tier = dataflow_stats_tier(stats_tier, stats_delivery) try: with telemetry.firstpull_load(session_id, tele_sink, "load_expr", build_dir=build_dir): @@ -554,7 +596,7 @@ async def post(self): with perf_log.perf_span("firstpull.dataflow_construct", session=session_id): xorq_dataflow = xorq_loading.XorqServerDataflow( expr, skip_main_serial=True, extra_klasses=extra_klasses, - **dataflow_kwargs) + stats_tier=dataflow_tier, **dataflow_kwargs) # Spanning metadata too leaves only the small klass-load step # unmeasured inside the outer firstpull.load_expr total. with perf_log.perf_span("firstpull.metadata", session=session_id): @@ -570,6 +612,15 @@ async def post(self): self.write(resp) return + # The dataflow and the count exist: this is where the size policy is + # resolved. A dataflow built at the schema tier has run no stat query, so + # the policy can still keep it there; one that ran its stats in the + # constructor has nothing left to resolve (None). + stats_policy = resolve_session_policy(stats_tier, dataflow_tier, metadata["rows"], len(metadata["columns"])) + # The guards for huge sources ride with the policy: a session a host opened + # with one has them, any other has none. + source_guards = resolve_session_guards(dataflow_tier, metadata["rows"], len(metadata["columns"])) + sessions = self.application.settings["sessions"] session = sessions.get_or_create(session_id, build_dir) session.mode = "buckaroo" @@ -579,6 +630,11 @@ 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.stats_policy = stats_policy + # Before the snapshot below: it holds the config a client sorts by. + apply_source_guards(session, xorq_dataflow, source_guards) session.xorq_dataflow = xorq_dataflow # Clear pandas-side state left by a prior /load on the same # session so WS dispatch can no longer reach a stale dataflow. @@ -616,17 +672,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" @@ -770,18 +822,13 @@ async def post(self): session.df_meta = display_state["df_meta"] session.mode = "viewer" telemetry.arm_session(session, tele_sink) + # A viewer session has no dataflow and so no deferred stats. + session.stats_tier, session.stats_delivery, session.stats_policy = "full", "inline", None + session.source_guards = None + 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): @@ -811,8 +858,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. The stats policy is resolved again for the pair. + Returns 404 when the session does not exist, 400 when it is not a xorq - session 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"] @@ -835,6 +896,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: @@ -849,9 +917,16 @@ async def post(self, session_id): xorq_loading.load_project_stat_klasses(session.project_root) + xorq_loading.load_project_post_processing_klasses(session.project_root) + xorq_loading.load_project_display_klasses(session.project_root)) + dataflow_tier = dataflow_stats_tier(stats_tier, stats_delivery) xorq_dataflow = xorq_loading.XorqServerDataflow( session.expr, skip_main_serial=True, extra_klasses=extra_klasses, - **session.dataflow_kwargs) + stats_tier=dataflow_tier, **session.dataflow_kwargs) + # Resolved again, with the count the load took: the expression is + # reused, so there is no new one to take. + stats_policy = resolve_session_policy(stats_tier, dataflow_tier, session.metadata["rows"], + len(session.metadata["columns"])) + source_guards = resolve_session_guards(dataflow_tier, session.metadata["rows"], + len(session.metadata["columns"])) except Exception: tb = traceback.format_exc() log.error("reload_expr error session=%s: %s", session_id, tb) @@ -876,29 +951,15 @@ async def post(self, session_id): if bs.get("quick_command_args"): xorq_dataflow.quick_command_args = bs["quick_command_args"] - refreshed = get_buckaroo_display_state(xorq_dataflow) session.xorq_dataflow = xorq_dataflow - session.df_display_args = refreshed["df_display_args"] - session.df_data_dict = refreshed["df_data_dict"] - session.df_meta = refreshed["df_meta"] - session.buckaroo_options = refreshed["buckaroo_options"] - session.command_config = refreshed["command_config"] - - if session.component_config and session.df_display_args: - for key in session.df_display_args: - dvc = session.df_display_args[key].get("df_viewer_config") - if dvc is not None: - dvc["component_config"] = { - **dvc.get("component_config", {}), - **session.component_config} - - for client in list(session.ws_clients): - try: - msg = build_state_message(session, - search_string=getattr(client, "search_string", "")) - client.write_message(json.dumps(msg)) - except Exception: - session.ws_clients.discard(client) + session.stats_tier = stats_tier + session.stats_delivery = stats_delivery + session.stats_policy = stats_policy + apply_source_guards(session, xorq_dataflow, source_guards) + refresh_session_snapshot(session, xorq_dataflow) + begin_stats_generation(session) + + broadcast_state(session) klass_count = len(extra_klasses) log.info("reload_expr session=%s project_root=%s klasses=%d", diff --git a/buckaroo/server/session.py b/buckaroo/server/session.py index 563ce28fc..1ecb623d8 100644 --- a/buckaroo/server/session.py +++ b/buckaroo/server/session.py @@ -1,3 +1,4 @@ +import copy import logging import time from dataclasses import dataclass, field @@ -11,6 +12,71 @@ _DEFAULT_SESSION_TTL_S = 3600.0 # 1 hour idle before eviction _DEFAULT_EVICTION_INTERVAL_S = 300.0 # check every 5 minutes +# How a session reaches its stats tier: ``inline`` computes them in the +# dataflow constructor (today's behaviour); ``deferred`` builds the schema tier +# first and leaves the rest to later requests. +STATS_DELIVERIES = ("inline", "deferred") + +# What a host may send as ``stats_tier`` on /load_expr and /reload_expr. ``auto`` +# lets the server decide from the size of the entry (``stats_policy.py``); the +# others name a tier, which the server can lower (the ceiling) but never raise. +# Not the dataflow's own tiers (``dataflow.STATS_TIERS``: ``full`` and ``schema``), +# which are what a dataflow is built at. +STATS_TIER_REQUESTS = ("auto", "full", "scalar", "schema") + +# Why a session is ``not_computed``, in ``df_meta.stats.reason``: the server's size +# rule, the host naming a tier, the cost guard pausing a run (phase 6a), or the +# ceiling lowering a request. (``error`` has its own reason, ``stats_failed``.) +STATS_REASONS = ("size", "host", "cost", "ceiling") + +# What a client assumes for a ``df_meta.stats`` field the server leaves out. The +# server omits a field that equals its default, so a message with no policy to +# report is the one it always was. +STATS_FIELD_DEFAULTS: Dict[str, Any] = {"auto_request": True, "requestable": ["full"], "omitted_keys": [], + "approx_keys": [], "demand_columns": []} + +# What a client assumes for the ``df_meta.sort`` and ``df_meta.search`` flags +# (rows-first p37) a message leaves out. The server sends a flag only when a +# source is too large for that request, so a session under the thresholds sends +# the message it always has. +SOURCE_GUARD_DEFAULTS: Dict[str, str] = {"sort": "enabled", "search": "enabled"} + + +def dataflow_stats_tier(stats_tier: str, stats_delivery: str) -> str: + """The tier a session's dataflow is constructed at. Only a session whose + stats run inline and whose host named no lower tier builds at the full tier; + a deferred one starts at the schema tier whatever tier it is headed for, and + so does one headed for ``scalar`` (no scalar-tier units exist yet) or + ``schema``. ``auto`` is resolved against a schema-tier dataflow, which an + inline session never has, so inline ``auto`` is ``full``.""" + if stats_delivery == "deferred" or stats_tier in ("scalar", "schema"): + return "schema" + return "full" + + +# What ``df_meta.stats.status`` says about a session's stats for its current +# ``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, policy: Optional[dict] = None) -> tuple[str, Optional[str]]: + """The status, and its reason, of a session whose stats generation has just + started, from the policy pair and the policy resolved at load. A target below + ``full`` is ``not_computed`` for the reason the policy gives; with no policy + (a dataflow built with its stats inline) a tier named below ``full`` is the + host's.""" + if policy is not None and policy["tier_target"] != "full": + return "not_computed", policy["reason"] + if policy is None and stats_tier in ("scalar", "schema"): + return "not_computed", "host" + if stats_delivery == "deferred": + return "pending", None + return "complete", None + @dataclass class SessionState: @@ -43,6 +109,38 @@ class SessionState: # extra_grid_config, init_sd, skip_stat_columns), replayed by /reload_expr so # a reload keeps stat caching and column config (#957). dataflow_kwargs: dict = field(default_factory=dict) + # The /load_expr stats policy. ``stats_tier`` is what the host asked for (see + # STATS_TIER_REQUESTS) and ``stats_delivery`` how the stats get there (see + # STATS_DELIVERIES). Kept apart from dataflow_kwargs, which is splatted into + # the dataflow constructor, and replayed by /reload_expr. The tier the + # dataflow is built at follows from the pair (dataflow_stats_tier). + stats_tier: str = "full" + stats_delivery: str = "inline" + # What ``stats_policy.resolve_stats_policy`` made of the pair and the size of + # the entry (``tier_target``, ``auto_request``, ``requestable``, ``reason``, + # ``estimate``), resolved by /load_expr and /reload_expr once the schema-tier + # dataflow and the count exist. ``None`` for a dataflow built with its stats + # inline, and for anything /load serves. Later phases add ``omitted_keys``, + # ``approx_keys`` and ``demand_columns``, which ``stats_meta`` reports. + stats_policy: Optional[dict] = None + # Whether sort and search stay on for the entry (``stats_policy.resolve_source_guards``: + # ``{"sort": "enabled" | "disabled", "search": ...}``), resolved by /load_expr and + # /reload_expr for a session a host opened with a stats policy. ``None`` for a + # dataflow built with its stats inline, and for anything /load serves, which are + # not guarded. ``build_state_message`` reports a disabled flag in ``df_meta``, + # and a sorted window of a session whose sort is disabled is refused + # (``sort_refusal``). + source_guards: Optional[dict] = None + # The stats generation: a counter the server owns, bumped whenever the + # dataflow state the stats describe changes (/load, /load_expr, /load_compare, + # /reload_expr, and a buckaroo_state_change that touches a dataflow field). + # It rides on df_meta.stats and on every stats message, so a client can drop + # a reply for a state it has left. Independent of any client-owned request + # sequence (state_seq). ``stats_status`` and ``stats_reason`` live here, not + # in ``df_meta``, because the dataflow rebuilds df_meta wholesale. + stats_gen: int = 0 + stats_status: str = "complete" + stats_reason: Optional[str] = None # Companion telemetry sink (#943): a fire-and-forget POST callable, built # once from the /load_expr payload's telemetry_url on the IOLoop (where # make_http_sink captures AsyncHTTPClient/IOLoop.current()). Stored here so @@ -80,8 +178,117 @@ def touch(self) -> None: this field is the runtime escape hatch.""" +def begin_stats_generation(session: "SessionState") -> None: + """Start a new stats generation: the session's dataflow state has changed, + so stats for the previous one no longer describe it. The status restarts + from the session's policy pair.""" + session.stats_gen += 1 + session.stats_status, session.stats_reason = initial_stats_status( + session.stats_tier, session.stats_delivery, session.stats_policy) + + +def stats_deferred_by_policy(session: "SessionState") -> bool: + """Whether the server, not the host, put the session below ``full``: the host + asked for ``auto`` or ``full`` and the policy resolved lower (a size rule, or + the ceiling). A tier the host named (``scalar``, ``schema``) is its own + choice. A client that cannot take the ``not_computed`` state is served such a + session as one headed for ``full`` (see ``stats_wire.stats_to_pull``).""" + return session.stats_status == "not_computed" and session.stats_tier in ("auto", "full") + + +def _policy_fields(policy: dict) -> dict: + """The part of a resolved policy ``df_meta.stats`` reports: ``tier_target`` and + ``estimate`` always, and the other fields only where they differ from + ``STATS_FIELD_DEFAULTS``.""" + fields: dict = {"tier_target": policy["tier_target"], "estimate": dict(policy["estimate"])} + for name in ("auto_request", "requestable"): + if policy[name] != STATS_FIELD_DEFAULTS[name]: + fields[name] = list(policy[name]) if isinstance(policy[name], list) else policy[name] + for name in ("omitted_keys", "approx_keys", "demand_columns"): + if policy.get(name): + fields[name] = list(policy[name]) + return fields + + +def stats_meta(session: "SessionState", ondemand: bool = True) -> Optional[dict]: + """The ``df_meta.stats`` value for a session: ``{status, tier, gen}`` plus + ``reason`` when there is one. ``tier`` is the tier reached so far, which is + ``full`` only once the session is complete. + + ``ondemand`` says whether the recipient advertised ``stats_ondemand``, so it + can take a ``not_computed`` state the server chose and the policy fields. Any + other client is told a session the server put below ``full`` is ``pending``, + headed for ``full``, with no policy fields. For an ondemand client a session + that is ``pending`` or ``not_computed`` adds the policy it resolved + (``_policy_fields``), unless that is an explicit ``full`` within the ceiling, + which sends the message it always has. + + ``None`` for a session on the default policy (inline delivery, full tier, + complete), which sends the message it always has; a client reads a missing + ``stats`` as complete (``stats_with_defaults``).""" + if session.stats_status == "complete" and session.stats_delivery != "deferred": + return None + status, reason = session.stats_status, session.stats_reason + if not ondemand and stats_deferred_by_policy(session): + status, reason = "pending", None + stats: dict = {"status": status, "tier": "full" if status == "complete" else "schema", "gen": session.stats_gen} + if reason: + stats["reason"] = reason + policy = session.stats_policy + reported = policy is not None and not (session.stats_tier == "full" and policy["tier_target"] == "full") + if ondemand and reported and status in ("pending", "not_computed"): + stats.update(_policy_fields(policy)) + return stats + + +def stats_with_defaults(df_meta: Optional[dict]) -> dict: + """``df_meta.stats`` as a client reads it, with every field filled in: what a + message leaves out takes its documented default. No ``stats`` at all is what + an old server sends and means ``complete``. ``tier_target`` defaults to the + tier reached when the status is ``complete`` or ``not_computed`` and to + ``full`` otherwise (a ``pending`` or ``error`` session is headed for + ``full``). ``gen``, ``reason`` and ``estimate`` are ``None`` when absent.""" + stats = (df_meta or {}).get("stats") or {"status": "complete", "tier": "full"} + status, tier = stats["status"], stats["tier"] + out: dict = {"status": status, "tier": tier, + "tier_target": stats.get("tier_target", tier if status in ("complete", "not_computed") else "full"), + "gen": stats.get("gen"), "reason": stats.get("reason"), "estimate": stats.get("estimate")} + for name, default in STATS_FIELD_DEFAULTS.items(): + out[name] = copy.deepcopy(stats[name] if name in stats else default) + return out + + +def guards_with_defaults(df_meta: Optional[dict]) -> dict: + """The ``sort`` and ``search`` flags of a ``df_meta`` as a client reads them: + one the message leaves out is ``enabled``, which is also what an old server + sends.""" + meta = df_meta or {} + return {name: meta.get(name, default) for name, default in SOURCE_GUARD_DEFAULTS.items()} + + +def _guard_flags(session: "SessionState") -> dict: + """The flags ``df_meta`` reports for the session: the ones that are not at their default.""" + return {name: value for name, value in (session.source_guards or {}).items() + if name in SOURCE_GUARD_DEFAULTS and value != SOURCE_GUARD_DEFAULTS[name]} + + +def sort_refusal(session: "SessionState", payload_args: Any) -> Optional[dict]: + """The ``infinite_resp`` that refuses a window with a ``sort`` on a session + whose sort is disabled, or ``None`` for any other request. It carries + ``error_code: "sort_disabled"`` beside the ``error_info`` every error response + has, and no rows, so the grid sees an error for that window and no parquet frame + follows.""" + guards = session.source_guards + if not guards or guards.get("sort") != "disabled": + return None + if not isinstance(payload_args, dict) or not payload_args.get("sort"): + return None + return {"type": "infinite_resp", "key": payload_args, "length": 0, "error_code": "sort_disabled", + "error_info": "Sorting is turned off for this source: it has too many rows to sort in a request."} + + def build_state_message(session: "SessionState", metadata: dict | None = None, - search_string: str = "") -> dict: + search_string: str = "", reply_seq: int | None = None, ondemand: bool = True) -> dict: """Build the full ``initial_state`` WebSocket payload from a session. Args: @@ -94,14 +301,33 @@ def build_state_message(session: "SessionState", metadata: dict | None = None, search box. The session itself never owns this value; callers must pass the right value per recipient (typically ``handler.search_string``). + reply_seq: The ``state_seq`` of the ``buckaroo_state_change`` this + message answers (#998). Set only on the copy sent to the client + that made the change; that client drops a reply older than its + latest change so an overlapping rerun can't put its + ``buckaroo_state`` back. Omitted (``None``) for the broadcast + copies other clients get, the ``/load`` push and a fresh + connection, which the client applies unconditionally. Optional + on the wire, so it doesn't bump ``PROTOCOL_VERSION``. + ondemand: Whether the recipient advertised ``stats_ondemand`` (see + ``stats_meta``); ``stats_wire.build_state_message_for`` passes it. Returns: A dict ready to be JSON-serialised and sent to WebSocket clients. """ + # The dataflow rebuilds df_meta wholesale, so the stats status and the guard + # flags are injected here, into a copy, rather than stored in it. + df_meta = session.df_meta + stats = stats_meta(session, ondemand=ondemand) + flags = _guard_flags(session) + if stats is not None or flags: + df_meta = {**df_meta, **flags, **({"stats": stats} if stats is not None else {})} msg: dict = {"type": "initial_state", "protocol_version": PROTOCOL_VERSION, "metadata": metadata if metadata is not None else session.metadata, "prompt": session.prompt, "df_display_args": session.df_display_args, "df_data_dict": session.df_data_dict, - "df_meta": session.df_meta, "mode": session.mode} + "df_meta": df_meta, "mode": session.mode} + if reply_seq is not None: + msg["reply_seq"] = reply_seq if session.mode == "buckaroo": # Per-client search_string overlay (#851 Codex P1): the snapshot # on the session is search-agnostic; we re-inject the recipient's diff --git a/buckaroo/server/stats_policy.py b/buckaroo/server/stats_policy.py new file mode 100644 index 000000000..d34f29d6c --- /dev/null +++ b/buckaroo/server/stats_policy.py @@ -0,0 +1,402 @@ +"""Which stats tier an entry should reach, decided from its size. + +Pure functions. Nothing in the handlers, the session or the dataflow imports +this module yet; it exists so the policy can be tested without a backend +before it is wired in. + +Tiers, lowest to highest: + +* ``schema``: dtype, ``_type``, identity keys and the row count. Never + requested automatically. +* ``scalar``: adds ``null_count``, ``min``, ``max``, ``mean``, ``std``, + ``empty_count`` and ``histogram_bins``. No histogram query, + ``value_counts``, median or (on xorq) ``distinct_count``. +* ``full``: today's stats. + +:func:`resolve_stats_policy` picks the tier. A hard ceiling is applied inside +it, so the result is the lower of the requested tier and the ceiling tier +whoever asks (a load, ``/reload_expr`` or a ``stats_request`` with ``force``). +A host tier can lower the result freely and raise it only up to the ceiling. +Eager pandas and polars resolve to ``full``: their stats sample 50,000 rows, +so the cost does not grow with the file. + +:func:`route_polars_entry` is the separate size rule for polars: an entry +above R rows belongs on xorq's ``/load_expr``, not eager ``/load``. + +:func:`resolve_source_guards` is the rule for the two requests that stay expensive +whether or not stats are computed, a sorted window and a searched one (rows-first +p37). Its thresholds are apart from the stats ones (:class:`GuardLimits`), because a +sort or a search scales with rows where the stats tiers scale with cells, and +:func:`disable_sorting` is the pass that turns sorting off in a display config. + +The thresholds are PROVISIONAL. They are the values the phase-0 measurements +proposed (one Apple M4 Pro, 115 tallyman telemetry loads, xorq stats on +parquet and CSV slices, eager polars RSS), and they rest on gaps: the +telemetry has nothing between 11.8M and 42.3M rows, nothing was measured on +the xorq side above 12M rows, and the scalar ceiling is an extrapolation. +They live in the ``DEFAULT_*`` constants below, each with the measurement it +comes from, and each has an environment override that is read on every call, +not at import. A value that is not a non-negative integer is ignored with a +warning. + +* ``BUCKAROO_STATS_FULL_AUTO_ROWS`` (12,000,000 rows) and + ``BUCKAROO_STATS_FULL_AUTO_CELLS`` (520,000,000 cells): ``full`` is chosen + automatically while both hold. +* ``BUCKAROO_STATS_SCALAR_AUTO_CELLS`` (1,000,000,000 cells): otherwise + ``scalar`` is chosen automatically up to this many cells (rows x columns). + Above it the entry stays at ``schema``. +* ``BUCKAROO_STATS_CEILING_FULL_ROWS`` (25,000,000 rows) and + ``BUCKAROO_STATS_CEILING_FULL_CELLS`` (1,000,000,000 cells): hard ceiling + on ``full``, refused above either bound, including for a forced request. +* ``BUCKAROO_STATS_CEILING_SCALAR_CELLS`` (4,000,000,000 cells): hard ceiling + on ``scalar``, in cells. Above it only ``schema`` is allowed. This number + is an extrapolation with no measurement behind it. +* ``BUCKAROO_POLARS_ROUTE_ROWS`` (8,000,000 rows): R for + :func:`route_polars_entry`. A memory threshold, set apart from the stats + ones, and sized for ``pre_limit`` False (see the constant below). + +The guard thresholds are PROVISIONAL too and read the same way (on every call, a +bad value ignored with a warning): + +* ``BUCKAROO_SORT_DISABLE_ROWS`` (25,000,000 rows): a source with more rows has + sorting turned off. +* ``BUCKAROO_SEARCH_DISABLE_ROWS`` (none): a source with more rows reports its + search as disabled. The flag is advice for a client; the server serves a search + either way. +""" +from __future__ import annotations + +import logging +import operator +import os +from dataclasses import dataclass +from typing import Any + +import pyarrow.parquet as pq + +log = logging.getLogger("buckaroo.server.stats_policy") + +# Provisional values from the phase-0 measurements; see the module docstring. +# 105 of 115 telemetry loads are at or below 12M rows (stats p90 4.4 s). Every +# load above it is one of three CSV unions (42.3M, 54.1M, 78.0M rows) at 75, +# 114 and 215 s. +DEFAULT_FULL_AUTO_ROWS: int = 12_000_000 +# Parquet 12M x 43 (516M cells) full stats took 2.9 s over 8 repetitions. The +# bound only bites on an entry wider than 43 columns at 12M rows. +DEFAULT_FULL_AUTO_CELLS: int = 520_000_000 +# The scalar class costs 1.1-1.3 ms per Mcell on a parquet scan and 3-6 ms on a +# CSV union. At the worst telemetry batch rates (30-48 ms per Mcell, scalar +# being 0.44-0.71 of the batch) that is about 20-34 s at 1.0B cells. +DEFAULT_SCALAR_AUTO_CELLS: int = 1_000_000_000 +# The ceiling on full separates the slowest in-limit load (11.8M rows, 22 s) +# from the fastest one above it (42.3M rows, 75 s). 25M rows and 1.0B cells are +# near the geometric midpoint of that gap (22M rows, 1.0B cells), and it is a +# gap in the telemetry, so the true limit could sit anywhere inside it. +DEFAULT_CEILING_FULL_ROWS: int = 25_000_000 +DEFAULT_CEILING_FULL_CELLS: int = 1_000_000_000 +# An extrapolation with no measurement behind it. The only evidence above 1.8B +# cells is the 3.36B-cell entry whose batch took 162.5 s, and scalar was not +# measured there. ``None`` means no ceiling. +DEFAULT_CEILING_SCALAR_CELLS: int | None = 4_000_000_000 +# Eager polars RSS is 0.28 GB + 0.543 GB per Mrow, and with ``pre_limit`` False +# a sorted window peaks at 0.23 + 1.62 GB per Mrow (3.0x the frame). With a +# 32 GB budget, four files resident and one window in flight, 8M rows is 27 GB +# and 10M is 33.5 GB. That figure assumes ``pre_limit`` False; main's 1M +# ``pre_limit`` windows do not copy the frame and the resident-only bound would +# be about 14M rows. +DEFAULT_POLARS_ROUTE_ROWS: int = 8_000_000 + +# Sorting a window sorts the whole source on the loop thread. At 10.8M rows x 43 +# columns that took 261-287 ms warm on parquet, and at 3M rows of CSV 326-337 ms +# (phase-0 measurements, item 4). Extrapolated linearly (nothing above 12M rows was +# measured) a 25M-row source costs about 0.7 s on parquet and 2.8 s on CSV, and the +# 42.3M, 54.1M and 78.0M entries 1.1 to 2.1 s and 4.6 to 8.5 s. 25M is the ceiling +# on full stats (the geometric midpoint of the telemetry gap between 11.8M and 42.3M +# rows), so an entry whose full stats the server refuses is one whose sort it +# refuses. A guess in the same gap, not a measurement. +DEFAULT_SORT_DISABLE_ROWS: int | None = 25_000_000 +# The server never refuses a search (plan 3 question 7), so there is no number for +# this one: the flag stays "enabled" until a host sets a threshold. +DEFAULT_SEARCH_DISABLE_ROWS: int | None = None + +# Environment variable -> StatsLimits field. +_ENV_FIELDS = {"BUCKAROO_STATS_FULL_AUTO_ROWS": "full_auto_rows", + "BUCKAROO_STATS_SCALAR_AUTO_CELLS": "scalar_auto_cells", "BUCKAROO_STATS_CEILING_FULL_ROWS": "ceiling_full_rows", + "BUCKAROO_STATS_CEILING_SCALAR_CELLS": "ceiling_scalar_cells", "BUCKAROO_POLARS_ROUTE_ROWS": "polars_route_rows", + "BUCKAROO_STATS_FULL_AUTO_CELLS": "full_auto_cells", "BUCKAROO_STATS_CEILING_FULL_CELLS": "ceiling_full_cells"} + +# Environment variable -> GuardLimits field. +_GUARD_ENV_FIELDS = {"BUCKAROO_SORT_DISABLE_ROWS": "sort_disable_rows", + "BUCKAROO_SEARCH_DISABLE_ROWS": "search_disable_rows"} + +# Lowest to highest. +TIERS = ("schema", "scalar", "full") + +# Backends whose stats cost does not grow with the file. +EAGER_BACKENDS = frozenset({"pandas", "polars"}) + + +def _env_int(name: str) -> int | None: + raw = os.environ.get(name, "").strip() + if not raw: + return None + try: + value = int(raw) + except ValueError: + value = -1 + if value < 0: + log.warning("%s=%r is not a non-negative integer; using the default", name, raw) + return None + return value + + +def _env_overrides(env_fields: dict[str, str]) -> dict[str, int]: + """The ``{field: value}`` of every variable in ``env_fields`` that is set to a + non-negative integer.""" + overrides = {} + for name, field in env_fields.items(): + value = _env_int(name) + if value is not None: + overrides[field] = value + return overrides + + +@dataclass(frozen=True) +class StatsLimits: + """The thresholds :func:`resolve_stats_policy` and + :func:`route_polars_entry` read. ``ceiling_scalar_cells=None`` means + ``scalar`` has no ceiling.""" + full_auto_rows: int = DEFAULT_FULL_AUTO_ROWS + scalar_auto_cells: int = DEFAULT_SCALAR_AUTO_CELLS + ceiling_full_rows: int = DEFAULT_CEILING_FULL_ROWS + ceiling_scalar_cells: int | None = DEFAULT_CEILING_SCALAR_CELLS + polars_route_rows: int = DEFAULT_POLARS_ROUTE_ROWS + # Added after the first five fields, so a positional construction keeps its meaning. + full_auto_cells: int = DEFAULT_FULL_AUTO_CELLS + ceiling_full_cells: int = DEFAULT_CEILING_FULL_CELLS + + @classmethod + def from_env(cls) -> StatsLimits: + """The defaults with any ``BUCKAROO_*`` overrides applied.""" + return cls(**_env_overrides(_ENV_FIELDS)) + + +@dataclass(frozen=True) +class GuardLimits: + """The thresholds :func:`resolve_source_guards` reads, in rows. ``None`` means + the flag is never turned off.""" + sort_disable_rows: int | None = DEFAULT_SORT_DISABLE_ROWS + search_disable_rows: int | None = DEFAULT_SEARCH_DISABLE_ROWS + + @classmethod + def from_env(cls) -> GuardLimits: + """The defaults with any ``BUCKAROO_*`` overrides applied.""" + return cls(**_env_overrides(_GUARD_ENV_FIELDS)) + + +def _count(name: str, value: Any) -> int: + try: + count = operator.index(value) + except TypeError: + raise TypeError(f"{name} must be an integer, got {value!r}") from None + if count < 0: + raise ValueError(f"{name} must not be negative, got {count}") + return count + + +def _size_tier(rows: int, cells: int, limits: StatsLimits) -> str: + """The tier the server picks on its own for an entry of this size.""" + if rows <= limits.full_auto_rows and cells <= limits.full_auto_cells: + return "full" + if cells <= limits.scalar_auto_cells: + return "scalar" + return "schema" + + +def _ceiling_tier(rows: int, cells: int, limits: StatsLimits) -> str: + """The highest tier any caller may reach for an entry of this size.""" + if limits.ceiling_scalar_cells is not None and cells > limits.ceiling_scalar_cells: + return "schema" + if rows <= limits.ceiling_full_rows and cells <= limits.ceiling_full_cells: + return "full" + return "scalar" + + +def resolve_stats_policy(backend: str, source_kind: str, rows: int, cols: int, bytes: int | None = None, + host_tier: str | None = None, limits: StatsLimits | None = None) -> dict[str, Any]: + """Resolve the stats tier for an entry. + + ``rows`` is the count load already took (the parquet footer's, or xorq's + cached ``_expr_count``); this function never counts anything. ``host_tier`` + is the request's ``stats_tier`` (``None`` or ``"auto"`` for the server's + own choice); a ``stats_request`` with ``force`` passes the tier it asks for + here. ``limits`` defaults to :meth:`StatsLimits.from_env`. + + Returns a JSON-serialisable dict: + + * ``tier_target``: the tier the session should reach. + * ``auto_request``: whether the client should request stats up to + ``tier_target`` on its own. False only for ``schema``. + * ``requestable``: tiers above ``tier_target`` that an on-demand request + may still be granted, in ascending order, up to the ceiling. + * ``reason``: ``None`` when ``tier_target`` is ``full``; otherwise + ``"ceiling"`` if the ceiling clipped the request, ``"host"`` if the host + asked for this tier, ``"size"`` if the server chose it. The ``"cost"`` + reason belongs to the cost guard and is never returned here. + * ``estimate``: the inputs the decision rests on (``rows``, ``cols`` and + ``bytes`` when given). + + Eager backends ignore ``host_tier``. ``source_kind`` and ``bytes`` are + accepted for the rules that will read them (a CSV size rule, and a bytes + threshold if calibration wants one); no threshold reads them yet. + ``source_kind`` must be a ``str``, its vocabulary being the caller's. + ``bytes`` is ``None`` or a non-negative integer, echoed as a plain ``int``. + A value of the wrong type raises ``TypeError`` and a negative count + ``ValueError``. + """ + if not isinstance(source_kind, str): + raise TypeError(f"source_kind must be a str, got {source_kind!r}") + rows = _count("rows", rows) + cols = _count("cols", cols) + if bytes is not None: + bytes = _count("bytes", bytes) + if host_tier == "auto": + host_tier = None + if host_tier is not None and host_tier not in TIERS: + raise ValueError(f"host_tier must be one of {TIERS} or 'auto', got {host_tier!r}") + if limits is None: + limits = StatsLimits.from_env() + + if backend in EAGER_BACKENDS: + requested = ceiling = "full" + from_host = False + else: + cells = rows * cols + ceiling = _ceiling_tier(rows, cells, limits) + from_host = host_tier is not None + requested = host_tier or _size_tier(rows, cells, limits) + + target = min(requested, ceiling, key=TIERS.index) + if target == "full": + reason = None + elif target != requested: + reason = "ceiling" + elif from_host: + reason = "host" + else: + reason = "size" + + requestable = list(TIERS[TIERS.index(target) + 1:TIERS.index(ceiling) + 1]) + estimate: dict[str, int] = {"rows": rows, "cols": cols} + if bytes is not None: + estimate["bytes"] = bytes + return {"tier_target": target, "auto_request": target != "schema", "requestable": requestable, "reason": reason, + "estimate": estimate} + + +def route_polars_entry(rows: int, cols: int, limits: StatsLimits | None = None) -> str: + """``"xorq"`` for a polars entry above R rows, ``"eager"`` otherwise. + + Eager polars holds about 6.8 GB at 12M rows and cannot hold the 78M-row + entry, so a host sends entries above R to ``/load_expr`` instead of + ``/load`` with ``backend="polars"``. R is ``limits.polars_route_rows``, + 8M rows by default. That figure assumes ``PolarsServerSampling.pre_limit`` + is False, so a sorted window copies the frame (3.0x its size at the peak); + with the 1M ``pre_limit`` on main windows do not copy it and the same + memory budget would allow about 14M rows. ``cols`` is accepted so the rule + can become a cell budget without a signature change; only ``rows`` decides + today. + """ + rows = _count("rows", rows) + _count("cols", cols) + if limits is None: + limits = StatsLimits.from_env() + return "xorq" if rows > limits.polars_route_rows else "eager" + + +def resolve_source_guards(backend: str, source_kind: str, rows: int, cols: int, + limits: GuardLimits | None = None) -> dict[str, str]: + """Whether sort and search stay on for a source of this size. + + Returns ``{"sort": ..., "search": ...}``, each ``"enabled"`` or ``"disabled"``: + a flag is ``"disabled"`` when the source has more rows than its threshold in + ``limits`` (default :meth:`GuardLimits.from_env`) and the threshold is set. + ``rows`` is the count load already took; nothing is counted here. Eager + backends always keep both, since their windows come from a bounded frame. + ``cols`` and ``source_kind`` are accepted so a rule can read them without a + signature change; only ``rows`` decides today. Inputs are checked as + :func:`resolve_stats_policy` checks them. + """ + if not isinstance(source_kind, str): + raise TypeError(f"source_kind must be a str, got {source_kind!r}") + rows = _count("rows", rows) + _count("cols", cols) + if limits is None: + limits = GuardLimits.from_env() + if backend in EAGER_BACKENDS: + return {"sort": "enabled", "search": "enabled"} + + def flag(threshold: int | None) -> str: + return "disabled" if threshold is not None and rows > threshold else "enabled" + + return {"sort": flag(limits.sort_disable_rows), "search": flag(limits.search_disable_rows)} + + +def _unsortable(column: Any) -> Any: + """A column config with ``ag_grid_specs.sortable`` false, its other specs kept.""" + if not isinstance(column, dict): + return column + specs = column.get("ag_grid_specs") + return {**column, "ag_grid_specs": {**(specs if isinstance(specs, dict) else {}), "sortable": False}} + + +def disable_sorting(display_args: Any) -> Any: + """A copy of ``df_display_args`` in which no column of a display served by + ``infinite_request`` can sort: every ``column_config`` and ``left_col_configs`` + entry of a display whose ``data_key`` is ``main`` gets ``ag_grid_specs.sortable`` + false, whatever a klass or an override set. A display with another + ``data_key`` (the summary grid) is a client-side grid and is left as it is. + + It is the last step of building the config (``XorqServerDataflow``), after the + klasses styled it and ``column_config_overrides`` were merged, so neither can + bypass it. ``display_args`` is not changed and may have any shape.""" + if not isinstance(display_args, dict): + return display_args + out = {} + for name, display in display_args.items(): + config = display.get("df_viewer_config") if isinstance(display, dict) else None + if not isinstance(config, dict) or display.get("data_key") != "main": + out[name] = display + continue + columns = {key: [_unsortable(column) for column in config[key]] + for key in ("column_config", "left_col_configs") if isinstance(config.get(key), list)} + out[name] = {**display, "df_viewer_config": {**config, **columns}} + return out + + +def probe_dtypes(obj: Any) -> dict[str, str]: + """Column name to dtype string, from the schema alone. + + Takes a pandas or polars frame (a ``LazyFrame`` is not collected) or a xorq + expression. No row is read and no query runs. + """ + # Polars first: reading ``LazyFrame.dtypes`` resolves the plan and warns. + if hasattr(obj, "collect_schema"): + pairs = obj.collect_schema().items() + elif callable(getattr(obj, "schema", None)): + pairs = obj.schema().items() + elif hasattr(getattr(obj, "dtypes", None), "items"): + pairs = obj.dtypes.items() + else: + raise TypeError(f"cannot read a schema from {type(obj).__name__}") + return {str(name): str(dtype) for name, dtype in pairs} + + +def probe_parquet_rows(source: Any) -> int: + """A parquet file's row count, from its footer. + + ``source`` is a path or a binary file object. Only the footer metadata is + read, so the cost does not grow with the file (0.4 ms on 4M rows). A + single file only; a dataset directory is not handled. + """ + return int(pq.read_metadata(source).num_rows) diff --git a/buckaroo/server/stats_wire.py b/buckaroo/server/stats_wire.py new file mode 100644 index 000000000..b7fdcd7e7 --- /dev/null +++ b/buckaroo/server/stats_wire.py @@ -0,0 +1,325 @@ +"""Server side of the stats wire protocol (rows-first s3). + +A deferred session (``stats_delivery="deferred"``) publishes its dataflow at the +schema tier and leaves the rest of the stats to later. Two kinds of client +reach them: + +* A client that advertised ``?caps=stats_update`` on its WebSocket URL gets a + stats-free ``initial_state`` (``df_meta.stats.status == "pending"``), pulls + the stats with ``stats_request {stats_gen, scope}`` and merges the + ``stats_update`` that answers it. A request for a generation the session has + left gets ``stats_aborted``. +* Any other client gets complete messages: ``build_state_message_for`` runs the + missing stats synchronously before it builds a message for one, which is + today's cost, paid on the loop. + +A session also carries the policy resolved at load (``stats_policy.py``). A +target the server chose below ``full`` (a size rule, or the ceiling on a host +that asked for ``auto`` or ``full``) is applied only for a client that +advertised ``?caps=stats_update,stats_ondemand``: it gets the schema tier, +``df_meta.stats.status == "not_computed"`` and the policy fields. A client with +``stats_update`` only is told the session is ``pending`` and pulls the stats as +above, and one with no caps gets them synchronously, as for any deferred +session. A tier the host named below ``full`` (``scalar``, ``schema``) is its +choice and reaches every client as ``not_computed``. ``stats_to_pull`` decides, +per client, whether a session has stats left to compute. + +A session opened with a stats policy also resolves the guards for huge sources +(``resolve_session_guards``, rows-first p37): above the sort threshold the dataflow +turns sorting off in its display config (``apply_source_guards``) and a sorted +window is refused (``session.sort_refusal``). Neither depends on the client or on +what stats the session has. + +Every send site goes through ``build_state_message_for`` (or ``broadcast_state``, +which calls it per client), because the session holds one shared snapshot and +the client is known only to the handler that owns the connection. + +The request here is the whole run: one synchronous call that computes the full +stats and applies the final assignment (``assign_full_stats``, +``refresh_session_snapshot``). Resumable units generalize it later. +""" +import json +import logging +import time +import traceback +from typing import Any, Callable, Optional + +from buckaroo.dataflow.sd_cache import split_chain_by_scope +from buckaroo.pluggable_analysis_framework import perf_log +from buckaroo.server.data_loading import get_buckaroo_display_state +from buckaroo.server.session import SessionState, build_state_message, stats_deferred_by_policy +from buckaroo.server.stats_policy import resolve_source_guards, resolve_stats_policy + +log = logging.getLogger("buckaroo.server.stats_wire") + +# The capability a client advertises, as one of the comma-separated values of +# ``?caps=`` on the WebSocket URL: it merges ``stats_update`` messages. The +# connection is the only place it can be recorded, because ``open()`` sends the +# first message before the client has said anything. +STATS_UPDATE_CAP = "stats_update" + +# The second capability bit: the client also renders ``not_computed``, honours +# ``auto_request`` and sends ``tier`` and ``force`` on ``stats_request``. It counts +# only together with ``stats_update``, which it builds on. +STATS_ONDEMAND_CAP = "stats_ondemand" + +# The scopes a ``stats_request`` can name. ``raw`` is the unfiltered table of +# the session's current generation; filtered stats exist only through a +# dataflow-field change (``quick_command_args.search``), which bumps the +# generation. +STATS_SCOPES = ("raw",) + + +def parse_caps(raw: str) -> frozenset: + """The capabilities in a ``?caps=a,b`` value.""" + return frozenset(cap for cap in (part.strip() for part in raw.split(",")) if cap) + + +def client_has_cap(client: Any, cap: str) -> bool: + """Whether a WebSocket handler recorded ``cap`` at open. Anything else with + no ``caps`` (a test double, a future non-WebSocket client) has none.""" + return cap in getattr(client, "caps", ()) + + +def client_has_ondemand(client: Any) -> bool: + """Whether a WebSocket handler advertised ``stats_update`` and + ``stats_ondemand``, so a session the server put below ``full`` is applied + for it as such.""" + return client_has_cap(client, STATS_UPDATE_CAP) and client_has_cap(client, STATS_ONDEMAND_CAP) + + +def stats_to_pull(session: SessionState, client: Any) -> bool: + """Whether the session has stats left to compute from this client's side: it + is ``pending``, or the server put it below ``full`` and the client cannot + take that (``stats_deferred_by_policy``). A client with ``stats_update`` + pulls them with ``stats_request``, one with no caps gets them at connect.""" + return session.stats_status == "pending" or (stats_deferred_by_policy(session) and not client_has_ondemand(client)) + + +def resolve_session_policy(stats_tier: str, dataflow_tier: str, rows: int, cols: int) -> Optional[dict]: + """The policy for a session the handler has just built: ``stats_tier`` is the + host's request, ``dataflow_tier`` the tier its dataflow was built at, and + ``rows`` and ``cols`` the size load already took (the cached count, so no + query runs). ``None`` when the dataflow ran its stats in the constructor + (``full``): there is nothing left to defer or refuse.""" + if dataflow_tier != "schema": + return None + return resolve_stats_policy("xorq", "xorq_build", rows, cols, host_tier=stats_tier) + + +def resolve_session_guards(dataflow_tier: str, rows: int, cols: int) -> Optional[dict]: + """The guards for a session the handler has just built, from the size load took + (``rows`` and ``cols``, so no query runs): ``{"sort": ..., "search": ...}``, each + ``"enabled"`` or ``"disabled"`` (``stats_policy.resolve_source_guards``). ``None`` + unless the dataflow was built at the schema tier, which is to say unless the host + opened the session with a stats policy, as ``resolve_session_policy`` does: a + session built with its stats inline is not guarded.""" + if dataflow_tier != "schema": + return None + return resolve_source_guards("xorq", "xorq_build", rows, cols) + + +def apply_source_guards(session: SessionState, dataflow: Any, guards: Optional[dict]) -> None: + """Record ``guards`` on the session and set the dataflow's sort switch to match, + which rebuilds its display config when the switch moves. Called before the + session snapshot is taken from the dataflow, since the config in the snapshot is + the one a client sorts by.""" + session.source_guards = guards + dataflow.set_sort_enabled(guards is None or guards["sort"] != "disabled") + + +def session_dataflow(session: Optional[SessionState]) -> Any: + """The dataflow behind a buckaroo-mode session, or ``None`` (viewer and lazy + sessions have none).""" + if session is None or session.mode != "buckaroo": + return None + return session.xorq_dataflow if session.backend == "xorq" else session.dataflow + + +def refresh_session_snapshot(session: SessionState, dataflow: Any) -> None: + """Copy the dataflow's display state onto the session snapshot that new + clients and every push read, and re-apply ``component_config`` so theme + settings survive. The dataflow reaches no client until this runs.""" + refreshed = get_buckaroo_display_state(dataflow) + session.df_display_args = refreshed["df_display_args"] + session.df_data_dict = refreshed["df_data_dict"] + session.df_meta = refreshed["df_meta"] + session.buckaroo_options = refreshed["buckaroo_options"] + session.command_config = refreshed["command_config"] + if session.component_config and session.df_display_args: + for key in session.df_display_args: + dvc = session.df_display_args[key].get("df_viewer_config") + if dvc is not None: + dvc["component_config"] = {**dvc.get("component_config", {}), **session.component_config} + + +def assign_full_stats(dataflow: Any) -> None: + """Take a schema-tier dataflow to the full tier: compute the current + state's full stats (or find them in ``summary_stats_cache``, where an + earlier visit to the same state left them), write them under the full-tier + key, then assign ``summary_sd``. Assigning runs the cascade, which fills the + other scopes' entries it still lacks and rebuilds ``merged_sd``, + ``df_data_dict`` and ``df_display_args``. + + The cache entry goes in first because ``_populate_sd_cache`` skips a key it + finds, so the filt scope is not computed a second time. The tier is a + plain attribute, flipped before the compute so ``_get_summary_sd`` runs the + full pipeline, and put back if anything raises.""" + tier = dataflow.stats_tier + filt_chain = split_chain_by_scope(dataflow.operations)["filt"] + key = dataflow._scope_cache_key(filt_chain, tier="full") + dataflow.stats_tier = "full" + try: + sd = dataflow.summary_stats_cache.get(key) + errs = {} + if sd is None: + sd, errs = dataflow._get_summary_sd(dataflow.processed_df) + dataflow.summary_stats_cache = {**dataflow.summary_stats_cache, key: sd} + dataflow.summary_sd = sd + dataflow.errs = errs + except Exception: + dataflow.stats_tier = tier + raise + + +def complete_stats(session: SessionState) -> bool: + """Run the stats a pending session is missing (or one the server put below + ``full``, for a client that cannot take that), in one synchronous call, and + publish them: the final assignment (``assign_full_stats``), then the session + snapshot refreshed and the status set to ``complete`` in the same step, so + ``all_stats`` and ``df_meta.stats`` cannot disagree. Returns whether the + session is complete. + + A failure is the session's state for this generation (``error``, reason + ``stats_failed``) and is not retried by the next request; the next + generation starts clean.""" + if session.stats_status == "complete": + return True + dataflow = session_dataflow(session) + if not (session.stats_status == "pending" or stats_deferred_by_policy(session)) or dataflow is None: + return False + with ( + perf_log.telemetry_context(session.session_id, session.tele_sink), + perf_log.perf_span("firstpull.stats_total", session=session.session_id, stats_gen=session.stats_gen), + ): + try: + assign_full_stats(dataflow) + refresh_session_snapshot(session, dataflow) + except Exception: + log.error("stats run failed session=%s stats_gen=%s: %s", session.session_id, session.stats_gen, + traceback.format_exc()) + session.stats_status, session.stats_reason = "error", "stats_failed" + return False + session.stats_status, session.stats_reason = "complete", None + return True + + +def build_state_message_for(session: SessionState, client: Any, metadata: Optional[dict] = None, + reply_seq: Optional[int] = None) -> dict: + """The ``initial_state`` message for one client. + + A client without the ``stats_update`` capability on a deferred session that + has stats to compute for it (``stats_to_pull``) gets them run first, so its + message is complete; a capable client gets the session snapshot as it is + (stats-free while the stats are to be pulled) and pulls the rest. A client + that also advertised ``stats_ondemand`` is told the session as the server + resolved it (``not_computed`` with the policy fields); the others are told + ``pending`` where the server chose a target below ``full``. The search term + is the recipient's own (#851), and ``reply_seq`` is passed through to + ``build_state_message`` (#998).""" + if (session.stats_delivery == "deferred" and stats_to_pull(session, client) + and not client_has_cap(client, STATS_UPDATE_CAP)): + complete_stats(session) + return build_state_message(session, metadata=metadata, search_string=getattr(client, "search_string", ""), + reply_seq=reply_seq, ondemand=client_has_ondemand(client)) + + +def broadcast_state(session: SessionState, metadata: Optional[dict] = None, reset_search: bool = False, + reply_to: Any = None, reply_seq: Optional[int] = None, + highlight: Optional[Callable[[Any, str], Any]] = None) -> 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. + + ``reply_seq`` goes on the copy for ``reply_to`` only, the client whose + ``buckaroo_state_change`` this answers (#998); the others made no change, + so every copy is current for them. ``highlight(df_display_args, term)``, + when given, puts each client's own live-search highlight on its copy.""" + 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 = "" + msg = build_state_message_for(session, client, metadata=metadata, + reply_seq=reply_seq if client is reply_to else None) + if highlight is not None: + msg["df_display_args"] = highlight(msg["df_display_args"], getattr(client, "search_string", "")) + client.write_message(json.dumps(msg)) + except Exception: + session.ws_clients.discard(client) + + +def _aborted(stats_gen: Any, scope: Any, reason: str, session: Optional[SessionState] = None) -> dict: + """``stats_aborted``: the request was not run. ``stats_gen`` echoes the + request so the client can pair them; ``current_gen`` is the session's, so a + stale client can ask again without waiting for the next ``initial_state``.""" + msg: dict = {"type": "stats_aborted", "stats_gen": stats_gen, "scope": scope, "reason": reason} + if session is not None: + msg["current_gen"] = session.stats_gen + return msg + + +def _answer_stats_request(session: Optional[SessionState], stats_gen: Any, scope: Any, client: Any, + started: float) -> dict: + dataflow = session_dataflow(session) + if session is None or dataflow is None: + return _aborted(stats_gen, scope, "no_data") + if stats_gen != session.stats_gen: + return _aborted(stats_gen, scope, "stale", session) + if scope not in STATS_SCOPES: + return _aborted(stats_gen, scope, "unsupported_scope", session) + to_pull = stats_to_pull(session, client) + if session.stats_status == "not_computed" and not to_pull: + return _aborted(stats_gen, scope, "not_requestable", session) + if to_pull: + complete_stats(session) + if session.stats_status != "complete": + return _aborted(stats_gen, scope, "error", session) + # Complete: from here the answer is the dataflow's own all_stats, with no + # query, whether this request ran the stats or an earlier one did. + return {"type": "stats_update", "stats_gen": stats_gen, "scope": scope, "tier": dataflow.stats_tier, + "final": True, "payload": session.df_data_dict["all_stats"], + "elapsed_ms": round((time.perf_counter() - started) * 1000, 1)} + + +def handle_stats_request(session: Optional[SessionState], msg: dict, client: Any = None) -> dict: + """Answer a ``stats_request {stats_gen, scope, columns?}``: a + ``stats_update`` carrying the complete ``all_stats`` as an inline wide + ``DFEnvelope`` (self-contained, so binary pairing stays single-slot), or a + ``stats_aborted``. ``columns`` is a hint a whole-run request ignores. + + The ``stats.request`` span records the request and its ``outcome`` + (``update`` or the abort reason), which is where updates sent, requests + dropped as stale and errors are counted. The caller binds the session's + telemetry sink around this call. ``client`` is the handler that sent the + request: whether the session has stats to pull depends on its caps + (``stats_to_pull``).""" + stats_gen, scope, columns = msg.get("stats_gen"), msg.get("scope", "raw"), msg.get("columns") + started = time.perf_counter() + with perf_log.perf_span("stats.request", session=session.session_id if session else None, stats_gen=stats_gen, + scope=scope, columns=len(columns) if isinstance(columns, list) else None) as span: + try: + reply = _answer_stats_request(session, stats_gen, scope, client, started) + except Exception: + log.error("stats_request error session=%s: %s", session.session_id if session else None, + traceback.format_exc()) + reply = _aborted(stats_gen, scope, "error", session) + if reply["type"] == "stats_update": + span.set_attr(outcome="update", tier=reply["tier"]) + else: + span.set_attr(outcome=reply["reason"]) + return reply diff --git a/buckaroo/server/websocket_handler.py b/buckaroo/server/websocket_handler.py index 2232211c7..5c5c21ff7 100644 --- a/buckaroo/server/websocket_handler.py +++ b/buckaroo/server/websocket_handler.py @@ -9,8 +9,9 @@ import tornado.websocket from buckaroo.pluggable_analysis_framework import perf_log -from buckaroo.server.data_loading import (handle_infinite_request, handle_infinite_request_buckaroo, handle_infinite_request_lazy, get_buckaroo_display_state) -from buckaroo.server.session import build_state_message +from buckaroo.server.data_loading import (handle_infinite_request, handle_infinite_request_buckaroo, handle_infinite_request_lazy) +from buckaroo.server.session import begin_stats_generation, dataflow_stats_tier, sort_refusal +from buckaroo.server.stats_wire import (broadcast_state, build_state_message_for, handle_stats_request, parse_caps, refresh_session_snapshot) def _handle_infinite_request_xorq(xorq_dataflow, payload_args, search_string=""): @@ -38,6 +39,11 @@ def open(self, session_id): # highlight overlay below — never broadcast, never stored on the # session. self.search_string = "" + # Capabilities the client advertised on the URL (``?caps=a,b``). Recorded + # per connection because this method sends the first message before the + # client can say anything, and the other send sites push one shared + # snapshot (stats_wire.build_state_message_for reads it per client). + self.caps = parse_caps(self.get_query_argument("caps", "")) sessions = self.application.settings["sessions"] sessions.add_ws_client(session_id, self) @@ -45,7 +51,7 @@ def open(self, session_id): # search_string="" — fresh connection, no per-client typing yet. session = sessions.get(session_id) if session and (session.df is not None or session.ldf is not None or session.xorq_dataflow is not None): - self.write_message(json.dumps(build_state_message(session, search_string=self.search_string))) + self.write_message(json.dumps(build_state_message_for(session, self))) def on_message(self, message): try: @@ -58,9 +64,29 @@ def on_message(self, message): if msg_type == "infinite_request": self._handle_infinite_request(msg.get("payload_args", {})) elif msg_type == "buckaroo_state_change": - self._handle_buckaroo_state_change(msg.get("new_state") or {}) + # state_seq (#998): the client's own counter for this change, + # echoed back as reply_seq on the initial_state it gets so it + # can drop a reply that an overlapping later change superseded. + # Optional; a client that sends none gets none back. + self._handle_buckaroo_state_change(msg.get("new_state") or {}, state_seq=msg.get("state_seq")) + elif msg_type == "stats_request": + self._handle_stats_request(msg) + + def _handle_stats_request(self, msg): + """Answer a client's ``stats_request`` with a ``stats_update`` or a + ``stats_aborted``. Synchronous, like ``infinite_request``: the request + runs the whole stats computation in this call (see ``stats_wire``). + + This branch is its own async context, so the session's telemetry sink + is bound here for the ``stats.request`` span and the stats spans under + it.""" + sessions = self.application.settings["sessions"] + session = sessions.get(self.session_id) + with perf_log.telemetry_context(self.session_id, session.tele_sink if session else None): + reply = handle_stats_request(session, msg, self) + self.write_message(json.dumps(reply)) - def _handle_buckaroo_state_change(self, new_state): + def _handle_buckaroo_state_change(self, new_state, state_seq=None): sessions = self.application.settings["sessions"] session = sessions.get(self.session_id) if not session or session.mode != "buckaroo": @@ -78,27 +104,39 @@ def _handle_buckaroo_state_change(self, new_state): # Decide up front whether the dataflow path will run, so we # don't both (a) send a per-client overlay AND (b) broadcast - # an initial_state in the same turn. The broadcast already - # carries any highlight_phrase that ``quick_command_args.search`` - # produced via cleaning_sd, so the overlay would be redundant - # (and tests/clients that read one message back would see two). + # an initial_state in the same turn. Each client's copy of the + # broadcast carries its own highlight, so the overlay would be + # redundant (and tests/clients that read one message back would + # see two). dataflow_changed = any(old_state.get(f) != new_state.get(f) for f in _DATAFLOW_FIELDS) - # Per-client live search (#838 / #851). Update self; if no - # dataflow change is coming, send a targeted highlight overlay - # to this client only. Never touches the session, never broadcasts. + # Per-client live search (#838 / #851). Recorded on self only: + # never on the session, never broadcast. new_search = new_state.get("search_string", "") new_search = new_search if isinstance(new_search, str) else "" - if self.search_string != new_search: + search_changed = self.search_string != new_search + if search_changed: self.search_string = new_search - if not dataflow_changed: - self._send_highlight_overlay(session) # Skip if no effective change to the fields that drive the dataflow. if not dataflow_changed: + # A new search term gets this client its highlight overlay. + # A numbered change gets an answer whatever it changed: the + # client drops a reply older than its latest change, so a + # no-op change that follows a dataflow change would otherwise + # leave the earlier reply dropped and nothing newer to apply + # (#998). + if search_changed or state_seq is not None: + self._send_client_state(session, new_state, reply_seq=state_seq) 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 +146,24 @@ 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) + # search box on another (or on the typing client itself), and + # its own live-search highlight, so the rebuilt display config + # doesn't drop it. + # Only the originating client gets reply_seq (#998): the + # others made no change, so every copy is current for them. + broadcast_state(session, reply_to=self, reply_seq=state_seq, highlight=self._with_highlight) except Exception: tb = traceback.format_exc() log.error("buckaroo_state_change error session=%s: %s", self.session_id, tb) @@ -148,48 +171,61 @@ def _handle_buckaroo_state_change(self, new_state): if _BUCKAROO_DEBUG: err["details"] = tb self.write_message(json.dumps(err)) - - def _send_highlight_overlay(self, session): - """Send this client an ``initial_state`` with highlight_phrase - injected into every string-column ``displayer_args`` so live-typed - matches highlight in the grid (#851). - - Per-client because ``search_string`` is per-client — another - client's term must not bleed into this one's highlight, and a - broadcast initial_state would clobber every other input box. We - deep-copy ``session.df_display_args`` so the overlay never - mutates the shared session snapshot. - - Empty term clears any prior highlight by sending the pristine - df_display_args back. Skipped when no display config is loaded - yet — the upcoming ``initial_state`` on first load will be - unhighlighted, which is correct (search starts empty). - """ - if not session.df_display_args: - return - term = self.search_string - overlay = copy.deepcopy(session.df_display_args) + # A numbered change is answered even when it fails (#998): the + # client has dropped the reply to its previous change as stale, + # so the error alone would leave it on the data from before + # both. The reply is the session's current data with the + # change's own buckaroo_state, as every numbered reply is: the + # session's would revert a search box that holds the failed + # term, and the box would send it again. + if state_seq is not None: + self._send_client_state(session, new_state if isinstance(new_state, dict) else old_state, + reply_seq=state_seq) + + @staticmethod + def _with_highlight(df_display_args, term): + """``df_display_args`` with ``term`` as the ``highlight_phrase`` of + every string column, as a deep copy so the shared session snapshot + is never mutated. With no ``term`` it is returned as it is, so the + highlight a committed ``quick_command_args.search`` put there + stays.""" + if not term or not df_display_args: + return df_display_args + overlay = copy.deepcopy(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: + if isinstance(disp, dict) and disp.get("displayer") == "string": disp["highlight_phrase"] = [term] - else: - disp.pop("highlight_phrase", None) - - # Pass self.search_string so the overlay's buckaroo_state - # round-trips the typed term back to this client (Codex P1 on - # #854 — without it the JS clears the search box on every - # keystroke). - msg = build_state_message(session, search_string=self.search_string) - msg["df_display_args"] = overlay + return overlay + + def _send_client_state(self, session, buckaroo_state, reply_seq=None): + """Send this client alone an ``initial_state``: the session's + current data, ``buckaroo_state`` with this client's search term, and + its live-search highlight (#851). It answers a change that reran no + dataflow (a new search term, or a numbered change that touched + nothing) and a numbered change that failed, so it never touches + the session or reaches another client. + + ``buckaroo_state`` is the change's own. The session records one + only on a dataflow change, so the session's would put back + ``show_commands``, ``df_display`` and ``sampled``. The search term + is re-injected so the reply doesn't clear the search box (Codex P1 + on #854). ``reply_seq`` is the change's ``state_seq`` (#998), so the + client drops this reply if it has sent a later change. + + The message is built first: for a client without the + ``stats_update`` capability it completes the session's stats, which + replaces the display config, and the highlight goes on that one. + """ + msg = build_state_message_for(session, self, reply_seq=reply_seq) + msg["buckaroo_state"] = {**buckaroo_state, "search_string": self.search_string} + msg["df_display_args"] = self._with_highlight(msg["df_display_args"], self.search_string) try: self.write_message(json.dumps(msg)) except Exception: - log.debug("highlight overlay write failed for session=%s", self.session_id) + log.debug("client state write failed for session=%s", self.session_id) def _handle_infinite_request(self, payload_args): sessions = self.application.settings["sessions"] @@ -209,6 +245,11 @@ def _dispatch(pa): # each backend can apply the filter in its native # expression layer. search = self.search_string or "" + # A source too large to sort in a request refuses the window before any + # query runs, the second window of a request included. + refusal = sort_refusal(session, pa) + if refusal is not None: + return refusal, b"" if session.mode == "lazy" and session.ldf is not None: return handle_infinite_request_lazy(session.ldf, session.orig_to_rw, session.rw_to_orig, session.metadata.get("rows", 0), pa) diff --git a/buckaroo/server/xorq_loading.py b/buckaroo/server/xorq_loading.py index 60a9226cf..585d53f9e 100644 --- a/buckaroo/server/xorq_loading.py +++ b/buckaroo/server/xorq_loading.py @@ -14,11 +14,15 @@ import os import traceback import uuid +import weakref +from collections import OrderedDict from pathlib import Path import pyarrow.parquet as pq +from buckaroo.customizations.xorq_commands import search_expr from buckaroo.server.git_state_guard import install_git_state_guard +from buckaroo.server.stats_policy import disable_sorting from buckaroo.server.window import clamp_window from buckaroo.serialization_utils import make_infinite_resp from buckaroo.xorq_buckaroo import ( @@ -64,8 +68,15 @@ 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. + + ``sort_enabled`` is whether the grid served by ``infinite_request`` may sort. + The server turns it off (``set_sort_enabled``) for a source with too many rows + to sort in a request (``stats_policy.resolve_source_guards``); the config is + then built with ``stats_policy.disable_sorting`` as its last step, so no + klass or override can turn sorting back on, and every rebuild keeps it. """ + sort_enabled = True sampling_klass = XorqInfiniteSampling autocleaning_klass = XorqAutocleaning autoclean_conf = (NoCleaningConfXorq,) @@ -81,6 +92,20 @@ def __init__(self, expr, *args, extra_klasses=None, cache_storage_path=None, **k self.cache_storage = _make_cache_storage(cache_storage_path) super().__init__(expr, *args, **kwargs) + def set_sort_enabled(self, enabled: bool) -> None: + """Turn sorting on or off in the display config, rebuilding it now when + the switch moves.""" + if enabled == self.sort_enabled: + return + self.sort_enabled = enabled + _unused, processed_df, merged_sd = self.widget_args_tuple + if processed_df is not None: + self.df_display_args = self._build_df_display_args(processed_df, merged_sd) + + def _build_df_display_args(self, processed_df, merged_sd): + display_args = super()._build_df_display_args(processed_df, merged_sd) + return display_args if self.sort_enabled else disable_sorting(display_args) + def load_expr_build_dir(build_dir: str, cache_dir=None): """Rehydrate an ibis expression from a xorq build directory. @@ -556,6 +581,40 @@ def _compile_project_display(path: Path, ColAnalysis, *extra_bases): return found +# The searched expression of each (base expression, term). ``search_expr`` builds +# a new expression on every call and ``_expr_count`` caches by expression object, so +# every request for a searched window counted its rows again (0.5 to 1 s on 10.8M +# rows). Holding the expression makes the second request a hit in ``_expr_count``'s +# cache, which still does not cache a count that failed. Keyed weakly by the base +# expression, the dataflow's ``processed_df``, which a state change replaces, so an +# entry lasts as long as its base; a base keeps at most this many terms, the least +# recently used going first. +_SEARCHED_EXPRS_PER_BASE = 16 +_searched_exprs: weakref.WeakKeyDictionary = weakref.WeakKeyDictionary() + + +def _searched_expr(base, term): + """``search_expr(base, term)``: the same object each time for the same pair.""" + if not term: + return base + terms = _searched_exprs.get(base) + expr = terms.get(term) if terms is not None else None + if expr is not None: + terms.move_to_end(term) + return expr + expr = search_expr(base, term) + if expr is base: + # No string column to search: there is nothing to hold, and the base held as + # its own value would never be freed. + return base + if terms is None: + terms = _searched_exprs[base] = OrderedDict() + terms[term] = expr + if len(terms) > _SEARCHED_EXPRS_PER_BASE: + terms.popitem(last=False) + return expr + + def handle_infinite_request_xorq(xorq_dataflow: XorqServerDataflow, payload_args: dict, search_string: str = "") -> tuple[dict, bytes]: """Drive one infinite_request window against a xorq expression. @@ -567,13 +626,12 @@ def handle_infinite_request_xorq(xorq_dataflow: XorqServerDataflow, delegates to ``window_to_parquet`` and returns the ``(json_msg, parquet_bytes)`` pair the WebSocket handler ships as a text + binary frame pair (matching the pandas/polars paths).""" - from buckaroo.customizations.xorq_commands import search_expr _unused, processed_df, merged_sd = xorq_dataflow.widget_args_tuple if processed_df is None: return ({"type": "infinite_resp", "key": payload_args, "length": 0}, b"") try: - filtered_expr = search_expr(processed_df, search_string) if search_string else processed_df + filtered_expr = _searched_expr(processed_df, search_string) sort = payload_args.get("sort") sort_col = None diff --git a/buckaroo/xorq_buckaroo.py b/buckaroo/xorq_buckaroo.py index 6b9149c40..e1aad2315 100644 --- a/buckaroo/xorq_buckaroo.py +++ b/buckaroo/xorq_buckaroo.py @@ -27,7 +27,7 @@ from .buckaroo_widget import BuckarooInfiniteWidget, BuckarooWidget from .customizations.styling import DefaultMainStyling, DefaultSummaryStatsStyling from .customizations.xorq_autoclean_conf import NoCleaningConfXorq -from .customizations.xorq_stats_v2 import XORQ_STATS_V2 +from .customizations.xorq_stats_v2 import XORQ_STATS_V2, schema_stats from .dataflow.autocleaning import PandasAutocleaning from .dataflow.dataflow import CustomizableDataflow from .dataflow.dataflow_extras import Sampling @@ -173,6 +173,8 @@ def _get_summary_sd(self, processed_df: "XorqExpr | pd.DataFrame"): 'orig_col_name': orig_col, 'rewritten_col_name': rewritten_col} return empty, {} + if self.stats_tier == "schema": + return self._get_schema_sd(processed_df), {} cache_storage = getattr(self, 'cache_storage', None) # The owning widget injects XorqDfStatsV2 as DFStatsClass (via its # InnerDataFlow subclass); the cast exposes cache_run_stats() below. @@ -213,6 +215,27 @@ def _get_summary_sd(self, processed_df: "XorqExpr | pd.DataFrame"): rewritten[rewritten_col] = col_meta return rewritten, errs + def _get_schema_sd(self, processed_df: "XorqExpr | pd.DataFrame") -> dict: + """Identity, dtype, typing flags and row count for every column of an + expression, with no query beyond ``_expr_count`` (cached per + expression). Carries the same values full stats give those keys, so + styling that reads only them renders the same. 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/packages/buckaroo-js-core/src/server/WebSocketModel.test.ts b/packages/buckaroo-js-core/src/server/WebSocketModel.test.ts new file mode 100644 index 000000000..386523bc7 --- /dev/null +++ b/packages/buckaroo-js-core/src/server/WebSocketModel.test.ts @@ -0,0 +1,104 @@ +/** + * WebSocketModel — state-change sequencing (#998). + * + * Each buckaroo_state_change the model sends carries an incrementing + * `state_seq`. The server echoes it as `reply_seq` on the `initial_state` + * it sends back to the originating client. A reply whose `reply_seq` is + * older than the latest change this model sent is stale and must not be + * applied; otherwise the reply for "alle" lands after "allen" was sent + * and sets the client's buckaroo_state back. Replies with no `reply_seq` + * (another tab's change, a /load push, a fresh connection) still apply. + */ +import { WebSocketModel } from "./WebSocketModel"; + +class FakeWebSocket { + readyState = WebSocket.OPEN; + sent: any[] = []; + onmessage: ((event: MessageEvent) => void) | null = null; + send(data: string) { + this.sent.push(JSON.parse(data)); + } + deliver(msg: Record) { + this.onmessage?.({ data: JSON.stringify(msg) } as MessageEvent); + } +} + +const stateWithSearch = (search: string) => ({ + post_processing: "", + cleaning_method: "", + quick_command_args: { search }, + df_display: "main", + show_commands: false, + sampled: false, + search_string: search, +}); + +function makeModel() { + const ws = new FakeWebSocket(); + const model = new WebSocketModel(ws as unknown as WebSocket, { + buckaroo_state: stateWithSearch(""), + df_meta: { total_rows: 1 }, + }); + const changes: any[] = []; + model.on("change:buckaroo_state", (v: any) => changes.push(v)); + return { ws, model, changes }; +} + +function sendSearch(model: WebSocketModel, search: string) { + model.set("buckaroo_state", stateWithSearch(search)); + model.save_changes(); +} + +describe("WebSocketModel state_seq (#998)", () => { + it("attaches an incrementing state_seq to each buckaroo_state_change", () => { + const { ws, model } = makeModel(); + sendSearch(model, "a"); + sendSearch(model, "al"); + expect(ws.sent.map((m) => m.type)).toEqual(["buckaroo_state_change", "buckaroo_state_change"]); + expect(ws.sent.map((m) => m.state_seq)).toEqual([1, 2]); + expect(ws.sent[1].new_state.quick_command_args.search).toBe("al"); + }); + + it("drops an initial_state whose reply_seq is older than the latest sent seq", () => { + const { ws, model, changes } = makeModel(); + sendSearch(model, "a"); + sendSearch(model, "al"); + changes.length = 0; + + ws.deliver({ + type: "initial_state", + reply_seq: 1, + buckaroo_state: stateWithSearch("a"), + df_meta: { total_rows: 10 }, + }); + expect(model.get("buckaroo_state").quick_command_args.search).toBe("al"); + expect(model.get("df_meta")).toEqual({ total_rows: 1 }); + expect(changes).toEqual([]); + + ws.deliver({ + type: "initial_state", + reply_seq: 2, + buckaroo_state: stateWithSearch("al"), + df_meta: { total_rows: 3 }, + }); + expect(model.get("buckaroo_state").quick_command_args.search).toBe("al"); + expect(model.get("df_meta")).toEqual({ total_rows: 3 }); + expect(changes).toHaveLength(1); + }); + + it("applies an initial_state that carries no reply_seq", () => { + const { ws, model, changes } = makeModel(); + sendSearch(model, "a"); + sendSearch(model, "al"); + changes.length = 0; + + ws.deliver({ + type: "initial_state", + buckaroo_state: stateWithSearch(""), + df_meta: { total_rows: 5 }, + }); + expect(model.get("buckaroo_state").quick_command_args.search).toBe(""); + expect(model.get("df_meta")).toEqual({ total_rows: 5 }); + expect(changes).toHaveLength(1); + }); +}); diff --git a/packages/buckaroo-js-core/src/server/WebSocketModel.ts b/packages/buckaroo-js-core/src/server/WebSocketModel.ts index 1e1fd4a2d..8c73b557c 100644 --- a/packages/buckaroo-js-core/src/server/WebSocketModel.ts +++ b/packages/buckaroo-js-core/src/server/WebSocketModel.ts @@ -13,6 +13,16 @@ * Binary protocol (matching anywidget's msg + buffers pattern): * Server sends a JSON text frame (infinite_resp), then a binary frame (Parquet). * This class pairs them and emits "msg:custom" with (msg, [DataView]). + * + * State-change sequencing (#998): + * Each buckaroo_state_change carries an incrementing `state_seq`. The + * server echoes it as `reply_seq` on the initial_state it sends back to + * this client. A reply whose `reply_seq` is older than the latest + * `state_seq` sent answers a change that a later one has superseded, so + * it is dropped; applying it would set buckaroo_state back and make the + * grid purge and re-request rows for the old state. An initial_state + * with no `reply_seq` (another tab's change, a /load push, a fresh + * connection) always applies. */ export class WebSocketModel { private ws: WebSocket; @@ -20,6 +30,7 @@ export class WebSocketModel { private handlers: Map> = new Map(); private state: Record; private pendingChanges: Set = new Set(); + private stateSeq = 0; constructor(ws: WebSocket, initialState: Record) { this.state = { ...initialState }; @@ -37,9 +48,13 @@ export class WebSocketModel { this.state._metadata = msg; this.emit("metadata", msg); } else if (msg.type === "initial_state") { + if (typeof msg.reply_seq === "number" && msg.reply_seq < this.stateSeq) { + // Stale reply to a change this model has since superseded (#998). + return; + } // Bulk state update from server for (const [k, v] of Object.entries(msg)) { - if (k === "type") continue; + if (k === "type" || k === "reply_seq") continue; this.state[k] = v; this.emit(`change:${k}`, v); } @@ -82,9 +97,11 @@ export class WebSocketModel { if (this.ws.readyState !== WebSocket.OPEN) return; // Sync buckaroo_state changes back to the server if (this.pendingChanges.has("buckaroo_state")) { + this.stateSeq += 1; this.ws.send(JSON.stringify({ type: "buckaroo_state_change", new_state: this.state["buckaroo_state"], + state_seq: this.stateSeq, })); } this.pendingChanges.clear(); 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..a8b0a5339 100644 --- a/tests/unit/dataflow/scoped_summary_stats_test.py +++ b/tests/unit/dataflow/scoped_summary_stats_test.py @@ -15,6 +15,8 @@ observable in ``merged_sd``. """ +import copy + import pandas as pd from typing import Any, TypedDict @@ -23,6 +25,7 @@ from buckaroo.customizations.pd_stats_v2 import PD_ANALYSIS_V2, PD_AUTOCLEAN_DEFAULT_V2, cleaning_gen_ops from buckaroo.dataflow.autocleaning import (AutocleaningConfig, PandasAutocleaning) from buckaroo.dataflow.dataflow import CustomizableDataflow, StylingAnalysis +from buckaroo.dataflow.sd_cache import split_chain_by_scope from buckaroo.pluggable_analysis_framework.stat_func import stat @@ -249,3 +252,67 @@ def test_analysis_klasses_change_invalidates_scoped_sd(): f"got the same key {key1} for both — likely the cache key still " f"omits analysis_klasses identity (Codex P2)" ) + + +def _scope_inputs(dfc): + """The pieces ``_merged_sd`` assembles, read back from the dataflow.""" + cache = dfc.summary_stats_cache + return dict( + init_sd=dfc.init_sd, cleaned_sd=dfc.cleaned_sd, raw_sd=cache[dfc.raw_sd_key], + processed_sd=dfc.processed_sd, processed_df=dfc.processed_df, + chains=split_chain_by_scope(dfc.operations), + clean_sd=cache[dfc.clean_sd_key], filt_sd=cache[dfc.filt_sd_key]) + + +def _assembly_dataflow(): + df = pd.DataFrame({'amount': ['10', '20', '30', '40', '50', '60', '70'], + 'word': ['foo', 'bar', 'foo', 'baz', 'foo', 'bar', 'foo']}) + # init_sd is keyed by the original column name; 'amount' is rewritten to 'a'. + dfc = ScopedDataflow(df, init_sd={ + 'amount': {'displayer_args': {'displayer': 'string', 'max_length': 99}, 'init_only': 1}}) + dfc.cleaning_method = 'default' + dfc.quick_command_args = {'search': ['foo']} + return dfc + + +def test_assemble_merged_sd_equals_merged_sd_with_init_sd_and_cleaning(): + """``assemble_merged_sd`` is the body of the ``merged_sd`` observer as a + pure function, so a server can assemble wire stats from an accumulating + raw sd without going through the traitlets cascade.""" + from buckaroo.dataflow.dataflow import assemble_merged_sd + dfc = _assembly_dataflow() + sd = dfc.merged_sd['a'] + assert sd['init_only'] == 1, "init_sd override must reach merged_sd under the rewritten name" + assert 'cleaned_mean' in sd and 'filtered_mean' in sd, "cleaning and filter layers must be active" + + assembled = assemble_merged_sd(**_scope_inputs(dfc)) + + assert assembled == dfc.merged_sd + + +def test_assemble_merged_sd_without_a_processed_df_merges_the_unrewritten_sds(): + from buckaroo.dataflow.dataflow import assemble_merged_sd + assembled = assemble_merged_sd( + init_sd={'x': {'init': 1}}, cleaned_sd={'x': {'cleaned': 2}}, raw_sd={'x': {'raw': 3}}, + processed_sd={'x': {'processed': 4}}, processed_df=None, + chains={'raw': [], 'clean': [], 'filt': []}) + assert assembled == {'x': {'init': 1, 'cleaned': 2, 'raw': 3, 'processed': 4}} + + +def test_assemble_merged_sd_does_not_mutate_its_inputs(): + from buckaroo.dataflow.dataflow import assemble_merged_sd + dfc = _assembly_dataflow() + inputs = _scope_inputs(dfc) + sd_names = ('init_sd', 'cleaned_sd', 'raw_sd', 'processed_sd', 'clean_sd', 'filt_sd') + + def key_sets(): + return {name: {col: sorted(stats) for col, stats in inputs[name].items()} for name in sd_names} + + before, chains_before = key_sets(), copy.deepcopy(inputs['chains']) + + assemble_merged_sd(**inputs) + + # Layering cleaned_* / filtered_* into a result column that shared its dict + # with an input would add keys to that input. + assert key_sets() == before + assert inputs['chains'] == chains_before diff --git a/tests/unit/server/test_data_loading_polars.py b/tests/unit/server/test_data_loading_polars.py index 58f06f5bb..6ed441642 100644 --- a/tests/unit/server/test_data_loading_polars.py +++ b/tests/unit/server/test_data_loading_polars.py @@ -9,15 +9,34 @@ * ``.json`` files are read as standard JSON arrays (matches ``pd.read_json`` default ``lines=False`` so the same file loads under either backend). + +It also pins the schema stats tier (rows-first s2) for the pandas and +polars server dataflows side by side: ``stats_tier="schema"`` publishes +the display state full stats give, without reading a value. """ +import base64 +import datetime +import decimal +import io import json import os import tempfile +import pandas as pd import polars as pl +import pyarrow.parquet as pq import pytest -from buckaroo.server.data_loading_polars import (create_polars_dataflow, handle_infinite_request_buckaroo_polars, load_file_polars) +from buckaroo.dataflow.dataflow import assemble_merged_sd +from buckaroo.dataflow.sd_cache import split_chain_by_scope +from buckaroo.jlisp.lisp_utils import s as lisp_sym +from buckaroo.pluggable_analysis_framework.col_analysis import ColAnalysis +from buckaroo.pluggable_analysis_framework.stat_pipeline import StatPipeline +from buckaroo.pluggable_analysis_framework.utils import PERVERSE_DF +from buckaroo.server.data_loading import ServerDataflow, handle_infinite_request_buckaroo +from buckaroo.server.data_loading_polars import (PolarsServerDataflow, create_polars_dataflow, handle_infinite_request_buckaroo_polars, load_file_polars) +from buckaroo.styling_helpers import float_, obj_ +from tests.unit.dataflow.scoped_summary_stats_test import _scope_inputs def _payload(start=0, end=100): @@ -89,3 +108,319 @@ def test_load_file_polars_ndjson_still_works(): def test_load_file_polars_unsupported_extension(): with pytest.raises(ValueError, match="Unsupported file format"): load_file_polars("/tmp/foo.xyz") + + +# --------------------------------------------------------------------------- +# stats_tier="schema" on the pandas and polars server dataflows (rows-first s2) +# --------------------------------------------------------------------------- + +_BACKENDS = { + "pandas": (ServerDataflow, pd.DataFrame, handle_infinite_request_buckaroo), + "polars": (PolarsServerDataflow, pl.DataFrame, handle_infinite_request_buckaroo_polars)} + +# price has a 1e9 maximum, so its estimated column width depends on the min and +# max that only full stats supply. +_TIER_DATA = { + "price": [12.5, 18.9, 7.4, 22.1, 1e9], + "qty": [1, 2, 1, 3, 2], + "category": ["a", "b", "a", "c", "b"]} + +# The same columns with a different number of rows per value in every column. +# polars does not keep the order of equal counts, so on _TIER_DATA the mode of +# two full-stats runs over the same frame can differ. +_UNTIED_DATA = { + "price": [12.5, 12.5, 12.5, 18.9, 18.9, 1e9], + "qty": [1, 1, 1, 2, 2, 3], + "category": ["a", "a", "a", "b", "b", "c"]} + +# Stats that read values. The schema tier must not publish any of them. +_DATA_STAT_KEYS = { + "mean", "std", "median", "min", "max", "null_count", "non_null_count", + "distinct_count", "unique_count", "value_counts", "mode", "histogram", + "histogram_bins", "memory_usage"} + + +@pytest.fixture(params=list(_BACKENDS)) +def backend(request): + return request.param + + +def _frame(backend, data=None): + return _BACKENDS[backend][1](_TIER_DATA if data is None else data) + + +def _typed_frame(backend): + """One column per kind of dtype the typing stats tell apart.""" + days = [datetime.date(2020, 1, 1), datetime.date(2020, 1, 2), datetime.date(2020, 1, 3)] + if backend == "pandas": + return pd.DataFrame({ + "float": [1.5, 2.5, 3.5], "int": [1, 2, 3], + "nullable_int": pd.array([1, None, 3], dtype="Int64"), + "bool": [True, False, True], "str": ["a", "b", "c"], + "datetime": pd.to_datetime(days), "datetime_utc": pd.to_datetime(days, utc=True), + "timedelta": pd.to_timedelta([1, 2, 3], unit="D"), + "category": pd.Categorical(["a", "b", "a"])}) + return pl.DataFrame({ + "float": [1.5, 2.5, 3.5], "int32": pl.Series([1, 2, 3], dtype=pl.Int32), + "bool": [True, False, True], "str": ["a", "b", "c"], "date": days, + "datetime": [datetime.datetime(2020, 1, d) for d in (1, 2, 3)], + "duration": [datetime.timedelta(days=d) for d in (1, 2, 3)], + "time": [datetime.time(h) for h in (1, 2, 3)], + "decimal": pl.Series([decimal.Decimal("1.5")] * 3, dtype=pl.Decimal(10, 2)), + "binary": [b"a", b"b", b"c"], + "category": pl.Series(["a", "b", "a"], dtype=pl.Categorical)}) + + +def _build_dataflow(backend, frame=None, **kwargs): + frame = _frame(backend) if frame is None else frame + return _BACKENDS[backend][0](frame, skip_main_serial=True, **kwargs) + + +def _three_scope_dataflow(backend, data=None, **kwargs): + """A dataflow with a search (filt scope) active and, on pandas, a user op + (clean scope), so each scope holds its own summary-stats cache entry. The + polars conf ships no command to run outside a search, so there raw and + clean are one scope.""" + dataflow = _build_dataflow(backend, _frame(backend, data), **kwargs) + if backend == "pandas": + dataflow.operations = [[lisp_sym("fillna"), {"symbol": "df"}, "qty", 0]] + dataflow.quick_command_args = {"search": ["a"]} + return dataflow + + +def _scope_keys(dataflow, tier): + chains = split_chain_by_scope(dataflow.operations) + return {scope: dataflow._scope_cache_key(chain, tier=tier) for scope, chain in chains.items()} + + +def _spy_stat_pipeline(monkeypatch): + """Record the columns of every frame ``StatPipeline.process_df`` runs on, + apart from the ``PERVERSE_DF`` its DAG self-check uses, which has no part + of the loaded data in it.""" + frames = [] + original = StatPipeline.process_df + + def spy(self, df, *args, **kwargs): + if list(df.columns) != list(PERVERSE_DF.columns): + frames.append(list(df.columns)) + return original(self, df, *args, **kwargs) + + monkeypatch.setattr(StatPipeline, "process_df", spy) + return frames + + +def _without_min_width(column_config): + return [{**cc, "ag_grid_specs": {k: v for k, v in cc["ag_grid_specs"].items() if k != "minWidth"}} + for cc in column_config] + + +def _as_json(sd): + """Comparable form of an sd: json equates NaN with NaN where == would not.""" + return json.dumps(sd, sort_keys=True, default=str) + + +def _wire_stat_names(dataflow): + """The stat names in the ``all_stats`` payload sent to the client.""" + envelope = dataflow.df_data_dict["all_stats"] + table = pq.read_table(io.BytesIO(base64.b64decode(envelope["data"]))) + return {name.split("__", 1)[1] for name in table.column_names} + + +def _viewer_config(dataflow, display="main"): + return dataflow.df_display_args[display]["df_viewer_config"] + + +class _NoopPostProcessing(ColAnalysis): + provides_defaults = {} + post_processing_method = "noop_post" + + @classmethod + def post_process_df(cls, df): + return [df, {}] + + +class TestStatsTierSchema: + """``stats_tier="schema"`` on ``ServerDataflow`` and + ``PolarsServerDataflow``: every column is typed from its dtype alone, so + the dataflow publishes the display state full stats give with no stat + computed on the data.""" + + def test_matches_full_stats_display_state(self, backend): + full = _build_dataflow(backend) + schema = _build_dataflow(backend, stats_tier="schema") + assert schema.df_display_args.keys() == full.df_display_args.keys() + for name, full_arg in full.df_display_args.items(): + schema_arg = schema.df_display_args[name] + assert schema_arg["data_key"] == full_arg["data_key"] + assert schema_arg["summary_stats_key"] == full_arg["summary_stats_key"] + full_cfg, schema_cfg = full_arg["df_viewer_config"], schema_arg["df_viewer_config"] + assert schema_cfg["pinned_rows"] == full_cfg["pinned_rows"] + assert (_without_min_width(schema_cfg["column_config"]) + == _without_min_width(full_cfg["column_config"])) + + def test_host_supplied_pinned_rows_match_full_stats(self, backend): + """pinned_rows is configuration, so it is the same in every tier even + when it names a stat the schema tier does not compute.""" + pinned = [obj_("dtype"), float_("mean")] + full = _build_dataflow(backend, pinned_rows=pinned) + schema = _build_dataflow(backend, stats_tier="schema", pinned_rows=pinned) + assert _viewer_config(schema)["pinned_rows"] == _viewer_config(full)["pinned_rows"] == pinned + + def test_min_width_is_the_stats_derived_difference(self, backend): + widths = {} + for tier in ("full", "schema"): + cfg = _viewer_config(_build_dataflow(backend, stats_tier=tier)) + widths[tier] = {cc["header_name"]: cc["ag_grid_specs"]["minWidth"] + for cc in cfg["column_config"]} + # price's 1e9 maximum widens it under full stats; without min and max + # the estimate falls back to a one-digit value. + assert widths["schema"]["price"] < widths["full"]["price"] + + def test_typing_matches_full_stats_for_every_dtype(self, backend): + full = _build_dataflow(backend, _typed_frame(backend)) + schema = _build_dataflow(backend, _typed_frame(backend), stats_tier="schema") + assert schema.merged_sd.keys() == full.merged_sd.keys() + for col, stats in schema.merged_sd.items(): + assert {"orig_col_name", "rewritten_col_name", "dtype", "_type", "length"} <= stats.keys() + assert stats == {k: full.merged_sd[col][k] for k in stats} + assert {"duration", "categorical", "datetime"} <= { + stats["_type"] for stats in schema.merged_sd.values()} + + def test_publishes_no_data_derived_stat(self, backend): + full = _build_dataflow(backend) + schema = _build_dataflow(backend, stats_tier="schema") + for stats in schema.merged_sd.values(): + assert not _DATA_STAT_KEYS & stats.keys() + assert _wire_stat_names(schema) == {"dtype"} + # The same reads on the full tier see the stats, so they would notice. + assert _DATA_STAT_KEYS & full.merged_sd["a"].keys() + assert {"dtype", "histogram_bins"} <= _wire_stat_names(full) + + def test_computes_no_stat_on_the_data(self, backend, monkeypatch): + frames = _spy_stat_pipeline(monkeypatch) + dataflow = _build_dataflow(backend, stats_tier="schema") + dataflow.quick_command_args = {"search": ["a"]} + dataflow.add_analysis(_NoopPostProcessing) + assert frames == [] + assert "noop_post" in dataflow.buckaroo_options["post_processing"] + + def test_the_spy_sees_the_pipeline_at_the_full_tier(self, backend, monkeypatch): + frames = _spy_stat_pipeline(monkeypatch) + _build_dataflow(backend) + assert frames, "the spy would not notice a stat computed on the data" + + def test_init_sd_hints_and_overrides_still_apply(self, backend): + dataflow = _build_dataflow( + backend, stats_tier="schema", + init_sd={"qty": {"displayer_args": {"displayer": "string", "max_length": 200}}}, + column_config_overrides={ + "category": {"displayer_args": {"displayer": "string", "max_length": 5000}}}) + by_header = {cc["header_name"]: cc for cc in _viewer_config(dataflow)["column_config"]} + assert by_header["qty"]["displayer_args"]["max_length"] == 200 + assert by_header["category"]["displayer_args"]["max_length"] == 5000 + + def test_skipped_column_keeps_init_sd_typing(self, backend): + """A column in ``skip_stat_columns`` gets only its names from the + pipeline, so its ``_type`` comes from ``init_sd``. The schema tier + must not layer the dtype-derived typing keys over it.""" + kwargs = { + "init_sd": {"qty": {"_type": "float", "mean": 1.8, "min": 1, "max": 3}}, + "skip_stat_columns": ["qty"]} + full = _build_dataflow(backend, **kwargs) + schema = _build_dataflow(backend, stats_tier="schema", **kwargs) + + def merged(dataflow, orig_col): + return next(v for v in dataflow.merged_sd.values() if v["orig_col_name"] == orig_col) + + assert merged(full, "qty")["_type"] == "float" + assert merged(schema, "qty")["_type"] == "float" + assert not {"is_numeric", "is_integer", "is_float", "dtype"} & merged(schema, "qty").keys() + assert merged(schema, "price")["_type"] == "float" + assert merged(schema, "price")["is_float"] is True + assert (_without_min_width(_viewer_config(schema)["column_config"]) + == _without_min_width(_viewer_config(full)["column_config"])) + + def test_empty_frame_matches_full_stats(self, backend): + """Full stats publish nothing for a frame with no rows (the pipeline + returns an empty summary), so the schema tier does the same.""" + empty = {"x": [], "y": []} + full = _build_dataflow(backend, _frame(backend, empty)) + schema = _build_dataflow(backend, _frame(backend, empty), stats_tier="schema") + assert full.merged_sd == {} + assert schema.merged_sd == {} + assert schema.df_display_args == full.df_display_args + + def test_sorted_infinite_request_works(self, backend): + dataflow = _build_dataflow(backend, stats_tier="schema") + qty = next(k for k, v in dataflow.merged_sd.items() if v["orig_col_name"] == "qty") + handler = _BACKENDS[backend][2] + resp, parquet = handler(dataflow, {"start": 0, "end": 5, "sourceName": "default", + "sort": qty, "sort_direction": "desc"}) + assert "error_info" not in resp + assert resp["length"] == 5 + assert pq.read_table(io.BytesIO(parquet)).column(qty).to_pylist() == [3, 2, 2, 1, 1] + + def test_pending_state_writes_no_full_tier_cache_key(self, backend): + dataflow = _three_scope_dataflow(backend, stats_tier="schema") + full_keys = set(_scope_keys(dataflow, "full").values()) + assert len(full_keys) == (3 if backend == "pandas" else 2) + assert dataflow.summary_stats_cache + assert not full_keys & dataflow.summary_stats_cache.keys() + + def test_later_full_assignment_reaches_merged_sd_for_all_scopes(self, backend): + full = _three_scope_dataflow(backend, _UNTIED_DATA) + dataflow = _three_scope_dataflow(backend, _UNTIED_DATA, stats_tier="schema") + assert "mean" not in dataflow.merged_sd["a"] + assert {"mean", "filtered_mean"} <= full.merged_sd["a"].keys() + + dataflow.stats_tier = "full" + dataflow.summary_sd = full.summary_sd + + assert _as_json(dataflow.merged_sd) == _as_json(full.merged_sd) + + def test_full_scope_sds_cached_first_are_used_without_the_pipeline(self, backend, monkeypatch): + full = _three_scope_dataflow(backend) + dataflow = _three_scope_dataflow(backend, stats_tier="schema") + cache = dict(dataflow.summary_stats_cache) + for scope, key in _scope_keys(dataflow, "full").items(): + cache[key] = full.summary_stats_cache[getattr(full, f"{scope}_sd_key")] + dataflow.summary_stats_cache = cache + frames = _spy_stat_pipeline(monkeypatch) + + dataflow.stats_tier = "full" + dataflow.summary_sd = full.summary_sd + + assert frames == [] + assert _as_json(dataflow.merged_sd) == _as_json(full.merged_sd) + + def test_assembled_sd_equals_merged_sd_with_init_sd_and_a_user_op(self, backend): + # init_sd is keyed by the original column name; price is rewritten to "a". + dataflow = _three_scope_dataflow(backend, stats_tier="schema", init_sd={ + "price": {"displayer_args": {"displayer": "string", "max_length": 99}, "init_only": 1}}) + sd = dataflow.merged_sd["a"] + assert sd["init_only"] == 1, "init_sd must reach merged_sd under the rewritten name" + assert "filtered_length" in sd, "the filter layer must be active" + if backend == "pandas": + assert "cleaned_length" in sd, "the cleaning layer must be active" + + assembled = assemble_merged_sd(**_scope_inputs(dataflow)) + + assert _as_json(assembled) == _as_json(dataflow.merged_sd) + + +def test_object_column_is_typed_as_a_string_at_the_schema_tier(): + """pandas gives an object column no dtype to tell strings from other + values, and calls an empty one a string, so the schema tier types every + object column as one. Full stats read the values and call the mixed column + ``obj``; the gap closes when full stats arrive.""" + def frame(): + return pd.DataFrame({ + "words": pd.Series(["a", "b", "c"], dtype=object), + "mixed": pd.Series([1, "b", 2.5], dtype=object)}) + + types = {} + for tier in ("full", "schema"): + sd = ServerDataflow(frame(), skip_main_serial=True, stats_tier=tier).merged_sd + types[tier] = {v["orig_col_name"]: v["_type"] for v in sd.values()} + assert types["full"] == {"words": "string", "mixed": "obj"} + assert types["schema"] == {"words": "string", "mixed": "string"} diff --git a/tests/unit/server/test_load_compare.py b/tests/unit/server/test_load_compare.py index 4b14fd8c0..9f9a810d0 100644 --- a/tests/unit/server/test_load_compare.py +++ b/tests/unit/server/test_load_compare.py @@ -10,6 +10,7 @@ from buckaroo.server import telemetry from buckaroo.server.app import make_app as _make_app +from buckaroo.server.session import build_state_message pytestmark = pytest.mark.skipif(sys.platform == "win32", reason="Temp file locking prevents cleanup on Windows") @@ -116,6 +117,32 @@ def test_eqs_reports_diffs(self): os.unlink(f1.name) os.unlink(f2.name) + def test_load_compare_resolves_to_full_and_clears_a_stats_policy(self): + """A compare session is a viewer with no dataflow, so no stats are + deferred or refused: a policy that an earlier /load_expr left on the + session is dropped (rows-first p33).""" + session = self._app.settings["sessions"].create("cmp-policy", "") + session.stats_tier, session.stats_delivery = "auto", "deferred" + session.stats_policy = {"tier_target": "schema", "reason": "size"} + df1 = pd.DataFrame({"id": [1, 2], "score": [100, 200]}) + df2 = pd.DataFrame({"id": [1, 2], "score": [100, 999]}) + + with tempfile.NamedTemporaryFile(suffix=".csv", delete=False) as f1, \ + tempfile.NamedTemporaryFile(suffix=".csv", delete=False) as f2: + _write_df(df1, f1.name) + _write_df(df2, f2.name) + try: + resp = self._post_compare({"session": "cmp-policy", "path1": f1.name, "path2": f2.name, + "join_columns": ["id"]}) + self.assertEqual(resp.code, 200) + finally: + os.unlink(f1.name) + os.unlink(f2.name) + self.assertEqual((session.stats_tier, session.stats_delivery), ("full", "inline")) + self.assertIsNone(session.stats_policy) + self.assertEqual(session.stats_status, "complete") + self.assertNotIn("stats", build_state_message(session)["df_meta"]) + def test_missing_fields(self): resp = self._post_compare({"session": "x", "path1": "/a"}) self.assertEqual(resp.code, 400) diff --git a/tests/unit/server/test_load_expr.py b/tests/unit/server/test_load_expr.py index bdbb47db6..8f5de964f 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 datetime +import gc import io import json import os import shutil import sys import tempfile +import weakref +from contextlib import contextmanager from pathlib import Path +from types import SimpleNamespace from unittest.mock import MagicMock, patch import pandas as pd import pyarrow.parquet as pq import pytest +import tornado.gen import tornado.httpclient import tornado.testing import tornado.websocket @@ -24,9 +30,17 @@ from xorq.common.utils.graph_utils import replace_nodes, walk_nodes # noqa: E402 from xorq.common.utils.provenance_utils import read_parquet_provenance # noqa: E402 from xorq.expr.relations import CachedNode # noqa: E402 - -from buckaroo.server import telemetry, xorq_loading # noqa: E402 +from xorq.vendor.ibis.expr.types.core import Expr # 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.xorq_stat_pipeline import XorqStatPipeline # noqa: E402 +from buckaroo.serialization_utils import resolve_summary_stats_payload # noqa: E402 +from buckaroo.server import session as session_mod # noqa: E402 +from buckaroo.server import stats_wire, telemetry, xorq_loading # noqa: E402 from buckaroo.server.app import make_app as _make_app # noqa: E402 +from buckaroo.server.websocket_handler import DataStreamHandler # noqa: E402 pytestmark = pytest.mark.skipif( sys.platform == "win32", @@ -910,6 +924,1989 @@ 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 + + +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_numbered_change_on_a_deferred_session_echoes_reply_seq_to_its_sender_only(self): + """#998's ``reply_seq`` rides the stats-aware broadcast: the caps client + that sent a numbered dataflow change gets it on its pending frame, and + the legacy client gets a complete frame without it.""" + sid = "sw-reply-seq" + a, a_open, b, _ = await self._pair(sid) + gen = self._stats(a_open)["gen"] + change = json.loads(_state_change(post_processing="first_three")) + a.write_message(json.dumps({**change, "state_seq": 5})) + a_frame, b_frame = await _read_json(a), await _read_json(b) + self.assertEqual(a_frame.get("reply_seq"), 5) + self.assertNotIn("reply_seq", b_frame) + self._assert_pending(a_frame, gen + 1) + self.assertEqual(self._stats(b_frame)["status"], "complete") + + @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, _with_highlight=DataStreamHandler._with_highlight) + self.assertEqual(getattr(session, "stats_status", None), "pending") + + DataStreamHandler._send_client_state(legacy, session, session.buckaroo_state) + + 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"])) + + +# The table _build_stats_wire_dir builds: five rows, three columns. +_WIRE_ESTIMATE = {"rows": 5, "cols": 3} +_ONDEMAND = "stats_update,stats_ondemand" +# Limits that put the table in the scalar or the schema tier by size. +_SCALAR_BY_SIZE = {"full_auto_rows": 3, "scalar_auto_cells": 1_000} +_SCHEMA_BY_SIZE = {"full_auto_rows": 3, "scalar_auto_cells": 10} + + +@contextmanager +def _stats_limits(**limits): + """Server stats thresholds for the block, as ``BUCKAROO_STATS_*`` + environment overrides, which ``resolve_stats_policy`` reads on every call.""" + env = {f"BUCKAROO_STATS_{name.upper()}": str(value) for name, value in limits.items()} + with patch.dict(os.environ, env): + yield + + +# (host tier, limits, df_meta.stats an ondemand client sees, minus gen). The +# five-by-three table is full by default; each set of limits moves it down. +_POLICY_CASES = { + "auto, within every threshold": ("auto", {}, {"status": "pending", "tier": "schema", "tier_target": "full", + "requestable": [], "estimate": _WIRE_ESTIMATE}), + "auto, scalar by size": ("auto", {"full_auto_rows": 3}, {"status": "not_computed", "tier": "schema", + "reason": "size", "tier_target": "scalar", "estimate": _WIRE_ESTIMATE}), + "auto, schema by size": ("auto", {"full_auto_rows": 3, "scalar_auto_cells": 10}, {"status": "not_computed", + "tier": "schema", "reason": "size", "tier_target": "schema", "auto_request": False, + "requestable": ["scalar", "full"], "estimate": _WIRE_ESTIMATE}), + "full, over the ceiling": ("full", {"ceiling_full_rows": 3}, {"status": "not_computed", "tier": "schema", + "reason": "ceiling", "tier_target": "scalar", "requestable": [], "estimate": _WIRE_ESTIMATE}), + "scalar, named by the host": ("scalar", {}, {"status": "not_computed", "tier": "schema", "reason": "host", + "tier_target": "scalar", "estimate": _WIRE_ESTIMATE}), + "schema, named by the host": ("schema", {}, {"status": "not_computed", "tier": "schema", "reason": "host", + "tier_target": "schema", "auto_request": False, "requestable": ["scalar", "full"], + "estimate": _WIRE_ESTIMATE}), + "scalar, over the scalar ceiling": ("scalar", {"ceiling_scalar_cells": 10}, {"status": "not_computed", + "tier": "schema", "reason": "ceiling", "tier_target": "schema", "auto_request": False, "requestable": [], + "estimate": _WIRE_ESTIMATE}), +} + + +class TestStatsPolicyWire(tornado.testing.AsyncHTTPTestCase): + """The stats policy on a ``/load_expr`` session (rows-first p33): + ``stats_tier`` ``auto | full | scalar | schema`` stored with the pair, the + policy resolved after the schema-tier dataflow and the count exist, reported + in ``df_meta.stats`` and applied at WebSocket open for a client that + advertises ``stats_ondemand`` alone, with the version-skew cases.""" + + def get_app(self): + return make_app() + + def setUp(self): + super().setUp() + self.builds_root = tempfile.mkdtemp() + self.project_root = tempfile.mkdtemp() + self.addCleanup(shutil.rmtree, self.builds_root, ignore_errors=True) + self.addCleanup(shutil.rmtree, self.project_root, ignore_errors=True) + self.build_path = _build_stats_wire_dir(self.builds_root) + self.clients = [] + + def tearDown(self): + for ws in self.clients: + ws.close() + super().tearDown() + + def _session(self, sid): + return self._app.settings["sessions"].get(sid) + + async def _load(self, sid, limits=None, **body): + with _stats_limits(**(limits or {})): + resp = await _post(self.get_http_port(), "/load_expr", + {"session": sid, "build_dir": self.build_path, "project_root": self.project_root, **body}) + self.assertEqual(resp.code, 200, resp.body) + + async def _reload(self, sid, limits=None, **body): + with _stats_limits(**(limits or {})): + resp = await _post(self.get_http_port(), f"/reload_expr/{sid}", body) + self.assertEqual(resp.code, 200, resp.body) + + async def _connect(self, sid, caps=None): + suffix = f"?caps={caps}" if caps else "" + ws = await tornado.websocket.websocket_connect( + f"ws://localhost:{self.get_http_port()}/ws/{sid}{suffix}") + self.clients.append(ws) + return ws, await _read_json(ws) + + def _stats(self, frame): + stats = frame["df_meta"].get("stats") + self.assertIsNotNone(stats, "initial_state carries no df_meta.stats") + return stats + + def _schema_rows(self, frame): + return list(_rows_by_stat(frame["df_data_dict"]["all_stats"])) + + async def _inline_frame(self, sid): + """The complete message an inline session of the same build sends.""" + await self._load(sid) + return (await self._connect(sid))[1] + + def _assert_complete(self, frame, inline_frame): + stats = self._stats(frame) + self.assertEqual((stats["status"], stats["tier"]), ("complete", "full")) + self.assertIn("histogram_bins", _rows_by_stat(frame["df_data_dict"]["all_stats"])) + self.assertEqual(_comparable(frame), _comparable(inline_frame)) + + @tornado.testing.gen_test + async def test_every_stats_tier_value_is_accepted_and_stored(self): + for tier in ("auto", "full", "scalar", "schema"): + sid = f"pw-accept-{tier}" + await self._load(sid, stats_tier=tier, stats_delivery="deferred") + session = self._session(sid) + self.assertEqual((session.stats_tier, session.stats_delivery), (tier, "deferred")) + self.assertEqual(session.xorq_dataflow.stats_tier, "schema") + resp = await _post(self.get_http_port(), f"/reload_expr/{sid}", {"stats_tier": "auto"}) + self.assertEqual(resp.code, 200, resp.body) + self.assertEqual(session.stats_tier, "auto") + resp = await _post(self.get_http_port(), "/load_expr", + {"session": "pw-bad", "build_dir": self.build_path, "stats_tier": "none"}) + self.assertEqual((resp.code, json.loads(resp.body)["error_code"]), (400, "invalid_stats_tier")) + + @tornado.testing.gen_test + async def test_a_tier_named_below_full_builds_a_schema_dataflow_with_inline_delivery(self): + """No scalar-tier units exist yet, so a scalar target is built and held at + the schema tier.""" + await self._load("pw-scalar-inline", stats_tier="scalar") + session = self._session("pw-scalar-inline") + self.assertEqual((session.xorq_dataflow.stats_tier, session.stats_status, session.stats_reason), + ("schema", "not_computed", "host")) + self.assertEqual(session.stats_policy["tier_target"], "scalar") + + @tornado.testing.gen_test + async def test_the_default_tier_stays_full(self): + await self._load("pw-default", stats_delivery="deferred") + session = self._session("pw-default") + self.assertEqual(session.stats_tier, "full") + self.assertEqual(session.stats_policy["tier_target"], "full") + + @tornado.testing.gen_test + async def test_an_inline_session_has_no_schema_dataflow_to_resolve_against(self): + """The constructor runs the stats of an inline session, so auto is full + there and the session sends the message it always has.""" + with _count_stat_queries() as queries: + await self._load("pw-auto-inline", limits={"full_auto_rows": 1, "ceiling_full_rows": 1}, + stats_tier="auto") + session = self._session("pw-auto-inline") + self.assertEqual((session.stats_tier, session.stats_delivery, session.stats_policy), ("auto", "inline", None)) + self.assertEqual(session.xorq_dataflow.stats_tier, "full") + self.assertTrue(queries, "an inline session runs its stats in the constructor") + _, frame = await self._connect("pw-auto-inline", caps=_ONDEMAND) + self.assertNotIn("stats", frame["df_meta"]) + + @tornado.testing.gen_test + async def test_the_policy_resolves_once_the_schema_dataflow_and_the_count_exist(self): + events, calls = [], [] + original_metadata = xorq_loading.get_xorq_metadata + original_resolve = stats_wire.resolve_stats_policy + + def metadata(*args, **kwargs): + events.append("dataflow and count") + return original_metadata(*args, **kwargs) + + def resolve(*args, **kwargs): + events.append("policy") + calls.append((args, kwargs)) + return original_resolve(*args, **kwargs) + + with ( + patch.object(xorq_loading, "get_xorq_metadata", metadata), + patch.object(stats_wire, "resolve_stats_policy", resolve), + _count_stat_queries() as queries, + ): + await self._load("pw-order", stats_tier="auto", stats_delivery="deferred") + self.assertEqual(events, ["dataflow and count", "policy"]) + (args, kwargs), = calls + self.assertEqual((args[0], args[2], args[3], kwargs["host_tier"]), ("xorq", 5, 3, "auto")) + self.assertEqual(queries, [], "resolving the policy runs no data stat query") + + @tornado.testing.gen_test + async def test_the_resolved_policy_is_stored_on_the_session(self): + limits = {"full_auto_rows": 3, "scalar_auto_cells": 10} + await self._load("pw-store", limits=limits, stats_tier="auto", stats_delivery="deferred") + session = self._session("pw-store") + self.assertEqual(session.stats_policy, {"tier_target": "schema", "auto_request": False, + "requestable": ["scalar", "full"], "reason": "size", "estimate": _WIRE_ESTIMATE}) + self.assertEqual((session.stats_status, session.stats_reason), ("not_computed", "size")) + + @tornado.testing.gen_test + async def test_df_meta_stats_carries_the_policy_fields_for_an_ondemand_client(self): + for index, (name, (tier, limits, expected)) in enumerate(_POLICY_CASES.items()): + sid = f"pw-fields-{index}" + await self._load(sid, limits=limits, stats_tier=tier, stats_delivery="deferred") + _, frame = await self._connect(sid, caps=_ONDEMAND) + stats = self._stats(frame) + self.assertEqual(stats, {**expected, "gen": stats["gen"]}, name) + self.assertEqual(self._schema_rows(frame), ["dtype"], f"{name}: the first message is the schema tier") + + @tornado.testing.gen_test + async def test_a_warm_repost_with_the_same_auto_tier_short_circuits(self): + body = {"session": "pw-warm", "build_dir": self.build_path, "stats_tier": "auto", "stats_delivery": "deferred"} + self.assertEqual((await _post(self.get_http_port(), "/load_expr", body)).code, 200) + with patch.object(xorq_loading, "load_expr_build_dir", + side_effect=AssertionError("an unchanged pair must take the warm exit")): + same = await _post(self.get_http_port(), "/load_expr", body) + omitted = await _post(self.get_http_port(), "/load_expr", + {"session": "pw-warm", "build_dir": self.build_path}) + self.assertEqual((same.code, omitted.code), (200, 200)) + session = self._session("pw-warm") + self.assertEqual((session.stats_tier, session.stats_delivery), ("auto", "deferred")) + + @tornado.testing.gen_test + async def test_a_changed_tier_rebuilds_and_resolves_again(self): + sid = "pw-change" + await self._load(sid, stats_tier="auto", stats_delivery="deferred") + self.assertEqual(self._session(sid).stats_policy["tier_target"], "full") + calls = [] + original = xorq_loading.load_expr_build_dir + + def counting_loader(bd, **kwargs): + calls.append(bd) + return original(bd, **kwargs) + + with patch.object(xorq_loading, "load_expr_build_dir", side_effect=counting_loader): + await self._load(sid, stats_tier="scalar", stats_delivery="deferred") + self.assertEqual(len(calls), 1, "a changed tier must rebuild") + self.assertEqual(self._session(sid).stats_policy["tier_target"], "scalar") + await self._load(sid, stats_tier="scalar", stats_delivery="deferred") + self.assertEqual(len(calls), 1, "the same tier takes the warm exit") + + @tornado.testing.gen_test + async def test_reload_expr_resolves_the_policy_again(self): + sid = "pw-reload" + await self._load(sid, stats_tier="auto", stats_delivery="deferred") + ws, first = await self._connect(sid, caps=_ONDEMAND) + gen = self._stats(first)["gen"] + self.assertEqual(self._stats(first)["tier_target"], "full") + + await self._reload(sid, limits={"full_auto_rows": 3}) + pushed = self._stats(await _read_json(ws)) + self.assertEqual((pushed["status"], pushed["reason"], pushed["tier_target"]), + ("not_computed", "size", "scalar")) + self.assertGreater(pushed["gen"], gen) + self.assertEqual(self._session(sid).stats_policy["tier_target"], "scalar") + + await self._reload(sid, stats_tier="schema") + pushed = self._stats(await _read_json(ws)) + self.assertEqual((pushed["reason"], pushed["tier_target"], pushed["auto_request"]), ("host", "schema", False)) + self.assertEqual(self._session(sid).stats_tier, "schema") + + await self._reload(sid, stats_tier="auto") + pushed = self._stats(await _read_json(ws)) + self.assertEqual((pushed["status"], pushed["tier_target"]), ("pending", "full")) + + @tornado.testing.gen_test + async def test_a_reload_that_omits_the_tier_keeps_the_session_s_and_resolves_it_against_the_stored_count(self): + """The expression is reused, so its count is the one the load took + (``session.metadata``), and the reload runs no second count.""" + sid = "pw-reload-keep" + await self._load(sid, stats_tier="auto", stats_delivery="deferred") + session = self._session(sid) + session.metadata = {**session.metadata, "rows": 20_000_000} + await self._reload(sid) + self.assertEqual(session.stats_tier, "auto") + self.assertEqual((session.stats_policy["tier_target"], session.stats_policy["estimate"]), + ("scalar", {"rows": 20_000_000, "cols": 3})) + + @tornado.testing.gen_test + async def test_the_policy_survives_a_dataflow_field_change(self): + await self._load("pw-keep", limits=_SCALAR_BY_SIZE, stats_tier="auto", stats_delivery="deferred") + ws, first = await self._connect("pw-keep", caps=_ONDEMAND) + gen = self._stats(first)["gen"] + ws.write_message(_state_change(quick_command_args={"search": ["a"]})) + stats = self._stats(await _read_json(ws)) + self.assertEqual((stats["status"], stats["reason"], stats["tier_target"], stats["gen"]), + ("not_computed", "size", "scalar", gen + 1)) + + # -- version skew ------------------------------------------------------- + + async def _policy_session(self, sid, **limits): + """A deferred ``auto`` session that resolved to the schema tier by size.""" + await self._load(sid, limits=limits or _SCHEMA_BY_SIZE, stats_tier="auto", stats_delivery="deferred") + + @tornado.testing.gen_test + async def test_no_caps_gets_a_complete_state_at_connect(self): + inline = await self._inline_frame("pw-skew-inline") + await self._policy_session("pw-skew-none") + _, legacy = await self._connect("pw-skew-none") + self._assert_complete(legacy, inline) + self.assertEqual(self._session("pw-skew-none").stats_status, "complete") + + @tornado.testing.gen_test + async def test_stats_ondemand_without_stats_update_is_a_client_with_no_caps(self): + inline = await self._inline_frame("pw-skew-inline-2") + await self._policy_session("pw-skew-odonly") + _, frame = await self._connect("pw-skew-odonly", caps="stats_ondemand") + self._assert_complete(frame, inline) + + @tornado.testing.gen_test + async def test_stats_update_only_is_served_as_a_session_headed_for_full(self): + await self._policy_session("pw-skew-update") + ws, first = await self._connect("pw-skew-update", caps="stats_update") + stats = self._stats(first) + self.assertEqual(stats, {"status": "pending", "tier": "schema", "gen": stats["gen"]}) + self.assertEqual(self._schema_rows(first), ["dtype"]) + + ws.write_message(_stats_request(stats["gen"])) + update = await _read_json(ws) + self.assertEqual((update["type"], update["final"], update["tier"]), ("stats_update", True, "full")) + self.assertIn("histogram_bins", _rows_by_stat(update["payload"])) + self.assertEqual(self._session("pw-skew-update").stats_status, "complete") + + @tornado.testing.gen_test + async def test_both_bits_get_the_schema_tier_not_computed_and_run_no_stat_query(self): + await self._policy_session("pw-skew-both") + with _count_stat_queries() as queries: + ws, first = await self._connect("pw-skew-both", caps=_ONDEMAND) + stats = self._stats(first) + self.assertEqual((stats["status"], stats["tier"], stats["reason"], stats["tier_target"]), + ("not_computed", "schema", "size", "schema")) + self.assertEqual((stats["auto_request"], stats["requestable"]), (False, ["scalar", "full"])) + self.assertEqual(self._schema_rows(first), ["dtype"]) + self.assertEqual(queries, []) + # The server enforces it: nothing is requestable until the request path serves a tier. + ws.write_message(_stats_request(stats["gen"])) + aborted = await _read_json(ws) + self.assertEqual((aborted["type"], aborted["reason"]), ("stats_aborted", "not_requestable")) + self.assertEqual(self._session("pw-skew-both").stats_status, "not_computed") + + @tornado.testing.gen_test + async def test_the_three_client_kinds_see_one_session_each_their_own_way(self): + inline = await self._inline_frame("pw-skew-inline-3") + await self._policy_session("pw-skew-mixed") + _, ondemand = await self._connect("pw-skew-mixed", caps=_ONDEMAND) + _, update = await self._connect("pw-skew-mixed", caps="stats_update") + self.assertEqual(self._stats(ondemand)["status"], "not_computed") + self.assertEqual(self._stats(update)["status"], "pending") + _, legacy = await self._connect("pw-skew-mixed") + self._assert_complete(legacy, inline) + + @tornado.testing.gen_test + async def test_an_ondemand_client_is_sent_its_frame_before_a_legacy_client_completes_the_session(self): + await self._policy_session("pw-skew-order", **_SCALAR_BY_SIZE) + a, a_open = await self._connect("pw-skew-order", caps=_ONDEMAND) + gen = self._stats(a_open)["gen"] + b, b_open = await self._connect("pw-skew-order") + self.assertEqual(self._stats(b_open)["status"], "complete") + + a.write_message(_state_change(quick_command_args={"search": ["a"]})) + a_frame, b_frame = await _read_json(a), await _read_json(b) + self.assertEqual((self._stats(a_frame)["status"], self._stats(a_frame)["gen"]), ("not_computed", gen + 1)) + self.assertEqual((self._stats(b_frame)["status"], self._stats(b_frame)["gen"]), ("complete", gen + 1)) + + @tornado.testing.gen_test + async def test_a_tier_the_host_named_reaches_every_client_as_the_host_chose_it(self): + await self._load("pw-skew-host", stats_tier="scalar", stats_delivery="deferred") + _, legacy = await self._connect("pw-skew-host") + self.assertEqual(self._schema_rows(legacy), ["dtype"]) + self.assertEqual(self._stats(legacy)["status"], "not_computed") + ws, update = await self._connect("pw-skew-host", caps="stats_update") + self.assertEqual(self._stats(update), {"status": "not_computed", "tier": "schema", + "gen": self._stats(update)["gen"], "reason": "host"}) + ws.write_message(_stats_request(self._stats(update)["gen"])) + self.assertEqual((await _read_json(ws))["reason"], "not_requestable") + _, ondemand = await self._connect("pw-skew-host", caps=_ONDEMAND) + self.assertEqual(self._stats(ondemand)["tier_target"], "scalar") + + @tornado.testing.gen_test + async def test_a_client_of_an_old_server_reads_a_frame_with_no_stats_as_complete(self): + """An old server sends no ``df_meta.stats`` and no policy field; a new + client falls back to the documented defaults and waits for nothing.""" + await self._load("pw-old") + _, frame = await self._connect("pw-old", caps=_ONDEMAND) + stats = session_mod.stats_with_defaults(frame["df_meta"]) + self.assertEqual((stats["status"], stats["tier"], stats["tier_target"]), ("complete", "full", "full")) + self.assertEqual((stats["auto_request"], stats["requestable"]), (True, ["full"])) + self.assertEqual(stats["demand_columns"], []) + + @tornado.testing.gen_test + async def test_an_explicit_full_changes_no_frame(self): + """With ``stats_tier`` full the message is the one sent without the + field, and the policy is recorded but not reported.""" + for delivery in ("inline", "deferred"): + await self._load(f"pw-same-{delivery}-a", stats_delivery=delivery) + await self._load(f"pw-same-{delivery}-b", stats_tier="full", stats_delivery=delivery) + # A client with no caps completes a deferred session, so it goes last. + for caps in ("stats_update", _ONDEMAND, None): + _, before = await self._connect(f"pw-same-{delivery}-a", caps=caps) + _, after = await self._connect(f"pw-same-{delivery}-b", caps=caps) + self.assertEqual(_as_json(before["df_meta"]), _as_json(after["df_meta"]), (delivery, caps)) + self.assertEqual(_comparable(before), _comparable(after), (delivery, caps)) + self.assertEqual(list(before), list(after)) + if delivery == "deferred" and caps: + stats = self._stats(after) + self.assertEqual(stats, {"status": "pending", "tier": "schema", "gen": stats["gen"]}, caps) + self.assertEqual(self._session("pw-same-deferred-b").stats_policy["tier_target"], "full") + self.assertIsNone(self._session("pw-same-inline-b").stats_policy) + + +class _LimitsWire(tornado.testing.AsyncHTTPTestCase): + """What the policy wire tests share: a ``/load_expr`` session over the + five-row table (deferred unless a test says otherwise) and clients that + advertise both capability bits. A request is judged against the thresholds + again, so a test that sets them holds them in the environment for the whole + block (``_stats_limits``), not only around the load.""" + + DISPLAY_FILES: dict = {} + + def get_app(self): + return make_app() + + def setUp(self): + super().setUp() + self.builds_root = tempfile.mkdtemp() + self.project_root = tempfile.mkdtemp() + self.addCleanup(shutil.rmtree, self.builds_root, ignore_errors=True) + self.addCleanup(shutil.rmtree, self.project_root, ignore_errors=True) + pp_dir = os.path.join(self.project_root, "post_processing") + os.makedirs(pp_dir) + with open(os.path.join(pp_dir, "first_three.py"), "w") as f: + f.write("def process(expr):\n return expr.limit(3)\n") + if self.DISPLAY_FILES: + display_dir = os.path.join(self.project_root, "display") + os.makedirs(display_dir) + for name, source in self.DISPLAY_FILES.items(): + with open(os.path.join(display_dir, name), "w") as f: + f.write(source) + self.build_path = _build_stats_wire_dir(self.builds_root) + self.clients = [] + + def tearDown(self): + for ws in self.clients: + ws.close() + super().tearDown() + + def _session(self, sid): + return self._app.settings["sessions"].get(sid) + + async def _load(self, sid, **body): + resp = await _post(self.get_http_port(), "/load_expr", {"session": sid, "build_dir": self.build_path, + "project_root": self.project_root, "stats_delivery": "deferred", **body}) + self.assertEqual(resp.code, 200, resp.body) + + async def _reload(self, sid, **body): + resp = await _post(self.get_http_port(), f"/reload_expr/{sid}", body) + self.assertEqual(resp.code, 200, resp.body) + + async def _connect(self, sid, caps=_ONDEMAND): + suffix = f"?caps={caps}" if caps else "" + ws = await tornado.websocket.websocket_connect( + f"ws://localhost:{self.get_http_port()}/ws/{sid}{suffix}") + self.clients.append(ws) + return ws, await _read_json(ws) + + def _stats(self, frame): + stats = frame["df_meta"].get("stats") + self.assertIsNotNone(stats, "initial_state carries no df_meta.stats") + return stats + + async def _inline_frame(self, sid): + """The complete message an inline session of the same build sends.""" + await self._load(sid, stats_delivery="inline") + return (await self._connect(sid, caps=None))[1] + + async def _state(self, ws, **changes): + """A state change from ``ws``, and the frame it broadcasts.""" + ws.write_message(_state_change(**changes)) + return await _read_json(ws) + + async def _ask_to_the_end(self, ws, gen, limit=12, **fields): + """``stats_request``s, the same each time, until a reply is final. A + request is the whole run, so that is one.""" + replies = [] + for _ in range(limit): + ws.write_message(_stats_request(gen, **fields)) + replies.append(await _read_json(ws)) + if replies[-1]["type"] != "stats_update" or replies[-1]["final"]: + break + return replies + + +# --------------------------------------------------------------------------- +# Guards for huge sources (rows-first p37) +# --------------------------------------------------------------------------- + +# A display klass that turns sorting back on for every column, as a klass that +# sets its own ag_grid_specs might. +_SORTABLE_STYLING = ''' +class SortableStyling(DefaultMainStyling): + df_display_name = "main" + + @classmethod + def style_column(cls, col, column_metadata): + base_config = super().style_column(col, column_metadata) + base_config["ag_grid_specs"] = {**base_config.get("ag_grid_specs", {}), "sortable": True, "minWidth": 90} + return base_config +''' + +# The same from the host's side: an override applied after the klasses have styled. +_SORTABLE_OVERRIDE = {"qty": {"ag_grid_specs": {"sortable": True}}} + + +@contextmanager +def _guard_limits(**limits): + """Sort and search thresholds for the block, as ``BUCKAROO_*`` environment + overrides, which are read each time a load resolves its guards. The table + ``_build_stats_wire_dir`` builds has five rows, so a ``sort_disable_rows`` of + 3 puts it above the sort threshold.""" + env = {f"BUCKAROO_{name.upper()}": str(value) for name, value in limits.items()} + with patch.dict(os.environ, env): + yield + + +@contextmanager +def _count_backend_queries(): + """Record the outermost op of every query sent to the backend (``execute`` + and ``to_pyarrow``) while the block runs: a count is ``CountStar``.""" + queries = [] + original_execute, original_to_pyarrow = Expr.execute, Expr.to_pyarrow + + def execute(self, *args, **kwargs): + queries.append(type(self.op()).__name__) + return original_execute(self, *args, **kwargs) + + def to_pyarrow(self, *args, **kwargs): + queries.append(type(self.op()).__name__) + return original_to_pyarrow(self, *args, **kwargs) + + with patch.object(Expr, "execute", execute), patch.object(Expr, "to_pyarrow", to_pyarrow): + yield queries + + +def _grid_columns(frame, display="main"): + """Every column config the grid of ``display`` is built from, the index + columns included.""" + config = frame["df_display_args"][display]["df_viewer_config"] + return config["column_config"] + config["left_col_configs"] + + +def _sortable(column): + return column.get("ag_grid_specs", {}).get("sortable") + + +class TestSortGuardDataflow: + """The sort guard on the dataflow (rows-first p37): ``set_sort_enabled`` turns + sorting off in the display config as a final pass, after the klasses and the + host's overrides, and the state of the dataflow survives a rebuild.""" + + @staticmethod + def _main(dataflow): + config = dataflow.df_display_args["main"]["df_viewer_config"] + return config["column_config"] + config["left_col_configs"] + + def test_a_dataflow_sorts_by_default(self): + dataflow = _build_dataflow() + assert dataflow.sort_enabled is True + assert all(_sortable(c) is not False for c in self._main(dataflow)) + + def test_disabling_sorting_covers_every_column_including_one_an_override_sets(self): + dataflow = _build_dataflow(column_config_overrides=_SORTABLE_OVERRIDE) + assert [_sortable(c) for c in self._main(dataflow) if c.get("header_name") == "qty"] == [True] + dataflow.set_sort_enabled(False) + columns = self._main(dataflow) + assert len(columns) == 4 + assert [_sortable(c) for c in columns] == [False] * 4 + + def test_the_summary_display_is_not_changed(self): + dataflow = _build_dataflow() + before = json.dumps(dataflow.df_display_args["summary"], sort_keys=True) + dataflow.set_sort_enabled(False) + assert json.dumps(dataflow.df_display_args["summary"], sort_keys=True) == before + + def test_enabling_sorting_again_restores_the_config(self): + dataflow = _build_dataflow() + before = json.dumps(dataflow.df_display_args, sort_keys=True) + dataflow.set_sort_enabled(False) + assert json.dumps(dataflow.df_display_args, sort_keys=True) != before + dataflow.set_sort_enabled(True) + assert json.dumps(dataflow.df_display_args, sort_keys=True) == before + + def test_the_guard_survives_a_state_change_and_new_stats(self): + dataflow = _build_dataflow(stats_tier="schema") + dataflow.set_sort_enabled(False) + dataflow.quick_command_args = {"search": ["a"]} + assert [_sortable(c) for c in self._main(dataflow)] == [False] * 4 + stats_wire.assign_full_stats(dataflow) + assert [_sortable(c) for c in self._main(dataflow)] == [False] * 4 + + +class TestSortGuardWire(_LimitsWire): + """A session a host opened with a stats policy, over the sort threshold + (rows-first p37): sorting is off in its config, a sorted window is refused + with ``sort_disabled``, and ``df_meta`` says so. The table has five rows, so a + ``sort_disable_rows`` of 3 is over.""" + + DISPLAY_FILES = {"sortable_styling.py": _SORTABLE_STYLING} + + async def _window(self, ws, **fields): + ws.write_message(json.dumps({"type": "infinite_request", "payload_args": {"start": 0, "end": 5, + "sourceName": "default", "origEnd": 5, **fields}})) + return await _read_json(ws) + + async def _served(self, ws, **fields): + """A window the server serves: the ``infinite_resp``, then its parquet frame.""" + resp = await self._window(ws, **fields) + self.assertEqual((resp["type"], resp.get("error_info")), ("infinite_resp", None), resp) + self.assertIsInstance(await ws.read_message(), bytes) + return resp + + async def _refused(self, ws, rows=5, **fields): + """A window the server refuses, with nothing after it: the next frame + read is the answer to the next request, not a parquet frame. ``rows`` is + the length of that next answer.""" + resp = await self._refusal(ws, **fields) + again = await self._window(ws) + self.assertEqual((again["type"], again["length"]), ("infinite_resp", rows), again) + self.assertIsInstance(await ws.read_message(), bytes) + return resp + + async def _refusal(self, ws, **fields): + resp = await self._window(ws, **fields) + self.assertEqual((resp["type"], resp.get("error_code")), ("infinite_resp", "sort_disabled"), resp) + return resp + + @tornado.testing.gen_test + async def test_every_column_stops_sorting_whatever_a_klass_or_an_override_says(self): + with _guard_limits(sort_disable_rows=3): + await self._load("sg-columns", column_config_overrides=_SORTABLE_OVERRIDE) + _, frame = await self._connect("sg-columns") + columns = _grid_columns(frame) + self.assertEqual(len(columns), 4) + self.assertEqual([_sortable(c) for c in columns], [False] * 4, columns) + by_name = {c["header_name"]: c["ag_grid_specs"] for c in columns if "header_name" in c and c["col_name"] != "index"} + self.assertEqual(by_name["price"], {"minWidth": 90, "sortable": False}, "the klass's other specs are kept") + self.assertEqual(by_name["qty"], {"sortable": False}, + "the override replaced the specs, and the guard has the last word") + self.assertEqual(frame["df_display_args"], self._session("sg-columns").df_display_args) + + @tornado.testing.gen_test + async def test_a_session_built_the_old_way_is_not_guarded(self): + """No stats policy, no guard: inline stats is what every host gets today.""" + with _guard_limits(sort_disable_rows=3): + await self._load("sg-inline", stats_delivery="inline") + ws, frame = await self._connect("sg-inline") + self.assertNotIn("sort", frame["df_meta"]) + self.assertTrue(all(_sortable(c) is not False for c in _grid_columns(frame))) + await self._served(ws, sort="a", sort_direction="asc") + + @tornado.testing.gen_test + async def test_a_session_under_both_thresholds_sends_no_flags(self): + with _guard_limits(sort_disable_rows=100, search_disable_rows=100): + await self._load("sg-none") + _, frame = await self._connect("sg-none") + self.assertFalse({"sort", "search"} & set(frame["df_meta"]), frame["df_meta"]) + + @tornado.testing.gen_test + async def test_a_session_within_the_threshold_is_as_it_was(self): + with _guard_limits(sort_disable_rows=5): + await self._load("sg-within") + ws, frame = await self._connect("sg-within") + self.assertNotIn("sort", frame["df_meta"]) + self.assertTrue(all(_sortable(c) is not False for c in _grid_columns(frame))) + self.assertEqual([_sortable(c) for c in frame["df_display_args"]["main"]["df_viewer_config"]["column_config"]], + [True] * 3, "the klass's own setting stands") + await self._served(ws, sort="a", sort_direction="asc") + + @tornado.testing.gen_test + async def test_the_summary_display_still_sorts(self): + with _guard_limits(sort_disable_rows=3): + await self._load("sg-summary") + _, frame = await self._connect("sg-summary") + self.assertTrue(_grid_columns(frame, "summary")) + self.assertTrue(all(_sortable(c) is not False for c in _grid_columns(frame, "summary"))) + + @tornado.testing.gen_test + async def test_unsorted_and_searched_windows_are_served(self): + with _guard_limits(sort_disable_rows=3): + await self._load("sg-serve") + ws, _ = await self._connect("sg-serve") + self.assertEqual((await self._served(ws))["length"], 5) + await self._state(ws, search_string="a") + self.assertEqual((await self._served(ws))["length"], 2) + + + @tornado.testing.gen_test + async def test_a_sorted_window_is_refused_and_no_query_runs(self): + with _guard_limits(sort_disable_rows=3): + await self._load("sg-refuse") + ws, _ = await self._connect("sg-refuse") + with _count_backend_queries() as queries: + refusal = await self._refusal(ws, sort="a", sort_direction="asc") + self.assertEqual(queries, []) + self.assertEqual(refusal["length"], 0) + self.assertEqual(refusal["key"]["sort"], "a") + self.assertTrue(isinstance(refusal["error_info"], str) and refusal["error_info"]) + again = await self._window(ws) + self.assertEqual((again["type"], again["length"]), ("infinite_resp", 5), again) + self.assertIsInstance(await ws.read_message(), bytes, "no frame followed the refusal") + + @tornado.testing.gen_test + async def test_the_second_window_of_a_request_is_judged_too(self): + with _guard_limits(sort_disable_rows=3): + await self._load("sg-second") + ws, _ = await self._connect("sg-second") + second = {"start": 5, "end": 10, "sourceName": "default", "origEnd": 10, "sort": "a", "sort_direction": "asc"} + first = await self._window(ws, second_request=second) + self.assertEqual((first["type"], first.get("error_info")), ("infinite_resp", None), first) + self.assertIsInstance(await ws.read_message(), bytes) + refusal = await _read_json(ws) + self.assertEqual((refusal["type"], refusal["error_code"], refusal["key"]), ("infinite_resp", "sort_disabled", + second)) + + + @tornado.testing.gen_test + async def test_df_meta_carries_the_sort_flag_for_every_client(self): + with _guard_limits(sort_disable_rows=3): + await self._load("sg-meta") + for caps in (_ONDEMAND, "stats_update", None): + _, frame = await self._connect("sg-meta", caps=caps) + self.assertEqual(frame["df_meta"].get("sort"), "disabled", caps) + self.assertNotIn("search", frame["df_meta"], caps) + + @tornado.testing.gen_test + async def test_df_meta_carries_the_search_flag_when_its_threshold_is_passed_and_search_is_still_served(self): + with _guard_limits(sort_disable_rows=100, search_disable_rows=3): + await self._load("sg-search") + ws, frame = await self._connect("sg-search") + self.assertEqual(frame["df_meta"].get("search"), "disabled") + self.assertNotIn("sort", frame["df_meta"]) + self.assertTrue(all(_sortable(c) is not False for c in _grid_columns(frame))) + await self._state(ws, search_string="a") + self.assertEqual((await self._served(ws))["length"], 2, "the server has no switch for search") + + + @tornado.testing.gen_test + async def test_the_guard_follows_a_state_change(self): + with _guard_limits(sort_disable_rows=3): + await self._load("sg-state") + ws, _ = await self._connect("sg-state") + frame = await self._state(ws, post_processing="first_three") + self.assertEqual(frame["df_meta"].get("sort"), "disabled") + self.assertEqual([_sortable(c) for c in _grid_columns(frame)], [False] * 4) + await self._refused(ws, rows=3, sort="a", sort_direction="desc") + + @tornado.testing.gen_test + async def test_the_guard_holds_once_the_stats_are_complete(self): + with _guard_limits(sort_disable_rows=3): + await self._load("sg-complete") + ws, first = await self._connect("sg-complete") + replies = await self._ask_to_the_end(ws, self._stats(first)["gen"]) + self.assertTrue(replies[-1]["final"]) + session = self._session("sg-complete") + self.assertEqual(session.stats_status, "complete") + config = session.df_display_args["main"]["df_viewer_config"] + self.assertEqual([_sortable(c) for c in config["column_config"] + config["left_col_configs"]], [False] * 4) + if "df_display_args" in replies[-1]: + held = replies[-1]["df_display_args"]["main"]["df_viewer_config"] + self.assertEqual([_sortable(c) for c in held["column_config"] + held["left_col_configs"]], [False] * 4) + await self._refused(ws, sort="b", sort_direction="asc") + + @tornado.testing.gen_test + async def test_a_reload_resolves_the_guard_again(self): + with _guard_limits(sort_disable_rows=100): + await self._load("sg-reload") + ws, first = await self._connect("sg-reload") + self.assertNotIn("sort", first["df_meta"]) + with _guard_limits(sort_disable_rows=3): + await self._reload("sg-reload") + frame = await _read_json(ws) + self.assertEqual(frame["df_meta"].get("sort"), "disabled") + self.assertEqual([_sortable(c) for c in _grid_columns(frame)], [False] * 4) + await self._refused(ws, sort="a", sort_direction="asc") + with _guard_limits(sort_disable_rows=100): + await self._reload("sg-reload") + frame = await _read_json(ws) + self.assertNotIn("sort", frame["df_meta"]) + self.assertTrue(all(_sortable(c) is not False for c in _grid_columns(frame))) + await self._served(ws, sort="a", sort_direction="asc") + + @tornado.testing.gen_test + async def test_a_load_of_a_file_ends_the_guard(self): + sid = "sg-load" + with _guard_limits(sort_disable_rows=3): + await self._load(sid) + ws, _ = await self._connect(sid) + csv_fd, csv_path = tempfile.mkstemp(suffix=".csv") + os.close(csv_fd) + try: + pd.DataFrame({"x": [1, 2, 3], "y": ["p", "q", "r"]}).to_csv(csv_path, index=False) + resp = await _post(self.get_http_port(), "/load", {"session": sid, "path": csv_path, "mode": "buckaroo"}) + self.assertEqual(resp.code, 200, resp.body) + finally: + os.unlink(csv_path) + frame = await _read_json(ws) + self.assertNotIn("sort", frame["df_meta"]) + self.assertIsNone(getattr(self._session(sid), "source_guards", "missing")) + self.assertEqual((await self._served(ws, sort="a", sort_direction="asc"))["length"], 3) + + @tornado.testing.gen_test + async def test_a_comparison_ends_the_guard(self): + sid = "sg-compare" + with _guard_limits(sort_disable_rows=3): + await self._load(sid) + ws, _ = await self._connect(sid) + paths = [] + try: + for frame in (pd.DataFrame({"id": [1, 2], "v": [10, 20]}), pd.DataFrame({"id": [1, 3], "v": [10, 30]})): + fd, path = tempfile.mkstemp(suffix=".csv") + os.close(fd) + frame.to_csv(path, index=False) + paths.append(path) + resp = await _post(self.get_http_port(), "/load_compare", + {"session": sid, "path1": paths[0], "path2": paths[1], "join_columns": ["id"]}) + self.assertEqual(resp.code, 200, resp.body) + finally: + for path in paths: + os.unlink(path) + frame = await _read_json(ws) + self.assertNotIn("sort", frame["df_meta"]) + self.assertIsNone(getattr(self._session(sid), "source_guards", "missing")) + + +class TestSearchedCountMemo: + """A searched xorq window counts the filtered rows, and ``search_expr`` builds + a new expression each time, which ``_expr_count`` caches by object, so every + request counted again. The count is now held for (base expression, term) + (rows-first p37, plan 2 section 4.3).""" + + WINDOW = {"start": 0, "end": 5, "sourceName": "default", "origEnd": 5} + + @staticmethod + def _counts(queries): + return queries.count("CountStar") + + def test_a_repeated_searched_window_issues_one_count(self): + dataflow = _build_dataflow() + with _count_backend_queries() as queries: + first, _ = xorq_loading.handle_infinite_request_xorq(dataflow, self.WINDOW, search_string="a") + second, _ = xorq_loading.handle_infinite_request_xorq(dataflow, self.WINDOW, search_string="a") + assert (first["length"], second["length"]) == (2, 2) + assert self._counts(queries) == 1 + + def test_another_window_of_the_same_search_reuses_the_count(self): + dataflow = _build_dataflow() + with _count_backend_queries() as queries: + for start, end in ((0, 1), (1, 2), (0, 2)): + resp, parquet = xorq_loading.handle_infinite_request_xorq(dataflow, + {**self.WINDOW, "start": start, "end": end}, search_string="a") + assert (resp["length"], pq.read_table(io.BytesIO(parquet)).num_rows) == (2, end - start) + assert self._counts(queries) == 1 + + def test_each_term_is_counted_once(self): + dataflow = _build_dataflow() + with _count_backend_queries() as queries: + lengths = [xorq_loading.handle_infinite_request_xorq(dataflow, self.WINDOW, search_string=term)[0]["length"] + for term in ("a", "b", "a", "b", "c", "a")] + assert lengths == [2, 2, 2, 2, 1, 2] + assert self._counts(queries) == 3 + + def test_a_new_base_expression_is_counted_again(self): + dataflow = _build_dataflow() + before, _ = xorq_loading.handle_infinite_request_xorq(dataflow, self.WINDOW, search_string="a") + dataflow.quick_command_args = {"search": ["b"]} + with _count_backend_queries() as queries: + after, _ = xorq_loading.handle_infinite_request_xorq(dataflow, self.WINDOW, search_string="a") + xorq_loading.handle_infinite_request_xorq(dataflow, self.WINDOW, search_string="a") + assert (before["length"], after["length"]) == (2, 0), "the rows the quick search left hold no a" + assert self._counts(queries) == 1 + + @staticmethod + def _terms_per_base(): + per_base = getattr(xorq_loading, "_SEARCHED_EXPRS_PER_BASE", None) + assert per_base is not None, "xorq_loading._SEARCHED_EXPRS_PER_BASE does not exist" + return per_base + + def test_a_base_keeps_a_bounded_number_of_terms(self): + dataflow = _build_dataflow() + terms = [f"t{i}" for i in range(self._terms_per_base() + 5)] + for term in terms: + xorq_loading.handle_infinite_request_xorq(dataflow, self.WINDOW, search_string=term) + with _count_backend_queries() as queries: + xorq_loading.handle_infinite_request_xorq(dataflow, self.WINDOW, search_string=terms[-1]) + assert self._counts(queries) == 0, "the latest term is held" + xorq_loading.handle_infinite_request_xorq(dataflow, self.WINDOW, search_string=terms[0]) + assert self._counts(queries) == 1, "the oldest was dropped to make room" + + def test_a_term_asked_again_is_the_last_to_go(self): + dataflow = _build_dataflow() + terms = [f"t{i}" for i in range(self._terms_per_base())] + for term in terms: + xorq_loading.handle_infinite_request_xorq(dataflow, self.WINDOW, search_string=term) + # t0 is the oldest of a full base; asking for it again makes it the newest, + # so the next new term takes t1's place. + xorq_loading.handle_infinite_request_xorq(dataflow, self.WINDOW, search_string=terms[0]) + xorq_loading.handle_infinite_request_xorq(dataflow, self.WINDOW, search_string="newcomer") + with _count_backend_queries() as queries: + xorq_loading.handle_infinite_request_xorq(dataflow, self.WINDOW, search_string=terms[0]) + assert self._counts(queries) == 0, "a term that was asked again is still held" + xorq_loading.handle_infinite_request_xorq(dataflow, self.WINDOW, search_string=terms[1]) + assert self._counts(queries) == 1, "the least recently asked term was dropped" + + + def test_a_window_with_no_term_counts_nothing_new(self): + dataflow = _build_dataflow() + with _count_backend_queries() as queries: + for _ in range(3): + resp, _ = xorq_loading.handle_infinite_request_xorq(dataflow, self.WINDOW, search_string="") + assert resp["length"] == 5 + assert self._counts(queries) == 0 + + def test_a_count_that_failed_is_not_remembered(self): + dataflow = _build_dataflow() + failures = [RuntimeError("the backend is down")] + original = Expr.execute + + def execute(expr, *args, **kwargs): + if type(expr.op()).__name__ == "CountStar" and failures: + raise failures.pop() + return original(expr, *args, **kwargs) + + with patch.object(Expr, "execute", execute): + first, _ = xorq_loading.handle_infinite_request_xorq(dataflow, self.WINDOW, search_string="a") + second, _ = xorq_loading.handle_infinite_request_xorq(dataflow, self.WINDOW, search_string="a") + assert (first["length"], second["length"]) == (0, 2), "the next request counts again and gets the real count" + @pytest.mark.parametrize("table", [{"x": [1, 2, 3], "name": ["a", "b", "a"]}, {"x": [1, 2, 3]}], + ids=["with a string column", "no string column"]) + def test_the_memo_does_not_keep_the_base_expression_alive(self, table): + dataflow = _build_dataflow(xo.memtable(table, name="t")) + resp, _ = xorq_loading.handle_infinite_request_xorq(dataflow, self.WINDOW, search_string="a") + assert resp["length"] == (2 if "name" in table else 3) + base = weakref.ref(dataflow.widget_args_tuple[1]) + del dataflow + gc.collect() + assert base() is None + + def test_a_source_with_no_string_column_is_not_searched_and_counts_nothing_new(self): + dataflow = _build_dataflow(xo.memtable({"x": [1, 2, 3]}, name="t")) + with _count_backend_queries() as queries: + lengths = [xorq_loading.handle_infinite_request_xorq(dataflow, self.WINDOW, search_string="a")[0]["length"] + for _ in range(2)] + assert lengths == [3, 3] + assert self._counts(queries) == 0 + + +class TestSearchedWindowCountsWire(_LimitsWire): + """The same through the WebSocket: the term a client types is per client, and + the count it costs is shared.""" + + async def _window(self, ws, **fields): + ws.write_message(json.dumps({"type": "infinite_request", "payload_args": {"start": 0, "end": 5, + "sourceName": "default", "origEnd": 5, **fields}})) + resp = await _read_json(ws) + self.assertIsInstance(await ws.read_message(), bytes) + return resp + + @tornado.testing.gen_test + async def test_a_repeated_searched_window_issues_one_count(self): + await self._load("sc-repeat", stats_delivery="inline") + ws, _ = await self._connect("sc-repeat", caps=None) + await self._state(ws, search_string="a") + with _count_backend_queries() as queries: + lengths = [(await self._window(ws, start=start, end=end))["length"] for start, end in ((0, 5), (0, 5), (1, 2))] + self.assertEqual(lengths, [2, 2, 2]) + self.assertEqual(queries.count("CountStar"), 1) + + @tornado.testing.gen_test + async def test_two_clients_searching_the_same_term_share_the_count(self): + await self._load("sc-share", stats_delivery="inline") + a, _ = await self._connect("sc-share", caps=None) + b, _ = await self._connect("sc-share", caps=None) + for ws in (a, b): + await self._state(ws, search_string="a") + with _count_backend_queries() as queries: + self.assertEqual([(await self._window(ws))["length"] for ws in (a, b, a, b)], [2, 2, 2, 2]) + self.assertEqual(queries.count("CountStar"), 1) + + class TestReloadExpr(tornado.testing.AsyncHTTPTestCase): def get_app(self): return make_app() diff --git a/tests/unit/server/test_server.py b/tests/unit/server/test_server.py index 3026bb36b..0ce356b3f 100644 --- a/tests/unit/server/test_server.py +++ b/tests/unit/server/test_server.py @@ -1,3 +1,4 @@ +import asyncio import json import os import sys @@ -12,6 +13,7 @@ from buckaroo.server import telemetry from buckaroo.server.app import make_app as _make_app +from buckaroo.server.session import build_state_message # Temp file cleanup fails on Windows due to file locking (WinError 32) pytestmark = pytest.mark.skipif(sys.platform == "win32", reason="Temp file locking prevents cleanup on Windows") @@ -28,6 +30,19 @@ def _write_test_csv(path): return df +def _string_highlights(df_display_args): + """The ``highlight_phrase`` of every string column across every display, as + a set of tuples; ``None`` stands for a column with no highlight.""" + phrases = set() + for dva in (df_display_args or {}).values(): + for col in ((dva or {}).get("df_viewer_config") or {}).get("column_config", []): + disp = col.get("displayer_args") or {} + if disp.get("displayer") == "string": + phrase = disp.get("highlight_phrase") + phrases.add(None if phrase is None else tuple(phrase)) + return phrases + + async def _async_fetch(port, path, method="GET", body=None): """Async HTTP fetch for use inside @gen_test methods.""" client = tornado.httpclient.AsyncHTTPClient() @@ -243,6 +258,30 @@ def test_load_buckaroo_without_optional_configs_keeps_defaults(self): finally: os.unlink(f.name) + def test_load_does_not_take_stats_tier_and_resolves_to_full(self): + """/load serves the eager backends, whose stats cost does not grow with + the file, so they resolve to full: the request fields are not read and + a policy left on the session by an earlier /load_expr is dropped + (rows-first p33).""" + with tempfile.NamedTemporaryFile(suffix=".csv", delete=False) as f: + _write_test_csv(f.name) + try: + body = {"session": "stats-tier-1", "path": f.name, "mode": "buckaroo", "stats_tier": "schema", + "stats_delivery": "deferred"} + for _ in range(2): + resp = self.fetch("/load", method="POST", body=json.dumps(body), + headers={"Content-Type": "application/json"}) + self.assertEqual(resp.code, 200) + session = self._app.settings["sessions"].get("stats-tier-1") + self.assertEqual((session.stats_tier, session.stats_delivery), ("full", "inline")) + self.assertIsNone(session.stats_policy) + self.assertEqual(session.stats_status, "complete") + self.assertNotIn("stats", build_state_message(session)["df_meta"]) + # A /load_expr that came first would have left a policy behind. + session.stats_policy = {"tier_target": "schema", "reason": "size"} + finally: + os.unlink(f.name) + class TestSessionPage(tornado.testing.AsyncHTTPTestCase): def get_app(self): @@ -665,6 +704,308 @@ async def test_search_string_per_client_in_dataflow_broadcast(self): finally: os.unlink(f.name) + @tornado.testing.gen_test + async def test_state_change_seq_echoed_to_originating_client_only(self): + """#998: a ``buckaroo_state_change`` that carries ``state_seq`` gets + that value back as ``reply_seq`` on the ``initial_state`` sent to + the client that made the change. The broadcast copy another + client receives carries no ``reply_seq``, so that client applies + it as it does any unsolicited state push.""" + with tempfile.NamedTemporaryFile(suffix=".csv", delete=False) as f: + _write_test_csv(f.name) + try: + sid = "ws-state-seq" + await _async_fetch(self.get_http_port(), "/load", + method="POST", + body=json.dumps({"session": sid, "path": f.name, "mode": "buckaroo"})) + + ws_a = await tornado.websocket.websocket_connect( + f"ws://localhost:{self.get_http_port()}/ws/{sid}") + first = json.loads(await ws_a.read_message()) + self.assertNotIn("reply_seq", first, "a fresh connection's initial_state is unsolicited") + ws_b = await tornado.websocket.websocket_connect( + f"ws://localhost:{self.get_http_port()}/ws/{sid}") + await ws_b.read_message() + + ws_a.write_message(json.dumps({ + "type": "buckaroo_state_change", + "state_seq": 7, + "new_state": { + "post_processing": "", "cleaning_method": "", + "quick_command_args": {"sort": "name"}, + "df_display": "main", + "show_commands": False, "sampled": False, + "search_string": ""}})) + msg_a = json.loads(await ws_a.read_message()) + msg_b = json.loads(await ws_b.read_message()) + self.assertEqual(msg_a["type"], "initial_state") + self.assertEqual(msg_a.get("reply_seq"), 7) + self.assertEqual(msg_b["type"], "initial_state") + self.assertNotIn("reply_seq", msg_b) + ws_a.close() + ws_b.close() + finally: + os.unlink(f.name) + + @tornado.testing.gen_test + async def test_state_change_seq_echoed_on_highlight_overlay(self): + """#998: the per-client highlight overlay answers a state change + too, so it carries ``reply_seq`` like a dataflow broadcast does. + Otherwise an overlay for an earlier term could land after a later + dataflow change and put the client's ``buckaroo_state`` back.""" + with tempfile.NamedTemporaryFile(suffix=".csv", delete=False) as f: + _write_test_csv(f.name) + try: + sid = "ws-state-seq-overlay" + await _async_fetch(self.get_http_port(), "/load", + method="POST", + body=json.dumps({"session": sid, "path": f.name, "mode": "buckaroo"})) + + ws = await tornado.websocket.websocket_connect( + f"ws://localhost:{self.get_http_port()}/ws/{sid}") + await ws.read_message() + + ws.write_message(json.dumps({ + "type": "buckaroo_state_change", + "state_seq": 3, + "new_state": { + "post_processing": "", "cleaning_method": "", + "quick_command_args": {}, "df_display": "main", + "show_commands": False, "sampled": False, + "search_string": "Alice"}})) + overlay = json.loads(await ws.read_message()) + self.assertEqual(overlay["type"], "initial_state") + self.assertEqual(overlay.get("reply_seq"), 3) + ws.close() + finally: + os.unlink(f.name) + + @tornado.testing.gen_test + async def test_state_change_seq_answered_for_no_op_change(self): + """#998: a ``buckaroo_state_change`` with a ``state_seq`` is always + answered, even when it changes no dataflow field. The client drops + a reply older than its latest seq, so when a dataflow change (seq 1) + is followed at once by a ``show_commands`` toggle (seq 2) the seq-1 + reply is dropped as stale. Without a reply for seq 2 the new + df_display_args / df_data_dict never reach the client. The reply + carries the change's own ``buckaroo_state`` so the toggle isn't undone.""" + with tempfile.NamedTemporaryFile(suffix=".csv", delete=False) as f: + _write_test_csv(f.name) + try: + sid = "ws-state-seq-noop" + await _async_fetch(self.get_http_port(), "/load", + method="POST", + body=json.dumps({"session": sid, "path": f.name, "mode": "buckaroo"})) + + ws = await tornado.websocket.websocket_connect( + f"ws://localhost:{self.get_http_port()}/ws/{sid}") + await ws.read_message() + + base = {"post_processing": "", "cleaning_method": "", + "quick_command_args": {"search": ["Al"]}, "df_display": "main", + "show_commands": False, "sampled": False, "search_string": ""} + ws.write_message(json.dumps({"type": "buckaroo_state_change", + "state_seq": 1, "new_state": base})) + ws.write_message(json.dumps({"type": "buckaroo_state_change", + "state_seq": 2, "new_state": {**base, "show_commands": True}})) + + first = json.loads(await ws.read_message()) + second = json.loads(await asyncio.wait_for(ws.read_message(), timeout=5)) + self.assertEqual(first.get("reply_seq"), 1) + self.assertEqual(second["type"], "initial_state") + self.assertEqual(second.get("reply_seq"), 2) + self.assertTrue(second["buckaroo_state"]["show_commands"]) + self.assertEqual(second["df_display_args"], first["df_display_args"]) + ws.close() + finally: + os.unlink(f.name) + + @tornado.testing.gen_test + async def test_state_change_without_seq_stays_silent_on_no_op(self): + """A client that sends no ``state_seq`` gets no reply for a change + that touches no dataflow field, as before #998.""" + with tempfile.NamedTemporaryFile(suffix=".csv", delete=False) as f: + _write_test_csv(f.name) + try: + sid = "ws-state-noseq-noop" + await _async_fetch(self.get_http_port(), "/load", + method="POST", + body=json.dumps({"session": sid, "path": f.name, "mode": "buckaroo"})) + + ws = await tornado.websocket.websocket_connect( + f"ws://localhost:{self.get_http_port()}/ws/{sid}") + await ws.read_message() + + ws.write_message(json.dumps({"type": "buckaroo_state_change", + "new_state": {"post_processing": "", "cleaning_method": "", + "quick_command_args": {}, "df_display": "main", + "show_commands": True, "sampled": False, "search_string": ""}})) + with self.assertRaises(asyncio.TimeoutError): + await asyncio.wait_for(ws.read_message(), timeout=0.5) + ws.close() + finally: + os.unlink(f.name) + + @tornado.testing.gen_test + async def test_state_change_seq_answered_when_change_fails(self): + """#998: a numbered change that fails still gets a numbered reply. + By the time seq 2 fails the client has dropped seq 1's reply as + stale, so an error frame alone leaves it on the data from before + both changes, with a df_meta that doesn't match the rows the server + serves. The reply carries the session's current data (seq 1's) and + the failed change's own ``buckaroo_state``, as the no-op ack does, + so a search box holding the failed term doesn't see it reverted + and send it again.""" + with tempfile.NamedTemporaryFile(suffix=".csv", delete=False) as f: + _write_test_csv(f.name) + try: + sid = "ws-state-seq-error" + await _async_fetch(self.get_http_port(), "/load", + method="POST", + body=json.dumps({"session": sid, "path": f.name, "mode": "buckaroo"})) + + ws = await tornado.websocket.websocket_connect( + f"ws://localhost:{self.get_http_port()}/ws/{sid}") + await ws.read_message() + + base = {"post_processing": "", "cleaning_method": "", + "quick_command_args": {"search": ["Al"]}, "df_display": "main", + "show_commands": False, "sampled": False, "search_string": ""} + ws.write_message(json.dumps({"type": "buckaroo_state_change", + "state_seq": 1, "new_state": base})) + ws.write_message(json.dumps({"type": "buckaroo_state_change", + "state_seq": 2, "new_state": {**base, "post_processing": "nonexistent_pp"}})) + + first = json.loads(await ws.read_message()) + error = json.loads(await asyncio.wait_for(ws.read_message(), timeout=5)) + reply = json.loads(await asyncio.wait_for(ws.read_message(), timeout=2)) + self.assertEqual(first.get("reply_seq"), 1) + self.assertEqual(error["type"], "error") + self.assertEqual(reply["type"], "initial_state") + self.assertEqual(reply.get("reply_seq"), 2) + self.assertEqual(reply["df_meta"], first["df_meta"]) + self.assertEqual(reply["df_data_dict"], first["df_data_dict"]) + self.assertEqual(reply["df_display_args"], first["df_display_args"]) + self.assertEqual(reply["buckaroo_state"]["post_processing"], "nonexistent_pp") + ws.close() + finally: + os.unlink(f.name) + + @tornado.testing.gen_test + async def test_state_change_seq_overlay_keeps_change_ui_fields(self): + """#998: the highlight overlay answers a numbered change, so the + client applies its ``buckaroo_state``. That has to be the change's + own, as the no-op ack's is. Built from the session's, it puts back + ``show_commands`` and ``df_display``, which only a dataflow change + records on the session.""" + with tempfile.NamedTemporaryFile(suffix=".csv", delete=False) as f: + _write_test_csv(f.name) + try: + sid = "ws-state-seq-overlay-ui" + await _async_fetch(self.get_http_port(), "/load", + method="POST", + body=json.dumps({"session": sid, "path": f.name, "mode": "buckaroo"})) + + ws = await tornado.websocket.websocket_connect( + f"ws://localhost:{self.get_http_port()}/ws/{sid}") + await ws.read_message() + + ws.write_message(json.dumps({"type": "buckaroo_state_change", "state_seq": 1, + "new_state": {"post_processing": "", "cleaning_method": "", + "quick_command_args": {}, "df_display": "summary", + "show_commands": True, "sampled": False, "search_string": "Al"}})) + overlay = json.loads(await ws.read_message()) + self.assertEqual(overlay.get("reply_seq"), 1) + self.assertEqual(overlay["buckaroo_state"]["search_string"], "Al") + self.assertTrue(overlay["buckaroo_state"]["show_commands"]) + self.assertEqual(overlay["buckaroo_state"]["df_display"], "summary") + ws.close() + finally: + os.unlink(f.name) + + @tornado.testing.gen_test + async def test_dataflow_broadcast_keeps_each_clients_live_highlight(self): + """A dataflow change rebuilds the display config, and each client's + copy of the broadcast keeps that client's live-search highlight, as + the overlay and the no-op ack do. Without it the highlight comes and + goes with whichever reply path ran last.""" + with tempfile.NamedTemporaryFile(suffix=".csv", delete=False) as f: + _write_test_csv(f.name) + try: + sid = "ws-broadcast-live-highlight" + await _async_fetch(self.get_http_port(), "/load", + method="POST", + body=json.dumps({"session": sid, "path": f.name, "mode": "buckaroo"})) + + ws_a = await tornado.websocket.websocket_connect( + f"ws://localhost:{self.get_http_port()}/ws/{sid}") + await ws_a.read_message() + ws_b = await tornado.websocket.websocket_connect( + f"ws://localhost:{self.get_http_port()}/ws/{sid}") + await ws_b.read_message() + + base = {"post_processing": "", "cleaning_method": "", + "quick_command_args": {}, "df_display": "main", + "show_commands": False, "sampled": False} + ws_a.write_message(json.dumps({"type": "buckaroo_state_change", "state_seq": 1, + "new_state": {**base, "search_string": "Al"}})) + overlay_a = json.loads(await ws_a.read_message()) + self.assertEqual(_string_highlights(overlay_a["df_display_args"]), {("Al",)}) + ws_b.write_message(json.dumps({"type": "buckaroo_state_change", + "new_state": {**base, "search_string": "Bo"}})) + await ws_b.read_message() + + ws_a.write_message(json.dumps({"type": "buckaroo_state_change", "state_seq": 2, + "new_state": {**base, "quick_command_args": {"sort": "name"}, "search_string": "Al"}})) + msg_a = json.loads(await ws_a.read_message()) + msg_b = json.loads(await ws_b.read_message()) + self.assertEqual(msg_a.get("reply_seq"), 2) + self.assertEqual(_string_highlights(msg_a["df_display_args"]), {("Al",)}) + self.assertEqual(_string_highlights(msg_b["df_display_args"]), {("Bo",)}) + ws_a.close() + ws_b.close() + finally: + os.unlink(f.name) + + @tornado.testing.gen_test + async def test_clearing_live_search_keeps_committed_highlight(self): + """Clearing the live search term takes away its own highlight, not + the committed search's: rows are still filtered by + ``quick_command_args.search``, and the broadcast that applied it + highlighted it.""" + with tempfile.NamedTemporaryFile(suffix=".csv", delete=False) as f: + _write_test_csv(f.name) + try: + sid = "ws-clear-live-highlight" + await _async_fetch(self.get_http_port(), "/load", + method="POST", + body=json.dumps({"session": sid, "path": f.name, "mode": "buckaroo"})) + + ws = await tornado.websocket.websocket_connect( + f"ws://localhost:{self.get_http_port()}/ws/{sid}") + await ws.read_message() + + base = {"post_processing": "", "cleaning_method": "", + "quick_command_args": {"search": ["Bo"]}, "df_display": "main", + "show_commands": False, "sampled": False, "search_string": ""} + ws.write_message(json.dumps({"type": "buckaroo_state_change", + "state_seq": 1, "new_state": base})) + committed = json.loads(await ws.read_message()) + self.assertEqual(_string_highlights(committed["df_display_args"]), {("Bo",)}) + ws.write_message(json.dumps({"type": "buckaroo_state_change", + "state_seq": 2, "new_state": {**base, "search_string": "Al"}})) + live = json.loads(await ws.read_message()) + self.assertEqual(_string_highlights(live["df_display_args"]), {("Al",)}) + + ws.write_message(json.dumps({"type": "buckaroo_state_change", + "state_seq": 3, "new_state": base})) + cleared = json.loads(await ws.read_message()) + self.assertEqual(cleared.get("reply_seq"), 3) + self.assertEqual(_string_highlights(cleared["df_display_args"]), {("Bo",)}) + ws.close() + finally: + os.unlink(f.name) + @tornado.testing.gen_test async def test_ws_request_no_data_loaded(self): ws = await tornado.websocket.websocket_connect( diff --git a/tests/unit/server/test_stats_policy.py b/tests/unit/server/test_stats_policy.py new file mode 100644 index 000000000..3a2ede73e --- /dev/null +++ b/tests/unit/server/test_stats_policy.py @@ -0,0 +1,1237 @@ +"""Unit tests for ``buckaroo.server.stats_policy`` (rows-first p31, p31b). + +Everything here is pure Python: the policy is a function of numbers, the +probes read a schema or a parquet footer, and nothing needs a server. The last +sections (rows-first p33) test how the session carries the result and how +``df_meta.stats`` reports it, still without a server; the HTTP and WebSocket +side is in test_load_expr.py. The last section (rows-first p37) tests the +guards for huge sources: the sort and search thresholds, the flags a session +reports, the refusal of a sorted window, and the pass that turns sorting off in a +display config. + +The thresholds are the provisional values proposed by the phase-0 +measurements (p31b). The boundary tables run on the module defaults, so a +change to a ``DEFAULT_*`` constant fails them and the table is updated +together with the constant. ``TestCalibratedDefaults`` pins each default to +its literal value. +""" +import dataclasses +import importlib +import importlib.util +import io +import json +import logging +from types import SimpleNamespace + +import numpy as np +import pandas as pd +import polars as pl +import pyarrow as pa +import pyarrow.parquet as pq +import pytest + +from buckaroo.server import session as session_mod +from buckaroo.server import stats_wire + +ENV_FIELDS = {"BUCKAROO_STATS_FULL_AUTO_ROWS": "full_auto_rows", "BUCKAROO_STATS_FULL_AUTO_CELLS": "full_auto_cells", + "BUCKAROO_STATS_SCALAR_AUTO_CELLS": "scalar_auto_cells", "BUCKAROO_STATS_CEILING_FULL_ROWS": "ceiling_full_rows", + "BUCKAROO_STATS_CEILING_FULL_CELLS": "ceiling_full_cells", + "BUCKAROO_STATS_CEILING_SCALAR_CELLS": "ceiling_scalar_cells", "BUCKAROO_POLARS_ROUTE_ROWS": "polars_route_rows"} + +# The phase-0 values (measurements-phase0.md, "Proposed values"), written out +# so a change to a module default has to change a test too. +CALIBRATED = {"full_auto_rows": 12_000_000, "full_auto_cells": 520_000_000, "scalar_auto_cells": 1_000_000_000, + "ceiling_full_rows": 25_000_000, "ceiling_full_cells": 1_000_000_000, "ceiling_scalar_cells": 4_000_000_000, + "polars_route_rows": 8_000_000} +CALIBRATED_CONSTANTS = {"DEFAULT_FULL_AUTO_ROWS": 12_000_000, "DEFAULT_FULL_AUTO_CELLS": 520_000_000, + "DEFAULT_SCALAR_AUTO_CELLS": 1_000_000_000, "DEFAULT_CEILING_FULL_ROWS": 25_000_000, + "DEFAULT_CEILING_FULL_CELLS": 1_000_000_000, "DEFAULT_CEILING_SCALAR_CELLS": 4_000_000_000, + "DEFAULT_POLARS_ROUTE_ROWS": 8_000_000} + +TIER_RANK = {"schema": 0, "scalar": 1, "full": 2} + + +@pytest.fixture +def sp(): + spec = importlib.util.find_spec("buckaroo.server.stats_policy") + assert spec is not None, "buckaroo.server.stats_policy is not importable" + return importlib.import_module("buckaroo.server.stats_policy") + + +@pytest.fixture(autouse=True) +def clean_env(monkeypatch): + """A developer's own BUCKAROO_STATS_* overrides must not leak in.""" + for name in ENV_FIELDS: + monkeypatch.delenv(name, raising=False) + + +@pytest.fixture +def lim(sp): + """The module defaults. The tables are written against the calibrated + values, so moving a ``DEFAULT_*`` constant fails them.""" + return sp.StatsLimits() + + +def _resolve(sp, lim, backend, source_kind, rows, cols, host_tier=None, bytes=None): + return sp.resolve_stats_policy( + backend, source_kind, rows, cols, bytes=bytes, host_tier=host_tier, limits=lim) + + +# (backend, source_kind, rows, cols, host_tier) -> (tier_target, auto_request, +# requestable, reason). The row counts of 42.3M, 54.1M and 78.0M and the 44 +# columns are the three tallyman entries over 70 s; 10.8M x 43 is parking_2017. +# Each threshold has an equal, a one-below and a one-above row, in rows and, +# where the rule reads cells, in cells (a row count or a column count moves +# the product). The comments give the cell count where it decides the row. +RESOLVE_TABLE = [ + # Eager backends resolve to full and ignore the host tier and the cell bounds. + ("pandas", "memory", 52_814, 27, None, ("full", True, [], None)), + ("pandas", "parquet", 1_000_000, 43, None, ("full", True, [], None)), + ("polars", "parquet", 10_800_000, 43, None, ("full", True, [], None)), + ("polars", "parquet", 10_800_000, 43, "schema", ("full", True, [], None)), + ("pandas", "csv", 52_814, 27, "scalar", ("full", True, [], None)), + ("pandas", "memory", 30_000_000, 100, None, ("full", True, [], None)), + # xorq sized by the policy alone: full up to 12M rows. + ("xorq", "parquet", 0, 0, None, ("full", True, [], None)), + ("xorq", "parquet", 52_814, 27, None, ("full", True, [], None)), + ("xorq", "parquet", 10_800_000, 43, None, ("full", True, [], None)), + ("xorq", "parquet", 11_999_999, 43, None, ("full", True, [], None)), + ("xorq", "parquet", 12_000_000, 43, None, ("full", True, [], None)), + ("xorq", "parquet", 12_000_001, 43, None, ("scalar", True, ["full"], "size")), + ("xorq", "parquet", 12_000_000, 10, None, ("full", True, [], None)), + ("xorq", "parquet", 12_000_001, 10, None, ("scalar", True, ["full"], "size")), + # The 11.7M-row CSV entries: 515.8M cells at 44 columns, 527.5M at 45. + ("xorq", "csv", 11_721_603, 44, None, ("full", True, [], None)), + ("xorq", "csv", 11_721_603, 45, None, ("scalar", True, ["full"], "size")), + # full auto also stops at 520M cells, reached by rows or by columns. + ("xorq", "parquet", 9_999_999, 52, None, ("full", True, [], None)), + ("xorq", "parquet", 10_000_000, 52, None, ("full", True, [], None)), + ("xorq", "parquet", 10_000_001, 52, None, ("scalar", True, ["full"], "size")), + ("xorq", "parquet", 8_000_000, 65, None, ("full", True, [], None)), + ("xorq", "parquet", 8_000_000, 66, None, ("scalar", True, ["full"], "size")), + # scalar auto up to 1.0B cells; full stays requestable up to its own ceiling. + ("xorq", "parquet", 11_000_000, 90, None, ("scalar", True, ["full"], "size")), + ("xorq", "parquet", 11_000_000, 100, None, ("schema", False, ["scalar"], "size")), + ("xorq", "parquet", 19_999_999, 50, None, ("scalar", True, ["full"], "size")), + ("xorq", "parquet", 20_000_000, 50, None, ("scalar", True, ["full"], "size")), + ("xorq", "parquet", 20_000_001, 50, None, ("schema", False, ["scalar"], "size")), + ("xorq", "parquet", 22_727_272, 44, None, ("scalar", True, ["full"], "size")), + ("xorq", "parquet", 22_727_273, 44, None, ("schema", False, ["scalar"], "size")), + ("xorq", "parquet", 12_500_000, 80, None, ("scalar", True, ["full"], "size")), + ("xorq", "parquet", 12_500_000, 81, None, ("schema", False, ["scalar"], "size")), + # The three slow tallyman entries are above 1.8B cells and start at schema. + ("xorq", "csv", 42_300_000, 44, None, ("schema", False, ["scalar"], "size")), + ("xorq", "csv", 50_000_000, 44, None, ("schema", False, ["scalar"], "size")), + ("xorq", "csv", 54_100_000, 44, None, ("schema", False, ["scalar"], "size")), + ("xorq", "csv", 78_000_000, 44, None, ("schema", False, ["scalar"], "size")), + ("xorq", "csv", 78_000_000, 44, "auto", ("schema", False, ["scalar"], "size")), + # Past the scalar ceiling (4.0B cells) nothing is requestable. + ("xorq", "csv", 100_000_000, 40, None, ("schema", False, ["scalar"], "size")), + ("xorq", "csv", 100_000_001, 40, None, ("schema", False, [], "size")), + # The host lowers freely. + ("xorq", "parquet", 1_000_000, 43, "schema", ("schema", False, ["scalar", "full"], "host")), + ("xorq", "parquet", 1_000_000, 43, "scalar", ("scalar", True, ["full"], "host")), + ("xorq", "parquet", 1_000_000, 43, "full", ("full", True, [], None)), + # The host raises only as far as the ceiling on full: 25M rows ... + ("xorq", "parquet", 12_000_000, 43, "full", ("full", True, [], None)), + ("xorq", "parquet", 24_999_999, 10, "full", ("full", True, [], None)), + ("xorq", "parquet", 25_000_000, 10, "full", ("full", True, [], None)), + ("xorq", "parquet", 25_000_001, 10, "full", ("scalar", True, [], "ceiling")), + ("xorq", "parquet", 25_000_000, 10, None, ("scalar", True, ["full"], "size")), + ("xorq", "parquet", 25_000_001, 10, None, ("scalar", True, [], "size")), + # ... or 1.0B cells, reached by rows or by columns. + ("xorq", "parquet", 24_999_999, 40, "full", ("full", True, [], None)), + ("xorq", "parquet", 25_000_000, 40, "full", ("full", True, [], None)), + ("xorq", "parquet", 25_000_000, 41, "full", ("scalar", True, [], "ceiling")), + ("xorq", "parquet", 20_000_000, 50, "full", ("full", True, [], None)), + ("xorq", "parquet", 20_000_001, 50, "full", ("scalar", True, [], "ceiling")), + ("xorq", "parquet", 11_000_000, 90, "full", ("full", True, [], None)), + ("xorq", "parquet", 11_000_000, 91, "full", ("scalar", True, [], "ceiling")), + ("xorq", "csv", 42_300_000, 44, "full", ("scalar", True, [], "ceiling")), + ("xorq", "csv", 50_000_000, 44, "full", ("scalar", True, [], "ceiling")), + ("xorq", "csv", 50_000_001, 44, "full", ("scalar", True, [], "ceiling")), + ("xorq", "csv", 78_000_000, 44, "full", ("scalar", True, [], "ceiling")), + ("xorq", "csv", 78_000_000, 44, "scalar", ("scalar", True, [], "host")), + ("xorq", "csv", 78_000_000, 44, "schema", ("schema", False, ["scalar"], "host")), + # The ceiling on scalar is 4.0B cells, reached by rows or by columns. + ("xorq", "csv", 99_999_999, 40, "scalar", ("scalar", True, [], "host")), + ("xorq", "csv", 100_000_000, 40, "scalar", ("scalar", True, [], "host")), + ("xorq", "csv", 100_000_001, 40, "scalar", ("schema", False, [], "ceiling")), + ("xorq", "csv", 50_000_000, 80, "scalar", ("scalar", True, [], "host")), + ("xorq", "csv", 50_000_000, 81, "scalar", ("schema", False, [], "ceiling")), + ("xorq", "csv", 100_000_000, 40, "full", ("scalar", True, [], "ceiling")), + ("xorq", "csv", 100_000_001, 40, "full", ("schema", False, [], "ceiling")), + ("xorq", "csv", 100_000_001, 40, "schema", ("schema", False, [], "host")), +] + + +def _table_id(case): + backend, _kind, rows, cols, host_tier, _expected = case + return f"{backend}-{rows}x{cols}-{host_tier}" + + +class TestResolveStatsPolicy: + @pytest.mark.parametrize( + "backend, source_kind, rows, cols, host_tier, expected", RESOLVE_TABLE, + ids=[_table_id(c) for c in RESOLVE_TABLE]) + def test_table(self, sp, lim, backend, source_kind, rows, cols, host_tier, expected): + out = _resolve(sp, lim, backend, source_kind, rows, cols, host_tier) + got = (out["tier_target"], out["auto_request"], out["requestable"], out["reason"]) + assert got == expected + + def test_result_has_exactly_the_planned_keys(self, sp, lim): + out = _resolve(sp, lim, "xorq", "parquet", 12_400_000, 43) + assert set(out) == {"tier_target", "auto_request", "requestable", "reason", "estimate"} + + def test_result_is_json_serialisable(self, sp, lim): + """It becomes df_meta.stats, so it has to survive json.dumps.""" + out = _resolve(sp, lim, "xorq", "csv", 78_000_000, 44, "full") + assert json.loads(json.dumps(out)) == out + + def test_estimate_reports_the_inputs(self, sp, lim): + out = _resolve(sp, lim, "xorq", "parquet", 12_400_000, 43) + assert out["estimate"] == {"rows": 12_400_000, "cols": 43} + + def test_estimate_includes_bytes_when_given(self, sp, lim): + out = _resolve(sp, lim, "xorq", "parquet", 12_400_000, 43, bytes=987_654_321) + assert out["estimate"] == {"rows": 12_400_000, "cols": 43, "bytes": 987_654_321} + + def test_numpy_integer_counts_are_accepted(self, sp, lim): + out = _resolve(sp, lim, "xorq", "parquet", np.int64(13_000_000), np.int64(43)) + assert out["tier_target"] == "scalar" + assert out["estimate"] == {"rows": 13_000_000, "cols": 43} + + @pytest.mark.parametrize("rows, cols", [(-1, 3), (3, -1), (None, 3), (3, None), (1.5, 3)]) + def test_bad_counts_are_rejected(self, sp, lim, rows, cols): + with pytest.raises((TypeError, ValueError)): + _resolve(sp, lim, "xorq", "parquet", rows, cols) + + def test_numpy_integer_bytes_are_echoed_as_plain_ints(self, sp, lim): + """df.memory_usage(deep=True).sum() is a np.int64, and the result has + to survive json.dumps.""" + out = _resolve(sp, lim, "xorq", "parquet", 12_400_000, 43, bytes=np.int64(987_654_321)) + assert type(out["estimate"]["bytes"]) is int + assert json.loads(json.dumps(out)) == out + + @pytest.mark.parametrize("bad", [-1, "lots", 1.5, object()]) + def test_bad_bytes_are_rejected(self, sp, lim, bad): + with pytest.raises((TypeError, ValueError)): + _resolve(sp, lim, "xorq", "parquet", 12_400_000, 43, bytes=bad) + + def test_a_host_tier_in_the_bytes_slot_is_rejected(self, sp): + """bytes sits before host_tier, so a positional slip must not be taken + for a byte count.""" + with pytest.raises((TypeError, ValueError)): + sp.resolve_stats_policy("xorq", "csv", 12_000_000, 43, "full") + + @pytest.mark.parametrize("bad", [None, 5, b"csv"]) + def test_bad_source_kind_is_rejected(self, sp, lim, bad): + with pytest.raises(TypeError): + _resolve(sp, lim, "xorq", bad, 1_000, 3) + + def test_unknown_host_tier_is_rejected(self, sp, lim): + with pytest.raises(ValueError): + _resolve(sp, lim, "xorq", "parquet", 1_000, 3, "everything") + + def test_unknown_host_tier_is_rejected_for_eager_backends_too(self, sp, lim): + with pytest.raises(ValueError): + _resolve(sp, lim, "pandas", "memory", 1_000, 3, "everything") + + def test_defaults_refuse_full_for_the_78m_row_entry(self, sp): + """No calibration may let a host request full on the 215 s entry.""" + out = sp.resolve_stats_policy("xorq", "csv", 78_000_000, 44, host_tier="full") + assert TIER_RANK[out["tier_target"]] < TIER_RANK["full"] + assert out["reason"] == "ceiling" + assert "full" not in out["requestable"] + + def test_defaults_let_the_control_recover_parking_2017(self, sp): + """10.8M x 43 columns took 3.7 s of stats and was fine, so a default + that withholds it must offer full on demand.""" + out = sp.resolve_stats_policy("xorq", "parquet", 10_800_000, 43) + assert out["tier_target"] == "full" or "full" in out["requestable"] + + +class TestCalibratedDefaults: + """The module defaults are the phase-0 proposals, one literal per value.""" + + def test_limits_defaults(self, sp): + assert dataclasses.asdict(sp.StatsLimits()) == CALIBRATED + + def test_module_constants(self, sp): + assert {name: getattr(sp, name, None) for name in CALIBRATED_CONSTANTS} == CALIBRATED_CONSTANTS + + def test_from_env_with_a_clean_environment_is_the_calibrated_set(self, sp): + assert dataclasses.asdict(sp.StatsLimits.from_env()) == CALIBRATED + + def test_every_limit_has_an_environment_override(self, sp): + """A new limit without an override could not be tuned on a server.""" + assert set(ENV_FIELDS.values()) == {f.name for f in dataclasses.fields(sp.StatsLimits)} + + +class TestCeiling: + """The ceiling is applied inside the function, so every caller gets it: + load, /reload_expr and a stats_request {force} all end in this call.""" + + HOST_TIERS = [None, "auto", "schema", "scalar", "full"] + + def test_host_full_above_the_ceiling_resolves_lower(self, sp, lim): + out = _resolve(sp, lim, "xorq", "csv", 78_000_000, 44, "full") + assert out["tier_target"] == "scalar" + assert out["reason"] == "ceiling" + + def test_reload_re_resolves_against_the_new_row_count(self, sp, lim): + """/reload_expr keeps the stored host tier and re-resolves. An entry + that fit under the ceiling and then grew past it loses full.""" + before = _resolve(sp, lim, "xorq", "parquet", 20_000_000, 44, "full") + assert before["tier_target"] == "full" + after = _resolve(sp, lim, "xorq", "parquet", 78_000_000, 44, "full") + assert after["tier_target"] == "scalar" + assert after["reason"] == "ceiling" + + def test_force_cannot_exceed_the_ceiling(self, sp, lim): + """A force request passes the requested tier as host_tier.""" + loaded = _resolve(sp, lim, "xorq", "csv", 78_000_000, 44) + assert loaded["tier_target"] == "schema" + assert "full" not in loaded["requestable"] + forced = _resolve(sp, lim, "xorq", "csv", 78_000_000, 44, "full") + assert TIER_RANK[forced["tier_target"]] < TIER_RANK["full"] + assert forced["reason"] == "ceiling" + + def test_ceiling_applies_to_the_automatic_choice_too(self, sp): + """A server whose ceiling sits below its own auto threshold is + clamped, with no host tier involved.""" + lim = sp.StatsLimits(full_auto_rows=10_000_000, ceiling_full_rows=5_000_000) + out = sp.resolve_stats_policy("xorq", "parquet", 7_000_000, 10, limits=lim) + assert out["tier_target"] == "scalar" + assert out["reason"] == "ceiling" + + def test_scalar_ceiling_can_leave_nothing_requestable(self, sp): + lim = sp.StatsLimits(ceiling_full_rows=50_000_000, ceiling_scalar_cells=1_000_000_000) + out = sp.resolve_stats_policy("xorq", "csv", 78_000_000, 44, host_tier="full", limits=lim) + assert out["tier_target"] == "schema" + assert out["reason"] == "ceiling" + assert out["auto_request"] is False + assert out["requestable"] == [] + + def test_scalar_ceiling_does_not_touch_smaller_entries(self, sp): + lim = sp.StatsLimits(ceiling_full_rows=50_000_000, ceiling_scalar_cells=1_000_000_000) + out = sp.resolve_stats_policy("xorq", "parquet", 1_000_000, 44, host_tier="full", limits=lim) + assert out["tier_target"] == "full" + + @pytest.mark.parametrize("rows", + [0, 1_000, 10_000_000, 11_000_000, 50_000_000, 50_000_001, 78_000_000, 500_000_000]) + @pytest.mark.parametrize("cols", [1, 44]) + @pytest.mark.parametrize("host_tier", HOST_TIERS) + def test_full_never_resolves_above_the_ceiling(self, sp, lim, rows, cols, host_tier): + out = _resolve(sp, lim, "xorq", "parquet", rows, cols, host_tier) + if rows > lim.ceiling_full_rows: + assert out["tier_target"] != "full" + assert "full" not in out["requestable"] + + @pytest.mark.parametrize("rows", [1_000, 10_800_000, 42_300_000, 50_000_001, 78_000_000]) + def test_requestable_is_exactly_what_a_force_would_grant(self, sp, lim, rows): + """Every tier the result offers is granted when asked for, and every + tier above the target that it does not offer is clipped.""" + out = _resolve(sp, lim, "xorq", "parquet", rows, 44) + for tier in ("scalar", "full"): + if TIER_RANK[tier] <= TIER_RANK[out["tier_target"]]: + continue + forced = _resolve(sp, lim, "xorq", "parquet", rows, 44, tier) + if tier in out["requestable"]: + assert forced["tier_target"] == tier + else: + assert forced["tier_target"] != tier + assert forced["reason"] == "ceiling" + + + @pytest.mark.parametrize("rows, cols", + [(20_000_001, 50), (25_000_000, 41), (11_000_000, 91), (24_999_999, 41), (42_300_000, 44), (78_000_000, 44)]) + @pytest.mark.parametrize("host_tier", HOST_TIERS) + def test_full_never_resolves_above_the_cell_ceiling(self, sp, lim, rows, cols, host_tier): + """1.0B cells, whatever the row count and the host asks for.""" + assert rows * cols > 1_000_000_000 + out = _resolve(sp, lim, "xorq", "parquet", rows, cols, host_tier) + assert out["tier_target"] != "full" + assert "full" not in out["requestable"] + + @pytest.mark.parametrize("rows, cols", + [(25_000_001, 10), (25_000_001, 1), (30_000_000, 5), (42_300_000, 4), (78_000_000, 1)]) + @pytest.mark.parametrize("host_tier", HOST_TIERS) + def test_full_never_resolves_above_25m_rows_however_few_the_cells(self, sp, lim, rows, cols, host_tier): + """The row ceiling holds on its own: these entries are under 1.0B cells.""" + assert rows * cols < 1_000_000_000 + out = _resolve(sp, lim, "xorq", "parquet", rows, cols, host_tier) + assert out["tier_target"] != "full" + assert "full" not in out["requestable"] + + @pytest.mark.parametrize("rows, cols", + [(100_000_001, 40), (50_000_000, 81), (1_000_000, 4_001), (500_000_000, 10)]) + @pytest.mark.parametrize("host_tier", HOST_TIERS) + def test_nothing_above_the_scalar_ceiling_resolves_past_schema(self, sp, lim, rows, cols, host_tier): + """4.0B cells, the placeholder scalar ceiling.""" + assert rows * cols > 4_000_000_000 + out = _resolve(sp, lim, "xorq", "parquet", rows, cols, host_tier) + assert out["tier_target"] == "schema" + assert out["requestable"] == [] + + def test_a_scalar_ceiling_is_on_by_default(self, sp): + out = sp.resolve_stats_policy("xorq", "csv", 100_000_001, 40, host_tier="scalar", limits=sp.StatsLimits()) + assert out["tier_target"] == "schema" + assert out["reason"] == "ceiling" + + + @pytest.mark.parametrize("rows, cols", + [(12_000_001, 43), (10_000_001, 52), (11_000_000, 91), (20_000_000, 50), (20_000_001, 50), (22_727_273, 44), + (25_000_000, 40), (25_000_000, 41), (25_000_001, 10), (50_000_000, 81), (100_000_000, 40), + (100_000_001, 40)]) + def test_requestable_matches_a_force_at_the_row_and_cell_boundaries(self, sp, lim, rows, cols): + """With the cell bounds, the tiers a result offers are still exactly + the ones a force request is granted.""" + out = _resolve(sp, lim, "xorq", "parquet", rows, cols) + for tier in ("scalar", "full"): + if TIER_RANK[tier] <= TIER_RANK[out["tier_target"]]: + continue + forced = _resolve(sp, lim, "xorq", "parquet", rows, cols, tier) + if tier in out["requestable"]: + assert forced["tier_target"] == tier + else: + assert forced["tier_target"] != tier + assert forced["reason"] == "ceiling" + + +class TestLimitsFromEnv: + def test_defaults_when_the_environment_is_empty(self, sp): + assert sp.StatsLimits.from_env() == sp.StatsLimits() + + @pytest.mark.parametrize("name, field", ENV_FIELDS.items()) + def test_each_variable_overrides_its_field(self, sp, monkeypatch, name, field): + monkeypatch.setenv(name, "12345") + assert getattr(sp.StatsLimits.from_env(), field) == 12345 + + def test_underscored_integers_are_accepted(self, sp, monkeypatch): + monkeypatch.setenv("BUCKAROO_STATS_FULL_AUTO_ROWS", "2_000_000") + assert sp.StatsLimits.from_env().full_auto_rows == 2_000_000 + + @pytest.mark.parametrize("name, field", ENV_FIELDS.items()) + @pytest.mark.parametrize("bad", ["lots", "-5", "1.5"]) + def test_invalid_value_falls_back_to_the_default_and_warns(self, sp, monkeypatch, caplog, name, field, bad): + monkeypatch.setenv(name, bad) + with caplog.at_level(logging.WARNING): + limits = sp.StatsLimits.from_env() + assert getattr(limits, field) == getattr(sp.StatsLimits(), field) + assert name in caplog.text + + def test_empty_value_means_unset(self, sp, monkeypatch): + monkeypatch.setenv("BUCKAROO_STATS_CEILING_SCALAR_CELLS", "") + assert sp.StatsLimits.from_env().ceiling_scalar_cells == sp.StatsLimits().ceiling_scalar_cells + + def test_resolve_reads_the_environment_when_no_limits_are_passed(self, sp, monkeypatch): + monkeypatch.setenv("BUCKAROO_STATS_CEILING_FULL_ROWS", "1000") + out = sp.resolve_stats_policy("xorq", "parquet", 5_000, 10, host_tier="full") + assert out["tier_target"] == "scalar" + assert out["reason"] == "ceiling" + + def test_the_full_auto_cell_bound_comes_from_the_environment(self, sp, monkeypatch): + args = ("xorq", "parquet", 5_000, 10) + assert sp.resolve_stats_policy(*args)["tier_target"] == "full" + monkeypatch.setenv("BUCKAROO_STATS_FULL_AUTO_CELLS", "1000") + out = sp.resolve_stats_policy(*args) + assert out["tier_target"] == "scalar" + assert out["reason"] == "size" + + def test_the_full_ceiling_cell_bound_comes_from_the_environment(self, sp, monkeypatch): + args = ("xorq", "parquet", 5_000, 10) + assert sp.resolve_stats_policy(*args, host_tier="full")["tier_target"] == "full" + monkeypatch.setenv("BUCKAROO_STATS_CEILING_FULL_CELLS", "1000") + out = sp.resolve_stats_policy(*args, host_tier="full") + assert out["tier_target"] == "scalar" + assert out["reason"] == "ceiling" + + def test_the_scalar_ceiling_comes_from_the_environment(self, sp, monkeypatch): + args = ("xorq", "parquet", 5_000, 10) + monkeypatch.setenv("BUCKAROO_STATS_CEILING_SCALAR_CELLS", "1000") + out = sp.resolve_stats_policy(*args, host_tier="scalar") + assert out["tier_target"] == "schema" + assert out["reason"] == "ceiling" + + def test_the_environment_is_read_on_each_call(self, sp, monkeypatch): + args = ("xorq", "parquet", 5_000, 10) + assert sp.resolve_stats_policy(*args, host_tier="full")["tier_target"] == "full" + monkeypatch.setenv("BUCKAROO_STATS_CEILING_FULL_ROWS", "1000") + assert sp.resolve_stats_policy(*args, host_tier="full")["tier_target"] == "scalar" + + +class _CountingFile(io.BytesIO): + """A binary file that counts the bytes read out of it.""" + + def __init__(self, data): + super().__init__(data) + self.bytes_read = 0 + + def read(self, size=-1): + out = super().read(size) + self.bytes_read += len(out) + return out + + def read1(self, size=-1): + out = super().read1(size) + self.bytes_read += len(out) + return out + + def readinto(self, b): + n = super().readinto(b) + self.bytes_read += n + return n + + +def _parquet_bytes(rows, row_group_size): + """Incompressible float columns, so the file is mostly row-group data.""" + rng = np.random.default_rng(0) + table = pa.table({f"c{i}": rng.random(rows) for i in range(4)}) + buf = io.BytesIO() + pq.write_table(table, buf, row_group_size=row_group_size, compression="none") + return buf.getvalue() + + +class TestProbes: + def test_parquet_rows_reads_the_footer_and_no_row_group(self, sp): + data = _parquet_bytes(200_000, 50_000) + source = _CountingFile(data) + assert sp.probe_parquet_rows(source) == 200_000 + assert len(data) > 5_000_000 + assert source.bytes_read * 20 < len(data) + + def test_parquet_rows_from_a_path_decodes_no_data(self, sp, tmp_path, monkeypatch): + path = tmp_path / "t.parquet" + path.write_bytes(_parquet_bytes(120_000, 40_000)) + + def boom(*args, **kwargs): + raise AssertionError("the footer probe decoded row-group data") + + monkeypatch.setattr(pq.ParquetFile, "read", boom) + monkeypatch.setattr(pq.ParquetFile, "read_row_group", boom) + monkeypatch.setattr(pq.ParquetFile, "iter_batches", boom) + monkeypatch.setattr(pq, "read_table", boom) + assert sp.probe_parquet_rows(str(path)) == 120_000 + assert sp.probe_parquet_rows(path) == 120_000 + + def test_parquet_rows_of_an_empty_file(self, sp): + buf = io.BytesIO() + pq.write_table(pa.table({"a": pa.array([], type=pa.int64())}), buf) + assert sp.probe_parquet_rows(io.BytesIO(buf.getvalue())) == 0 + + def test_dtypes_of_a_pandas_frame(self, sp): + df = pd.DataFrame({"a": [1, 2], "b": ["x", "y"]}) + assert sp.probe_dtypes(df) == {"a": "int64", "b": str(df["b"].dtype)} + + def test_dtypes_of_a_polars_frame(self, sp): + df = pl.DataFrame({"a": [1, 2], "b": ["x", "y"]}) + assert sp.probe_dtypes(df) == {"a": "Int64", "b": "String"} + + @pytest.mark.filterwarnings("error") + def test_dtypes_of_a_polars_scan_run_no_collect(self, sp, tmp_path, monkeypatch): + """Reading ``LazyFrame.dtypes`` resolves the plan and warns, so the + error filter also pins that the probe goes through ``collect_schema``.""" + path = tmp_path / "t.parquet" + pl.DataFrame({"a": [1, 2], "b": ["x", "y"]}).write_parquet(path) + lf = pl.scan_parquet(path) + + def boom(*args, **kwargs): + raise AssertionError("the dtype probe collected a LazyFrame") + + monkeypatch.setattr(pl.LazyFrame, "collect", boom) + assert sp.probe_dtypes(lf) == {"a": "Int64", "b": "String"} + + def test_dtypes_keep_column_order(self, sp): + df = pd.DataFrame({"z": [1], "a": [2], "m": [3]}) + assert list(sp.probe_dtypes(df)) == ["z", "a", "m"] + + def test_dtypes_of_something_without_a_schema_raise(self, sp): + with pytest.raises(TypeError): + sp.probe_dtypes(object()) + + def test_xorq_probes_and_a_known_count_execute_nothing(self, sp, monkeypatch): + """The count is an input: the caller passes the one load already took + (``_expr_count``) and the policy never runs one itself.""" + xo = pytest.importorskip("xorq.api") + from xorq.vendor.ibis.expr.types.core import Expr + + expr = xo.memtable({"a": [1, 2, 3], "b": ["x", "y", "z"]}, name="t") + + def boom(*args, **kwargs): + raise AssertionError("a probe or the policy executed a xorq query") + + for name in ("execute", "to_pyarrow", "to_pyarrow_batches", "to_pandas"): + if hasattr(Expr, name): + monkeypatch.setattr(Expr, name, boom) + + dtypes = sp.probe_dtypes(expr) + assert dtypes == {"a": "int64", "b": "string"} + out = sp.resolve_stats_policy("xorq", "expr", 78_000_000, len(dtypes), host_tier="full") + assert out["reason"] == "ceiling" + + +class TestRoutePolarsEntry: + def test_above_r_goes_to_xorq(self, sp, lim): + r = lim.polars_route_rows + assert sp.route_polars_entry(r + 1, 43, limits=lim) == "xorq" + assert sp.route_polars_entry(78_000_000, 44, limits=lim) == "xorq" + + def test_at_and_below_r_stay_eager(self, sp, lim): + r = lim.polars_route_rows + assert sp.route_polars_entry(r, 43, limits=lim) == "eager" + assert sp.route_polars_entry(52_814, 27, limits=lim) == "eager" + assert sp.route_polars_entry(0, 0, limits=lim) == "eager" + + def test_r_comes_from_the_limits_object(self, sp): + lim = sp.StatsLimits(polars_route_rows=1_000) + assert sp.route_polars_entry(1_001, 3, limits=lim) == "xorq" + assert sp.route_polars_entry(1_000, 3, limits=lim) == "eager" + + def test_r_comes_from_the_environment_by_default(self, sp, monkeypatch): + assert sp.route_polars_entry(5_000, 3) == "eager" + monkeypatch.setenv("BUCKAROO_POLARS_ROUTE_ROWS", "1000") + assert sp.route_polars_entry(5_000, 3) == "xorq" + + @pytest.mark.parametrize("rows, expected", + [(0, "eager"), (7_999_999, "eager"), (8_000_000, "eager"), (8_000_001, "xorq"), (10_000_000, "xorq"), + (10_800_000, "xorq")]) + def test_default_r_is_8m_rows(self, sp, rows, expected): + """Equal, one below and one above R, on the module defaults; 10.8M + rows (parking_2017) was eager under the old 10M threshold.""" + assert sp.route_polars_entry(rows, 43) == expected + + def test_defaults_route_the_78m_row_entry_to_xorq(self, sp): + assert sp.route_polars_entry(78_000_000, 44) == "xorq" + + def test_the_result_is_one_of_two_strings(self, sp, lim): + for rows in (0, 10_000_000, 10_000_001): + assert sp.route_polars_entry(rows, 43, limits=lim) in ("xorq", "eager") + + @pytest.mark.parametrize("rows, cols", [(-1, 3), (3, -1), (None, 3), (1.5, 3)]) + def test_bad_counts_are_rejected(self, sp, lim, rows, cols): + with pytest.raises((TypeError, ValueError)): + sp.route_polars_entry(rows, cols, limits=lim) + + +# --------------------------------------------------------------------------- +# The policy on the session and the wire (rows-first p33) +# --------------------------------------------------------------------------- + +# The table the wire tests load: five rows and three columns. +ESTIMATE = {"rows": 5, "cols": 3} +ONDEMAND_CAPS = frozenset({"stats_update", "stats_ondemand"}) +UPDATE_CAPS = frozenset({"stats_update"}) +LEGACY_CAPS = frozenset() + +# Limits that put a five-by-three table in each tier. +SCALAR_BY_SIZE = dict(full_auto_rows=3, scalar_auto_cells=1_000) +SCHEMA_BY_SIZE = dict(full_auto_rows=3, scalar_auto_cells=10) +FULL_OVER_CEILING = dict(ceiling_full_rows=3) +SCALAR_OVER_CEILING = dict(ceiling_scalar_cells=10) + + +def _session(sp, stats_tier="auto", delivery="deferred", limits=None, resolve=True): + """A buckaroo-mode session after a load that resolved its policy and began + its first stats generation, the way ``/load_expr`` leaves one.""" + session = session_mod.SessionState(session_id="s", path="p") + session.mode = "buckaroo" + session.df_meta = {"total_rows": 5} + session.stats_tier, session.stats_delivery = stats_tier, delivery + session.stats_policy = ( + sp.resolve_stats_policy("xorq", "xorq_build", 5, 3, host_tier=stats_tier, + limits=sp.StatsLimits(**(limits or {}))) + if resolve else None) + session_mod.begin_stats_generation(session) + return session + + +def _client(caps): + return SimpleNamespace(caps=caps, search_string="") + + +# (host tier, delivery, limits) -> what an ``stats_ondemand`` client is told. +# A field equal to its documented default is left out, and ``tier_target`` and +# ``estimate`` are always there once a policy is reported. +POLICY_FRAMES = [ + ("auto", "deferred", {}, {"status": "pending", "tier": "schema", "gen": 1, "tier_target": "full", + "requestable": [], "estimate": ESTIMATE}), + ("auto", "deferred", SCALAR_BY_SIZE, {"status": "not_computed", "tier": "schema", "gen": 1, "reason": "size", + "tier_target": "scalar", "estimate": ESTIMATE}), + ("auto", "deferred", SCHEMA_BY_SIZE, {"status": "not_computed", "tier": "schema", "gen": 1, "reason": "size", + "tier_target": "schema", "auto_request": False, "requestable": ["scalar", "full"], "estimate": ESTIMATE}), + ("full", "deferred", FULL_OVER_CEILING, {"status": "not_computed", "tier": "schema", "gen": 1, + "reason": "ceiling", "tier_target": "scalar", "requestable": [], "estimate": ESTIMATE}), + ("scalar", "deferred", {}, {"status": "not_computed", "tier": "schema", "gen": 1, "reason": "host", + "tier_target": "scalar", "estimate": ESTIMATE}), + ("schema", "deferred", {}, {"status": "not_computed", "tier": "schema", "gen": 1, "reason": "host", + "tier_target": "schema", "auto_request": False, "requestable": ["scalar", "full"], "estimate": ESTIMATE}), + ("scalar", "inline", SCALAR_OVER_CEILING, {"status": "not_computed", "tier": "schema", "gen": 1, + "reason": "ceiling", "tier_target": "schema", "auto_request": False, "requestable": [], "estimate": ESTIMATE}), + ("schema", "inline", {}, {"status": "not_computed", "tier": "schema", "gen": 1, "reason": "host", + "tier_target": "schema", "auto_request": False, "requestable": ["scalar", "full"], "estimate": ESTIMATE}), +] + + +class TestStatsFieldDefaults: + """What a client assumes for a ``df_meta.stats`` field the server leaves out.""" + + def test_a_frame_with_no_stats_is_complete(self): + """An old server sends no ``df_meta.stats``, and a new client reads it as complete.""" + assert session_mod.stats_with_defaults({"total_rows": 3}) == {"status": "complete", "tier": "full", + "tier_target": "full", "gen": None, "reason": None, "auto_request": True, "requestable": ["full"], + "estimate": None, "omitted_keys": [], "approx_keys": [], "demand_columns": []} + + def test_a_pending_frame_is_headed_for_full(self): + out = session_mod.stats_with_defaults({"stats": {"status": "pending", "tier": "schema", "gen": 3}}) + assert (out["status"], out["tier"], out["tier_target"], out["gen"]) == ("pending", "schema", "full", 3) + assert (out["auto_request"], out["requestable"]) == (True, ["full"]) + + def test_a_not_computed_frame_without_a_target_targets_the_tier_it_has(self): + out = session_mod.stats_with_defaults({"stats": {"status": "not_computed", "tier": "schema", "gen": 2, + "reason": "host"}}) + assert (out["tier_target"], out["reason"]) == ("schema", "host") + + def test_fields_the_frame_carries_win(self): + stats = {"status": "not_computed", "tier": "schema", "gen": 2, "reason": "size", "tier_target": "scalar", + "auto_request": False, "requestable": ["scalar", "full"], "estimate": ESTIMATE, "omitted_keys": ["a"], + "approx_keys": ["b"], "demand_columns": ["c"]} + out = session_mod.stats_with_defaults({"stats": stats}) + assert {key: out[key] for key in stats} == stats + + def test_the_defaults_are_not_shared_between_calls(self): + first = session_mod.stats_with_defaults({}) + first["requestable"].append("scalar") + first["omitted_keys"].append("x") + second = session_mod.stats_with_defaults({}) + assert (second["requestable"], second["omitted_keys"]) == (["full"], []) + + def test_the_documented_defaults(self): + assert session_mod.STATS_FIELD_DEFAULTS == {"auto_request": True, "requestable": ["full"], "omitted_keys": [], + "approx_keys": [], "demand_columns": []} + assert session_mod.STATS_TIER_REQUESTS == ("auto", "full", "scalar", "schema") + assert session_mod.STATS_REASONS == ("size", "host", "cost", "ceiling") + + +class TestInitialStatsStatus: + @pytest.mark.parametrize("stats_tier, delivery, policy, expected", [ + ("full", "inline", None, ("complete", None)), + ("full", "deferred", None, ("pending", None)), + ("schema", "inline", None, ("not_computed", "host")), + ("scalar", "deferred", None, ("not_computed", "host")), + # auto resolves against a schema-tier dataflow, which an inline session never has. + ("auto", "inline", None, ("complete", None)), + ("auto", "deferred", None, ("pending", None)), + ("auto", "deferred", {"tier_target": "full", "reason": None}, ("pending", None)), + ("auto", "deferred", {"tier_target": "scalar", "reason": "size"}, ("not_computed", "size")), + ("auto", "deferred", {"tier_target": "schema", "reason": "size"}, ("not_computed", "size")), + ("full", "deferred", {"tier_target": "scalar", "reason": "ceiling"}, ("not_computed", "ceiling")), + ("schema", "deferred", {"tier_target": "schema", "reason": "host"}, ("not_computed", "host")), + ("scalar", "inline", {"tier_target": "schema", "reason": "ceiling"}, ("not_computed", "ceiling"))]) + def test_the_status_follows_the_pair_and_the_policy(self, stats_tier, delivery, policy, expected): + assert session_mod.initial_stats_status(stats_tier, delivery, policy) == expected + + @pytest.mark.parametrize("stats_tier, delivery, expected", [ + ("full", "inline", "full"), ("auto", "inline", "full"), ("full", "deferred", "schema"), + ("auto", "deferred", "schema"), ("scalar", "inline", "schema"), ("schema", "inline", "schema"), + ("scalar", "deferred", "schema")]) + def test_the_dataflow_is_built_at_the_schema_tier_unless_stats_run_inline(self, stats_tier, delivery, expected): + """No scalar-tier units exist yet, so a scalar target builds at the schema tier.""" + assert session_mod.dataflow_stats_tier(stats_tier, delivery) == expected + + +class TestStatsMeta: + @pytest.mark.parametrize("stats_tier, delivery, limits, expected", POLICY_FRAMES) + def test_a_stats_ondemand_client_is_told_the_policy(self, sp, stats_tier, delivery, limits, expected): + session = _session(sp, stats_tier, delivery, limits) + assert session_mod.stats_meta(session) == expected + message = session_mod.build_state_message(session) + assert message["df_meta"] == {"total_rows": 5, "stats": expected} + + @pytest.mark.parametrize("stats_tier, limits", [("auto", SCALAR_BY_SIZE), ("auto", SCHEMA_BY_SIZE), + ("full", FULL_OVER_CEILING)]) + def test_a_client_without_stats_ondemand_is_told_the_stats_are_on_their_way(self, sp, stats_tier, limits): + """A target the server chose below full is applied only for a client that + can take it; the others are served as a deferred session headed for full.""" + session = _session(sp, stats_tier, "deferred", limits) + pending = {"status": "pending", "tier": "schema", "gen": 1} + assert session_mod.stats_meta(session, ondemand=False) == pending + assert session_mod.build_state_message(session, ondemand=False)["df_meta"]["stats"] == pending + + @pytest.mark.parametrize("stats_tier, delivery, limits, reason", [ + ("scalar", "deferred", {}, "host"), ("schema", "deferred", {}, "host"), ("schema", "inline", {}, "host"), + ("scalar", "inline", SCALAR_OVER_CEILING, "ceiling")]) + def test_a_tier_the_host_named_reaches_every_client(self, sp, stats_tier, delivery, limits, reason): + """The host chose it, so there is nothing for an older client to complete, + and it gets the status without the policy fields.""" + session = _session(sp, stats_tier, delivery, limits) + assert session_mod.stats_meta(session, ondemand=False) == {"status": "not_computed", "tier": "schema", + "gen": 1, "reason": reason} + + @pytest.mark.parametrize("ondemand", [True, False]) + def test_an_explicit_full_within_the_ceiling_reports_no_policy(self, sp, ondemand): + """The message is the one sent before the policy existed, though the + policy is resolved and recorded.""" + session = _session(sp, "full", "deferred") + assert session.stats_policy["tier_target"] == "full" + assert session_mod.stats_meta(session, ondemand=ondemand) == {"status": "pending", "tier": "schema", "gen": 1} + inline = _session(sp, "full", "inline", resolve=False) + assert session_mod.stats_meta(inline, ondemand=ondemand) is None + assert session_mod.build_state_message(inline, ondemand=ondemand)["df_meta"] == {"total_rows": 5} + + @pytest.mark.parametrize("stats_tier, limits", + [("auto", {}), ("auto", SCALAR_BY_SIZE), ("full", FULL_OVER_CEILING)]) + def test_a_completed_session_reports_no_policy(self, sp, stats_tier, limits): + session = _session(sp, stats_tier, "deferred", limits) + session.stats_status, session.stats_reason = "complete", None + assert session_mod.stats_meta(session) == {"status": "complete", "tier": "full", "gen": 1} + + def test_the_gen_follows_the_generation_and_the_status_restarts_from_the_policy(self, sp): + session = _session(sp, "auto", "deferred", SCALAR_BY_SIZE) + session.stats_status, session.stats_reason = "complete", None + session_mod.begin_stats_generation(session) + assert session_mod.stats_meta(session) == {"status": "not_computed", "tier": "schema", "gen": 2, + "reason": "size", "tier_target": "scalar", "estimate": ESTIMATE} + + def test_omitted_approx_and_demand_keys_ride_along_when_there_are_some(self, sp): + session = _session(sp, "auto", "deferred", SCALAR_BY_SIZE) + assert not {"omitted_keys", "approx_keys", "demand_columns"} & set(session_mod.stats_meta(session)) + session.stats_policy.update(omitted_keys=["value_counts"], approx_keys=["distinct_count"], + demand_columns=["a", "b"]) + stats = session_mod.stats_meta(session) + assert (stats["omitted_keys"], stats["approx_keys"], stats["demand_columns"]) == ( + ["value_counts"], ["distinct_count"], ["a", "b"]) + # A client reads what is left out as empty. + session.stats_policy.update(omitted_keys=[], approx_keys=[], demand_columns=[]) + assert not {"omitted_keys", "approx_keys", "demand_columns"} & set(session_mod.stats_meta(session)) + + def test_every_reason_the_server_sends_is_a_documented_one(self): + for *_, expected in POLICY_FRAMES: + assert expected.get("reason", "size") in session_mod.STATS_REASONS + + +class TestPolicyForTheSession: + """``resolve_session_policy``: what a load handler resolves, and when.""" + + @pytest.mark.parametrize("stats_tier", ["auto", "full", "scalar", "schema"]) + def test_a_schema_dataflow_resolves_against_the_count_it_has(self, stats_tier): + out = stats_wire.resolve_session_policy(stats_tier, "schema", 5, 3) + assert out["estimate"] == ESTIMATE + assert out["tier_target"] == {"auto": "full", "full": "full", "scalar": "scalar", "schema": "schema"}[stats_tier] + + def test_a_dataflow_that_ran_its_stats_has_nothing_to_resolve(self): + assert stats_wire.resolve_session_policy("full", "full", 5, 3) is None + assert stats_wire.resolve_session_policy("auto", "full", 5_000_000_000, 40) is None + + def test_the_ceiling_clips_a_host_that_asks_for_full(self, sp): + out = stats_wire.resolve_session_policy("full", "schema", 78_000_000, 44) + assert (out["tier_target"], out["reason"]) == ("scalar", "ceiling") + + def test_the_limits_are_read_on_each_call(self, monkeypatch): + assert stats_wire.resolve_session_policy("auto", "schema", 5, 3)["tier_target"] == "full" + monkeypatch.setenv("BUCKAROO_STATS_FULL_AUTO_ROWS", "3") + assert stats_wire.resolve_session_policy("auto", "schema", 5, 3)["tier_target"] == "scalar" + + +class TestCapabilities: + def test_the_second_bit_is_named(self): + assert stats_wire.STATS_UPDATE_CAP == "stats_update" + assert stats_wire.STATS_ONDEMAND_CAP == "stats_ondemand" + assert stats_wire.parse_caps("stats_update,stats_ondemand") == ONDEMAND_CAPS + + @pytest.mark.parametrize("caps, expected", [ + (ONDEMAND_CAPS, True), (UPDATE_CAPS, False), (LEGACY_CAPS, False), + # The bit means the client also merges stats_update; alone it is not enough. + (frozenset({"stats_ondemand"}), False)]) + def test_only_stats_update_with_stats_ondemand_counts_as_ondemand(self, caps, expected): + assert stats_wire.client_has_ondemand(_client(caps)) is expected + + def test_a_client_with_no_caps_attribute_is_not_ondemand(self): + assert stats_wire.client_has_ondemand(SimpleNamespace()) is False + + @pytest.mark.parametrize("caps, expected", [(ONDEMAND_CAPS, False), (UPDATE_CAPS, True), (LEGACY_CAPS, True)]) + def test_a_policy_target_below_full_is_missing_stats_to_every_client_but_an_ondemand_one(self, sp, caps, expected): + session = _session(sp, "auto", "deferred", SCHEMA_BY_SIZE) + assert stats_wire.stats_to_pull(session, _client(caps)) is expected + + @pytest.mark.parametrize("caps", [ONDEMAND_CAPS, UPDATE_CAPS, LEGACY_CAPS]) + def test_a_pending_session_has_stats_to_pull_for_everyone(self, sp, caps): + assert stats_wire.stats_to_pull(_session(sp, "auto", "deferred"), _client(caps)) is True + + @pytest.mark.parametrize("caps", [ONDEMAND_CAPS, UPDATE_CAPS, LEGACY_CAPS]) + def test_a_tier_the_host_named_has_nothing_to_pull(self, sp, caps): + assert stats_wire.stats_to_pull(_session(sp, "schema", "deferred"), _client(caps)) is False + assert stats_wire.stats_to_pull(_session(sp, "scalar", "deferred"), _client(caps)) is False + + +# --------------------------------------------------------------------------- +# Guards for huge sources (rows-first p37) +# --------------------------------------------------------------------------- + +GUARD_ENV_FIELDS = {"BUCKAROO_SORT_DISABLE_ROWS": "sort_disable_rows", + "BUCKAROO_SEARCH_DISABLE_ROWS": "search_disable_rows"} +# The provisional defaults, written out so a change to a module constant has to +# change a test too. Search has no default threshold: the server never refuses a +# search (plan 3 question 7), so the flag is only ever turned off by a host. +GUARD_DEFAULTS = {"sort_disable_rows": 25_000_000, "search_disable_rows": None} +ENABLED = {"sort": "enabled", "search": "enabled"} + + +@pytest.fixture(autouse=True) +def clean_guard_env(monkeypatch): + """A developer's own sort and search thresholds must not leak in.""" + for name in GUARD_ENV_FIELDS: + monkeypatch.delenv(name, raising=False) + + +def _guard_api(sp, name): + api = getattr(sp, name, None) + assert api is not None, f"buckaroo.server.stats_policy.{name} does not exist" + return api + + +def _guards(sp, backend, rows, cols=10, source_kind="parquet", **limits): + """``resolve_source_guards`` on the module defaults, or on the limits given.""" + resolve = _guard_api(sp, "resolve_source_guards") + if limits: + return resolve(backend, source_kind, rows, cols, limits=_guard_api(sp, "GuardLimits")(**limits)) + return resolve(backend, source_kind, rows, cols) + + +class TestGuardLimits: + def test_the_defaults_are_the_provisional_values(self, sp): + assert dataclasses.asdict(_guard_api(sp, "GuardLimits")()) == GUARD_DEFAULTS + assert getattr(sp, "DEFAULT_SORT_DISABLE_ROWS", None) == GUARD_DEFAULTS["sort_disable_rows"] + assert getattr(sp, "DEFAULT_SEARCH_DISABLE_ROWS", "missing") is None + + def test_from_env_with_a_clean_environment_is_the_default_set(self, sp): + assert dataclasses.asdict(_guard_api(sp, "GuardLimits").from_env()) == GUARD_DEFAULTS + + def test_the_stats_limits_do_not_carry_the_guard_thresholds(self, sp): + """They are separate numbers (plan 3 section 3.5): a threshold on stats + says nothing about the cost of a sorted window.""" + assert not {"sort_disable_rows", "search_disable_rows"} & {f.name for f in dataclasses.fields(sp.StatsLimits)} + + @pytest.mark.parametrize("name, field", GUARD_ENV_FIELDS.items()) + def test_each_variable_overrides_its_field(self, sp, monkeypatch, name, field): + monkeypatch.setenv(name, "12345") + assert getattr(_guard_api(sp, "GuardLimits").from_env(), field) == 12345 + + def test_underscored_integers_are_accepted(self, sp, monkeypatch): + monkeypatch.setenv("BUCKAROO_SORT_DISABLE_ROWS", "2_000_000") + assert _guard_api(sp, "GuardLimits").from_env().sort_disable_rows == 2_000_000 + + @pytest.mark.parametrize("name, field", GUARD_ENV_FIELDS.items()) + @pytest.mark.parametrize("bad", ["lots", "-5", "1.5"]) + def test_invalid_value_falls_back_to_the_default_and_warns(self, sp, monkeypatch, caplog, name, field, bad): + monkeypatch.setenv(name, bad) + with caplog.at_level(logging.WARNING): + limits = _guard_api(sp, "GuardLimits").from_env() + assert getattr(limits, field) == GUARD_DEFAULTS[field] + assert name in caplog.text + + def test_empty_value_means_unset(self, sp, monkeypatch): + monkeypatch.setenv("BUCKAROO_SORT_DISABLE_ROWS", "") + assert _guard_api(sp, "GuardLimits").from_env().sort_disable_rows == GUARD_DEFAULTS["sort_disable_rows"] + + +# (backend, rows, limits) -> (sort, search). 10.8M x 43 is parking_2017; the +# 42.3M, 54.1M and 78.0M entries are the three tallyman loads over 70 s. +SOURCE_GUARD_TABLE = [ + ("xorq", 0, {}, ("enabled", "enabled")), + ("xorq", 52_814, {}, ("enabled", "enabled")), + ("xorq", 10_800_000, {}, ("enabled", "enabled")), + ("xorq", 24_999_999, {}, ("enabled", "enabled")), + ("xorq", 25_000_000, {}, ("enabled", "enabled")), + ("xorq", 25_000_001, {}, ("disabled", "enabled")), + ("xorq", 42_300_000, {}, ("disabled", "enabled")), + ("xorq", 54_100_000, {}, ("disabled", "enabled")), + ("xorq", 78_000_000, {}, ("disabled", "enabled")), + ("xorq", 10**12, {}, ("disabled", "enabled")), + # Eager backends keep both: their windows come from a bounded frame. + ("pandas", 78_000_000, {}, ("enabled", "enabled")), + ("polars", 78_000_000, {}, ("enabled", "enabled")), + ("polars", 78_000_000, {"sort_disable_rows": 1, "search_disable_rows": 1}, ("enabled", "enabled")), + # The sort threshold: equal, one above, and none at all. + ("xorq", 1000, {"sort_disable_rows": 1000}, ("enabled", "enabled")), + ("xorq", 1001, {"sort_disable_rows": 1000}, ("disabled", "enabled")), + ("xorq", 10**12, {"sort_disable_rows": None}, ("enabled", "enabled")), + # The search threshold is its own number. + ("xorq", 1000, {"search_disable_rows": 1000}, ("enabled", "enabled")), + ("xorq", 1001, {"search_disable_rows": 1000}, ("enabled", "disabled")), + ("xorq", 1001, {"sort_disable_rows": 5000, "search_disable_rows": 1000}, ("enabled", "disabled")), + ("xorq", 5001, {"sort_disable_rows": 5000, "search_disable_rows": 1000}, ("disabled", "disabled")), +] + + +def _source_guard_id(case): + backend, rows, limits, _expected = case + return f"{backend}-{rows}-{'-'.join(f'{k}={v}' for k, v in limits.items()) or 'defaults'}" + + +class TestResolveSourceGuards: + @pytest.mark.parametrize("backend, rows, limits, expected", SOURCE_GUARD_TABLE, + ids=[_source_guard_id(case) for case in SOURCE_GUARD_TABLE]) + def test_the_table(self, sp, backend, rows, limits, expected): + assert _guards(sp, backend, rows, **limits) == dict(zip(("sort", "search"), expected)) + + def test_the_result_has_exactly_the_two_flags(self, sp): + assert set(_guards(sp, "xorq", 5)) == {"sort", "search"} + + def test_the_column_count_does_not_move_a_row_threshold(self, sp): + assert _guards(sp, "xorq", 25_000_000, cols=1) == _guards(sp, "xorq", 25_000_000, cols=500) == ENABLED + + def test_the_guards_ignore_the_stats_thresholds(self, sp, monkeypatch): + """A ceiling that refuses full stats for the entry leaves its sort alone.""" + monkeypatch.setenv("BUCKAROO_STATS_FULL_AUTO_ROWS", "1") + monkeypatch.setenv("BUCKAROO_STATS_CEILING_FULL_ROWS", "1") + monkeypatch.setenv("BUCKAROO_STATS_CEILING_SCALAR_CELLS", "1") + assert _guards(sp, "xorq", 5_000) == ENABLED + + def test_the_stats_policy_ignores_the_guard_thresholds(self, sp, monkeypatch): + monkeypatch.setenv("BUCKAROO_SORT_DISABLE_ROWS", "1") + monkeypatch.setenv("BUCKAROO_SEARCH_DISABLE_ROWS", "1") + assert sp.resolve_stats_policy("xorq", "parquet", 5_000, 10)["tier_target"] == "full" + + def test_the_environment_is_read_on_each_call(self, sp): + resolve = _guard_api(sp, "resolve_source_guards") + assert resolve("xorq", "parquet", 5_000, 10) == ENABLED + with pytest.MonkeyPatch.context() as env: + env.setenv("BUCKAROO_SORT_DISABLE_ROWS", "1000") + assert resolve("xorq", "parquet", 5_000, 10)["sort"] == "disabled" + assert resolve("xorq", "parquet", 5_000, 10) == ENABLED + + def test_the_environment_sets_the_search_flag(self, sp, monkeypatch): + monkeypatch.setenv("BUCKAROO_SEARCH_DISABLE_ROWS", "1000") + assert _guard_api(sp, "resolve_source_guards")("xorq", "parquet", 5_000, 10) == {"sort": "enabled", + "search": "disabled"} + + @pytest.mark.parametrize("rows, cols, error", [(-1, 3, ValueError), (5, -1, ValueError), (1.5, 3, TypeError), + ("5", 3, TypeError), (5, None, TypeError)]) + def test_a_bad_count_is_refused_as_the_stats_policy_refuses_it(self, sp, rows, cols, error): + with pytest.raises(error): + _guard_api(sp, "resolve_source_guards")("xorq", "parquet", rows, cols) + + def test_a_source_kind_that_is_not_a_string_is_refused(self, sp): + with pytest.raises(TypeError): + _guard_api(sp, "resolve_source_guards")("xorq", None, 5, 3) + + +def _guarded(sort="enabled", search="enabled"): + """A buckaroo-mode session whose load resolved these guard flags.""" + session = session_mod.SessionState(session_id="s", path="p") + session.mode = "buckaroo" + session.df_meta = {"total_rows": 5} + session.source_guards = {"sort": sort, "search": search} + return session + + +class TestGuardFlagsInDfMeta: + """``df_meta.sort`` and ``df_meta.search``: sent when a session has turned + one off, left out when it is on, which is what a client assumes.""" + + def test_a_session_with_no_guards_by_default(self): + assert getattr(session_mod.SessionState(session_id="s", path="p"), "source_guards", "missing") is None + + def test_the_documented_defaults(self): + assert getattr(session_mod, "SOURCE_GUARD_DEFAULTS", None) == ENABLED + + def test_a_disabled_sort_is_reported(self): + assert session_mod.build_state_message(_guarded(sort="disabled"))["df_meta"] == {"total_rows": 5, + "sort": "disabled"} + + def test_a_disabled_search_is_reported(self): + assert session_mod.build_state_message(_guarded(search="disabled"))["df_meta"] == {"total_rows": 5, + "search": "disabled"} + + def test_both_flags_are_reported_when_both_are_off(self): + assert session_mod.build_state_message(_guarded("disabled", "disabled"))["df_meta"] == {"total_rows": 5, + "sort": "disabled", "search": "disabled"} + + @pytest.mark.parametrize("session", [_guarded(), session_mod.SessionState(session_id="s", path="p")], + ids=["enabled", "unresolved"]) + def test_enabled_flags_leave_the_message_as_it_was(self, session): + session.df_meta = {"total_rows": 5} + assert session_mod.build_state_message(session)["df_meta"] == {"total_rows": 5} + + @pytest.mark.parametrize("ondemand", [True, False]) + def test_every_client_is_told(self, ondemand): + """The server refuses the sort whoever asks, so no capability gates the flag.""" + message = session_mod.build_state_message(_guarded(sort="disabled"), ondemand=ondemand) + assert message["df_meta"]["sort"] == "disabled" + + def test_the_flags_ride_with_the_stats_status(self, sp): + session = _session(sp, "auto", "deferred", SCALAR_BY_SIZE) + session.source_guards = {"sort": "disabled", "search": "enabled"} + df_meta = session_mod.build_state_message(session)["df_meta"] + assert df_meta == {"total_rows": 5, "sort": "disabled", "stats": session_mod.stats_meta(session)} + + def test_the_session_s_df_meta_is_not_changed(self): + session = _guarded("disabled", "disabled") + session_mod.build_state_message(session) + assert session.df_meta == {"total_rows": 5} + + def test_a_client_reads_what_is_left_out_as_enabled(self): + read = getattr(session_mod, "guards_with_defaults", None) + assert read is not None, "session.guards_with_defaults does not exist" + assert read({"total_rows": 5}) == ENABLED + assert read({}) == ENABLED + assert read(None) == ENABLED + assert read({"sort": "disabled"}) == {"sort": "disabled", "search": "enabled"} + assert read({"search": "disabled", "sort": "enabled"}) == {"sort": "enabled", "search": "disabled"} + + +class TestSortRefusal: + """``sort_refusal``: the ``infinite_resp`` a sorted window gets when the + session's sort is off, and nothing for any other request.""" + + WINDOW = {"start": 0, "end": 5, "sourceName": "default", "origEnd": 5} + + @staticmethod + def _refusal(session, payload): + refuse = getattr(session_mod, "sort_refusal", None) + assert refuse is not None, "session.sort_refusal does not exist" + return refuse(session, payload) + + def test_a_sorted_window_is_refused_with_a_code(self): + payload = {**self.WINDOW, "sort": "a", "sort_direction": "asc"} + refusal = self._refusal(_guarded(sort="disabled"), payload) + assert refusal is not None + assert (refusal["type"], refusal["key"], refusal["length"]) == ("infinite_resp", payload, 0) + assert refusal["error_code"] == "sort_disabled" + assert isinstance(refusal["error_info"], str) and refusal["error_info"] + assert "payload" not in refusal, "no rows travel with a refusal" + + @pytest.mark.parametrize("payload", [{}, {"sort": None}, {"sort": ""}, {"sort_direction": "asc"}]) + def test_a_window_with_no_sort_is_served(self, payload): + assert self._refusal(_guarded(sort="disabled"), {**self.WINDOW, **payload}) is None + + def test_a_session_with_the_sort_on_serves_a_sorted_window(self): + assert self._refusal(_guarded(), {**self.WINDOW, "sort": "a"}) is None + assert self._refusal(_guarded(search="disabled"), {**self.WINDOW, "sort": "a"}) is None + + def test_a_session_that_resolved_no_guards_serves_a_sorted_window(self): + session = session_mod.SessionState(session_id="s", path="p") + assert self._refusal(session, {**self.WINDOW, "sort": "a"}) is None + + @pytest.mark.parametrize("payload", [None, "sort", ["sort"], 5]) + def test_a_payload_that_is_not_a_dict_is_not_refused(self, payload): + assert self._refusal(_guarded(sort="disabled"), payload) is None + + +class TestResolveSessionGuards: + """``stats_wire.resolve_session_guards``: what a load handler resolves, and + for which sessions.""" + + @staticmethod + def _resolve(*args): + resolve = getattr(stats_wire, "resolve_session_guards", None) + assert resolve is not None, "stats_wire.resolve_session_guards does not exist" + return resolve(*args) + + def test_a_schema_dataflow_resolves_against_the_count_it_has(self, monkeypatch): + monkeypatch.setenv("BUCKAROO_SORT_DISABLE_ROWS", "3") + assert self._resolve("schema", 5, 3) == {"sort": "disabled", "search": "enabled"} + assert self._resolve("schema", 3, 3) == ENABLED + + def test_a_dataflow_that_ran_its_stats_has_no_guards(self, monkeypatch): + """Only a session a host opened with a stats policy is guarded. One built + the way it always was, with its stats inline, stays as it was.""" + monkeypatch.setenv("BUCKAROO_SORT_DISABLE_ROWS", "3") + assert self._resolve("full", 5, 3) is None + assert self._resolve("full", 5_000_000_000, 40) is None + + def test_the_defaults_guard_a_session_of_the_size_the_stats_ceiling_refuses(self): + assert self._resolve("schema", 78_000_000, 44)["sort"] == "disabled" + assert self._resolve("schema", 10_800_000, 43) == ENABLED + + def test_the_thresholds_are_read_on_each_call(self, monkeypatch): + assert self._resolve("schema", 5, 3) == ENABLED + monkeypatch.setenv("BUCKAROO_SORT_DISABLE_ROWS", "3") + assert self._resolve("schema", 5, 3)["sort"] == "disabled" + + +def _grid_display(main_columns, left=(), summary_columns=()): + """A ``df_display_args`` as the dataflow builds it: ``main`` is the display + the infinite grid serves from the server, ``summary`` is a client-side grid.""" + def display(data_key, columns, left_columns): + return {"data_key": data_key, "summary_stats_key": "all_stats", + "df_viewer_config": {"pinned_rows": [], "column_config": list(columns), + "left_col_configs": list(left_columns), "extra_grid_config": {}, "component_config": {}}} + return {"main": display("main", main_columns, left), "summary": display("empty", summary_columns, ())} + + +class TestDisableSorting: + """``disable_sorting``: the final pass that turns sorting off for every column + of a display served by ``infinite_request``, after the klasses and the + overrides have had their say.""" + + @staticmethod + def _disable(sp, display): + return _guard_api(sp, "disable_sorting")(display) + + @staticmethod + def _columns(display, key="main"): + config = display[key]["df_viewer_config"] + return config["column_config"] + config["left_col_configs"] + + def test_every_column_of_the_infinite_display_stops_sorting(self, sp): + display = _grid_display( + [{"col_name": "a", "header_name": "price"}, + {"col_name": "b", "header_name": "qty", "ag_grid_specs": {"minWidth": 90}}, + {"col_name": "c", "header_name": "category", "ag_grid_specs": {"sortable": True}}, + {"col_path": ("x", "y"), "field": "d", "ag_grid_specs": {"sortable": True, "minWidth": 70}}], + left=[{"col_name": "index", "header_name": "index"}]) + out = self._disable(sp, display) + assert [c["ag_grid_specs"]["sortable"] for c in self._columns(out)] == [False] * 5 + + def test_the_other_column_specs_are_kept(self, sp): + out = self._disable(sp, _grid_display([{"col_name": "b", "ag_grid_specs": {"minWidth": 90, "sortable": True}}])) + assert out["main"]["df_viewer_config"]["column_config"][0]["ag_grid_specs"] == {"minWidth": 90, + "sortable": False} + + def test_specs_that_are_not_a_dict_are_replaced(self, sp): + out = self._disable(sp, _grid_display([{"col_name": "a", "ag_grid_specs": None}])) + assert out["main"]["df_viewer_config"]["column_config"] == [{"col_name": "a", + "ag_grid_specs": {"sortable": False}}] + + def test_a_column_entry_that_is_not_a_dict_is_left_as_it_is(self, sp): + out = self._disable(sp, _grid_display([None, "a", {"col_name": "b"}])) + assert out["main"]["df_viewer_config"]["column_config"] == [None, "a", {"col_name": "b", + "ag_grid_specs": {"sortable": False}}] + + def test_a_display_the_client_sorts_itself_is_left_alone(self, sp): + summary = [{"col_name": "a", "header_name": "price", "ag_grid_specs": {"sortable": True}}, {"col_name": "b"}] + display = _grid_display([{"col_name": "a"}], summary_columns=summary) + out = self._disable(sp, display) + assert out["summary"] == display["summary"] + + def test_the_rest_of_the_display_is_unchanged(self, sp): + display = _grid_display([{"col_name": "a", "displayer_args": {"displayer": "obj"}}]) + out = self._disable(sp, display) + main, held = out["main"], display["main"] + assert (main["data_key"], main["summary_stats_key"]) == (held["data_key"], held["summary_stats_key"]) + for key in ("pinned_rows", "extra_grid_config", "component_config"): + assert main["df_viewer_config"][key] == held["df_viewer_config"][key] + assert main["df_viewer_config"]["column_config"][0]["displayer_args"] == {"displayer": "obj"} + + def test_the_display_given_is_not_changed(self, sp): + display = _grid_display([{"col_name": "a", "ag_grid_specs": {"sortable": True}}, {"col_name": "b"}], + left=[{"col_name": "index"}]) + before = json.dumps(display, sort_keys=True) + out = self._disable(sp, display) + assert json.dumps(display, sort_keys=True) == before + assert out is not display + + @pytest.mark.parametrize("display", [{}, {"main": None}, {"main": {}}, {"main": {"data_key": "main"}}, + {"main": {"data_key": "main", "df_viewer_config": None}}, + {"main": {"data_key": "main", "df_viewer_config": {"column_config": None}}}, + {"main": {"data_key": "main", "df_viewer_config": {"column_config": []}}}]) + def test_a_display_that_is_not_the_usual_shape_does_not_raise(self, sp, display): + assert isinstance(self._disable(sp, display), dict) + + def test_a_display_with_no_left_columns_is_handled(self, sp): + display = _grid_display([{"col_name": "a"}]) + del display["main"]["df_viewer_config"]["left_col_configs"] + out = self._disable(sp, display) + assert out["main"]["df_viewer_config"]["column_config"][0]["ag_grid_specs"]["sortable"] is False diff --git a/tests/unit/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