feat(xorq): stats tier machinery and the xorq schema tier (rows-first s1)

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)")
 
 

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
 # ============================================================

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)

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
@@ -602,18 +602,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)

buckaroo/server/handlers.py

@@ -14,7 +14,9 @@
 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, build_state_message, dataflow_stats_tier)
 from buckaroo.server import telemetry
 from buckaroo.pluggable_analysis_framework import perf_log
 
@@ -324,6 +326,9 @@ async def post(self):
         session.backend = backend
         session.xorq_dataflow = None
         session.expr = None
+        # 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
@@ -396,6 +401,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.
@@ -454,7 +481,10 @@ async def post(self):
         # Config-bearing fields that change how the result is computed or
         # rendered. If the caller passes any of these on a warm POST we must
         # re-run the pipeline — returning cached metadata would silently
-        # ignore the new config.
+        # ignore the new config. stats_tier and stats_delivery are not in this
+        # tuple: it tests truthiness, and a host that sends the pair on every
+        # POST would never get the warm exit (#944). They are compared with the
+        # session's stored pair below instead.
         has_config = any(body.get(k) for k in (
             "component_config", "column_config_overrides", "extra_grid_config",
             "init_sd", "skip_stat_columns"))
@@ -479,6 +509,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
@@ -492,7 +532,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)
@@ -564,6 +605,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 +631,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.
@@ -821,8 +865,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 +904,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 +927,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()
@@ -888,6 +955,8 @@ async def post(self, session_id):
 
         refreshed = get_buckaroo_display_state(xorq_dataflow)
         session.xorq_dataflow = xorq_dataflow
+        session.stats_tier = stats_tier
+        session.stats_delivery = stats_delivery
         session.df_display_args = refreshed["df_display_args"]
         session.df_data_dict = refreshed["df_data_dict"]
         session.df_meta = refreshed["df_meta"]

buckaroo/server/session.py

@@ -11,6 +11,20 @@
 _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
+
 
 @dataclass
 class SessionState:
@@ -43,6 +57,13 @@ 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
     # 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