Skip to content

perf(line): batch multi-trace streaming envelopes into one scan #8

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
jvdd wants to merge 1 commit into
base: main
Choose a base branch
from
Open

perf(line): batch multi-trace streaming envelopes into one scan #8

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
53 changes: 43 additions & 10 deletions flexviz/LF.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand Down Expand Up @@ -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]] = {}
Expand Down
129 changes: 68 additions & 61 deletions flexviz/trace/line.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand All @@ -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
Expand All @@ -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]
Expand All @@ -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
Expand All @@ -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

Expand Down Expand Up @@ -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,
Expand All @@ -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(
Expand Down
51 changes: 51 additions & 0 deletions tests/test_engine.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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 ---------------------------------------------

Expand Down