diff --git a/buckaroo/buckaroo_widget.py b/buckaroo/buckaroo_widget.py index 758b1b0f2..28c0225d9 100644 --- a/buckaroo/buckaroo_widget.py +++ b/buckaroo/buckaroo_widget.py @@ -31,7 +31,7 @@ from .serialization_utils import EMPTY_DF_WHOLE, check_and_fix_df, pd_to_obj, to_parquet, send_infinite_resp from .dataflow.dataflow import CustomizableDataflow from .dataflow.dataflow_extras import (Sampling, exception_protect) -from .dataflow.styling_core import (ComponentConfig, InitSD, OverrideColumnConfig, PinnedRowConfig, StylingAnalysis, merge_column_config, EMPTY_DFVIEWER_CONFIG) +from .dataflow.styling_core import (ComponentConfig, InitSD, OverrideColumnConfig, PinnedRowConfig, StylingAnalysis, EMPTY_DFVIEWER_CONFIG) from .dataflow.autocleaning import PandasAutocleaning from pathlib import Path @@ -359,39 +359,10 @@ def _handle_widget_change(self, change): _bk_flash("_handle_widget_change EXIT — processed_df is None") return - # 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 - - # to expedite processing maybe future provided dfs from - # postprcoessing could default to empty until that is - # selected, optionally - - #note this needs to be empty so that we can do the infinite stuff - self.df_data_dict = {'main': [], - 'all_stats': self._sd_to_jsondf(merged_sd), - 'empty': []} - - temp_display_args = {} - for display_name, A_Klass in self.dataflow.df_display_klasses.items(): - df_viewer_config = A_Klass.get_dfviewer_config(merged_sd, processed_df) - base_column_config = df_viewer_config['column_config'] - df_viewer_config['column_config'] = merge_column_config( - base_column_config, self.dataflow.processed_df, self.dataflow.column_config_overrides) - disp_arg = {'data_key': A_Klass.data_key, - #'df_viewer_config': json.loads(json.dumps(df_viewer_config)), - 'df_viewer_config': df_viewer_config, - 'summary_stats_key': A_Klass.summary_stats_key} - temp_display_args[display_name] = disp_arg - - if self.dataflow.pinned_rows is not None: - temp_display_args['main']['df_viewer_config']['pinned_rows'] = self.dataflow.pinned_rows - if self.dataflow.extra_grid_config: - temp_display_args['main']['df_viewer_config']['extra_grid_config'] = self.dataflow.extra_grid_config - if self.dataflow.component_config: - temp_display_args['main']['df_viewer_config']['component_config'] = self.dataflow.component_config - - self.df_display_args = temp_display_args + # The dataflow is built with skip_main_serial, so 'main' stays empty for + # the infinite rows. + self.df_data_dict = self.dataflow._build_df_data_dict(processed_df, merged_sd) + self.df_display_args = self.dataflow._build_df_display_args(processed_df, merged_sd) _bk_flash("_handle_widget_change EXIT (df_display_args → JS)") diff --git a/buckaroo/customizations/xorq_stats_v2.py b/buckaroo/customizations/xorq_stats_v2.py index 239a261f4..74a2a2b19 100644 --- a/buckaroo/customizations/xorq_stats_v2.py +++ b/buckaroo/customizations/xorq_stats_v2.py @@ -139,6 +139,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 8b6553621..61973f860 100644 --- a/buckaroo/dataflow/dataflow.py +++ b/buckaroo/dataflow/dataflow.py @@ -5,12 +5,13 @@ import six import warnings import pandas as pd -from traitlets import Unicode, Any, observe, Dict +from traitlets import Unicode, Any, observe, Dict, Enum from buckaroo.pluggable_analysis_framework.col_analysis import ColAnalysis, ErrDict, SDType from ..serialization_utils import pd_to_obj, sd_to_parquet_b64, project_sd from buckaroo.pluggable_analysis_framework.utils import (filter_analysis) from buckaroo.pluggable_analysis_framework.df_stats_v2 import DfStatsV2 +from buckaroo.pluggable_analysis_framework.stat_pipeline import with_stat from .autocleaning import SentinelAutocleaning from .dataflow_extras import (exception_protect, Sampling) from .styling_core import ( @@ -27,6 +28,63 @@ from .df_types import DataFrameT from .sd_cache import hash_chain, split_chain_by_scope +# Content levels a dataflow can be built at. ``full`` runs the summary-stats +# pipeline. ``schema`` publishes identity and typing keys derived from the +# schema alone and runs no data query (backends opt in, see ``_get_schema_sd``). +STATS_TIERS = ("full", "schema") + + +def assemble_merged_sd(init_sd: InitSD, cleaned_sd: SDType, raw_sd: SDType, processed_sd: SDType, + processed_df: TAny, chains: TDict[str, List[TAny]], + clean_sd: Optional[SDType] = None, filt_sd: Optional[SDType] = None) -> SDType: + """Build the wire ``merged_sd`` from the pieces ``_merged_sd`` reads. + + Bare keys come from the raw scope's sd (``raw_sd``, computed on + ``sampled_df``) merged under ``init_sd`` (keyed by original column name, + so rewritten here), ``cleaned_sd`` and ``processed_sd``. ``cleaned_*`` keys + are layered on top from ``clean_sd`` when cleaning is active, and + ``filtered_*`` keys from ``filt_sd`` when a search filter is active. + + Pure: the inputs are not mutated, so a caller holding an accumulating + ``raw_sd`` can assemble it repeatedly without the traitlets cascade. + ``chains`` is ``split_chain_by_scope(operations)``. + """ + if processed_df is None: + # on initial startup + return merge_sds(init_sd, cleaned_sd, raw_sd, processed_sd) + + # Gate each prefixed layer on a chain-shape diff between scopes. + # ``cleaned_*`` fires when the clean chain has ops the raw chain + # doesn't (cleaning is on); ``filtered_*`` fires when the filt + # chain has ops the clean chain doesn't (at least one + # quick-command op is present, i.e. a search filter is on). + # Keying off ``filt_sd_key != raw_sd_key`` would also fire + # ``filtered_*`` for cleaning-only states, mislabelling cleaned + # stats as filtered. + cleaning_active = chains['clean'] != chains['raw'] + filter_active = chains['filt'] != chains['clean'] + + #we do this to get rewrtten keys for init_sd + rewritten_init_sd = merge_sd_overrides({}, processed_df, init_sd) + intermediate_sd = merge_sds(rewritten_init_sd, cleaned_sd, raw_sd) + base = merge_sd_overrides(intermediate_sd, processed_df, processed_sd) + + # Layer ``cleaned_*`` keys on top when cleaning is active. + if cleaning_active and clean_sd: + for col, stats in clean_sd.items(): + col_dict = base.setdefault(col, {}) + for stat_name, val in stats.items(): + col_dict[f'cleaned_{stat_name}'] = val + + # Layer ``filtered_*`` keys on top when a filter is active. + if filter_active and filt_sd: + for col, stats in filt_sd.items(): + col_dict = base.setdefault(col, {}) + for stat_name, val in stats.items(): + col_dict[f'filtered_{stat_name}'] = val + + return base + class DfTrait(Any): """Any-trait for values that may contain a pandas/polars DataFrame. @@ -161,6 +219,11 @@ 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). A + # trait, so assigning it reruns _summary_sd at the new tier and an unknown + # value is refused. ``set_stats_tier`` switches without that compute when + # the new tier's stats are already in hand. + stats_tier = Enum(STATS_TIERS, default_value="full") summary_sd = Any() df_meta = Any() @@ -316,9 +379,9 @@ def _get_summary_sd(self, processed_df: DataFrameT, scope: str = 'filt') -> Tupl ret_summary[col] = {} return ret_summary, {} - _summary_sd_cache_key = (None, None) + _summary_sd_cache_key = (None, None, None) - @observe('processed_result', 'analysis_klasses') + @observe('processed_result', 'analysis_klasses', 'stats_tier') @exception_protect('summary_sd-protector') def _summary_sd(self, change): # Dedupe: the autocleaning operations cascade re-fires @@ -326,20 +389,25 @@ 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. # Matches the guards in _merged_sd and _populate_sd_cache, and # narrows Optional[DataFrameT] -> DataFrameT for _get_summary_sd. return - klasses = self.analysis_klasses - if (id(df), id(klasses)) == self._summary_sd_cache_key: + key = (id(df), id(self.analysis_klasses), self.stats_tier) + if key == self._summary_sd_cache_key: return - self._summary_sd_cache_key = (id(df), id(klasses)) result_summary_sd, errs = self._get_summary_sd(df, scope='filt') - self.summary_sd = result_summary_sd + # Recorded once the stats exist: after a run that raises, the key still + # names the frame summary_sd belongs to, which _populate_sd_cache checks. + self._summary_sd_cache_key = key + # errs first: assigning summary_sd runs _populate_sd_cache, which files + # them with the sd it caches (see _summary_errs_cache). self.errs = errs + self.summary_sd = result_summary_sd @observe('summary_sd', 'processed_result') @exception_protect('merged_sd-protector') @@ -394,7 +462,14 @@ 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 + # The errs of the runs behind summary_stats_cache entries that had any, + # by the same key (see set_stats_tier). + self._summary_errs_cache: TDict[str, ErrDict] = {} self.init_sd: InitSD if init_sd is None: self.init_sd = {} @@ -524,43 +599,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. @@ -602,14 +643,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. @@ -621,11 +662,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') @@ -657,9 +705,16 @@ 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. - if keys['filt'] not in new_cache: + # filt scope reuses the SD that _summary_sd just produced, but only if + # it was produced for this frame, klass list and tier. One left over + # from an earlier frame (its run raised) stays out of the cache: stored + # under this key it would never be repaired, since a present key is a + # hit. The next _summary_sd run for this state fills it. + current = (id(self.processed_df), id(self.analysis_klasses), self.stats_tier) + if keys['filt'] not in new_cache and self._summary_sd_cache_key == current: new_cache[keys['filt']] = dict(self.summary_sd or {}) + if getattr(self, 'errs', None): + self._summary_errs_cache[keys['filt']] = self.errs cache_grew = True # raw + clean: fresh compute, but only on cache miss. @@ -715,6 +770,15 @@ def _build_error_dataframe(self, e) -> pd.DataFrame: #TAny closer to some error type @override def _get_summary_sd(self, processed_df: DataFrameT, scope: str = 'filt') -> Tuple[SDType, ErrDict]: + """``processed_df``'s summary stats at the dataflow's tier: the stats + pipeline's (``_get_full_sd``) or the schema's (``_get_schema_sd``).""" + if self.stats_tier == "schema": + return self._get_schema_sd(processed_df, scope), {} + return self._get_full_sd(processed_df, scope) + + def _get_full_sd(self, processed_df: DataFrameT, scope: str) -> Tuple[SDType, ErrDict]: + """The stats pipeline's summary of ``processed_df``, the frame of + ``scope``. A backend overrides this to run its own pipeline.""" stats = self.DFStatsClass( processed_df, self.analysis_klasses, @@ -732,7 +796,61 @@ def _get_summary_sd(self, processed_df: DataFrameT, scope: str = 'filt') -> Tupl return sdf, {} - # ### end summary stats block + def _get_schema_sd(self, processed_df: DataFrameT, scope: str) -> 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. ``scope`` is the + frame's scope, for a backend whose row count is a query. A backend opts + in to ``stats_tier="schema"`` by overriding this. + """ + raise NotImplementedError(f"{type(self).__name__} has no schema stats tier") + + def set_stats_tier(self, tier: str, summary: Optional[Tuple[SDType, ErrDict]] = None) -> None: + """Switch to ``tier`` without computing stats that are already in hand. + + ``summary`` is the current frame's ``(sd, errs)`` at ``tier``, when the + caller computed it elsewhere (off the IOLoop, say). Without it, the + entry an earlier visit to this state left in ``summary_stats_cache`` is + used, and failing that this is plain assignment: the ``stats_tier`` + observer computes the tier's stats. The cascade then fills the raw and + clean scopes from the cache, computing the ones it lacks. + """ + filt_key = self._scope_cache_key(split_chain_by_scope(self.operations)['filt'], tier=tier) + prior_tier, prior_key = self.stats_tier, self._summary_sd_cache_key + try: + if summary is None: + cached = self.summary_stats_cache.get(filt_key) + if cached is None: + self.stats_tier = tier + return + # The errs of the run that made the entry, kept apart because + # the cache holds sds only. + summary = (cached, self._summary_errs_cache.get(filt_key, {})) + else: + # Stored under the state's key at the new tier, so merged_sd reads + # the stats given rather than an entry already there. + new_cache = dict(self.summary_stats_cache) + new_cache[filt_key] = summary[0] + self.summary_stats_cache = new_cache + if summary[1]: + self._summary_errs_cache[filt_key] = summary[1] + sd, errs = summary + # Recorded first, so the stats_tier observer finds this frame's stats + # present and leaves them to the assignments below. + self._summary_sd_cache_key = (id(self.processed_df), id(self.analysis_klasses), tier) + self.stats_tier = tier + self.errs = errs + self.summary_sd = sd + except Exception: + # The key goes back first, so the observer the tier change fires + # finds the frame's stats present and runs nothing. + self._summary_sd_cache_key = prior_key + self.stats_tier = prior_tier + raise + + # ### end summary stats block def _sd_to_jsondf(self, sd:SDType): """Serialize summary stats to the wire payload (parquet-b64 tagged dict). @@ -750,34 +868,25 @@ def _df_to_obj(self, df: DataFrameT) -> List[TDict[str, TAny]]: return pd_to_obj(self.sampling_klass.serialize_sample(df)) 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 - self.DFStatsClass.verify_analysis_objects(self.analysis_klasses) + """Add ``analysis_klass``, replacing a class of the same name, and rerun + the summary stats with it at the dataflow's tier.""" + klasses = with_stat(self.analysis_klasses, analysis_klass) + # Validate before assigning, so a klass that fails leaves the list as it was. + self.DFStatsClass.verify_analysis_objects(klasses) + self.analysis_klasses = klasses self.setup_options_from_analysis() - #force recomputation + # analysis_klasses is a plain attribute, so the cascade starts here; the + # new list misses _summary_sd's dedupe key. + self._summary_sd({}) + # A klass that changes only the displays (post-processing, styling) + # leaves summary_sd as it was, and nothing downstream would rebuild them. 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 @@ -785,15 +894,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) @@ -811,8 +920,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 9823b8b0f..3404c748f 100644 --- a/buckaroo/pluggable_analysis_framework/stat_pipeline.py +++ b/buckaroo/pluggable_analysis_framework/stat_pipeline.py @@ -23,6 +23,16 @@ from .utils import PERVERSE_DF +def with_stat(inputs: list, stat_func_or_class) -> list: + """``inputs`` with ``stat_func_or_class`` added: a class of the same name is + replaced, anything else is appended. The rule both pipelines' ``add_stat`` + and ``CustomizableDataflow.add_analysis`` follow.""" + if isinstance(stat_func_or_class, type): + inputs = [inp for inp in inputs + if not (isinstance(inp, type) and inp.__name__ == stat_func_or_class.__name__)] + return [*inputs, stat_func_or_class] + + def _normalize_inputs(inputs: list) -> List[StatFunc]: """Convert a mixed list of StatFunc, @stat functions, and stat-group classes to StatFuncs. @@ -349,15 +359,7 @@ def add_stat(self, stat_func_or_class) -> Tuple[bool, List[StatError]]: Validates the DAG and runs unit test against PERVERSE_DF. """ - new_inputs = list(self._original_inputs) - - # Remove existing with same name if re-adding - if isinstance(stat_func_or_class, type): - new_inputs = [ - inp for inp in new_inputs - if not (isinstance(inp, type) and inp.__name__ == stat_func_or_class.__name__) - ] - new_inputs.append(stat_func_or_class) + new_inputs = with_stat(self._original_inputs, stat_func_or_class) try: new_funcs = _normalize_inputs(new_inputs) diff --git a/buckaroo/pluggable_analysis_framework/xorq_stat_pipeline.py b/buckaroo/pluggable_analysis_framework/xorq_stat_pipeline.py index 80bce5dc1..30929533d 100644 --- a/buckaroo/pluggable_analysis_framework/xorq_stat_pipeline.py +++ b/buckaroo/pluggable_analysis_framework/xorq_stat_pipeline.py @@ -38,7 +38,7 @@ from .stat_cache import (CachedScope, CachedStatError, StatCache, length_stat_id, make_scope_id, stat_hashes) from .stat_func import XorqColumn, XorqExpr, XorqExecute, RAW_MARKER_TYPES, StatFunc -from .stat_pipeline import _execute_stat_func, _normalize_inputs, errors_to_errdict +from .stat_pipeline import _execute_stat_func, _normalize_inputs, errors_to_errdict, with_stat from .stat_result import Err, Ok, StatError, StatResult, resolve_accumulator from .typed_dag import build_column_dag, build_typed_dag from .utils import PERVERSE_DF @@ -667,18 +667,7 @@ def add_stat(self, stat_func_or_class) -> Tuple[bool, List[StatError]]: against). Validates the DAG; returns ``(True, [])`` on success or ``(False, [config_error])`` if the DAG can't be built. """ - new_inputs = list(self._original_inputs) - - if isinstance(stat_func_or_class, type): - new_inputs = [ - inp - for inp in new_inputs - if not ( - isinstance(inp, type) - and inp.__name__ == stat_func_or_class.__name__ - ) - ] - new_inputs.append(stat_func_or_class) + new_inputs = with_stat(self._original_inputs, stat_func_or_class) try: new_funcs = _normalize_inputs(new_inputs) diff --git a/buckaroo/server/handlers.py b/buckaroo/server/handlers.py index aa56b5f13..6a25b36ee 100644 --- a/buckaroo/server/handlers.py +++ b/buckaroo/server/handlers.py @@ -14,7 +14,10 @@ from buckaroo.compare import col_join_dfs from buckaroo.df_util import old_col_new_col from buckaroo.server.focus import find_or_create_session_window -from buckaroo.server.session import build_state_message +from buckaroo.dataflow.dataflow import STATS_TIERS +from buckaroo.server.session import ( + DEFAULT_STATS_DELIVERY, DEFAULT_STATS_TIER, STATS_DELIVERIES, begin_stats_generation, dataflow_stats_tier) +from buckaroo.server.stats_wire import apply_component_config, broadcast_state, refresh_session_snapshot from buckaroo.server import telemetry from buckaroo.pluggable_analysis_framework import perf_log @@ -192,17 +195,10 @@ def _push_state_to_clients(self, session, metadata: dict): if not session.ws_clients: return - for client in list(session.ws_clients): - try: - # Reset per-client live search first (dataset changed), - # then build the msg so the injected ``buckaroo_state`` - # mirrors the post-reset value. - client.search_string = "" - msg = build_state_message(session, metadata=metadata, - search_string=client.search_string) - client.write_message(json.dumps(msg)) - except Exception: - session.ws_clients.discard(client) + # Each client's search is reset first (dataset changed), then its + # message is built so the injected ``buckaroo_state`` mirrors the + # post-reset value. + broadcast_state(session, metadata=metadata, reset_search=True) def _handle_browser_window(self, session_id: str) -> str: """Handle browser window management.""" @@ -323,6 +319,9 @@ async def post(self): # to the new dataflow rather than a stale xorq one. session.backend = backend session.release_xorq_state() + # The stats policy is /load_expr's too, so a later /load_expr that + # omits it starts from the defaults rather than an earlier load's. + session.stats_tier, session.stats_delivery = DEFAULT_STATS_TIER, DEFAULT_STATS_DELIVERY session.prompt = prompt if component_config: session.component_config = component_config @@ -386,6 +385,11 @@ async def post(self): **component_config, } + # /load builds no deferred stats, whatever policy a prior /load_expr on + # this session left behind, and the generation moves on with the data. + session.stats_tier, session.stats_delivery = "full", "inline" + begin_stats_generation(session) + # Notify connected clients and open browser self._push_state_to_clients(session, metadata) browser_action = "skipped" if no_browser else self._handle_browser_window(session_id) @@ -395,6 +399,28 @@ async def post(self): self.write({"session": session_id, "server_pid": os.getpid(), "browser_action": browser_action, **metadata}) +def _stats_policy_from_body(body: dict, current_tier: str, current_delivery: str): + """Read ``stats_tier`` and ``stats_delivery`` from a request body. + + A field the body omits (or sends as null) keeps the current value (the + session's, or the default for a new session), as ``cache_dir`` does. Returns + ``(stats_tier, stats_delivery, None)``, or ``(None, None, error)`` with the + 400 response body for an unknown value.""" + stats_tier = body.get("stats_tier") + if stats_tier is None: + stats_tier = current_tier + if stats_tier not in STATS_TIERS: + return None, None, {"error_code": "invalid_stats_tier", + "message": f"stats_tier must be one of {list(STATS_TIERS)}, got {stats_tier!r}"} + stats_delivery = body.get("stats_delivery") + if stats_delivery is None: + stats_delivery = current_delivery + if stats_delivery not in STATS_DELIVERIES: + return None, None, {"error_code": "invalid_stats_delivery", + "message": f"stats_delivery must be one of {list(STATS_DELIVERIES)}, got {stats_delivery!r}"} + return stats_tier, stats_delivery, None + + class LoadExprHandler(tornado.web.RequestHandler): """POST /load_expr — load a xorq/ibis expression from a build dir and serve it via the xorq-backed buckaroo dataflow. @@ -453,7 +479,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")) @@ -478,6 +507,16 @@ 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 DEFAULT_STATS_TIER, + existing.stats_delivery if existing is not None else DEFAULT_STATS_DELIVERY) + if policy_error is not None: + self.set_status(400) + self.write(policy_error) + return # The identity of the rows the build reads (e.g. tallyman's snapshot # digest); it keys the summary-stat cache. Omitted, it carries over # like cache_dir, but only for the same build. A new value means the @@ -491,7 +530,8 @@ async def post(self): if (not force_reload and not has_config and existing and existing.backend == "xorq" and existing.build_dir == build_dir and existing.cache_dir == cache_dir and existing_kwargs.get("data_id") == data_id - and existing.metadata): + 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) @@ -563,6 +603,7 @@ async def post(self): with perf_log.perf_span("firstpull.dataflow_construct", session=session_id): xorq_dataflow = xorq_loading.XorqServerDataflow( expr, skip_main_serial=True, extra_klasses=extra_klasses, + stats_tier=dataflow_stats_tier(stats_tier, stats_delivery), **dataflow_kwargs) # Spanning metadata too leaves only the small klass-load step # unmeasured inside the outer firstpull.load_expr total. @@ -589,6 +630,8 @@ async def post(self): session.cache_dir = cache_dir session.project_root = project_root session.dataflow_kwargs = dataflow_kwargs + session.stats_tier = stats_tier + session.stats_delivery = stats_delivery session.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. @@ -618,25 +661,15 @@ async def post(self): "generated_py_code": "# server mode"} session.operations = [] - if 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", {}), - **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) + apply_component_config(session.df_display_args, component_config) + + # 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" @@ -780,18 +813,12 @@ 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 = DEFAULT_STATS_TIER, DEFAULT_STATS_DELIVERY + 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): @@ -821,8 +848,23 @@ class ReloadExprHandler(tornado.web.RequestHandler): ``cache_storage_path`` store are cache hits, and column overrides, extra grid config, init_sd and skip_stat_columns survive the reload. + The session's stored ``stats_tier`` / ``stats_delivery`` are replayed too. + The body is optional; a pair in it replaces the stored one (and is stored) + when the reload succeeds. + Returns 404 when the session does not exist, 400 when it is not a xorq - session or has no project_root recorded, 501 when xorq is not installed.""" + session, has no project_root recorded or carries an invalid stats policy, + 501 when xorq is not installed.""" + + def _optional_body(self) -> dict: + """The JSON object body, or ``{}`` when there is none or it is not an + object: the endpoint took no body before the stats policy fields.""" + try: + body = json.loads(self.request.body) + # ValueError covers JSONDecodeError and UnicodeDecodeError (non-UTF-8 bytes). + except (ValueError, TypeError): + return {} + return body if isinstance(body, dict) else {} async def post(self, session_id): sessions = self.application.settings["sessions"] @@ -845,6 +887,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: @@ -861,6 +910,7 @@ async def post(self, session_id): + xorq_loading.load_project_display_klasses(session.project_root)) xorq_dataflow = xorq_loading.XorqServerDataflow( session.expr, skip_main_serial=True, extra_klasses=extra_klasses, + stats_tier=dataflow_stats_tier(stats_tier, stats_delivery), **session.dataflow_kwargs) except Exception: tb = traceback.format_exc() @@ -886,29 +936,13 @@ async def post(self, session_id): if bs.get("quick_command_args"): xorq_dataflow.quick_command_args = bs["quick_command_args"] - refreshed = get_buckaroo_display_state(xorq_dataflow) session.xorq_dataflow = xorq_dataflow - session.df_display_args = refreshed["df_display_args"] - session.df_data_dict = refreshed["df_data_dict"] - session.df_meta = refreshed["df_meta"] - session.buckaroo_options = refreshed["buckaroo_options"] - session.command_config = refreshed["command_config"] + session.stats_tier = stats_tier + session.stats_delivery = stats_delivery + refresh_session_snapshot(session, xorq_dataflow) + begin_stats_generation(session) - if session.component_config and session.df_display_args: - for key in session.df_display_args: - dvc = session.df_display_args[key].get("df_viewer_config") - if dvc is not None: - dvc["component_config"] = { - **dvc.get("component_config", {}), - **session.component_config} - - for client in list(session.ws_clients): - try: - msg = build_state_message(session, - search_string=getattr(client, "search_string", "")) - client.write_message(json.dumps(msg)) - except Exception: - session.ws_clients.discard(client) + broadcast_state(session) klass_count = len(extra_klasses) log.info("reload_expr session=%s project_root=%s klasses=%d", diff --git a/buckaroo/server/session.py b/buckaroo/server/session.py index 4b4af41c3..787ce7f6a 100644 --- a/buckaroo/server/session.py +++ b/buckaroo/server/session.py @@ -11,6 +11,30 @@ _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") +# The policy of a session no /load_expr has set one for. +DEFAULT_STATS_TIER = "full" +DEFAULT_STATS_DELIVERY = "inline" + + +def dataflow_stats_tier(stats_tier: str, stats_delivery: str) -> str: + """The tier a session's dataflow is constructed at. A deferred session + starts at the schema tier whatever tier it is headed for.""" + return "schema" if stats_delivery == "deferred" else stats_tier + + +def initial_stats_status(stats_tier: str, stats_delivery: str) -> tuple[str, Optional[str]]: + """The status, and its reason, of a session whose stats generation has just + started, from the policy pair alone.""" + if stats_tier != "full": + return "not_computed", "host" + if stats_delivery == "deferred": + return "pending", None + return "complete", None + @dataclass class SessionState: @@ -43,6 +67,35 @@ class SessionState: # extra_grid_config, init_sd, skip_stat_columns), replayed by /reload_expr so # a reload keeps stat caching and column config (#957). dataflow_kwargs: dict = field(default_factory=dict) + # The /load_expr stats policy. ``stats_tier`` is the tier the session is + # headed for (see dataflow.STATS_TIERS) and ``stats_delivery`` how it gets + # there (see STATS_DELIVERIES). Kept apart from dataflow_kwargs, which is + # splatted into the dataflow constructor, and replayed by /reload_expr. The + # tier the dataflow is built at follows from the pair (dataflow_stats_tier). + stats_tier: str = DEFAULT_STATS_TIER + stats_delivery: str = DEFAULT_STATS_DELIVERY + # 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 + # What ``df_meta.stats.status`` says about the stats of the 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; the next connection to + # open retries it, see stats_wire.rearm_failed_stats). + stats_status: str = "complete" + stats_reason: Optional[str] = None + # The ``df_display_args`` the full stats changed, set when a run completes + # (a float column's minWidth reads its min and max) and carried by the + # ``stats_update`` that delivers them, so a client whose pending frame held + # the schema tier's config gets the stats-derived one. ``None`` while pending, + # and when the stats left the config as it was. + stats_display_args: Optional[dict] = 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 @@ -90,6 +143,34 @@ def release_xorq_state(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_display_args = None + + +def stats_meta(session: "SessionState") -> Optional[dict]: + """The ``df_meta.stats`` value for a session: ``{status, tier, gen}`` plus + ``reason`` when there is one. ``tier`` is the tier reached so far, which is + the one the session is headed for only once it is complete. + + ``None`` for a session on the default policy (inline delivery, full tier, + complete), which sends the message it always has; a client reads a missing + ``stats`` as complete.""" + if session.stats_status == "complete" and session.stats_delivery != "deferred": + return None + stats: dict = {"status": session.stats_status, + "tier": session.stats_tier if session.stats_status == "complete" else "schema", + "gen": session.stats_gen} + if session.stats_reason: + stats["reason"] = session.stats_reason + return stats + + def build_state_message(session: "SessionState", metadata: dict | None = None, search_string: str = "", reply_seq: int | None = None) -> dict: """Build the full ``initial_state`` WebSocket payload from a session. @@ -116,10 +197,16 @@ def build_state_message(session: "SessionState", metadata: dict | None = None, Returns: A dict ready to be JSON-serialised and sent to WebSocket clients. """ + # The dataflow rebuilds df_meta wholesale, so the stats status is injected + # here, into a copy, rather than stored in it. + df_meta = session.df_meta + stats = stats_meta(session) + if stats is not None: + df_meta = {**df_meta, "stats": stats} msg: dict = {"type": "initial_state", "protocol_version": PROTOCOL_VERSION, "metadata": metadata if metadata is not None else session.metadata, "prompt": session.prompt, "df_display_args": session.df_display_args, "df_data_dict": session.df_data_dict, - "df_meta": session.df_meta, "mode": session.mode} + "df_meta": df_meta, "mode": session.mode} if reply_seq is not None: msg["reply_seq"] = reply_seq if session.mode == "buckaroo": diff --git a/buckaroo/server/stats_wire.py b/buckaroo/server/stats_wire.py new file mode 100644 index 000000000..f6f0bac8c --- /dev/null +++ b/buckaroo/server/stats_wire.py @@ -0,0 +1,298 @@ +"""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"``) and is + owed the stats for that generation. The server pushes them as a + ``stats_update`` once it has finished sending the client's first row reply + (``push_stats``, ADR-002 D2), and the client merges it. The client may also + ask with ``stats_request {stats_gen, scope}``; a request for a generation the + session has left gets ``stats_aborted``. +* Any other client gets complete messages: ``build_state_message_for`` runs the + missing stats synchronously before it builds a message for one, which is + today's cost, paid on the loop. + +Every send site goes through ``build_state_message_for`` (or ``broadcast_state``, +which calls it per client), because the session holds one shared snapshot and +the client is known only to the handler that owns the connection. + +A push and a request are each the whole run: one synchronous call that computes +the full stats and applies the final assignment (``set_stats_tier``, +``refresh_session_snapshot``). Resumable units generalize it later. +""" +import json +import logging +import os +import time +import traceback +from contextlib import nullcontext +from typing import Any, Callable, Optional + +from buckaroo.pluggable_analysis_framework import perf_log +from buckaroo.server.data_loading import get_buckaroo_display_state +from buckaroo.server.session import SessionState, build_state_message + +log = logging.getLogger("buckaroo.server.stats_wire") + +# The capability a client advertises, as one of the comma-separated values of +# ``?caps=`` on the WebSocket URL: it merges ``stats_update`` messages. The +# connection is the only place it can be recorded, because ``open()`` sends the +# first message before the client has said anything. +STATS_UPDATE_CAP = "stats_update" + +# The scopes a ``stats_request`` can name. ``raw`` is the unfiltered table of +# the session's current generation; filtered stats exist only through a +# dataflow-field change (``quick_command_args.search``), which bumps the +# generation. +STATS_SCOPES = ("raw",) + +# Seconds ``complete_stats`` waits before it runs the full stats. For tests +# that must see a session between its rows and its stats (the Playwright suite +# sets it); unset or 0 in every other use. +TEST_STATS_DELAY_ENV = "BUCKAROO_TEST_STATS_DELAY_S" + + +def parse_caps(raw: str) -> frozenset: + """The capabilities in a ``?caps=a,b`` value.""" + return frozenset(cap for cap in (part.strip() for part in raw.split(",")) if cap) + + +def client_has_cap(client: Any, cap: str) -> bool: + """Whether a WebSocket handler recorded ``cap`` at open. Anything else with + no ``caps`` (a test double, a future non-WebSocket client) has none.""" + return cap in getattr(client, "caps", ()) + + +def session_dataflow(session: Optional[SessionState]) -> Any: + """The dataflow behind a buckaroo-mode session, or ``None`` (viewer and lazy + sessions have none).""" + if session is None or session.mode != "buckaroo": + return None + return session.xorq_dataflow if session.backend == "xorq" else session.dataflow + + +def apply_component_config(df_display_args: Optional[dict], component_config: Optional[dict]) -> None: + """Merge ``component_config`` (theme and layout settings from the load + request) into every display's ``df_viewer_config``, in place, over what the + dataflow put there.""" + if not component_config or not df_display_args: + return + for display in df_display_args.values(): + dvc = display.get("df_viewer_config") + if dvc is not None: + dvc["component_config"] = {**dvc.get("component_config", {}), **component_config} + + +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"] + apply_component_config(session.df_display_args, session.component_config) + + +def rearm_failed_stats(session: SessionState) -> None: + """Put a deferred session whose stats run failed back to pending, for the + connection that is about to open (a page reload, a second tab): its pending + frame owes it a push, which runs the stats again. A failure that persists + ends the same way (``error`` again, after one more run); a transient one, + such as a backend timeout, no longer outlasts the generation it hit. + Requests and a client's later frames do not retry, so a failing query runs + once per connection and not once per keystroke.""" + if session.stats_delivery == "deferred" and session.stats_status == "error": + session.stats_status, session.stats_reason = "pending", None + + +def complete_stats(session: SessionState) -> bool: + """Run the stats a pending session is missing, in one synchronous call, and + publish them: the final assignment (``set_stats_tier``), 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. + + The run is timed as ``stats.complete`` on the session's telemetry sink, or + on the sink already bound when the session has none. It is not a + ``firstpull.*`` span: a state change or a legacy client's connect completes + stats long after the load. + + A failure is the session's state (``error``, reason ``stats_failed``) and is + not retried by the next request; ``rearm_failed_stats`` retries it when the + next connection opens, and the next generation starts clean. + + ``session.stats_display_args`` is set to the refreshed display config when + the stats changed it (see ``SessionState``).""" + if session.stats_status == "complete": + return True + dataflow = session_dataflow(session) + if session.stats_status != "pending" or dataflow is None: + return False + with ( + perf_log.telemetry_context(session.session_id, session.tele_sink) if session.tele_sink else nullcontext(), + perf_log.perf_span("stats.complete", session=session.session_id, stats_gen=session.stats_gen), + ): + try: + schema_display_args = session.df_display_args + time.sleep(float(os.environ.get(TEST_STATS_DELAY_ENV) or 0)) + dataflow.set_stats_tier("full") + 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 + session.stats_display_args = session.df_display_args if session.df_display_args != schema_display_args else 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 pending deferred + session gets its missing stats run first, so its message is complete; a + capable client gets the session snapshot as it is (stats-free while the + session is pending) and is recorded as owed that generation's stats + (``client.stats_owed``), which ``push_stats`` delivers after its first row + reply. 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 session.stats_status == "pending" + and not client_has_cap(client, STATS_UPDATE_CAP)): + complete_stats(session) + elif session.stats_status == "pending" and client_has_cap(client, STATS_UPDATE_CAP): + client.stats_owed = session.stats_gen + return build_state_message(session, metadata=metadata, search_string=getattr(client, "search_string", ""), + reply_seq=reply_seq) + + +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, started: float) -> dict: + dataflow = session_dataflow(session) + if session is None or dataflow is None: + return _aborted(stats_gen, scope, "no_data") + if stats_gen != session.stats_gen: + return _aborted(stats_gen, scope, "stale", session) + if scope not in STATS_SCOPES: + return _aborted(stats_gen, scope, "unsupported_scope", session) + if session.stats_status == "not_computed": + return _aborted(stats_gen, scope, "not_requestable", session) + if session.stats_status == "error" or (session.stats_status == "pending" and not complete_stats(session)): + return _aborted(stats_gen, scope, "error", session) + # Complete: from here the answer is the session's own all_stats, with no + # query, whether this request ran the stats or an earlier one did. The tier + # is the session's, not the dataflow's: a state change that failed halfway + # leaves the dataflow at the schema tier while the snapshot still holds the + # full stats. + reply = {"type": "stats_update", "stats_gen": stats_gen, "scope": scope, "tier": session.stats_tier, + "final": True, "payload": session.df_data_dict["all_stats"], + "elapsed_ms": round((time.perf_counter() - started) * 1000, 1)} + if session.stats_display_args is not None: + reply["df_display_args"] = session.stats_display_args + return reply + + +def handle_stats_request(session: Optional[SessionState], msg: dict) -> dict: + """Answer a ``stats_request {stats_gen, scope, columns?}``: a + ``stats_update`` carrying the complete ``all_stats`` as an inline wide + ``DFEnvelope`` (self-contained, so binary pairing stays single-slot), or a + ``stats_aborted``. ``columns`` is a hint a whole-run request ignores. + + The ``stats.request`` span records the request and its ``outcome`` + (``update`` or the abort reason), which is where updates sent, requests + dropped as stale and errors are counted. The caller binds the session's + telemetry sink around this call.""" + stats_gen, scope, columns = msg.get("stats_gen"), msg.get("scope", "raw"), msg.get("columns") + started = time.perf_counter() + with perf_log.perf_span("stats.request", session=session.session_id if session else None, stats_gen=stats_gen, + scope=scope, columns=len(columns) if isinstance(columns, list) else None) as span: + try: + reply = _answer_stats_request(session, stats_gen, scope, started) + except Exception: + log.error("stats_request error session=%s: %s", session.session_id if session else None, + traceback.format_exc()) + reply = _aborted(stats_gen, scope, "error", session) + _record_outcome(span, reply) + return reply + + +def _record_outcome(span: Any, reply: dict) -> None: + if reply["type"] == "stats_update": + span.set_attr(outcome="update", tier=reply["tier"]) + else: + span.set_attr(outcome=reply["reason"]) + + +def push_stats(session: Optional[SessionState], client: Any, stats_gen: int, rows_done: float) -> Optional[dict]: + """The ``stats_update`` (or ``stats_aborted``) owed to ``client`` for + ``stats_gen``, or ``None`` when nothing is owed any more: the client was + already pushed this generation, or the session has moved to another one (the + ``initial_state`` that started it registered its own push). ``rows_done`` is + the ``perf_counter`` time the row reply finished sending. + + The caller runs this on the write future of the client's first row reply + and sends the result, so rows reach the socket before any stats work starts. + The stats run once per session: a later connection's push is answered from + the session. The ``stats.push`` span records the gap since the rows and the + ``outcome``; the caller binds the session's telemetry sink around this call.""" + if getattr(client, "stats_owed", None) != stats_gen: + return None + client.stats_owed = None + if session is None or session.stats_gen != stats_gen: + return None + started = time.perf_counter() + with perf_log.perf_span("stats.push", session=session.session_id, stats_gen=stats_gen, + gap_ms=round((started - rows_done) * 1000, 1)) as span: + try: + reply = _answer_stats_request(session, stats_gen, "raw", started) + except Exception: + log.error("stats push error session=%s: %s", session.session_id, traceback.format_exc()) + reply = _aborted(stats_gen, "raw", "error", session) + _record_outcome(span, reply) + return reply diff --git a/buckaroo/server/websocket_handler.py b/buckaroo/server/websocket_handler.py index 7ef51aeb5..66d1dc32c 100644 --- a/buckaroo/server/websocket_handler.py +++ b/buckaroo/server/websocket_handler.py @@ -2,6 +2,7 @@ import json import logging import os +import time import traceback from contextlib import nullcontext from urllib.parse import urlparse @@ -9,8 +10,9 @@ import tornado.websocket from buckaroo.pluggable_analysis_framework import perf_log -from buckaroo.server.data_loading import (handle_infinite_request, handle_infinite_request_buckaroo, handle_infinite_request_lazy, get_buckaroo_display_state) -from buckaroo.server.session import build_state_message +from buckaroo.server.data_loading import (handle_infinite_request, handle_infinite_request_buckaroo, handle_infinite_request_lazy) +from buckaroo.server.session import begin_stats_generation, dataflow_stats_tier +from buckaroo.server.stats_wire import (broadcast_state, build_state_message_for, handle_stats_request, parse_caps, push_stats, rearm_failed_stats, refresh_session_snapshot) def _handle_infinite_request_xorq(xorq_dataflow, payload_args, search_string=""): @@ -38,6 +40,15 @@ 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", "")) + # The stats_gen this connection is owed a ``stats_update`` for: set when + # it is sent a pending frame (stats_wire.build_state_message_for), + # cleared by the push that follows its next row reply. + self.stats_owed = None sessions = self.application.settings["sessions"] sessions.add_ws_client(session_id, self) @@ -45,7 +56,8 @@ 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))) + rearm_failed_stats(session) + self.write_message(json.dumps(build_state_message_for(session, self))) def on_message(self, message): try: @@ -56,13 +68,29 @@ def on_message(self, message): msg_type = msg.get("type") if msg_type == "infinite_request": - self._handle_infinite_request(msg.get("payload_args", {})) + self._handle_infinite_request(msg.get("payload_args", {}), stats_gen=msg.get("stats_gen")) elif msg_type == "buckaroo_state_change": # 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._write_stats_reply(reply) def _handle_buckaroo_state_change(self, new_state, state_seq=None): sessions = self.application.settings["sessions"] @@ -109,36 +137,36 @@ def _handle_buckaroo_state_change(self, new_state, state_seq=None): log.debug("buckaroo_state_change no-op session=%s — skipping rebroadcast", self.session_id) return - # 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", "") - if old_state.get("cleaning_method") != new_state.get("cleaning_method"): - dataflow.cleaning_method = new_state.get("cleaning_method", "") - if old_state.get("quick_command_args") != new_state.get("quick_command_args"): - dataflow.quick_command_args = new_state.get("quick_command_args", {}) + # 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 a push. The tier + # is assigned last and the notifications are held until every field + # is in, so the cascade runs once, for the new state: resetting the + # tier first would run it on the state being left, and a failed + # change leaves the tier alone (a rollback would run the full stats + # of a half-applied state; the session's own stats, and the tier it + # reports, are those of its snapshot, which the failure did not touch). + deferred = session.stats_delivery == "deferred" + with dataflow.hold_trait_notifications() if deferred else nullcontext(): + # 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", "") + if old_state.get("cleaning_method") != new_state.get("cleaning_method"): + dataflow.cleaning_method = new_state.get("cleaning_method", "") + if old_state.get("quick_command_args") != new_state.get("quick_command_args"): + dataflow.quick_command_args = new_state.get("quick_command_args", {}) + if deferred: + dataflow.stats_tier = dataflow_stats_tier(session.stats_tier, session.stats_delivery) # 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"] + refresh_session_snapshot(session, dataflow) + # The state the stats describe has changed, so the generation moves on. + 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 @@ -148,15 +176,7 @@ def _handle_buckaroo_state_change(self, new_state, state_seq=None): # doesn't drop it. # Only the originating client gets reply_seq (#998): the # others made no change, so every copy is current for them. - for client in list(session.ws_clients): - try: - search = getattr(client, "search_string", "") - msg = build_state_message(session, search_string=search, - reply_seq=state_seq if client is self else None) - msg["df_display_args"] = self._with_highlight(session.df_display_args, search) - client.write_message(json.dumps(msg)) - except Exception: - session.ws_clients.discard(client) + 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) @@ -207,16 +227,20 @@ def _send_client_state(self, session, buckaroo_state, reply_seq=None): 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(session, search_string=self.search_string, reply_seq=reply_seq) + 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(session.df_display_args, 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("client state write failed for session=%s", self.session_id) - def _handle_infinite_request(self, payload_args): + def _handle_infinite_request(self, payload_args, stats_gen=None): sessions = self.application.settings["sessions"] session = sessions.get(self.session_id) @@ -258,6 +282,8 @@ def _dispatch(pa): tele_sink = session.tele_sink if not_seen else None first_payload = not_seen and (perf_log.enabled() or tele_sink is not None) + last_write = None + def _dispatch_and_send(pa, span_label): # Dispatch one window and send its two-frame reply: a JSON text # frame, then the binary Parquet frame when non-empty. On the initial @@ -267,11 +293,12 @@ def _dispatch_and_send(pa, span_label): # v2. span = (perf_log.perf_span(span_label, session=self.session_id) if first_payload else nullcontext()) + nonlocal last_write with span: resp, parquet = _dispatch(pa) - self.write_message(json.dumps(resp)) + last_write = self.write_message(json.dumps(resp)) if parquet: - self.write_message(parquet, binary=True) + last_write = self.write_message(parquet, binary=True) try: with perf_log.telemetry_context(self.session_id, tele_sink): @@ -289,6 +316,42 @@ def _dispatch_and_send(pa, span_label): log.error("infinite_request error session=%s: %s", self.session_id, tb) self.write_message(json.dumps({"type": "infinite_resp", "key": payload_args, "length": 0, "error_info": tb if _BUCKAROO_DEBUG else "Request failed"})) + if self.stats_owed is not None and stats_gen == self.stats_owed and last_write is not None: + # D2: the stats follow the rows. Only the reply to a request made for + # the owed generation starts them: one made before the client saw + # that state (``stats_gen`` older, or absent) would run the stats + # ahead of the rows the client now needs. A sort or a scroll changes + # which rows are asked for, not the generation. The continuation runs + # once the last frame of this reply is written, so none of the stats + # work can delay it. + owed, rows_done = self.stats_owed, time.perf_counter() + last_write.add_done_callback(lambda _: self._push_stats(owed, rows_done)) + + def _push_stats(self, stats_gen, rows_done): + """Send the ``stats_update`` this connection is owed for ``stats_gen``. + Does nothing when the connection has closed since the row reply, or + when ``push_stats`` finds nothing owed (a newer generation, or already + pushed). A newer generation's own pending frame registered its own + push.""" + if self.ws_connection is None: + return + 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 = push_stats(session, self, stats_gen, rows_done) + if reply is not None: + try: + self._write_stats_reply(reply) + except tornado.websocket.WebSocketClosedError: + log.debug("stats push write failed for session=%s", self.session_id) + + def _write_stats_reply(self, reply): + """Send a ``stats_update`` or ``stats_aborted``. The display config a + final update carries replaces the client's, so it gets this client's + live-search highlight like every other frame that carries one (#851).""" + if "df_display_args" in reply: + reply = {**reply, "df_display_args": self._with_highlight(reply["df_display_args"], self.search_string)} + self.write_message(json.dumps(reply)) def on_close(self): sessions = self.application.settings["sessions"] diff --git a/buckaroo/xorq_buckaroo.py b/buckaroo/xorq_buckaroo.py index 74cfb73da..555f04aba 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 @@ -144,7 +144,7 @@ class XorqDataflow(CustomizableDataflow["XorqExpr | pd.DataFrame"]): 1. ``populate_df_meta`` can't ``len(expr)`` — it issues an ``expr.count().execute()`` for the row count. - 2. ``_get_summary_sd`` re-keys the summary dict from original column + 2. ``_get_full_sd`` re-keys the summary dict from original column names (what ``XorqStatPipeline`` produces) to the rewritten ``a, b, c`` names that ``pd_to_obj`` and the styling layer expect. @@ -177,7 +177,7 @@ def _scope_chain(self, scope: str) -> list: assigned, before ``operations`` catches up, so ``operations`` would still hold the previous search.""" if scope == 'filt': - return split_chain_by_scope(self.merged_operations)['filt'] + return split_chain_by_scope(self.merged_operations or [])['filt'] return split_chain_by_scope(self.operations)[scope] def _stat_scope_id(self, chain: list): @@ -220,6 +220,9 @@ def _get_summary_sd(self, processed_df: "XorqExpr | pd.DataFrame", scope: str = 'orig_col_name': orig_col, 'rewritten_col_name': rewritten_col} return empty, {} + return super()._get_summary_sd(processed_df, scope) + + def _get_full_sd(self, processed_df: "XorqExpr | pd.DataFrame", scope: str): cache_storage = getattr(self, 'cache_storage', None) chain = self._scope_chain(scope) scope_id = self._stat_scope_id(chain) if cache_storage is not None else None @@ -239,10 +242,10 @@ def _get_summary_sd(self, processed_df: "XorqExpr | pd.DataFrame", scope: str = # signal only the server observes (#943). Carries the write side too # (snapshots/bytes/write_errors, #951) so a cache that stops writing — # or a run with write_errors > 0 — is visible to a telemetry consumer - # rather than only to a server log no deployment reads. The pandas - # branch above already returned, so stats here is always an - # XorqDfStatsV2, which always exposes cache_run_stats(); call it - # unconditionally so a missing signal fails loudly rather than + # rather than only to a server log no deployment reads. + # _get_summary_sd keeps pandas frames from reaching here, so stats is + # always an XorqDfStatsV2, which always exposes cache_run_stats(); call + # it unconditionally so a missing signal fails loudly rather than # silently vanishing. cs = stats.cache_run_stats() span.set_attr( @@ -266,6 +269,31 @@ def _get_summary_sd(self, processed_df: "XorqExpr | pd.DataFrame", scope: str = rewritten[rewritten_col] = col_meta return rewritten, errs + def _get_schema_sd(self, processed_df: "XorqExpr | pd.DataFrame", scope: str) -> dict: + """Identity, dtype, typing flags and row count for every column of an + expression. The count (``_expr_count``, cached per expression) is the + only query: the filtered frame's is the one ``populate_df_meta`` takes + anyway, and the raw frame changes only with the source or the + post-processing. The clean scope's frame is rebuilt on every cache miss, + so counting it would be a query of its own, and its sd has no + ``length``. 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 = {} if scope == 'clean' else {'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} + return sd + _XORQ_ANALYSIS_KLASSES = list(XORQ_STATS_V2) + [DefaultSummaryStatsStyling, DefaultMainStyling] diff --git a/docs/plans/ADR-002-rows-first-stats-delivery.md b/docs/plans/ADR-002-rows-first-stats-delivery.md new file mode 100644 index 000000000..365ea5406 --- /dev/null +++ b/docs/plans/ADR-002-rows-first-stats-delivery.md @@ -0,0 +1,108 @@ +# ADR: Rows-first delivery with server-pushed summary stats + +- **Status:** Proposed (2026-10-06). Revised the same day: stats are pushed by the server after rows instead of pulled by the client, and the design assumes ADR-001's cache. Implemented in part as draft PRs that are parked on purpose: rows-first delivery is wanted eventually, behind the out-of-core `scan_parquet` backend. Nothing in this ADR changes a default. +- **Affected code:** `buckaroo/server/session.py` (`build_state_message`, the session snapshot), `buckaroo/server/websocket_handler.py` (`DataStreamHandler._handle_infinite_request`, `on_message`), `buckaroo/server/handlers.py` (`/load_expr`, `/reload_expr`), `buckaroo/dataflow/dataflow.py` (`_get_summary_sd`, `_populate_sd_cache`, `_scope_cache_key`), `buckaroo/xorq_buckaroo.py`, `buckaroo/pluggable_analysis_framework/xorq_stat_pipeline.py`, and on the client `WebSocketModel.ts`, `StateOrchestrator.ts`, `BuckarooServerView.tsx` and `packages/js/standalone.tsx`. +- **Implementing PRs:** #1021 (schema tier core, xorq), #1022 (schema tier for pandas and polars), #1024 (wire: `stats_update`, `df_meta.stats`; its `stats_request` branch narrows to the explicit request in D2) and the client PR #1025 (merge `stats_update`). #1027's pending, not-computed and error states still apply. #1020, the client hardening that makes a two-message protocol safe, is already on main. #1026 (resumable units and `StatRun`), #1028 (request runs units for a time budget), #1030 (incremental requests), #1035 (column-chunk split of the xorq batch) and #1027's scheduler implement the client-pull design that D2 replaced. As of 2026-10-06 #1027, #1028 and #1030 are closed and #1026 and #1035 are open. +- **Related:** #998 and #1014 (stale replies), #896 (eager second window), #923 (aggregate-backed xorq row cost), #829 (filtered stats), ADR-001 (the stat cache, which this ADR assumes), ADR-003 (the stats policy that decides what a push contains). + +## Terms + +- **schema tier**: a dataflow built from dtypes and the row count only: identity and typing keys (`orig_col_name`, `rewritten_col_name`, `dtype`, `typing_stats`, `_type`, `is_*`, `length`), no data query beyond the count. +- **stats_delivery**: `inline` (stats computed in the dataflow constructor, today's behaviour) or `deferred` (schema-tier first message, stats afterwards). +- **stats_gen**: a server-owned counter that identifies the state a stats message describes. Bumped by `/load`, `/reload_expr` and any `buckaroo_state_change` that changes a dataflow field. +- **row reply**: the frames that answer one `infinite_request`: a JSON `infinite_resp`, then a binary parquet frame when the window is non-empty, then the same pair for the eager `second_request` (#896) when the request carries one. +- **stats policy**: ADR-003's `resolve_stats_policy` result for a session (`tier_target`, `auto_request`, `requestable`, `reason`). +- **push**: the `stats_update` messages the server sends a connection without being asked. + +## Problem + +On `/load_expr`, `XorqServerDataflow` computes summary stats inside its constructor, and `session.xorq_dataflow` is assigned only afterwards. No row is served until the stats finish. Across 115 tallyman loads, `firstpull.load_expr` has p50 1.29 s, p90 4.5 s and a maximum of 215.8 s, and stats are a median 93% of it. The worst load spent 162.5 s in a single batch aggregate. Constructing the same dataflow at the schema tier takes about 1 ms against 393 ms with stats on a 27-column entry. + +The server is synchronous Tornado, so any stats query blocks every WebSocket message and HTTP request, `/health` included. Stats in a worker thread failed on xorq with `RuntimeError: Already borrowed` (two threads on one `SessionContext`). Stats stay on the loop, so what this ADR controls is their order relative to rows. + +This ADR assumes ADR-001's cache works. A warm load reads its stats in a few milliseconds (about 2 ms for keys, 1 to 3 ms to read and coalesce, 1 ms to decode), which is about what a row window costs. Cold stats, on the first load of an entry or after a stat's code changes, still run on the loop for as long as they take. Bounding that is left to ADR-001 and ADR-003 (see "Not decided here"). + +Three facts of the current code get in the way of a second message: + +1. A client that receives two `initial_state` frames can drop the second one while the first is decoding, and nothing in the protocol identifies which state a frame describes (#998). #1020, already on main, hardens the client for it. +2. `df_meta` is rebuilt wholesale on each state change, so a flag the dataflow writes there is lost at the next change. +3. A session snapshot is pushed from five places besides `open()`, so a check at connect time protects only the first message. + +## Decisions + +### D1. Stats travel in a typed `stats_update`, never a second `initial_state` + +The first `initial_state` carries the schema tier and `df_meta.stats = {status, tier, gen, reason?}`. `status` is `complete`, `pending`, `not_computed` or `error`; absent `stats` means `complete`, which is what old servers send. Stats arrive as `stats_update {stats_gen, scope, tier, final, payload, elapsed_ms}`, where `payload` is a self-contained wide `DFEnvelope` that the client key-merges into `all_stats`. The server may split one run across several updates, for example one per tier; the last has `final: true`. An explicit request the server will not run gets `stats_aborted`. `PROTOCOL_VERSION` stays 1. + +A second `initial_state` was rejected: existing clients can drop it, it reuses the message type whose lack of identity is #998, and it can say "not computed" only as a flag. + +### D2. The server pushes stats after it has finished sending rows + +On a deferred session the server sends a connection no summary stats until it has finished that connection's first row reply for the current `stats_gen`. Finished means every frame of the reply, the eager second window included, has been handed to the socket. The server waits on the `Future` that `write_message` returns for the reply's last frame, then reads or computes the stats and sends `stats_update`. A stats frame never lands inside a row reply, and stats are never computed while row bytes for that connection sit in Tornado's write buffer. An error reply counts as a finished reply, so a failed window does not leave the stats pending. + +Each connection holds the `stats_gen` it is owed stats for. It is set when the connection is sent an `initial_state` that starts a `stats_gen` (connect, `/load`, `/reload_expr`, a dataflow-field change, D7) and cleared by the push. Such an `initial_state` carries `status: "pending"` and no stats, even when the session already holds them; a later frame for the same `stats_gen`, sent after the push, carries them as today. + +What the push contains follows the session's stats policy (ADR-003). The server pushes the tier the policy runs automatically (`tier_target` when `auto_request` holds) and nothing above it. When the policy runs nothing, the push is a single `stats_update {final: true}` with `status: "not_computed"` and the policy fields, which ends the pending state. A tier above the pushed one runs only on an explicit `stats_request {stats_gen, scope, tier, columns?, force}`, which ADR-003's compute control sends. + +Stats for a `stats_gen` are computed once per session, at the first connection's first row reply, and kept on the session (D5). A connection whose first row reply comes later is pushed the stored result. A session nobody views computes nothing, because nothing starts a push until a client asks for rows. + +The push is one continuation registered on the write future. There is no timer, thread or chain of `add_callback` ticks. Before it runs it checks that the connection is still open and that `stats_gen` is still the one it was registered for. If either check fails it does nothing; a newer `stats_gen` gets its own push after that connection's next row reply. + +A request-and-reply transport such as `IpcDuckModel` cannot receive a frame it did not ask for. It does not advertise `stats_update`, so a server speaking to one uses D6. + +The earlier form of this decision had the client pull: after first paint the client sent `stats_request`, and the server ran resumable stat units for a bounded time per request (a unit plan and `StatRun` on the session, a cursor per connection, a 75 ms budget, a client scheduler). That machinery existed to keep a cold run from stalling row requests. With the cache, a warm push costs about what a row window costs, and on a cold run the units only spread the same work across requests. Pushing at load time, before any row request, was rejected because on a cold cache it puts stats ahead of rows again. + +### D3. `stats_gen` keys every stats message; `state_seq` does not + +The client drops a `stats_update` whose `stats_gen` differs from the one it holds, and advances its expected value on a broadcast frame. #1014's `state_seq` is client-owned, echoed to the originating client only, and not bumped by `/load` or `/reload_expr`, so it cannot guard a reply to a different client's state change. The two tokens are independent and #1014 does not conflict with this ADR. + +### D4. The schema tier is a dataflow level, and the tier is part of the cache scope + +`stats_tier` (`full` default, `schema`) is a dataflow attribute. At `schema`, `_get_summary_sd` returns the schema sd and runs no data query, and `add_analysis` gets its own tier check because it builds `DFStatsClass` itself. The tier joins `_scope_cache_key`, and a run assigns `summary_sd` only when it is a complete run. Without this, a schema-only construction followed by a full assignment never reaches `merged_sd`, because all three scopes hit the cache, and `_populate_sd_cache` can store a stale sd under a new filtered key (a probe stored `filtered_length` 1000 against a true 250). + +`pinned_rows` is identical at every tier, so nothing has to be restored when stats finish. A pinned key with no value renders a placeholder with a unique row id while `pending`. + +Pandas and polars dataflows accept the schema tier too, so one mechanism covers all three. Their `/load` path stays `full` and inline: eager stats cost 64 to 176 ms against 0.5 to 9 ms for a window, and deferring trades that for a second message and the layout problems of two messages. + +### D5. The end of a run does the final assignment + +When a run finishes, in this order: write the sd into `summary_stats_cache`, then assign `summary_sd`; refresh the session snapshot through one helper (replacing three copies) and set `status` to `complete` in the same step; attach the rebuilt `df_display_args` to the `final` update when its hash changed (a float `minWidth` was 114 with `max` 123456789.5 and 78 without stats). The snapshot is what D2 pushes to a connection that arrives later. + +### D6. Clients without the capability get complete stats synchronously + +Capability rides the connect URL, `?caps=stats_update`, because `open()` sends before any client hello. All six send sites and the highlight overlay go through one `build_state_message_for(session, client)`, which is also where a capable client's state-starting frame drops its stats (D2). A legacy client on a deferred session gets the stats computed synchronously before the first message built for it, which costs what today costs and, with the cache, a cache read. A permanent schema tier for them would strip stats at their first search, because stock clients search through `quick_command_args.search`. A new client against an old server sees no `df_meta.stats` and treats the session as complete. + +### D7. Search and state changes + +A dataflow-field change on a deferred session sets the traits at the schema tier, bumps `stats_gen` and broadcasts a `pending` frame. Each connection is then pushed the new stats after its next row reply (D2). This removes the double `_handle_widget_change` per search (596 ms to 12.7 ms on the 27-column entry). A `search_string` change touches no stats. Filtered stats are out of scope: after a search the client shows unfiltered stats with a note. #829 and #923 own the filtered case. + +### D8. Telemetry + +The push continuation and the `stats_request` branch enter `telemetry_context` with `session.tele_sink`, emit `stats.push` and `stats.request` spans, and count updates sent, pushes dropped as stale (closed connection or newer `stats_gen`) and errors. `stats.push` records the gap from the end of the row reply to the start of the push. `firstpull.load_expr` stops measuring stats, so tallyman must be told before the default flips, and `firstpull.stats_total` is emitted per completed run. + +## Consequences + +- The default stays `inline`. A host selects `deferred` per load with a body field. The new fields ride `dataflow_kwargs` and a stored pair, never the truthiness-based `has_config` tuple, or a host that sends them on every POST defeats the warm short-circuit. +- Tallyman pins `buckaroo-js-core` 0.15.8, so enabling deferred delivery for it needs a client release and a bump. +- `df_meta.stats` replaces per-cell sentinels and `stats_omitted` as the single flag for "stats are not here". +- A row request that arrives while a push is being built waits for it. With a warm cache that is the cache read. On a cold cache it is the whole stats run, which today happens before first paint and now happens after it. Flipping the default is gated on the cold case being acceptable, which ADR-001 and ADR-003 own. +- Time to complete stats grows by the time it takes to send the first rows, because stats start after them. The measure to hold is the end of the first row reply against the last `stats_update`. +- The pull-side PRs (#1026, #1028, #1030, #1035 and #1027's scheduler) are not needed under D2. #1027, #1028 and #1030 are closed. #1026 and #1035 are still open and are to be closed or cut down rather than rebased. The batch split in #1035 stays an option for cold stats (see "Not decided here"). + +## Alternatives considered + +- A second partial `initial_state` (the #787 shape). +- Client-pulled stats with resumable units, per-connection cursors and a per-request time budget (the earlier D2). +- A push started at load, before any row request. +- A server-pushed chain of `add_callback` units. +- A worker thread (fails on xorq) or a child process for the batch tail (worked at 2.69 s while the parent paged 47 windows at p50 17 ms, but needs `cache_storage_path`, an identical `build_dir`, and process lifecycle and fork handling, #885). +- Reusing `state_seq` as the stats key. +- Deferring eager pandas and polars. + +## Not decided here + +- **Which stats a session runs at a given size.** That is ADR-003. This ADR pushes what that policy runs automatically and carries `not_computed` and the `tier` field on the wire. +- **Cold stats.** On a cache miss the push runs the stats on the loop, after rows, for as long as they take. Splitting that work (ADR-001's column bisect, D10, with #1038), the xorq batch split of #1035, or a child process for the batch tail are where that would be decided. +- **A lazy or `scan_parquet` backend.** The destination for polars `/load` is an out-of-core scan. It would reuse this protocol, and its design is a separate document that does not exist yet. +- **A host `row_count` hint** for aggregate-backed xorq, where `count()` executes the plan (#923). +- **Whether capability negotiation stays in the URL query** or becomes a client hello, which costs a round trip on every connect. diff --git a/packages/buckaroo-js-core/playwright.config.server.ts b/packages/buckaroo-js-core/playwright.config.server.ts index 8d9d9d544..3ad2d1175 100644 --- a/packages/buckaroo-js-core/playwright.config.server.ts +++ b/packages/buckaroo-js-core/playwright.config.server.ts @@ -32,6 +32,9 @@ export default defineConfig({ command: `${PYTHON} -m buckaroo.server --no-browser --port ${PORT}`, cwd: '../..', url: `http://localhost:${PORT}/health`, + // Full stats on a deferred session wait this long, so a test can see the + // session between its rows and its stats (BUCKAROO_TEST_STATS_DELAY_S). + env: { BUCKAROO_TEST_STATS_DELAY_S: process.env.BUCKAROO_TEST_STATS_DELAY_S ?? '1' }, reuseExistingServer: !process.env.CI, timeout: 15_000, }, diff --git a/packages/buckaroo-js-core/pw-tests/server-load-expr.spec.ts b/packages/buckaroo-js-core/pw-tests/server-load-expr.spec.ts index 329b34b17..44e904d09 100644 --- a/packages/buckaroo-js-core/pw-tests/server-load-expr.spec.ts +++ b/packages/buckaroo-js-core/pw-tests/server-load-expr.spec.ts @@ -25,16 +25,23 @@ expr = xo.memtable({ print(xo.build_expr(expr, builds_dir=sys.argv[1])) `; -function buildExprDir(buildsRoot: string): string { +// A float column whose 1e9 maximum changes its estimated width once full stats arrive. +const BUILD_PRICE_EXPR_PY = ` +import sys, xorq.api as xo +expr = xo.memtable({'price': [12.5, 18.9, 7.4, 22.1, 1e9], 'qty': [1, 2, 1, 3, 2]}, name='t') +print(xo.build_expr(expr, builds_dir=sys.argv[1])) +`; + +function buildExprDir(buildsRoot: string, code: string = BUILD_EXPR_PY): string { const [python, ...pythonArgs] = (process.env.BUCKAROO_SERVER_PYTHON ?? 'uv run python').split(' '); - const out = execFileSync(python, [...pythonArgs, '-c', BUILD_EXPR_PY, buildsRoot], + const out = execFileSync(python, [...pythonArgs, '-c', code, buildsRoot], { cwd: ROOT_DIR, encoding: 'utf8' }); return out.trim().split('\n').pop()!; } -async function loadExpr(request: any, sessionId: string, buildDir: string) { +async function loadExpr(request: any, sessionId: string, buildDir: string, extra: Record = {}) { const resp = await request.post(`${BASE}/load_expr`, { - data: { session: sessionId, build_dir: buildDir, no_browser: true }, + data: { session: sessionId, build_dir: buildDir, no_browser: true, ...extra }, }); if (!resp.ok()) { throw new Error(`/load_expr failed (${resp.status()}): ${await resp.text()}`); @@ -46,6 +53,17 @@ async function getPinnedRowCount(page: any): Promise { return await page.locator('.df-viewer .ag-floating-top-container .ag-row').count(); } +// The text of every pinned cell in one column, top to bottom. In the summary +// view these are the stats: dtype, length, min, max, distinct count and so on. +async function getPinnedCellTexts(page: any, colId: string): Promise { + const texts = await page.locator(`.df-viewer .ag-floating-top-container [col-id="${colId}"]`).allInnerTexts(); + return texts.map((t: string) => t.trim()); +} + +async function showSummaryView(page: any) { + await page.locator('.status-bar').locator('select').first().selectOption('summary'); +} + test.describe('POST /load_expr', () => { let buildsRoot: string; let buildDir: string; @@ -113,4 +131,66 @@ test.describe('POST /load_expr', () => { await expect.poll(() => getCellText(page, COL.idx, 1), { timeout: 15_000 }).toBe('3'); expect(await getCellText(page, COL.name, 1)).toBe('alpha'); }); + + test('eager stats reach the DOM: the summary view shows the computed values', async ({ page, request }) => { + const session = `lx-eager-${Date.now()}`; + await loadExpr(request, session, buildDir, { stats_delivery: 'inline' }); + + await page.goto(`${BASE}/s/${session}`); + await waitForGrid(page); + await showSummaryView(page); + + // 'name' has 7 distinct values; 'idx' runs 0 to 9. + await expect.poll(() => getPinnedCellTexts(page, COL.name), { timeout: 15_000 }).toContain('7'); + expect(await getPinnedCellTexts(page, COL.idx)).toEqual(expect.arrayContaining(['0', '9'])); + }); + + test('deferred stats: the rows render first and the stats then reach the DOM', async ({ page, request }) => { + const session = `lx-deferred-${Date.now()}`; + await loadExpr(request, session, buildDir, { stats_delivery: 'deferred' }); + + await page.goto(`${BASE}/s/${session}`); + await waitForGrid(page); + expect(await getCellText(page, COL.idx, 0)).toBe('0'); + expect(await getCellText(page, COL.name, 0)).toBe('alpha'); + + await showSummaryView(page); + // The server holds the full stats back for a second (playwright.config.server.ts), + // so the summary view first shows the schema tier only: no distinct count, min or max. + expect(await getPinnedCellTexts(page, COL.name)).not.toContain('7'); + expect(await getPinnedCellTexts(page, COL.idx)).not.toEqual(expect.arrayContaining(['0', '9'])); + + await expect.poll(() => getPinnedCellTexts(page, COL.name), { timeout: 15_000 }).toContain('7'); + expect(await getPinnedCellTexts(page, COL.idx)).toEqual(expect.arrayContaining(['0', '9'])); + }); + + test('the display config the stats change arrives with them, and the grid does not ask for its first rows again', async ({ page, request }) => { + const session = `lx-display-${Date.now()}`; + await loadExpr(request, session, buildExprDir(buildsRoot, BUILD_PRICE_EXPR_PY), { stats_delivery: 'deferred' }); + + const events: { dir: string; type: string; start?: number; hasDisplayArgs?: boolean; minWidth?: number }[] = []; + page.on('websocket', (ws) => { + const record = (dir: string) => (e: { payload: string | Buffer }) => { + if (typeof e.payload !== 'string') return; + const msg = JSON.parse(e.payload); + const price = msg.df_display_args?.main?.df_viewer_config?.column_config?.find((c: any) => c.col_name === 'a'); + events.push({ dir, type: msg.type, start: msg.payload_args?.start, hasDisplayArgs: 'df_display_args' in msg, + minWidth: price?.ag_grid_specs?.minWidth }); + }; + ws.on('framesent', record('sent')); + ws.on('framereceived', record('received')); + }); + + await page.goto(`${BASE}/s/${session}`); + await waitForGrid(page); + await expect.poll(() => events.some((e) => e.type === 'stats_update'), { timeout: 15_000 }).toBe(true); + await page.waitForTimeout(1_000); + + const at = events.findIndex((e) => e.type === 'stats_update'); + expect(events[at].hasDisplayArgs).toBe(true); + // The price column's minWidth reads its 1e9 maximum, so it grows once the full stats are in. + const pending = events.find((e) => e.type === 'initial_state')!; + expect(events[at].minWidth).toBeGreaterThan(pending.minWidth!); + expect(events.slice(at).filter((e) => e.dir === 'sent' && e.type === 'infinite_request' && e.start === 0)).toEqual([]); + }); }); diff --git a/packages/buckaroo-js-core/pw-tests/server.spec.ts b/packages/buckaroo-js-core/pw-tests/server.spec.ts index ad9634a05..3327d908a 100644 --- a/packages/buckaroo-js-core/pw-tests/server.spec.ts +++ b/packages/buckaroo-js-core/pw-tests/server.spec.ts @@ -485,6 +485,22 @@ test.describe('WebSocket data flow', () => { expect(await getCellText(page, COL.name, 0)).toBe('Alice'); }); + test('standalone page advertises stats_update on its WebSocket URL', async ({ page, request }) => { + const session = `ws-caps-${Date.now()}`; + await loadSession(request, session, csvPath); + + // The server reads capabilities from ?caps= at open, before the client + // sends anything, so the page has to put them on the URL it connects to. + const socketUrls: string[] = []; + page.on('websocket', (socket) => socketUrls.push(socket.url())); + await page.goto(`${BASE}/s/${session}`); + await waitForGrid(page); + + const wsUrls = socketUrls.filter((u) => u.includes(`/ws/${session}`)); + expect(wsUrls).toHaveLength(1); + expect(new URL(wsUrls[0]).searchParams.get('caps')).toBe('stats_update'); + }); + test('WebSocket receives data for scrolled rows', async ({ page, request }) => { // Create a larger dataset (100 rows) to force infinite scrolling const rows = []; diff --git a/packages/buckaroo-js-core/pw-tests/stats-update-channel.spec.ts b/packages/buckaroo-js-core/pw-tests/stats-update-channel.spec.ts new file mode 100644 index 000000000..2a1cc6349 --- /dev/null +++ b/packages/buckaroo-js-core/pw-tests/stats-update-channel.spec.ts @@ -0,0 +1,84 @@ +/** + * Playwright tests for the client half of the stats wire (rows-first c2), on + * the StatsUpdateChannel "Manual" story: a real WebSocketModel and BuckarooView + * on a fake in-page server, with stats_update frames delivered by button. + * + * The pinned `mean` row starts as a placeholder because the first frame + * carries only the schema tier. Its stats come in two column chunks, `age` + * then `score` (final), each padding the other column with null. Each test + * checks the grid and the model state the story mirrors (df_meta.stats and the + * stat rows in all_stats). + */ +import { test, expect, Page } from "@playwright/test"; +import { waitForCells } from "./ag-pw-utils"; + +const STORY_URL = + "http://localhost:6006/iframe.html?viewMode=story&id=buckaroo-server-statsupdatechannel--manual&globals=&args="; + +const pinnedCell = (page: Page, rowId: string, colId: string) => + page.locator(`.df-viewer .ag-floating-top .ag-row[row-id="${rowId}"] [col-id="${colId}"]`); + +async function openStory(page: Page) { + await page.goto(STORY_URL); + await waitForCells(page); + await expect(page.locator(".df-viewer .ag-center-cols-container")).toContainText("Bob", { timeout: 10_000 }); + await expect(pinnedCell(page, "main-dtype", "age")).toHaveText("int64"); + await expect(pinnedCell(page, "main-mean", "index")).toHaveText("mean"); + await expect(pinnedCell(page, "main-mean", "age")).toHaveText(""); + await expect(page.getByTestId("model-stats")).toHaveText("pending schema gen 1"); +} + +test("column chunks of a stats_update merge into the pinned rows, the final one completes", async ({ page }) => { + const pageErrors: string[] = []; + page.on("pageerror", (e) => pageErrors.push(e.message)); + await openStory(page); + + // age chunk: age's mean arrives, score stays a placeholder, the gen stays pending. + await page.getByTestId("deliver-age-gen-1").click(); + await expect(pinnedCell(page, "main-mean", "age")).toHaveText("31.00"); + await expect(pinnedCell(page, "main-mean", "score")).toHaveText(""); + await expect(page.getByTestId("model-stats")).toHaveText("pending schema gen 1"); + + // score chunk (final): its null age leaves the merged age alone, and the + // schema-tier dtype row is kept. + await page.getByTestId("deliver-score-gen-1").click(); + await expect(pinnedCell(page, "main-mean", "score")).toHaveText("87.00"); + await expect(page.getByTestId("model-stats")).toHaveText("complete full gen 1"); + await expect(pinnedCell(page, "main-mean", "age")).toHaveText("31.00"); + await expect(pinnedCell(page, "main-dtype", "age")).toHaveText("int64"); + await expect(page.getByTestId("model-stat-rows")).toHaveText("dtype,mean"); + expect(pageErrors).toEqual([]); +}); + +test("stats for the state before a search are dropped, the search's own stats merge", async ({ page }) => { + const pageErrors: string[] = []; + page.on("pageerror", (e) => pageErrors.push(e.message)); + await openStory(page); + + const searchInput = page.locator(".FakeSearchEditor input[type='text']"); + await searchInput.fill("Alice"); + await searchInput.press("Enter"); + + // The server answers the search with gen 2: filtered rows, schema-only stats. + await expect(page.getByTestId("server-gen")).toHaveText("2"); + await expect(page.getByTestId("model-stats")).toHaveText("pending schema gen 2"); + const body = page.locator(".df-viewer .ag-center-cols-container"); + await expect(body).toContainText("Alice"); + await expect(body).not.toContainText("Bob"); + + // gen 1's stats arrive late, after gen 2's initial_state: dropped. + await page.getByTestId("deliver-age-gen-1").click(); + await page.getByTestId("deliver-score-gen-1").click(); + await expect(page.getByTestId("settled")).toHaveText("2"); + await expect(page.getByTestId("model-stats")).toHaveText("pending schema gen 2"); + await expect(page.getByTestId("model-stat-rows")).toHaveText("dtype"); + await expect(pinnedCell(page, "main-mean", "age")).toHaveText(""); + + // gen 2's stats describe Alice alone. + await page.getByTestId("deliver-age-gen-2").click(); + await page.getByTestId("deliver-score-gen-2").click(); + await expect(pinnedCell(page, "main-mean", "age")).toHaveText("25.00"); + await expect(pinnedCell(page, "main-mean", "score")).toHaveText("80.00"); + await expect(page.getByTestId("model-stats")).toHaveText("complete full gen 2"); + expect(pageErrors).toEqual([]); +}); diff --git a/packages/buckaroo-js-core/src/components/BuckarooWidgetInfinite.statsGen.test.ts b/packages/buckaroo-js-core/src/components/BuckarooWidgetInfinite.statsGen.test.ts new file mode 100644 index 000000000..6b4176165 --- /dev/null +++ b/packages/buckaroo-js-core/src/components/BuckarooWidgetInfinite.statsGen.test.ts @@ -0,0 +1,36 @@ +/** + * A row request names the stats_gen of the state it was made for, so the + * server can tell a reply to it from a reply to a request that predates the + * state on screen (the push that follows rows is keyed on it). + */ +import { getKeySmartRowCache } from "./BuckarooWidgetInfinite"; +import { PayloadArgs } from "./DFViewerParts/SmartRowCache"; + +const fakeModel = (expectedGen: number | undefined) => { + const send = jest.fn(); + return { send, model: { send, on: jest.fn(), stats: { expectedGen } } }; +}; + +const request = (sort?: string): PayloadArgs => ({ + sourceName: "default", start: 0, end: 50, origEnd: 50, sort, sort_direction: sort ? "desc" : undefined, +}); + +describe("infinite_request", () => { + it("carries the gen the client is showing, whatever the sort", () => { + const { send, model } = fakeModel(4); + const cache = getKeySmartRowCache(model, jest.fn()); + cache.getRequestRows(request(), jest.fn(), jest.fn()); + cache.getRequestRows(request("b"), jest.fn(), jest.fn()); + expect(send.mock.calls.map(([msg]) => [msg.type, msg.stats_gen])).toEqual([ + ["infinite_request", 4], + ["infinite_request", 4], + ]); + }); + + it("carries no gen when the model has no stats channel (a Jupyter widget)", () => { + const send = jest.fn(); + const cache = getKeySmartRowCache({ send, on: jest.fn() }, jest.fn()); + cache.getRequestRows(request(), jest.fn(), jest.fn()); + expect(send.mock.calls[0][0].stats_gen).toBeUndefined(); + }); +}); diff --git a/packages/buckaroo-js-core/src/components/BuckarooWidgetInfinite.tsx b/packages/buckaroo-js-core/src/components/BuckarooWidgetInfinite.tsx index d126a87c7..0e9d02536 100644 --- a/packages/buckaroo-js-core/src/components/BuckarooWidgetInfinite.tsx +++ b/packages/buckaroo-js-core/src/components/BuckarooWidgetInfinite.tsx @@ -95,7 +95,9 @@ export const getKeySmartRowCache = (model: any, setRespError:any) => { //const symNum = counter(); const reqFn: RequestFN = (pa: PayloadArgs) => { bkLog("model.send infinite_request → Python", { start: pa.start, end: pa.end }); - model.send({ type: 'infinite_request', payload_args: pa }) + // The gen of the state the client is showing: the server starts the stats it owes + // after the reply to a request for that gen, not one made before this state arrived. + model.send({ type: 'infinite_request', payload_args: pa, stats_gen: model.stats?.expectedGen }) } const src = new KeyAwareSmartRowCache(reqFn) diff --git a/packages/buckaroo-js-core/src/components/WidgetTypes.tsx b/packages/buckaroo-js-core/src/components/WidgetTypes.tsx index e7f63015c..678ca0290 100644 --- a/packages/buckaroo-js-core/src/components/WidgetTypes.tsx +++ b/packages/buckaroo-js-core/src/components/WidgetTypes.tsx @@ -1,9 +1,24 @@ +// Where the summary stats stand, as the server reports it in df_meta.stats. +// "pending": a stats_update is expected. "not_computed": none will be sent +// unless the user asks. A missing df_meta.stats means "complete", which is +// what servers that predate the field send. +export type StatsStatus = "complete" | "pending" | "not_computed" | "error"; + +export interface DFMetaStats { + status: StatsStatus; + tier?: string; + reason?: string; + gen?: number; +} + export interface DFMeta { // static, total_rows: number; columns: number; filtered_rows: number; rows_shown: number; + // Absent when the server predates the two-message protocol. + stats?: DFMetaStats; } export interface BuckarooOptions { diff --git a/packages/buckaroo-js-core/src/index.ts b/packages/buckaroo-js-core/src/index.ts index f26c4860d..11d10f022 100644 --- a/packages/buckaroo-js-core/src/index.ts +++ b/packages/buckaroo-js-core/src/index.ts @@ -19,6 +19,7 @@ import { BuckarooServerView, buckarooWsUrl } from './server/BuckarooServerView'; import { BuckarooView } from './server/BuckarooView'; import { WebSocketModel } from './server/WebSocketModel'; import { makeLatestDictDecoder } from './server/latestDictDecoder'; +import { withStatsCapability } from './server/StatsChannel'; import { HistogramCell } from "./components/DFViewerParts/HistogramCell"; import { InfiniteEx } from "./components/DFViewerParts/TableInfinite"; @@ -62,6 +63,7 @@ export default { buckarooWsUrl, WebSocketModel, makeLatestDictDecoder, + withStatsCapability, }; // Named exports for direct imports @@ -92,6 +94,7 @@ export { buckarooWsUrl, WebSocketModel, makeLatestDictDecoder, + withStatsCapability, }; export type { IModel } from './server/IModel'; diff --git a/packages/buckaroo-js-core/src/server/BuckarooServerView.caps.test.tsx b/packages/buckaroo-js-core/src/server/BuckarooServerView.caps.test.tsx new file mode 100644 index 000000000..21e8b8421 --- /dev/null +++ b/packages/buckaroo-js-core/src/server/BuckarooServerView.caps.test.tsx @@ -0,0 +1,84 @@ +/** + * BuckarooServerView — capability advertisement (rows-first c2). + * + * The server records a client's capabilities from `?caps=` on the WebSocket + * URL, because it sends the first message before the client says anything. A + * client that merges `stats_update` must put it there, whatever URL the host + * passes in. + */ +import { render, cleanup, waitFor } from "@testing-library/react"; +import { BuckarooServerView } from "./BuckarooServerView"; + +const capturedViewProps: any[] = []; + +jest.mock("./BuckarooView", () => ({ + BuckarooView: (props: any) => { + capturedViewProps.push(props); + return
; + }, + pickMode: (m: unknown) => (m === "buckaroo" ? "buckaroo" : "viewer"), +})); + +jest.mock("./WebSocketModel", () => ({ + WebSocketModel: class { constructor(_ws: any, _state: any) {} }, +})); + +class FakeWebSocket { + static instances: FakeWebSocket[] = []; + binaryType = "arraybuffer"; + onopen: (() => void) | null = null; + onerror: ((e: any) => void) | null = null; + private listeners: Record void>> = {}; + constructor(public url: string) { + FakeWebSocket.instances.push(this); + setTimeout(() => { + this.onopen?.(); + setTimeout(() => { + const initial = { + type: "initial_state", + df_meta: { total_rows: 4, columns: 2, filtered_rows: 4, rows_shown: 4 }, + df_data_dict: {}, + df_display_args: {}, + mode: "viewer", + }; + this.listeners["message"]?.forEach((h) => h({ data: JSON.stringify(initial) } as any)); + }, 0); + }, 0); + } + addEventListener(ev: string, h: (e: any) => void) { + (this.listeners[ev] ??= new Set()).add(h); + } + removeEventListener(ev: string, h: (e: any) => void) { + this.listeners[ev]?.delete(h); + } + close() {} +} + +const origWebSocket = (globalThis as any).WebSocket; + +beforeAll(() => { + (globalThis as any).WebSocket = FakeWebSocket; +}); +afterAll(() => { + (globalThis as any).WebSocket = origWebSocket; +}); +afterEach(() => { + capturedViewProps.length = 0; + FakeWebSocket.instances.length = 0; + cleanup(); +}); + +describe("BuckarooServerView advertises stats_update", () => { + it("opens the socket with ?caps=stats_update", async () => { + render(); + await waitFor(() => expect(capturedViewProps.length).toBeGreaterThan(0)); + expect(FakeWebSocket.instances).toHaveLength(1); + expect(FakeWebSocket.instances[0].url).toBe("ws://x/ws/s?caps=stats_update"); + }); + + it("keeps the query string the host passed", async () => { + render(); + await waitFor(() => expect(capturedViewProps.length).toBeGreaterThan(0)); + expect(FakeWebSocket.instances[0].url).toBe("ws://x/ws/s?token=abc&caps=stats_update"); + }); +}); diff --git a/packages/buckaroo-js-core/src/server/BuckarooServerView.tsx b/packages/buckaroo-js-core/src/server/BuckarooServerView.tsx index 49fedbe3a..1687dddaa 100644 --- a/packages/buckaroo-js-core/src/server/BuckarooServerView.tsx +++ b/packages/buckaroo-js-core/src/server/BuckarooServerView.tsx @@ -2,6 +2,7 @@ import * as React from "react"; import { decodeDFDataDict } from "../components/DFViewerParts/resolveDFData"; import { WebSocketModel } from "./WebSocketModel"; +import { withStatsCapability } from "./StatsChannel"; import { BuckarooView, BuckarooServerMetadata, @@ -106,7 +107,8 @@ export function BuckarooServerView({ (async () => { try { - ws = new WebSocket(wsUrl); + // The server reads capabilities from the URL at open. + ws = new WebSocket(withStatsCapability(wsUrl)); ws.binaryType = "arraybuffer"; await new Promise((resolve, reject) => { diff --git a/packages/buckaroo-js-core/src/server/BuckarooView.statsDisplay.test.tsx b/packages/buckaroo-js-core/src/server/BuckarooView.statsDisplay.test.tsx new file mode 100644 index 000000000..6ff5533d3 --- /dev/null +++ b/packages/buckaroo-js-core/src/server/BuckarooView.statsDisplay.test.tsx @@ -0,0 +1,67 @@ +/** + * The display config a final stats_update carries reaches the AG-Grid column + * definitions: a float column's minWidth reads its min and max, so the config + * built from the schema tier and the one built from the full stats differ there. + * A real WebSocketModel and the real widget, with AgGridReact stubbed by the spy. + */ +import { act, cleanup, render } from "@testing-library/react"; +import { BuckarooView } from "./BuckarooView"; +import { WebSocketModel } from "./WebSocketModel"; +import { getSpyCalls, resetSpy } from "../test-utils/agGridSpy"; + +jest.mock("ag-grid-react", () => require("../test-utils/agGridSpy").agGridReactMockFactory()); +jest.mock("../components/useColorScheme", () => ({ useColorScheme: () => "light" })); +jest.mock("../components/StatusBar", () => ({ StatusBar: () =>
})); + +class FakeSocket { + readyState = 1; + onmessage: ((e: MessageEvent) => void) | null = null; + send() {} + deliver(msg: object) { + this.onmessage?.({ data: JSON.stringify(msg) } as MessageEvent); + } +} + +const displayArgs = (minWidth: number) => ({ + main: { + data_key: "main", + summary_stats_key: "all_stats", + df_viewer_config: { + pinned_rows: [], + left_col_configs: [], + column_config: [ + { col_name: "a", header_name: "price", displayer_args: { displayer: "obj" }, ag_grid_specs: { minWidth } }, + ], + }, + }, +}); + +const settle = () => act(async () => { await new Promise((resolve) => setTimeout(resolve, 0)); }); + +afterEach(cleanup); + +it("a final stats_update's display config reaches the grid's column definitions without remounting it", async () => { + resetSpy(); + const ws = new FakeSocket(); + const meta = { total_rows: 5, columns: 1, filtered_rows: 5, rows_shown: 5, stats: { status: "pending", tier: "schema", gen: 1 } }; + const initialState = { df_meta: meta, df_data_dict: { main: [], all_stats: [] }, df_display_args: displayArgs(100), + buckaroo_options: { sampled: [], cleaning_method: [""], post_processing: [""], df_display: ["main"], show_commands: [] }, + buckaroo_state: { sampled: false, cleaning_method: false, quick_command_args: {}, post_processing: false, df_display: "main", show_commands: false }, + command_config: { argspecs: {}, defaultArgs: {} } }; + const model = new WebSocketModel(ws as unknown as WebSocket, initialState); + + await act(async () => { + render(); + }); + const minWidth = () => getSpyCalls().lastProps.columnDefs.find((c: any) => c.field === "a")?.minWidth; + expect(minWidth()).toBe(100); + const mounts = getSpyCalls().mountCount; + + ws.deliver({ type: "stats_update", stats_gen: 1, scope: "raw", tier: "full", final: true, + payload: { format: "json", layout: "wide", data: [{ index: "max", level_0: "max", a: 1e9 }] }, + df_display_args: displayArgs(135) }); + await settle(); + + expect(minWidth()).toBe(135); + expect(getSpyCalls().mountCount).toBe(mounts); +}); diff --git a/packages/buckaroo-js-core/src/server/StatsChannel.test.ts b/packages/buckaroo-js-core/src/server/StatsChannel.test.ts new file mode 100644 index 000000000..fb562745a --- /dev/null +++ b/packages/buckaroo-js-core/src/server/StatsChannel.test.ts @@ -0,0 +1,427 @@ +/** + * StatsChannel — the client half of the stats wire (rows-first c2). + * + * A capable client receives a stats-free `initial_state` (df_meta.stats.status + * "pending"), asks for the stats and merges the `stats_update` that answers, + * keyed by `stats_gen`. These tests drive a WebSocketModel with a fake socket, + * the way the server's frames reach it. + */ +import { WebSocketModel } from "./WebSocketModel"; +import { withStatsCapability } from "./StatsChannel"; +import { decodeDFData } from "../components/DFViewerParts/resolveDFData"; + +// A wide summary-stats envelope (parquet_b64, layout "wide") as the server +// sends one; the decoder tests use the same fixture. +// eslint-disable-next-line @typescript-eslint/no-var-requires +const wideFixture = require("../components/DFViewerParts/test-fixtures/summary_stats_parquet_b64.json"); + +// The real decoder, except that an envelope carrying `hold` waits on `gate`, +// so a test can deliver a frame while a decode is in flight. +let gate: Promise = Promise.resolve(); +jest.mock("../components/DFViewerParts/resolveDFData", () => { + const actual = jest.requireActual("../components/DFViewerParts/resolveDFData"); + return { + ...actual, + decodeDFData: jest.fn(async (env: any, buffers?: DataView[]) => { + if (env && env.hold) await gate; + return actual.decodeDFData(env, buffers); + }), + }; +}); + +class FakeSocket { + readyState = 1; // WebSocket.OPEN + onmessage: ((e: MessageEvent) => void) | null = null; + sent: any[] = []; + send(data: string) { + this.sent.push(JSON.parse(data)); + } + deliver(msg: object) { + this.onmessage?.({ data: JSON.stringify(msg) } as MessageEvent); + } +} + +// Lets every promise continuation (the payload decode) run. +const settle = () => new Promise((resolve) => setTimeout(resolve, 0)); + +const row = (stat: string, cells: Record) => ({ index: stat, level_0: stat, ...cells }); + +// What the schema tier ships: identity rows only. +const schemaStats = () => [ + row("dtype", { a: "int64", b: "float64", c: "object" }), + row("length", { a: 3, b: 3, c: 3 }), +]; + +const metaFor = (stats?: Record) => ({ + total_rows: 3, columns: 3, filtered_rows: 3, rows_shown: 3, + ...(stats === undefined ? {} : { stats }), +}); + +const pending = (gen: number) => ({ status: "pending", tier: "schema", gen }); + +const frame = (gen: number | undefined, allStats: any = schemaStats()) => ({ + type: "initial_state", + df_meta: metaFor(gen === undefined ? undefined : pending(gen)), + df_data_dict: { all_stats: allStats }, +}); + +// The server's payload is a wide DFEnvelope; the json format decodes to the +// same row shape without a parquet fixture. +const update = (gen: number, rows: any[], extra: object = {}) => ({ + type: "stats_update", + stats_gen: gen, + scope: "raw", + tier: "full", + final: true, + payload: { format: "json", layout: "wide", data: rows }, + elapsed_ms: 12.5, + ...extra, +}); + +// `gen` null builds a model whose df_meta carries no stats, as an old server's does. +function makeModel(gen: number | null = 3, allStats: any = schemaStats()) { + const ws = new FakeSocket(); + const model = new WebSocketModel(ws as unknown as WebSocket, { + df_meta: metaFor(gen === null ? undefined : pending(gen)), + df_data_dict: { all_stats: allStats }, + }); + const events: { event: string; value: any }[] = []; + for (const key of ["df_data_dict", "df_meta"]) { + model.on(`change:${key}`, (value: any) => events.push({ event: `change:${key}`, value })); + } + return { ws, model, events }; +} + +describe("withStatsCapability", () => { + it("adds ?caps=stats_update to a bare URL", () => { + expect(withStatsCapability("ws://localhost:8700/ws/sales")).toBe("ws://localhost:8700/ws/sales?caps=stats_update"); + }); + + it("appends to an existing query string", () => { + expect(withStatsCapability("ws://h/ws/s?token=abc")).toBe("ws://h/ws/s?token=abc&caps=stats_update"); + }); + + it("extends an existing caps value", () => { + expect(withStatsCapability("ws://h/ws/s?caps=other")).toBe("ws://h/ws/s?caps=other,stats_update"); + }); + + it("leaves a URL that already advertises the capability alone", () => { + expect(withStatsCapability("ws://h/ws/s?caps=stats_update")).toBe("ws://h/ws/s?caps=stats_update"); + expect(withStatsCapability("ws://h/ws/s?caps=a,stats_update")).toBe("ws://h/ws/s?caps=a,stats_update"); + }); + + it("keeps the fragment last", () => { + expect(withStatsCapability("ws://h/ws/s#frag")).toBe("ws://h/ws/s?caps=stats_update#frag"); + }); +}); + +describe("stats_update merge", () => { + it("key-merges the payload's columns into all_stats and keeps the other columns", async () => { + const { ws, model } = makeModel(3); + ws.deliver(update(3, [ + row("length", { a: 3, b: 5 }), + row("mean", { a: 2, b: 4.5 }), + ])); + await settle(); + expect(model.get("df_data_dict").all_stats).toEqual([ + row("dtype", { a: "int64", b: "float64", c: "object" }), + row("length", { a: 3, b: 5, c: 3 }), + row("mean", { a: 2, b: 4.5 }), + ]); + }); + + it("does not let a null in the payload erase a value already merged", async () => { + const { ws, model } = makeModel(3); + // The wide pivot pads a stat a column did not carry with null. + ws.deliver(update(3, [row("dtype", { a: null, b: "float32" })])); + await settle(); + const dtype = model.get("df_data_dict").all_stats.find((r: any) => r.index === "dtype"); + expect(dtype).toEqual(row("dtype", { a: "int64", b: "float32", c: "object" })); + }); + + it("assigns a new df_data_dict and leaves the previous objects untouched", async () => { + const { ws, model, events } = makeModel(3); + const before = model.get("df_data_dict"); + const beforeStats = before.all_stats; + const snapshot = JSON.parse(JSON.stringify(beforeStats)); + ws.deliver(update(3, [row("length", { a: 99 }), row("mean", { a: 2 })])); + await settle(); + const after = model.get("df_data_dict"); + expect(after).not.toBe(before); + expect(after.all_stats).not.toBe(beforeStats); + expect(beforeStats).toEqual(snapshot); + const dictEvents = events.filter((e) => e.event === "change:df_data_dict"); + expect(dictEvents).toHaveLength(1); + expect(dictEvents[0].value).toBe(after); + }); + + it("merges onto an all_stats that arrived as an undecoded envelope", async () => { + const { ws, model } = makeModel(3); + // A later initial_state hands the model the dict as the server sent it. + ws.deliver(frame(4, { format: "json", layout: "wide", data: schemaStats() })); + ws.deliver(update(4, [row("mean", { a: 2, b: 4.5, c: null })])); + await settle(); + const stats = model.get("df_data_dict").all_stats; + expect(Array.isArray(stats)).toBe(true); + expect(stats.map((r: any) => r.index)).toEqual(["dtype", "length", "mean"]); + }); + + it("keeps the other df_data_dict keys", async () => { + const ws = new FakeSocket(); + const model = new WebSocketModel(ws as unknown as WebSocket, { + df_meta: metaFor(pending(3)), + df_data_dict: { all_stats: schemaStats(), empty: [], main: [{ a: 1 }] }, + }); + ws.deliver(update(3, [row("mean", { a: 2 })])); + await settle(); + const dict = model.get("df_data_dict"); + expect(dict.empty).toEqual([]); + expect(dict.main).toEqual([{ a: 1 }]); + expect(dict.all_stats).toHaveLength(3); + }); + + it("merges a wide parquet_b64 payload as the server sends it", async () => { + const { ws, model } = makeModel(3, [row("orig_col_name", { a: "first" })]); + const decoded: any[] = await decodeDFData(wideFixture); + expect(decoded.length).toBeGreaterThan(1); + ws.deliver(update(3, [], { payload: wideFixture })); + await settle(); + const stats = model.get("df_data_dict").all_stats; + expect(stats[0]).toEqual(row("orig_col_name", { a: "first" })); + for (const decodedRow of decoded) { + expect(stats).toContainEqual(decodedRow); + } + }); + + it("builds all_stats when the model holds no dict yet", async () => { + const ws = new FakeSocket(); + const model = new WebSocketModel(ws as unknown as WebSocket, { df_meta: metaFor(pending(3)) }); + ws.deliver(update(3, [row("mean", { a: 2 })])); + await settle(); + expect(model.get("df_data_dict")).toEqual({ all_stats: [row("mean", { a: 2 })] }); + }); + + it("adds all_stats to a dict that has none", async () => { + const ws = new FakeSocket(); + const model = new WebSocketModel(ws as unknown as WebSocket, { + df_meta: metaFor(pending(3)), + df_data_dict: { main: [{ a: 1 }] }, + }); + ws.deliver(update(3, [row("mean", { a: 2 })])); + await settle(); + expect(model.get("df_data_dict")).toEqual({ main: [{ a: 1 }], all_stats: [row("mean", { a: 2 })] }); + }); + + it("applies updates that arrive back to back, in order", async () => { + const { ws, model } = makeModel(3); + ws.deliver(update(3, [row("mean", { a: 2 })], { final: false })); + ws.deliver(update(3, [row("mean", { b: 4.5 }), row("max", { a: 9 })])); + await settle(); + const stats = model.get("df_data_dict").all_stats; + expect(stats.find((r: any) => r.index === "mean")).toEqual(row("mean", { a: 2, b: 4.5 })); + expect(stats.find((r: any) => r.index === "max")).toEqual(row("max", { a: 9 })); + }); +}); + +describe("stats_update order", () => { + it("merges updates in arrival order when an earlier payload decodes slower", async () => { + let release: () => void = () => {}; + gate = new Promise((resolve) => { release = resolve; }); + try { + const { ws, model, events } = makeModel(3); + const held = { format: "json", layout: "wide", data: [row("mean", { a: 1 })], hold: true }; + ws.deliver(update(3, [], { final: false, payload: held })); + ws.deliver(update(3, [row("mean", { a: 2 })])); + await settle(); // the second payload has decoded; the first is held + release(); + await settle(); + // The final update goes last: its value stands and the status + // completes only after both merges. + const mean = model.get("df_data_dict").all_stats.find((r: any) => r.index === "mean"); + expect(mean.a).toBe(2); + expect(events.map((e) => e.event)).toEqual(["change:df_data_dict", "change:df_data_dict", "change:df_meta"]); + } finally { + release(); + gate = Promise.resolve(); + } + }); +}); + +describe("stats_update and df_meta.stats", () => { + it("a final update marks the stats complete at the update's tier", async () => { + const { ws, model, events } = makeModel(3); + const metaBefore = model.get("df_meta"); + ws.deliver(update(3, [row("mean", { a: 2 })])); + await settle(); + const meta = model.get("df_meta"); + expect(meta).not.toBe(metaBefore); + expect(meta.stats).toEqual({ status: "complete", tier: "full", gen: 3 }); + expect(meta.total_rows).toBe(3); + expect(events.filter((e) => e.event === "change:df_meta")).toHaveLength(1); + }); + + it("a non-final update merges but leaves the status pending", async () => { + const { ws, model, events } = makeModel(3); + ws.deliver(update(3, [row("mean", { a: 2 })], { final: false })); + await settle(); + expect(model.get("df_data_dict").all_stats).toHaveLength(3); + expect(model.get("df_meta").stats.status).toBe("pending"); + expect(events.filter((e) => e.event === "change:df_meta")).toHaveLength(0); + }); +}); + +describe("stats_update and df_display_args", () => { + it("a final update carrying df_display_args replaces the model's, and one without leaves them", async () => { + const { ws, model } = makeModel(3); + const before = { main: { df_viewer_config: { column_config: [{ col_name: "a" }] } } }; + const after = { main: { df_viewer_config: { column_config: [{ col_name: "a", ag_grid_specs: { minWidth: 114 } }] } } }; + model.set("df_display_args", before); + + ws.deliver(update(3, [row("mean", { a: 2 })])); + await settle(); + expect(model.get("df_display_args")).toBe(before); + + ws.deliver(update(3, [row("max", { a: 9 })], { df_display_args: after })); + await settle(); + expect(model.get("df_display_args")).toEqual(after); + }); + + it("drops the df_display_args of an update for a gen the client has left", async () => { + const { ws, model } = makeModel(3); + const before = { main: {} }; + model.set("df_display_args", before); + ws.deliver(update(2, [row("mean", { a: 2 })], { df_display_args: { main: { stale: true } } })); + await settle(); + expect(model.get("df_display_args")).toBe(before); + }); +}); + +describe("stats_gen", () => { + it("drops a stats_update whose stats_gen is not the expected one", async () => { + const { ws, model, events } = makeModel(3); + ws.deliver(update(2, [row("mean", { a: 2 })])); + ws.deliver(update(3, [row("max", { a: 9 })])); + await settle(); + expect(model.get("df_data_dict").all_stats.map((r: any) => r.index)).toEqual(["dtype", "length", "max"]); + expect(events.filter((e) => e.event === "change:df_data_dict")).toHaveLength(1); + }); + + it("drops every stats_update when the server reported no stats", async () => { + const { ws, model, events } = makeModel(null); + ws.deliver(update(0, [row("mean", { a: 2 })])); + await settle(); + expect(model.get("df_data_dict").all_stats).toHaveLength(2); + expect(events).toHaveLength(0); + }); + + it("clears the expectation when an initial_state carries no df_meta.stats", async () => { + const { ws, model } = makeModel(3); + ws.deliver(frame(undefined)); + ws.deliver(update(3, [row("mean", { a: 2 })])); + await settle(); + expect(model.get("df_data_dict").all_stats).toHaveLength(2); + }); + + it("advances the expected gen on a broadcast initial_state with no reply_seq", async () => { + const { ws, model } = makeModel(3); + ws.deliver(frame(4)); + ws.deliver(update(3, [row("mean", { a: 2 })])); + await settle(); + expect(model.get("df_data_dict").all_stats).toHaveLength(2); + ws.deliver(update(4, [row("mean", { a: 2 })])); + await settle(); + expect(model.get("df_data_dict").all_stats).toHaveLength(3); + }); + + it("discards a merge when an initial_state for a newer gen arrives while it decodes", async () => { + const { ws, model } = makeModel(3); + ws.deliver(update(3, [row("mean", { a: 2 })])); + // Still decoding: the model has not merged anything yet. + ws.deliver(frame(4, [row("dtype", { a: "int32" })])); + await settle(); + expect(model.get("df_data_dict").all_stats).toEqual([row("dtype", { a: "int32" })]); + expect(model.get("df_meta").stats).toEqual(pending(4)); + ws.deliver(update(4, [row("mean", { a: 2 })])); + await settle(); + expect(model.get("df_data_dict").all_stats).toEqual([row("dtype", { a: "int32" }), row("mean", { a: 2 })]); + }); + + it("merges onto the new dict when a same-gen initial_state replaces it while the old one decodes", async () => { + let release: () => void = () => {}; + gate = new Promise((resolve) => { release = resolve; }); + try { + const held = { format: "json", layout: "wide", data: schemaStats(), hold: true }; + const { ws, model } = makeModel(3, held); + ws.deliver(update(3, [row("mean", { a: 2 })])); + await settle(); // the update is now waiting on the held decode + expect((decodeDFData as jest.Mock).mock.calls.some(([env]) => env === held)).toBe(true); + ws.deliver(frame(3, [row("dtype", { a: "int32" })])); + release(); + await settle(); + expect(model.get("df_data_dict").all_stats).toEqual([ + row("dtype", { a: "int32" }), + row("mean", { a: 2 }), + ]); + expect(model.get("df_meta").stats).toEqual({ status: "complete", tier: "full", gen: 3 }); + } finally { + release(); + gate = Promise.resolve(); + } + }); +}); + +describe("stats_aborted", () => { + it("marks the stats failed when the run for the expected gen failed", async () => { + const { ws, model } = makeModel(3); + ws.deliver({ type: "stats_aborted", stats_gen: 3, current_gen: 3, scope: "raw", reason: "error" }); + await settle(); + expect(model.get("df_meta").stats).toEqual({ status: "error", tier: "schema", gen: 3, reason: "stats_failed" }); + }); + + it("marks the stats not computed when the server says they cannot be requested", async () => { + const { ws, model } = makeModel(3); + ws.deliver({ type: "stats_aborted", stats_gen: 3, current_gen: 3, scope: "raw", reason: "not_requestable" }); + await settle(); + expect(model.get("df_meta").stats.status).toBe("not_computed"); + }); + + it("changes nothing for a stale reply, an unsupported scope or a session with no data", async () => { + const { ws, events } = makeModel(3); + for (const reason of ["stale", "unsupported_scope", "no_data"]) { + ws.deliver({ type: "stats_aborted", stats_gen: 3, current_gen: 4, scope: "raw", reason }); + } + await settle(); + expect(events).toHaveLength(0); + }); + + it("ignores a reply for a gen the client has left", async () => { + const { ws, model } = makeModel(3); + ws.deliver({ type: "stats_aborted", stats_gen: 2, current_gen: 3, scope: "raw", reason: "error" }); + await settle(); + expect(model.get("df_meta").stats.status).toBe("pending"); + }); +}); + +describe("other messages", () => { + it("ignores a type it does not know and keeps applying frames", async () => { + const { ws, model, events } = makeModel(3); + expect(() => ws.deliver({ type: "something_new", stats_gen: 3 })).not.toThrow(); + ws.deliver({ type: "error", message: "boom" }); + await settle(); + expect(events).toHaveLength(0); + ws.deliver(frame(4, [row("dtype", { a: "int32" })])); + expect(model.get("df_data_dict").all_stats).toEqual([row("dtype", { a: "int32" })]); + expect(model.get("df_meta").stats).toEqual(pending(4)); + }); + + it("still pairs an infinite_resp with the binary frame that follows it", () => { + const { ws, model } = makeModel(3); + const seen: [any, DataView[]][] = []; + model.on("msg:custom", (msg: any, buffers: DataView[]) => seen.push([msg, buffers])); + ws.deliver({ type: "infinite_resp", key: { start: 0, end: 5 }, length: 5 }); + ws.onmessage?.({ data: new ArrayBuffer(8) } as MessageEvent); + expect(seen).toHaveLength(1); + expect(seen[0][0].type).toBe("infinite_resp"); + expect(seen[0][1][0].byteLength).toBe(8); + }); +}); diff --git a/packages/buckaroo-js-core/src/server/StatsChannel.ts b/packages/buckaroo-js-core/src/server/StatsChannel.ts new file mode 100644 index 000000000..677fecd3d --- /dev/null +++ b/packages/buckaroo-js-core/src/server/StatsChannel.ts @@ -0,0 +1,215 @@ +/** + * StatsChannel — the client half of the stats wire (rows-first c2). + * + * A client that advertises `?caps=stats_update` gets a first `initial_state` + * whose `df_meta.stats` says the stats are pending, asks for them with + * `stats_request {stats_gen, scope}`, and receives either + * + * stats_update {stats_gen, scope, tier, final, payload, elapsed_ms} + * stats_aborted {stats_gen, current_gen?, scope, reason} + * + * `payload` is an inline wide DFEnvelope holding `all_stats`. `stats_gen` is the + * server's counter for the state the stats describe; it rides on every + * `initial_state` as `df_meta.stats.gen`, and a reply for any other gen is for + * a state the client has left. Sending the request is the scheduler's job, not + * this module's. + * + * Merge semantics are WebSocket-only: `WebSocketModel` hands every frame to + * `handle()` first, and Jupyter's widget sets `df_data_dict` whole. + */ +import { decodeDFData } from "../components/DFViewerParts/resolveDFData"; +import { DFData, DFDataOrPayload } from "../components/DFViewerParts/DFWhole"; +import { DFMeta, DFMetaStats } from "../components/WidgetTypes"; +import { IModel } from "./IModel"; + +/** The capability this client advertises, as one value of `?caps=` on the + * WebSocket URL: it merges `stats_update` messages. The server records it per + * connection when the socket opens, since it sends the first message before + * the client can say anything. */ +export const STATS_UPDATE_CAP = "stats_update"; + +const decodeQueryValue = (value: string): string => { + try { + return decodeURIComponent(value); + } catch { + return value; + } +}; + +/** `wsUrl` with `caps=stats_update` added: a new query on a bare URL, a new + * parameter after an existing query, or a comma-joined value when the host + * already passes `caps`. The fragment stays last and other parameters are + * left as the host wrote them. */ +export function withStatsCapability(wsUrl: string): string { + const hashAt = wsUrl.indexOf("#"); + const fragment = hashAt === -1 ? "" : wsUrl.slice(hashAt); + const beforeFragment = hashAt === -1 ? wsUrl : wsUrl.slice(0, hashAt); + const queryAt = beforeFragment.indexOf("?"); + const path = queryAt === -1 ? beforeFragment : beforeFragment.slice(0, queryAt); + const params = queryAt === -1 ? [] : beforeFragment.slice(queryAt + 1).split("&").filter((p) => p !== ""); + + const capsAt = params.findIndex((p) => p === "caps" || p.startsWith("caps=")); + if (capsAt === -1) { + params.push(`caps=${STATS_UPDATE_CAP}`); + } else { + const caps = decodeQueryValue(params[capsAt].slice("caps=".length)) + .split(",") + .map((cap) => cap.trim()) + .filter((cap) => cap !== ""); + if (!caps.includes(STATS_UPDATE_CAP)) caps.push(STATS_UPDATE_CAP); + params[capsAt] = `caps=${caps.join(",")}`; + } + return `${path}?${params.join("&")}${fragment}`; +} + +export interface StatsUpdateMessage { + type: "stats_update"; + stats_gen: number; + scope?: string; + tier?: string; + final?: boolean; + payload?: DFDataOrPayload; + /** The display config the stats change (a float column's minWidth reads + * its min and max), present on a final update when it differs from the + * config the client was sent while the stats were pending. */ + df_display_args?: Record; + elapsed_ms?: number; +} + +export interface StatsAbortedMessage { + type: "stats_aborted"; + /** The request's gen. */ + stats_gen?: number; + /** The session's gen, omitted when the session has no data. */ + current_gen?: number; + scope?: string; + reason?: "stale" | "unsupported_scope" | "not_requestable" | "error" | "no_data"; +} + +/** + * Key-merge a stats payload into `all_stats`. Both are row-per-stat tables + * (`{index: , : , ...}`). For each stat row the payload + * names, every column it carries replaces that cell, columns it does not carry + * keep theirs, and a stat the table lacks is appended. + * + * A `null` in the payload never replaces a value already merged: the wide + * pivot fills a stat a column did not carry in this message with `null`. + * + * Returns new row objects and a new array. `base` may be the decoder's cached + * array, so nothing in it is touched. + */ +export function mergeStatRows(base: DFData, update: DFData): DFData { + const merged = base.slice(); + const at = new Map(); + merged.forEach((row, i) => at.set(row.index, i)); + + for (const updateRow of update) { + const i = at.get(updateRow.index); + if (i === undefined) { + at.set(updateRow.index, merged.length); + merged.push({ ...updateRow }); + continue; + } + const row = { ...merged[i] }; + for (const [column, value] of Object.entries(updateRow)) { + if (column === "index" || column === "level_0") continue; + if (value === null && row[column] != null) continue; + row[column] = value; + } + merged[i] = row; + } + return merged; +} + +const genOf = (meta: DFMeta | undefined): number | undefined => { + const gen = meta?.stats?.gen; + return typeof gen === "number" ? gen : undefined; +}; + +/** What the channel needs of a model. */ +type StatsModel = Pick; + +export class StatsChannel { + // Updates are applied one at a time: each reads the dict the previous one + // wrote, so a later update cannot overwrite an earlier one's merge. + private applying: Promise = Promise.resolve(); + + constructor(private model: StatsModel) {} + + /** + * The stats_gen of the state the client is showing, or `undefined` when the + * server reports none; a `stats_request` carries it. It is read off the + * model's `df_meta` each time, so it starts from the frame the model was + * built from and every applied `initial_state` (a broadcast frame with no + * `reply_seq` included) moves it. A `df_meta` with no `stats.gen` leaves + * nothing expected, so a server that stops reporting stats cannot have a + * late `stats_update` merged. Reading the model, not the message, keeps this + * correct under any scheme that drops stale `initial_state` frames. + */ + get expectedGen(): number | undefined { + return genOf(this.model.get("df_meta")); + } + + /** Consume a stats message. Returns false for every other message type, + * which the model handles (or ignores) as before. */ + handle(msg: { type?: string }): boolean { + if (msg.type === "stats_update") { + this.receiveUpdate(msg as StatsUpdateMessage); + return true; + } + if (msg.type === "stats_aborted") { + this.receiveAborted(msg as StatsAbortedMessage); + return true; + } + return false; + } + + private receiveUpdate(msg: StatsUpdateMessage): void { + if (msg.stats_gen !== this.expectedGen) return; + this.applying = this.applying + .then(() => this.applyUpdate(msg)) + .catch((e) => console.error("[StatsChannel] stats_update failed:", e)); + } + + private async applyUpdate(msg: StatsUpdateMessage): Promise { + const update = await decodeDFData(msg.payload); + for (;;) { + // A frame may have moved the gen on while something decoded. + if (msg.stats_gen !== this.expectedGen) return; + const dict: Record | null | undefined = this.model.get("df_data_dict"); + // The dict is decoded when the seed built it and raw when a later + // initial_state did. + const base = await decodeDFData(dict?.all_stats); + // A frame replaced the dict (or moved the gen) while it decoded: + // start over from the new one. + if (dict !== this.model.get("df_data_dict") || msg.stats_gen !== this.expectedGen) continue; + this.model.set("df_data_dict", { ...dict, all_stats: mergeStatRows(base, update) }); + if (msg.df_display_args) this.model.set("df_display_args", msg.df_display_args); + if (msg.final) { + this.replaceStats((stats) => ({ status: "complete", tier: msg.tier ?? stats.tier, gen: stats.gen })); + } + return; + } + } + + private receiveAborted(msg: StatsAbortedMessage): void { + // Only a reply to a request for the state on screen says anything about + // it. A `stale` reply is answered by the initial_state that carries the + // new gen; taking `current_gen` without that frame would merge stats + // for a state the client has not seen into the one it shows. + if (msg.stats_gen !== this.expectedGen) return; + if (msg.reason === "error") { + this.replaceStats((stats) => ({ ...stats, status: "error", reason: "stats_failed" })); + } else if (msg.reason === "not_requestable") { + this.replaceStats((stats) => ({ ...stats, status: "not_computed" })); + } + } + + /** Replace `df_meta.stats` in a new `df_meta`, the reference that c0a's + * `inFlight` rule and the pinned rows' placeholders key on. */ + private replaceStats(change: (stats: DFMetaStats) => DFMetaStats): void { + const meta: DFMeta | undefined = this.model.get("df_meta"); + if (!meta?.stats) return; + this.model.set("df_meta", { ...meta, stats: change(meta.stats) }); + } +} diff --git a/packages/buckaroo-js-core/src/server/WebSocketModel.ts b/packages/buckaroo-js-core/src/server/WebSocketModel.ts index 8c73b557c..9fbc7da8f 100644 --- a/packages/buckaroo-js-core/src/server/WebSocketModel.ts +++ b/packages/buckaroo-js-core/src/server/WebSocketModel.ts @@ -23,7 +23,11 @@ * 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. + * + * stats_update and stats_aborted frames go to `stats` (see StatsChannel). */ +import { StatsChannel } from "./StatsChannel"; + export class WebSocketModel { private ws: WebSocket; private pendingMsg: any = null; @@ -31,14 +35,17 @@ export class WebSocketModel { private state: Record; private pendingChanges: Set = new Set(); private stateSeq = 0; + readonly stats: StatsChannel; constructor(ws: WebSocket, initialState: Record) { this.state = { ...initialState }; this.ws = ws; + this.stats = new StatsChannel(this); this.ws.onmessage = (event: MessageEvent) => { if (typeof event.data === "string") { const msg = JSON.parse(event.data); + if (this.stats.handle(msg)) return; if (msg.type === "infinite_resp") { // Expect a following binary frame — stash this JSON diff --git a/packages/buckaroo-js-core/src/stories/StatsUpdateChannel.stories.tsx b/packages/buckaroo-js-core/src/stories/StatsUpdateChannel.stories.tsx new file mode 100644 index 000000000..2f44143ab --- /dev/null +++ b/packages/buckaroo-js-core/src/stories/StatsUpdateChannel.stories.tsx @@ -0,0 +1,278 @@ +/** + * The stats wire from the client side (rows-first c2), with no server. + * + * A real WebSocketModel renders through BuckarooView in buckaroo mode, but its + * socket is FakeStatsServer, which answers the way the Python server does: + * + * - infinite_request → infinite_resp (JSON frame, then a binary frame) after + * rowDelayMs, filtered by the server's current search. + * - buckaroo_state_change → initial_state after stateDelayMs, carrying a new + * stats gen whose df_meta.stats is pending and whose all_stats holds only + * the schema-tier dtype row. + * + * stats_update frames go out only when a "deliver" button is clicked, one pair + * of buttons per gen the server has issued, so the order of the stats against + * the other frames is up to the person (or test) clicking. Each gen's `mean` + * row comes in two column chunks, the way a column-chunked tier sends it: `age` + * (not final), then `score` (final). Each chunk pads the column it does not + * carry with null, as the wide pivot does, so the `score` chunk also checks + * that a null leaves an already-merged value alone. Used by + * stats-update-channel.spec.ts. + */ +import type { Meta, StoryObj } from "@storybook/react"; +import React, { useEffect, useMemo, useState } from "react"; +import { DFData, DFViewerConfig } from "../components/DFViewerParts/DFWhole"; +import { IDisplayArgs } from "../components/DFViewerParts/gridUtils"; +import { PayloadArgs } from "../components/DFViewerParts/SmartRowCache"; +import { BuckarooOptions, BuckarooState, DFMeta, DFMetaStats } from "../components/WidgetTypes"; +import { BuckarooView } from "../server/BuckarooView"; +import { WebSocketModel } from "../server/WebSocketModel"; + +const NAMES = ["Alice", "Bob", "Charlie", "Diana", "Eve"]; + +const allRows: DFData = NAMES.map((name, i) => ({ + index: i, + name, + age: 25 + i * 3, + score: 80 + i * 3.5, +})); + +const mean = (rows: DFData, col: "age" | "score") => + rows.length === 0 ? null : rows.reduce((acc, r) => acc + (r[col] as number), 0) / rows.length; + +const schemaStats = (): DFData => [{ index: "dtype", name: "object", age: "int64", score: "float64" }]; + +const floatArgs = { displayer: "float", min_fraction_digits: 2, max_fraction_digits: 2 } as const; + +const viewerConfig: DFViewerConfig = { + column_config: [ + { col_name: "name", header_name: "name", displayer_args: { displayer: "obj" } }, + { col_name: "age", header_name: "age", displayer_args: { displayer: "integer", min_digits: 1, max_digits: 5 } }, + { col_name: "score", header_name: "score", displayer_args: floatArgs }, + ], + left_col_configs: [{ col_name: "index", header_name: "index", displayer_args: { displayer: "obj" } }], + pinned_rows: [ + { primary_key_val: "dtype", displayer_args: { displayer: "obj" } }, + { primary_key_val: "mean", displayer_args: floatArgs }, + ], +}; + +const dfDisplayArgs: Record = { + main: { data_key: "main", df_viewer_config: viewerConfig, summary_stats_key: "all_stats" }, +}; + +const buckarooOptions: BuckarooOptions = { + sampled: [], + cleaning_method: [], + post_processing: [], + df_display: ["main"], + show_commands: [], +}; + +const initialBuckarooState: BuckarooState = { + sampled: false, + cleaning_method: false, + quick_command_args: {}, + post_processing: false, + df_display: "main", + show_commands: false, +}; + +const metaFor = (rows: DFData, gen: number): DFMeta => ({ + total_rows: allRows.length, + columns: 3, + filtered_rows: rows.length, + rows_shown: rows.length, + stats: { status: "pending", tier: "schema", gen }, +}); + +type StatsChunk = "age" | "score"; + +interface FakeServerOptions { + rowDelayMs: number; + stateDelayMs: number; + // Called whenever the server issues a new gen, so the controls can re-render. + onGen: (gen: number) => void; +} + +/** The socket WebSocketModel is built on: it takes the client's frames through + * send() and answers through the onmessage handler the model installs. */ +class FakeStatsServer { + readyState = 1; // WebSocket.OPEN + onmessage: ((e: MessageEvent) => void) | null = null; + gen = 1; + private rowsByGen = new Map([[1, allRows]]); + + constructor(private opts: FakeServerOptions) {} + + get rows(): DFData { + return this.rowsByGen.get(this.gen) ?? []; + } + + send(data: string): void { + const msg = JSON.parse(data); + if (msg.type === "infinite_request") { + this.answerRows(msg.payload_args as PayloadArgs, this.rows); + } else if (msg.type === "buckaroo_state_change") { + this.answerStateChange(msg.new_state as BuckarooState); + } + } + + /** Sends one column chunk of `gen`'s stats as a stats_update, the way the + * server answers a stats_request. The `score` chunk is the final one. */ + deliverStats(gen: number, column: StatsChunk): void { + const rows = this.rowsByGen.get(gen) ?? []; + const meanRow = { + index: "mean", + name: null, + age: column === "age" ? mean(rows, "age") : null, + score: column === "score" ? mean(rows, "score") : null, + }; + const final = column === "score"; + this.push({ + type: "stats_update", + stats_gen: gen, + scope: "raw", + tier: "full", + final, + payload: { format: "json", data: [meanRow] }, + elapsed_ms: 1, + }); + } + + private answerRows(pa: PayloadArgs, rows: DFData): void { + setTimeout(() => { + this.push({ + type: "infinite_resp", + key: pa, + length: rows.length, + payload: { format: "json", data: rows.slice(pa.start, Math.min(pa.end, rows.length)) }, + }); + // The model pairs every infinite_resp with the binary frame after it. + this.onmessage?.({ data: new ArrayBuffer(0) } as MessageEvent); + }, this.opts.rowDelayMs); + } + + private answerStateChange(state: BuckarooState): void { + const term = String(state.quick_command_args?.search?.[0] ?? "").toLowerCase(); + const rows = allRows.filter((r) => String(r.name).toLowerCase().includes(term)); + this.gen += 1; + const gen = this.gen; + this.rowsByGen.set(gen, rows); + this.opts.onGen(gen); + setTimeout(() => { + this.push({ + type: "initial_state", + df_meta: metaFor(rows, gen), + df_data_dict: { main: [], all_stats: schemaStats(), empty: [] }, + buckaroo_state: state, + }); + }, this.opts.stateDelayMs); + } + + private push(msg: object): void { + this.onmessage?.({ data: JSON.stringify(msg) } as MessageEvent); + } +} + +const describeStats = (stats?: DFMetaStats) => + stats ? `${stats.status} ${stats.tier ?? "-"} gen ${stats.gen ?? "-"}` : "none"; + +interface StatsUpdateChannelProps { + // How long the server takes to answer a row request. + rowDelayMs: number; + // How long the server takes to answer a state change (a search). + stateDelayMs: number; +} + +const StatsUpdateChannel: React.FC = ({ rowDelayMs, stateDelayMs }) => { + const [serverGen, setServerGen] = useState(1); + const [modelStats, setModelStats] = useState(undefined); + const [statRows, setStatRows] = useState([]); + const [settled, setSettled] = useState(0); + + const { server, model, initialState } = useMemo(() => { + const server = new FakeStatsServer({ rowDelayMs, stateDelayMs, onGen: setServerGen }); + const initialState = { + df_meta: metaFor(allRows, 1), + df_data_dict: { main: [], all_stats: schemaStats(), empty: [] }, + df_display_args: dfDisplayArgs, + buckaroo_state: initialBuckarooState, + buckaroo_options: buckarooOptions, + }; + const model = new WebSocketModel(server as unknown as WebSocket, initialState); + return { server, model, initialState }; + }, [rowDelayMs, stateDelayMs]); + + // Mirror what the model holds, not what the grid renders, so a dropped + // update can be told from one that has not rendered yet. + useEffect(() => { + const onMeta = (meta: DFMeta) => setModelStats(meta?.stats); + const onDict = (dict: Record) => { + const allStats = dict?.all_stats; + setStatRows(Array.isArray(allStats) ? allStats.map((r) => String(r.index)) : []); + }; + onMeta(model.get("df_meta")); + onDict(model.get("df_data_dict")); + model.on("change:df_meta", onMeta); + model.on("change:df_data_dict", onDict); + return () => { + model.off("change:df_meta", onMeta); + model.off("change:df_data_dict", onDict); + }; + }, [model]); + + // A json payload merges within microtasks, so once a macrotask has run after + // a delivery the model holds whatever the delivery did; "settled" counts + // those deliveries. + const deliver = (gen: number, column: StatsChunk) => { + server.deliverStats(gen, column); + setTimeout(() => setSettled((n) => n + 1), 0); + }; + + const gens = Array.from({ length: serverGen }, (_, i) => i + 1); + return ( +
+
+
+ server gen: {serverGen} + {" · "}model df_meta.stats: {describeStats(modelStats)} + {" · "}model all_stats rows: {statRows.join(",")} + {" · "}deliveries settled: {settled} +
+
+ {gens.map((gen) => ( + + + + + ))} +
+
+
+ +
+
+ ); +}; + +const meta = { + title: "Buckaroo/Server/StatsUpdateChannel", + component: StatsUpdateChannel, + parameters: { layout: "centered" }, + argTypes: { + rowDelayMs: { control: { type: "range", min: 0, max: 3000, step: 50 } }, + stateDelayMs: { control: { type: "range", min: 0, max: 3000, step: 50 } }, + }, +} satisfies Meta; + +export default meta; +type Story = StoryObj; + +export const Manual: Story = { + args: { rowDelayMs: 10, stateDelayMs: 10 }, +}; diff --git a/packages/js/standalone.tsx b/packages/js/standalone.tsx index 37f191c31..d7ccfe48f 100644 --- a/packages/js/standalone.tsx +++ b/packages/js/standalone.tsx @@ -287,7 +287,8 @@ async function main() { // Show connecting state rootEl.textContent = "Connecting..."; - const ws = new WebSocket(wsUrl); + // The server reads capabilities from the URL at open. + const ws = new WebSocket(srt.withStatsCapability(wsUrl)); ws.binaryType = "arraybuffer"; // Wait for connection diff --git a/scripts/test_playwright_storybook.sh b/scripts/test_playwright_storybook.sh index 6d835ab2a..fcae548ae 100755 --- a/scripts/test_playwright_storybook.sh +++ b/scripts/test_playwright_storybook.sh @@ -96,6 +96,7 @@ STORYBOOK_TESTS=( "pw-tests/outside-params.spec.ts" "pw-tests/fitcontent-height.spec.ts" "pw-tests/stats-pending-pinned-rows.spec.ts" + "pw-tests/stats-update-channel.spec.ts" # "pw-tests/example.spec.ts" # Has pre-existing failures, excluded for now ) 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..688155d81 100644 --- a/tests/unit/dataflow/customizable_dataflow_test.py +++ b/tests/unit/dataflow/customizable_dataflow_test.py @@ -21,6 +21,11 @@ def always_fail(ser: RawSeries) -> int: raise ZeroDivisionError("always fails") +@stat() +def double_length(length: int) -> int: + return length * 2 + + BASIC_DF = get_basic_df2() EMPTY_DF_JSON = { 'dfviewer_config': { @@ -75,6 +80,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 @@ -276,6 +300,14 @@ class PostDCFC(ACDFC): assert p_dfc.cleaned_sd == {} assert p_dfc.df_display_args['main']['df_viewer_config'] == DFVIEWER_CONFIG_DEFAULT +def test_add_analysis_recomputes_the_summary_stats(): + """The summary stats rerun with the new klass list, so the added stat + reaches merged_sd.""" + bw = BuckarooWidget(BASIC_DF) + bw.add_analysis(double_length) + for col in ('a', 'b'): + assert bw.dataflow.merged_sd[col]['double_length'] == 2 * bw.dataflow.merged_sd[col]['length'] + class HidePostProcessingAnalysis2(ColAnalysis): provides_defaults = {} diff --git a/tests/unit/dataflow/scoped_summary_stats_test.py b/tests/unit/dataflow/scoped_summary_stats_test.py index b5648ab56..1e990c58e 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 @@ -22,7 +24,8 @@ from buckaroo.customizations.pd_autoclean_conf import NoCleaningConf 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.dataflow import CustomizableDataflow, StylingAnalysis, assemble_merged_sd +from buckaroo.dataflow.sd_cache import split_chain_by_scope from buckaroo.pluggable_analysis_framework.stat_func import stat @@ -249,3 +252,64 @@ 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.""" + 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(): + 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(): + dfc = _assembly_dataflow() + inputs = _scope_inputs(dfc) + sd_names = ('init_sd', 'cleaned_sd', 'raw_sd', 'processed_sd', 'clean_sd', 'filt_sd') + + def key_sets(): + return {name: {col: sorted(stats) for col, stats in inputs[name].items()} for name in sd_names} + + before, chains_before = key_sets(), copy.deepcopy(inputs['chains']) + + assemble_merged_sd(**inputs) + + # Layering cleaned_* / filtered_* into a result column that shared its dict + # with an input would add keys to that input. + assert key_sets() == before + assert inputs['chains'] == chains_before diff --git a/tests/unit/server/test_load_expr.py b/tests/unit/server/test_load_expr.py index 1d6a2ea93..86dd46331 100644 --- a/tests/unit/server/test_load_expr.py +++ b/tests/unit/server/test_load_expr.py @@ -1,5 +1,6 @@ """End-to-end tests for POST /load_expr — server load path for XorqBuckarooInfiniteWidget over a xorq/ibis expression.""" +import datetime import gc import io import json @@ -8,15 +9,19 @@ 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 +from traitlets import TraitError xo = pytest.importorskip("xorq.api") @@ -28,9 +33,19 @@ from xorq.common.utils.provenance_utils import read_parquet_provenance # noqa: E402 from xorq.expr.relations import CachedNode # noqa: E402 from xorq.vendor.ibis.expr import operations as ops # noqa: E402 - +from xorq.vendor.ibis.expr.types.core import Expr # noqa: E402 + +from buckaroo.dataflow.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 import perf_log # noqa: E402 +from buckaroo.pluggable_analysis_framework.col_analysis import ColAnalysis # noqa: E402 +from buckaroo.pluggable_analysis_framework.xorq_stat_pipeline import XorqStatPipeline # noqa: E402 +from buckaroo.serialization_utils import resolve_summary_stats_payload # noqa: E402 from buckaroo.server import telemetry, xorq_loading # noqa: E402 from buckaroo.server.app import make_app as _make_app # noqa: E402 +from buckaroo.server.session import SessionState, begin_stats_generation # noqa: E402 +from buckaroo.server.stats_wire import complete_stats # noqa: E402 +from buckaroo.server.websocket_handler import DataStreamHandler # noqa: E402 pytestmark = pytest.mark.skipif( sys.platform == "win32", @@ -1065,6 +1080,1260 @@ async def test_reload_releases_previous_memtables(self): shutil.rmtree(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.""" + 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 _V1Analysis(ColAnalysis): + """Fails DAG validation: v1 ColAnalysis classes are no longer supported.""" + provides_defaults = {"foo": 1} + + +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_an_op_under_a_search_counts_only_the_filtered_frame(self, monkeypatch): + """The clean scope's frame is rebuilt on every cache miss, so counting + it is a query of its own. The filtered frame's count is one df_meta + runs anyway.""" + dataflow = _build_dataflow(stats_tier="schema") + dataflow.quick_command_args = {"search": ["a"]} + ops, _stat_queries = _spy_data_queries(monkeypatch) + dataflow.operations = [[lisp_sym("fillna"), {"symbol": "df"}, "qty", 0]] + assert ops == ["CountStar"] + + def test_add_analysis_that_fails_validation_keeps_the_klasses(self): + dataflow = _build_dataflow(stats_tier="schema") + before = list(dataflow.analysis_klasses) + with pytest.raises(TypeError): + dataflow.add_analysis(_V1Analysis) + assert dataflow.analysis_klasses == before + dataflow.add_analysis(_NoopPostProcessing) + assert dataflow.analysis_klasses == [*before, _NoopPostProcessing] + + 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_switching_the_tier_computes_its_stats_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" + + assert _as_json(dataflow.merged_sd) == _as_json(full.merged_sd) + + def test_an_unknown_tier_is_rejected(self): + dataflow = _build_dataflow(stats_tier="schema") + with pytest.raises(TraitError): + dataflow.stats_tier = "Full" + assert dataflow.stats_tier == "schema" + + def test_set_stats_tier_installs_the_summary_it_is_given(self, monkeypatch): + """Stats computed off the IOLoop go in as they are: the filtered scope + is not computed again, only the scopes the cascade still lacks.""" + expr = _stats_tier_expr() + full = _three_scope_dataflow(expr) + dataflow = _three_scope_dataflow(expr, stats_tier="schema") + scopes = [] + original = dataflow._get_summary_sd + monkeypatch.setattr(dataflow, "_get_summary_sd", + lambda processed_df, scope="filt": scopes.append(scope) or original(processed_df, scope=scope)) + + dataflow.set_stats_tier("full", (full.summary_sd, {})) + + assert scopes == ["raw", "clean"] + assert dataflow.stats_tier == "full" + 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.set_stats_tier("full") + + 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): + 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_a_load_in_between_resets_the_pair(self): + """/load replaces the session's xorq state, so a later /load_expr that + omits the pair starts from the defaults, not from an earlier load's.""" + builds_root = tempfile.mkdtemp() + csv_fd, csv_path = tempfile.mkstemp(suffix=".csv") + os.close(csv_fd) + try: + build_path = _build_expr_dir(builds_root) + pd.DataFrame({"x": [1, 2]}).to_csv(csv_path, index=False) + sid = "sp-load-between" + await _post(self.get_http_port(), "/load_expr", + {"session": sid, "build_dir": build_path, "stats_delivery": "deferred"}) + await _post(self.get_http_port(), "/load", + {"session": sid, "path": csv_path, "mode": "buckaroo"}) + resp = await _post(self.get_http_port(), "/load_expr", + {"session": sid, "build_dir": build_path}) + 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") + finally: + shutil.rmtree(builds_root, ignore_errors=True) + os.unlink(csv_path) + + @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 + + +@contextmanager +def _record_summary_tiers(dataflow): + """The ``stats_tier`` of every summary-stats run ``dataflow`` makes while + the block runs.""" + tiers = [] + original = dataflow._get_summary_sd + + def spy(df, scope="filt"): + tiers.append(dataflow.stats_tier) + return original(df, scope) + + with patch.object(dataflow, "_get_summary_sd", spy): + yield tiers + + +def _deferred_session(): + """A pending deferred session over a schema-tier dataflow, as /load_expr + leaves one.""" + session = SessionState(session_id="s", path="", mode="buckaroo", backend="xorq", + xorq_dataflow=_build_dataflow(stats_tier="schema"), stats_tier="full", stats_delivery="deferred") + begin_stats_generation(session) + return session + + +class TestCompleteStats: + """``complete_stats`` takes a deferred session's dataflow to the full tier.""" + + def test_a_cache_hit_keeps_the_errs_of_the_run_that_filled_the_cache(self): + session = _deferred_session() + dataflow = session.xorq_dataflow + errs = {"a": {"stat": "boom"}} + real = dataflow._get_summary_sd + dataflow._get_summary_sd = lambda df, scope="filt": (real(df, scope)[0], errs) if dataflow.stats_tier == "full" else real(df, + scope) + assert complete_stats(session) + + # A search and its clearing, each reset to the schema tier as a state + # change does, bring the dataflow back to a state whose stats are cached. + for args in ({"search": ["a"]}, {}): + dataflow.stats_tier = "schema" + dataflow.quick_command_args = args + begin_stats_generation(session) + assert complete_stats(session) + + assert dataflow.errs == errs + + def test_spans_reach_the_bound_sink_when_the_session_has_none_and_the_run_is_not_a_first_pull(self): + session = _deferred_session() + records = [] + + with perf_log.telemetry_context("outer", records.append): + assert complete_stats(session) + + names = {record["name"] for record in records} + assert "stats.complete" in names + assert "firstpull.stats_total" not in names + + +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) + frame = await _read_json(ws) + ws.stats_gen = (frame["df_meta"].get("stats") or {}).get("gen") + return ws, frame + + 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, "stats_gen": ws.stats_gen})) + self.assertEqual((await _read_json(ws))["length"], 5) + pq.read_table(io.BytesIO(await ws.read_message())) + self.assertEqual((await _read_json(ws))["type"], "stats_update", "the push follows the rows") + 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, "stats_gen": ws.stats_gen})) + self.assertEqual((await _read_json(ws))["type"], "infinite_resp", + "a stats_update must leave no stray binary frame in the stream") + + async def _first_rows(self, ws, stats_gen=None, **window): + """Ask for the first window, as made for ``stats_gen`` (the generation + ``ws`` last saw by default), and read its two frames.""" + ws.write_message(json.dumps({"type": "infinite_request", + "stats_gen": ws.stats_gen if stats_gen is None else stats_gen, + "payload_args": {"start": 0, "end": 5, "sourceName": "default", "origEnd": 5, **window}})) + resp = await _read_json(ws) + self.assertEqual(resp["type"], "infinite_resp") + pq.read_table(io.BytesIO(await ws.read_message())) + + @tornado.testing.gen_test + async def test_stats_are_pushed_after_the_first_row_reply_without_a_request(self): + await self._load("sw-push", stats_delivery="deferred") + inline = await self._inline_frame("sw-push-inline") + ws, first = await self._connect("sw-push", caps="stats_update") + gen = self._stats(first)["gen"] + + await self._first_rows(ws) + update = await _read_json(ws) + + self.assertEqual((update["type"], update["stats_gen"], update["final"]), ("stats_update", gen, True)) + self.assertEqual(_rows_by_stat(update["payload"]), _rows_by_stat(inline["df_data_dict"]["all_stats"])) + await self._first_rows(ws) # the push is owed once: a second row reply is followed by no second update + + @tornado.testing.gen_test + async def test_a_second_connection_is_pushed_the_stats_the_first_connection_ran(self): + await self._load("sw-push-two", stats_delivery="deferred") + a, _ = await self._connect("sw-push-two", caps="stats_update") + b, _ = await self._connect("sw-push-two", caps="stats_update") + await self._first_rows(a) + self.assertEqual((await _read_json(a))["type"], "stats_update") + + with _count_stat_queries() as queries: + await self._first_rows(b) + self.assertEqual((await _read_json(b))["type"], "stats_update") + self.assertEqual(queries, [], "stats are computed once per session, at the first connection's push") + + @tornado.testing.gen_test + async def test_a_push_emits_a_stats_push_span(self): + captured: list = [] + with patch.object(telemetry, "make_http_sink", lambda url, **kw: captured.append): + await self._load("sw-push-span", stats_delivery="deferred", + telemetry_url="http://companion.invalid/internal/telemetry") + ws, first = await self._connect("sw-push-span", caps="stats_update") + await self._first_rows(ws) + await _read_json(ws) + + (push,) = [r for r in captured if r["name"] == "stats.push"] + self.assertEqual(push["trace"], "sw-push-span") + self.assertEqual((push["attrs"]["stats_gen"], push["attrs"]["outcome"], push["attrs"]["tier"]), + (self._stats(first)["gen"], "update", "full")) + self.assertGreaterEqual(push["attrs"]["gap_ms"], 0) + + @tornado.testing.gen_test + async def test_a_stale_stats_gen_gets_stats_aborted_and_runs_nothing(self): + await self._load("sw-stale", stats_delivery="deferred") + ws, first = await self._connect("sw-stale", caps="stats_update") + gen = self._stats(first)["gen"] + + with _count_stat_queries() as queries: + ws.write_message(_stats_request(gen - 1)) + aborted = await _read_json(ws) + self.assertEqual(aborted["type"], "stats_aborted") + self.assertEqual((aborted["stats_gen"], aborted["current_gen"], aborted["reason"]), + (gen - 1, gen, "stale")) + self.assertEqual(queries, [], "a stale request must run no stat query") + + ws.write_message(_stats_request(gen)) + self.assertEqual((await _read_json(ws))["type"], "stats_update", + "the stale request must not have consumed or failed the session's stats") + + @tornado.testing.gen_test + async def test_an_unsupported_scope_gets_stats_aborted(self): + await self._load("sw-scope", stats_delivery="deferred") + ws, first = await self._connect("sw-scope", caps="stats_update") + ws.write_message(_stats_request(self._stats(first)["gen"], scope="filt")) + aborted = await _read_json(ws) + self.assertEqual((aborted["type"], aborted["reason"], aborted["scope"]), + ("stats_aborted", "unsupported_scope", "filt")) + + @tornado.testing.gen_test + async def test_load_expr_and_reload_expr_bump_stats_gen(self): + await self._load("sw-gen", stats_delivery="deferred") + ws, first = await self._connect("sw-gen", caps="stats_update") + gen = self._stats(first)["gen"] + + await self._load("sw-gen", stats_delivery="deferred", force_reload=True) + pushed = await _read_json(ws) + self.assertGreater(self._stats(pushed)["gen"], gen, "/load_expr must bump stats_gen") + gen = self._stats(pushed)["gen"] + + resp = await _post(self.get_http_port(), "/reload_expr/sw-gen", {}) + self.assertEqual(resp.code, 200, resp.body) + pushed = await _read_json(ws) + self.assertGreater(self._stats(pushed)["gen"], gen, "/reload_expr must bump stats_gen") + + @tornado.testing.gen_test + async def test_df_meta_stats_survives_a_dataflow_field_change(self): + await self._load("sw-change", stats_delivery="deferred") + ws, first = await self._connect("sw-change", caps="stats_update") + gen = self._stats(first)["gen"] + ws.write_message(_stats_request(gen)) + self.assertEqual((await _read_json(ws))["type"], "stats_update") + + with _count_stat_queries() as queries: + ws.write_message(_state_change(quick_command_args={"search": ["a"]})) + changed = await _read_json(ws) + self._assert_pending(changed, gen + 1) + self.assertEqual(queries, [], "a dataflow-field change on a deferred session must run no stats") + # The dataflow rebuilt df_meta for the filtered state; its keys are still there. + self.assertEqual((changed["df_meta"]["total_rows"], changed["df_meta"]["filtered_rows"]), (5, 2)) + + ws.write_message(_stats_request(gen)) + aborted = await _read_json(ws) + self.assertEqual((aborted["type"], aborted["current_gen"], aborted["reason"]), + ("stats_aborted", gen + 1, "stale")) + ws.write_message(_stats_request(gen + 1)) + update = await _read_json(ws) + self.assertEqual((update["type"], update["stats_gen"], update["final"]), + ("stats_update", gen + 1, True)) + + @tornado.testing.gen_test + async def test_a_legacy_client_stays_complete_while_a_caps_client_gets_a_stats_free_frame(self): + sid = "sw-ab" + inline = await self._inline_frame("sw-ab-inline", post_processing="first_three") + await self._load(sid, stats_delivery="deferred") + a1, a1_open = await self._connect(sid, caps="other,stats_update") + a2, _ = await self._connect(sid, caps="stats_update") + # Connecting runs the missing stats for a legacy client, so the session + # is complete when B1 changes it. + b1, b1_open = await self._connect(sid, caps="unknown") + b2, _ = await self._connect(sid) + gen = self._stats(a1_open)["gen"] + self.assertEqual(self._stats(b1_open)["status"], "complete") + + with _count_stat_queries() as queries: + b1.write_message(_state_change(post_processing="first_three")) + frames = {name: await _read_json(ws) + for name, ws in (("a1", a1), ("a2", a2), ("b1", b1), ("b2", b2))} + for name in ("a1", "a2"): + self._assert_pending(frames[name], gen + 1) + for name in ("b1", "b2"): + self._assert_complete(frames[name], gen + 1, inline) + ran = len(queries) + self.assertGreater(ran, 0, "B's frame must have run the stats it carries") + + a1.write_message(_stats_request(gen + 1)) + update = await _read_json(a1) + self.assertEqual((update["type"], update["stats_gen"]), ("stats_update", gen + 1)) + self.assertEqual(_rows_by_stat(update["payload"]), + _rows_by_stat(frames["b1"]["df_data_dict"]["all_stats"])) + self.assertEqual(len(queries), ran, + "A's request must be answered from the stats B's frame computed") + + @tornado.testing.gen_test + async def test_load_expr_push_keeps_a_legacy_client_complete(self): + sid = "sw-push-load-expr" + inline = await self._inline_frame("sw-push-load-expr-inline") + a, a_open, b, _ = await self._pair(sid) + gen = self._stats(a_open)["gen"] + await self._load(sid, stats_delivery="deferred", force_reload=True) + self._assert_pending(await _read_json(a), gen + 1) + self._assert_complete(await _read_json(b), gen + 1, inline) + + @tornado.testing.gen_test + async def test_reload_expr_push_keeps_a_legacy_client_complete(self): + sid = "sw-push-reload" + inline = await self._inline_frame("sw-push-reload-inline") + a, a_open, b, _ = await self._pair(sid) + gen = self._stats(a_open)["gen"] + resp = await _post(self.get_http_port(), f"/reload_expr/{sid}", {}) + self.assertEqual(resp.code, 200, resp.body) + self._assert_pending(await _read_json(a), gen + 1) + self._assert_complete(await _read_json(b), gen + 1, inline) + + @tornado.testing.gen_test + async def test_load_push_leaves_both_clients_complete(self): + """/load swaps the session to pandas, which has no deferred stats: the + stored policy must not leave it looking pending.""" + sid = "sw-push-load" + a, _, b, _ = await self._pair(sid) + csv_fd, csv_path = tempfile.mkstemp(suffix=".csv") + os.close(csv_fd) + try: + pd.DataFrame({"x": [1, 2, 3], "y": ["p", "q", "r"]}).to_csv(csv_path, index=False) + resp = await _post(self.get_http_port(), "/load", + {"session": sid, "path": csv_path, "mode": "buckaroo"}) + self.assertEqual(resp.code, 200, resp.body) + finally: + os.unlink(csv_path) + frames = [await _read_json(a), await _read_json(b)] + for frame in frames: + self.assertNotIn("stats", frame["df_meta"]) + self.assertIn("histogram_bins", _rows_by_stat(frame["df_data_dict"]["all_stats"])) + self.assertEqual(_comparable(frames[0]), _comparable(frames[1])) + + @tornado.testing.gen_test + async def test_load_compare_push_leaves_both_clients_complete(self): + sid = "sw-push-compare" + a, _, b, _ = await self._pair(sid) + paths = [] + try: + for frame in (pd.DataFrame({"id": [1, 2], "v": [10, 20]}), + pd.DataFrame({"id": [1, 3], "v": [10, 30]})): + fd, path = tempfile.mkstemp(suffix=".csv") + os.close(fd) + frame.to_csv(path, index=False) + paths.append(path) + resp = await _post(self.get_http_port(), "/load_compare", + {"session": sid, "path1": paths[0], "path2": paths[1], "join_columns": ["id"]}) + self.assertEqual(resp.code, 200, resp.body) + finally: + for path in paths: + os.unlink(path) + frames = [await _read_json(a), await _read_json(b)] + for frame in frames: + self.assertNotIn("stats", frame["df_meta"]) + self.assertEqual(_comparable(frames[0]), _comparable(frames[1])) + + @tornado.testing.gen_test + async def test_highlight_overlay_is_complete_for_a_legacy_client_and_stats_free_for_a_caps_client(self): + sid = "sw-overlay" + await self._load(sid, stats_delivery="deferred") + a, a_open = await self._connect(sid, caps="stats_update") + gen = self._stats(a_open)["gen"] + + a.write_message(_state_change(search_string="ca")) + a_overlay = await _read_json(a) + self._assert_pending(a_overlay, gen) + + b, _ = await self._connect(sid) + b.write_message(_state_change(search_string="ca")) + b_overlay = await _read_json(b) + self.assertEqual(self._stats(b_overlay)["status"], "complete") + self.assertIn("histogram_bins", _rows_by_stat(b_overlay["df_data_dict"]["all_stats"])) + # The overlay's display config is the complete one with the highlight on top. + session_args = self._session(sid).df_display_args["main"]["df_viewer_config"]["column_config"] + by_col = {cc["col_name"]: cc for cc in b_overlay["df_display_args"]["main"]["df_viewer_config"]["column_config"]} + for expected in session_args: + self.assertEqual(by_col[expected["col_name"]]["ag_grid_specs"], expected["ag_grid_specs"]) + highlighted = [cc for cc in by_col.values() + if cc.get("displayer_args", {}).get("highlight_phrase") == ["ca"]] + self.assertTrue(highlighted, "the overlay must still carry the highlight") + + @tornado.testing.gen_test + async def test_overlay_for_a_legacy_client_is_built_after_its_stats_are_completed(self): + """Completing the stats replaces the session's display config, so an + overlay that copied it first would send a legacy client the schema + tier's. Every send completes a session that has a legacy client + connected, so only a direct call reaches a pending session here.""" + await self._load("sw-overlay-order", stats_delivery="deferred") + session = self._session("sw-overlay-order") + sent: list = [] + legacy = SimpleNamespace(search_string="ca", caps=frozenset(), session_id="sw-overlay-order", + write_message=sent.append, _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_failed_state_change_runs_no_stats_and_the_session_keeps_the_stats_it_has(self): + """The change resets a completed session's dataflow to the schema tier + before it applies the change. When applying it raises, nothing may run + the full stats of the half-applied state, and the session still serves + the full stats of the snapshot it has.""" + await self._load("sw-failed-change", stats_delivery="deferred") + ws, _ = await self._connect("sw-failed-change") + session = self._session("sw-failed-change") + gen = session.stats_gen + self.assertEqual(session.stats_status, "complete") + + with _record_summary_tiers(session.xorq_dataflow) as tiers: + with patch("buckaroo.server.websocket_handler.refresh_session_snapshot", side_effect=RuntimeError("boom")): + ws.write_message(_state_change(quick_command_args={"search": ["a"]})) + self.assertEqual((await _read_json(ws))["error_code"], "state_change_error") + self.assertNotIn("full", tiers) + + self.assertEqual((session.stats_status, session.stats_gen), ("complete", gen)) + ws.write_message(_stats_request(gen)) + update = await _read_json(ws) + self.assertEqual((update["type"], update["tier"]), ("stats_update", "full")) + self.assertIn("histogram_bins", _rows_by_stat(update["payload"])) + + @tornado.testing.gen_test + async def test_a_state_change_on_a_completed_session_runs_the_summary_cascade_once(self): + await self._load("sw-cascade-once", stats_delivery="deferred") + ws, _ = await self._connect("sw-cascade-once", caps="stats_update") + session = self._session("sw-cascade-once") + await self._first_rows(ws) + self.assertEqual((await _read_json(ws))["type"], "stats_update") + self.assertEqual(session.stats_status, "complete") + + with _record_summary_tiers(session.xorq_dataflow) as tiers: + ws.write_message(_state_change(quick_command_args={"search": ["a"]})) + self.assertEqual((await _read_json(ws))["type"], "initial_state") + + self.assertEqual(tiers, ["schema"], "the cascade must run once, for the new state, at the schema tier") + + @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_is_not_retried_by_a_request(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 until a connection opens: a request + # does not retry it. + 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_a_failed_stats_run_is_retried_when_the_next_connection_opens(self): + await self._load("sw-retry", stats_delivery="deferred") + ws, first = await self._connect("sw-retry", 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)) + self.assertEqual((await _read_json(ws))["reason"], "error") + + # The cause is gone, as after a transient backend failure, and the page reloads. + reopened, frame = await self._connect("sw-retry", caps="stats_update") + self._assert_pending(frame, gen) + await self._first_rows(reopened) + update = await _read_json(reopened) + self.assertEqual((update["type"], update["stats_gen"], update["final"]), ("stats_update", gen, True)) + self.assertEqual(self._session("sw-retry").stats_status, "complete") + + @tornado.testing.gen_test + async def test_the_final_stats_update_carries_the_display_config_the_stats_change(self): + inline = await self._inline_frame("sw-display-inline") + await self._load("sw-display", stats_delivery="deferred") + a, a_open = await self._connect("sw-display", caps="stats_update") + b, _ = await self._connect("sw-display", caps="stats_update") + self.assertNotEqual(_as_json(a_open["df_display_args"]), _as_json(inline["df_display_args"]), + "the schema tier's display config must differ, or this test sees nothing") + + for ws in (a, b): # b's push is answered from the session, not from a run + await self._first_rows(ws) + update = await _read_json(ws) + self.assertEqual(_as_json(update["df_display_args"]), _as_json(inline["df_display_args"])) + + @tornado.testing.gen_test + async def test_only_a_row_reply_for_the_owed_generation_pushes_stats(self): + await self._load("sw-keyed", stats_delivery="deferred") + ws, first = await self._connect("sw-keyed", caps="stats_update") + gen = self._stats(first)["gen"] + + # _first_rows reads two frames and asserts the first is an infinite_resp, so + # a stats_update slipping in anywhere in this sequence fails it. + await self._first_rows(ws, stats_gen=gen - 1) # made before the client saw this state: no push + # A sort changes the frame's order, not its stats: the sorted request is + # the first reply for this generation, and the stats follow it once. + await self._first_rows(ws, sort="b", sort_direction="desc") + update = await _read_json(ws) + self.assertEqual((update["type"], update["stats_gen"]), ("stats_update", gen)) + await self._first_rows(ws, sort="b", sort_direction="asc") + await self._first_rows(ws, sort="b", sort_direction="desc") + self.assertEqual(self._session("sw-keyed").stats_gen, gen) + + @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("stats.complete", [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"])) + + class TestReloadExpr(tornado.testing.AsyncHTTPTestCase): def get_app(self): return make_app() @@ -1177,6 +2446,30 @@ async def test_reload_expr_returns_200_with_zero_klasses(self): shutil.rmtree(builds_root, ignore_errors=True) shutil.rmtree(project_root, ignore_errors=True) + @tornado.testing.gen_test + async def test_reload_expr_ignores_a_body_it_cannot_parse(self): + """The body is optional, so one that isn't JSON (or isn't UTF-8) is + treated as absent rather than failing the reload.""" + builds_root = tempfile.mkdtemp() + project_root = tempfile.mkdtemp() + try: + build_path = _build_expr_dir(builds_root) + sid = "re-bad-body" + await _post(self.get_http_port(), "/load_expr", + {"session": sid, "build_dir": build_path, + "project_root": project_root}) + client = tornado.httpclient.AsyncHTTPClient() + for raw in (b"not json", b"\xff"): + resp = await client.fetch( + f"http://localhost:{self.get_http_port()}/reload_expr/{sid}", + method="POST", body=raw, + headers={"Content-Type": "application/json"}, + raise_error=False) + self.assertEqual(resp.code, 200, raw) + finally: + shutil.rmtree(builds_root, ignore_errors=True) + shutil.rmtree(project_root, ignore_errors=True) + @tornado.testing.gen_test async def test_reload_expr_preserves_load_expr_config(self): """#957: /reload_expr must rebuild the dataflow with the same 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