diff --git a/flexviz/LF.py b/flexviz/LF.py index 4971be7..d1f5ee8 100644 --- a/flexviz/LF.py +++ b/flexviz/LF.py @@ -60,12 +60,17 @@ class AggregationSpec: uid: str = "" global_stats_cols: Tuple[str, ...] = () #: Optional escape hatch for an aggregation that cannot be a select - #: expression. Called as ``plan(filtered_ldf)`` and must return a one-row - #: DataFrame whose single column is named ``uid``, i.e. exactly the column - #: the batched ``select`` would have produced. Set only when a spec needs - #: its own plan — the out-of-core line envelope uses a streaming group_by - #: that cannot ride the shared select. - plan: "Callable[[pl.LazyFrame], pl.DataFrame] | None" = None + #: expression. Called as ``plan(filtered_ldf)`` or, when batched, + #: ``plan(filtered_ldf, batch_args_list)`` — see ``plan_batch_key``. + #: Must return a DataFrame whose columns are named after the uid(s). + plan: "Callable[..., pl.DataFrame] | None" = None + #: Hashable key for plan batching. Plans with the same non-None key are + #: fused into one call: ``plan(filtered_ldf, [spec.plan_batch_args ...])`` + #: returning one column per spec. Used by the streaming line envelope to + #: run one scan for N traces sharing the same x column and viewport. + plan_batch_key: Any = None + #: Per-trace arguments passed to a batched plan callback. + plan_batch_args: Any = None @dataclass(frozen=True) @@ -324,16 +329,44 @@ def aggregate( regular_df = filtered_ldf.select(*[s.expr for s in expr_specs]).collect() else: regular_df = pl.DataFrame() - for spec in plan_specs: + + # Group plan specs by batch key: specs sharing a key run as one call + # (one scan for N traces). Specs without a key run individually. + unbatched = [s for s in plan_specs if s.plan_batch_key is None] + plan_batches: dict[Any, list[AggregationSpec]] = {} + for s in plan_specs: + if s.plan_batch_key is not None: + plan_batches.setdefault(s.plan_batch_key, []).append(s) + + def _hstack(planned: pl.DataFrame) -> None: + nonlocal regular_df + regular_df = ( + planned if regular_df.is_empty() else regular_df.hstack(planned) + ) + + for spec in unbatched: planned = spec.plan(filtered_ldf) if planned.width != 1 or planned.columns[0] != spec.uid: raise ValueError( f"AggregationSpec.plan for {spec.uid!r} must return exactly " f"one column named {spec.uid!r}, got {planned.columns!r}" ) - regular_df = ( - planned if regular_df.is_empty() else regular_df.hstack(planned) - ) + _hstack(planned) + + for batch in plan_batches.values(): + if len(batch) == 1: + planned = batch[0].plan(filtered_ldf) + else: + planned = batch[0].plan( + filtered_ldf, [s.plan_batch_args for s in batch] + ) + expected = {s.uid for s in batch} + got = set(planned.columns) + if got != expected: + raise ValueError( + f"Batched plan must return columns {expected!r}, got {got!r}" + ) + _hstack(planned) grouped_dfs: dict[str, pl.DataFrame] = {} grouped_batches: dict[tuple, list[GroupedAggregationSpec]] = {} diff --git a/flexviz/trace/line.py b/flexviz/trace/line.py index 8f425c2..046883a 100644 --- a/flexviz/trace/line.py +++ b/flexviz/trace/line.py @@ -110,6 +110,35 @@ def _viewport_window( # --------------------------------------------------------------------------- +def _envelope_empty(x_col: str, y_col: str, uid: str) -> pl.DataFrame: + return pl.DataFrame( + {uid: [[]]}, + schema={uid: pl.List(pl.Struct({x_col: pl.Null, y_col: pl.Null}))}, + ) + + +def _envelope_points( + result: pl.DataFrame, x_col: str, y_col: str, uid: str +) -> pl.DataFrame: + """Reassemble a per-trace envelope from the group_by result columns.""" + lo_c, hi_c = f"__lo_{uid}", f"__hi_{uid}" + xlo_c, xhi_c = f"__xlo_{uid}", f"__xhi_{uid}" + pts = pl.concat( + [ + result.select(pl.col(xlo_c).alias("__x"), pl.col(lo_c).alias("__y")), + result.select(pl.col(xhi_c).alias("__x"), pl.col(hi_c).alias("__y")), + ] + ).drop_nulls("__x") + pts = pts.unique(subset=["__x", "__y"]).sort("__x") + return pts.select( + pl.struct( + **{x_col: pl.col("__x").alias(x_col), y_col: pl.col("__y").alias(y_col)} + ) + .implode() + .alias(uid) + ) + + def _streaming_envelope_plan( x_col: str, y_col: str, @@ -121,8 +150,10 @@ def _streaming_envelope_plan( ): """Streaming min-max envelope using equal-width buckets in x. - Replaces the two-pass ``_native_envelope_plan``. One streaming collect, no - intermediate collects. + One streaming collect per call, no intermediate collects. When called with + ``batch_args`` (a list of ``(uid, y_col)`` tuples from sibling traces), + all y columns are aggregated in a single ``group_by`` — one scan for N + traces. Buckets partition the x range into ``n_points // 2`` equal-width bins. ``min_by``/``max_by`` locate the x value at each y extremum in a single @@ -134,26 +165,22 @@ def _streaming_envelope_plan( On an exact y plateau, ``min_by`` picks an arbitrary member, and which member can vary with ``POLARS_MAX_THREADS``. Any member is a valid envelope - point. This is a deliberate trade: the previous two-pass plan paid 2.9x - runtime and 39x memory for deterministic tie-breaking. + point. """ import math n_out = max(n_points // 2, 1) - empty = pl.DataFrame( - {uid: [[]]}, - schema={uid: pl.List(pl.Struct({x_col: pl.Null, y_col: pl.Null}))}, - ) - def run(filtered_ldf: pl.LazyFrame) -> pl.DataFrame: + def run(filtered_ldf: pl.LazyFrame, batch_args=None) -> pl.DataFrame: + traces = batch_args if batch_args is not None else [(uid, y_col)] + y_cols = list(dict.fromkeys(yc for _, yc in traces)) + src = filtered_ldf if vp_filter is None else filtered_ldf.filter(vp_filter) - src = src.select(x_col, y_col) + src = src.select(x_col, *y_cols) dtype = schema.get(x_col) if schema else None is_temporal = dtype is not None and dtype.is_temporal() - # Bucket arithmetic runs on the physical representation so that - # temporal columns reduce to plain integer division. phys_expr = pl.col(x_col).to_physical() if is_temporal else pl.col(x_col) if x_range is not None: x_lo, x_hi = x_range[0], x_range[1] @@ -171,14 +198,12 @@ def run(filtered_ldf: pl.LazyFrame) -> pl.DataFrame: x_hi = domain["__hi"].item() if x_lo is None or x_hi is None: - return empty + return pl.concat( + [_envelope_empty(x_col, yc, u) for u, yc in traces], + how="horizontal_extend", + ) span = x_hi - x_lo - # Float columns need true division: integer ceiling division rounds a - # sub-1 width up to 1 (0.002 / 500 -> 1), collapsing every row into - # one bucket. Integer and temporal columns use ceiling division to keep - # the width whole. The check is on the column dtype, not the Python - # type of span, because JSON deserializes 100.0 as int. use_float_div = dtype is not None and dtype.is_float() if span <= 0: bsz = 1 @@ -189,53 +214,31 @@ def run(filtered_ldf: pl.LazyFrame) -> pl.DataFrame: lo_lit = pl.lit(x_lo) bsz_lit = pl.lit(bsz) - - x, y = pl.col(x_col), pl.col(y_col) - result = ( - src.group_by( - # True division lands x_hi exactly on n_out; fold that lone - # top row back into the last bucket instead of letting it open - # an n_out + 1-th one and overrun the n_points budget. - ((phys_expr - lo_lit) // bsz_lit) - .clip(upper_bound=n_out - 1) - .alias("__b") - ) - .agg( - y.min().alias("__lo"), - y.max().alias("__hi"), - x.min_by(y).alias("__xlo"), - x.max_by(y).alias("__xhi"), + bkt = ((phys_expr - lo_lit) // bsz_lit).clip(upper_bound=n_out - 1).alias("__b") + + agg_exprs = [] + x = pl.col(x_col) + for u, yc in traces: + y = pl.col(yc) + agg_exprs.extend( + [ + y.min().alias(f"__lo_{u}"), + y.max().alias(f"__hi_{u}"), + x.min_by(y).alias(f"__xlo_{u}"), + x.max_by(y).alias(f"__xhi_{u}"), + ] ) - .collect(engine="streaming") - ) - if result.is_empty(): - return empty + result = src.group_by(bkt).agg(*agg_exprs).collect(engine="streaming") - pts = pl.concat( - [ - result.select( - pl.col("__xlo").alias("__x"), - pl.col("__lo").alias("__y"), - ), - result.select( - pl.col("__xhi").alias("__x"), - pl.col("__hi").alias("__y"), - ), - ] - ).drop_nulls("__x") - pts = pts.unique(subset=["__x", "__y"]).sort("__x") - - return pts.select( - pl.struct( - **{ - x_col: pl.col("__x").alias(x_col), - y_col: pl.col("__y").alias(y_col), - } + if result.is_empty(): + return pl.concat( + [_envelope_empty(x_col, yc, u) for u, yc in traces], + how="horizontal_extend", ) - .implode() - .alias(uid) - ) + + parts = [_envelope_points(result, x_col, yc, u) for u, yc in traces] + return pl.concat(parts, how="horizontal_extend") return run @@ -628,6 +631,8 @@ def get_aggregation_spec( if scan_source and self.downsample == "minmax": # `nth` already streams (a stride needs no state) and `fpcs` has no # streaming formulation at all, so only minmax needs the swap. + # Traces sharing the same x column and viewport batch into one scan. + vp_key = tuple(x_range) if x_range is not None else None return AggregationSpec( expr=pl.lit(None).alias(self.uid), uid=self.uid, @@ -644,6 +649,8 @@ def get_aggregation_spec( x_range, schema, ), + plan_batch_key=("streaming_envelope", self.x_col, vp_key), + plan_batch_args=(self.uid, self.y_col), ) expr = _plugin_line_agg_expr( diff --git a/tests/test_engine.py b/tests/test_engine.py index 62352f1..df5c152 100644 --- a/tests/test_engine.py +++ b/tests/test_engine.py @@ -2038,6 +2038,57 @@ def test_only_minmax_swaps(self, downsample): line = LinePlot(x="ts", y="val", n_points=1000, downsample=downsample) assert line.get_aggregation_spec({}, scan_source=True).plan is None + def test_multi_trace_scan_batches_into_one_plan(self, tmp_path): + """Multiple minmax line traces sharing x on a scan source fuse into + one streaming group_by instead of N sequential scans.""" + n = 20_000 + ys1 = [((i * 7919) % 1000) / 7.0 for i in range(n)] + ys2 = [((i * 4001) % 1000) / 5.0 for i in range(n)] + ys3 = [float(i % 89) for i in range(n)] + df = pl.DataFrame( + {"ts": list(range(n)), "v1": ys1, "v2": ys2, "v3": ys3}, + schema={ + "ts": pl.Int64, + "v1": pl.Float64, + "v2": pl.Float64, + "v3": pl.Float64, + }, + ) + path = tmp_path / "multi.parquet" + df.write_parquet(path) + + traces = [ + LinePlot(x="ts", y="v1", n_points=1000), + LinePlot(x="ts", y="v2", n_points=1000), + LinePlot(x="ts", y="v3", n_points=1000), + ] + lf = LFQueryBuilder(pl.scan_parquet(path)) + engine = FlexEngine( + backend_lf=lf, + scalable_traces={t.uid: t for t in traces}, + ) + infos = [ + TraceInfo(uid=t.uid, axes=("x", "y"), trace_type="line") for t in traces + ] + deltas = engine.process(InteractionEvent(type="init", force_update=True), infos) + assert len(deltas) == 3 + + src_v1 = set(df["v1"].to_list()) + src_v2 = set(df["v2"].to_list()) + src_v3 = set(df["v3"].to_list()) + for delta, src_ys in zip(deltas, [src_v1, src_v2, src_v3]): + xs = list(delta.updates["x"]) + ys = list(delta.updates["y"]) + assert len(xs) > 0 + assert len(xs) <= 1000 + for yi in ys: + assert yi in src_ys + + # The specs must share a batch key so the engine fuses them. + specs = [t.get_aggregation_spec({}, scan_source=True) for t in traces] + keys = {s.plan_batch_key for s in specs} + assert len(keys) == 1, f"expected one batch key, got {keys}" + # ---- descending viewport ranges ---------------------------------------------