From 2b4558620e8ac562eccb0fdc32e9b81cb372905e Mon Sep 17 00:00:00 2001 From: Jeroen Van Der Donckt Date: Thu, 10 Sep 2026 11:36:16 +0200 Subject: [PATCH] perf: stream domain probes on resident frames --- .github/workflows/ci.yml | 6 ++- Makefile | 4 ++ flexviz/LF.py | 8 +-- pyproject.toml | 2 +- tests/test_domain_resolution.py | 12 +++-- tests/test_perf_choices.py | 87 +++++++++++++++++++++++++++++++++ 6 files changed, 109 insertions(+), 10 deletions(-) create mode 100644 tests/test_perf_choices.py diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index c03ef94..26596b2 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -168,7 +168,8 @@ jobs: run: uv run --no-sync pytest -m browser -p no:randomly --override-ini="addopts=" -n 6 # A separate job: it spawns one subprocess per trace and size (1-2 minutes), - # and its memory counter only exists on Linux and macOS. + # its memory counter only exists on Linux and macOS, and it also runs the + # performance-decision checks, which need the built plugin and several cores. ooc: needs: plugin runs-on: ubuntu-latest @@ -193,3 +194,6 @@ jobs: - name: Out-of-core memory matrix run: uv run --no-sync pytest -m ooc --override-ini="addopts=" -v tests/test_ooc.py + + - name: Performance choices + run: uv run --no-sync pytest -m benchmark --override-ini="addopts=" -v tests/test_perf_choices.py diff --git a/Makefile b/Makefile index 3cfca01..85e012c 100644 --- a/Makefile +++ b/Makefile @@ -20,6 +20,10 @@ test-browser: test-ooc: uv run pytest -m ooc --override-ini="addopts=" -v tests/test_ooc.py +.PHONY: test-perf +test-perf: + uv run pytest -m benchmark --override-ini="addopts=" -v tests/test_perf_choices.py + .PHONY: docs docs: uv run --group docs mkdocs serve diff --git a/flexviz/LF.py b/flexviz/LF.py index b03bfed..9e6cbe7 100644 --- a/flexviz/LF.py +++ b/flexviz/LF.py @@ -245,8 +245,8 @@ def collect_engine(self) -> str: Fixed per source kind rather than left to ``"auto"``: a file scan must stream, a resident frame must not pay the streaming machinery. The line - bucket plan and the grouped histogram plan are the exceptions: both - stream on both source kinds. + bucket plan, the grouped histogram plan and the domain probe are the + exceptions: all stream on both source kinds. """ return "streaming" if self.is_scan else "in-memory" @@ -300,7 +300,9 @@ def physical_minmax( val = val.to_physical() exprs.append(val.min().alias(f"__min_{c}__")) exprs.append(val.max().alias(f"__max_{c}__")) - stats = self._ldf.select(exprs).collect(engine=self.collect_engine) + # Always streaming: the min/max select is ~2x faster on the + # streaming engine than on the in-memory one, on both source kinds. + stats = self._ldf.select(exprs).collect(engine="streaming") for c in missing: memo[c] = (stats[f"__min_{c}__"].item(), stats[f"__max_{c}__"].item()) return {c: memo[c] for c in columns} diff --git a/pyproject.toml b/pyproject.toml index 1d03c8c..6924d7d 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -98,7 +98,7 @@ markers = [ "integration: tests that spin up a FastAPI TestClient", "slow: benchmarks and large-dataset tests", "browser: browser-based tests that require Playwright (run with make test-browser)", - "benchmark: performance benchmarks (reserved for the future benchmark suite)", + "benchmark: performance-decision checks in tests/test_perf_choices.py (run with make test-perf)", "ooc: out-of-core memory matrix, one subprocess per trace and size (run with make test-ooc)", ] diff --git a/tests/test_domain_resolution.py b/tests/test_domain_resolution.py index 4a8028f..26a9505 100644 --- a/tests/test_domain_resolution.py +++ b/tests/test_domain_resolution.py @@ -3,8 +3,8 @@ Bin edges come from the unfiltered frame, so a cross-filter never moves them. The engine resolves every column a request needs in one min/max pass: the Parquet footer on a single-file Parquet scan, one collect otherwise. Every -collect names its Polars engine: streaming for a scan, in-memory for a resident -frame. +collect names its Polars engine: streaming for a scan and for the probe, +in-memory for a resident frame's aggregation. """ from __future__ import annotations @@ -184,13 +184,15 @@ def test_scan_collect_counts(self, tmp_path, collects, n_traces): class TestEnginePinning: - def test_resident_frame_collects_in_memory(self, collects): + def test_resident_frame_streams_the_probe_and_aggregates_in_memory(self, collects): + """The probe streams on both source kinds; only the aggregation pins.""" df = pl.DataFrame({"a": [float(i) for i in range(50)]}) engine, infos = _engine(df, [Histogram(x="a", bins=10)]) _init(engine, infos) - assert len(collects.minmax) == 1 - assert set(collects.engines) == {"in-memory"} + assert [engine for engine, _ in collects.minmax] == ["streaming"] + rest = [engine for engine, plan in collects.calls if "__min_" not in plan] + assert set(rest) == {"in-memory"} def test_scan_collects_streaming(self, tmp_path, collects): # A line, not a histogram: a scanned histogram folds over batches and diff --git a/tests/test_perf_choices.py b/tests/test_perf_choices.py new file mode 100644 index 0000000..939faeb --- /dev/null +++ b/tests/test_perf_choices.py @@ -0,0 +1,87 @@ +"""Performance-decision checks: one test per decision. + +Each test names the decision it checks, the evidence that motivated it (with +numbers), and what the assertion verifies. Run with ``make test-perf``; +excluded from ``make test`` (marker ``benchmark``) because timings need a +quiet machine and at least 4 Polars threads. +""" + +from __future__ import annotations + +import statistics +import time + +import numpy as np +import polars as pl +import pytest +from ooc_child import PeakSampler + +from flexviz.LF import LFQueryBuilder + +pytestmark = pytest.mark.benchmark + +# 80 MB of Float64: past the streaming engine's fixed overhead (it already wins +# at 10M rows) and still ~2 ms per probe, so the whole check costs under a second. +PROBE_ROWS = 10_000_000 +# Measured on this data at 4 threads, the CI runner's core count: 2.47 ms +# in-memory against 1.53 ms streaming, a ratio of 0.62. 0.80 leaves headroom +# for a shared runner. Below 4 threads the streaming engine has no win, hence +# the skip. +PROBE_TIME_RATIO = 0.80 +MIN_PROBE_THREADS = 4 +# The streaming engine reduces morsel by morsel, so it must never copy the +# column. Measured 0.2 MB at 200M rows; 64 MB is well under the 80 MB column. +PROBE_PEAK_MB = 64 + + +def _median_ms(call) -> float: + call() # warm-up: first touch of the column + times = [] + for _ in range(5): + start = time.perf_counter() + call() + times.append(time.perf_counter() - start) + return statistics.median(times) * 1e3 + + +@pytest.mark.skipif( + pl.thread_pool_size() < MIN_PROBE_THREADS, + reason=f"the streaming min/max needs >= {MIN_PROBE_THREADS} threads to win", +) +def test_domain_probe_streams_on_resident_frames() -> None: + """Decision: `LFQueryBuilder.physical_minmax` collects its min/max select + with the streaming engine even on a resident frame, although every other + resident collect is pinned to the in-memory engine. + + Evidence: the in-memory engine reads the column twice (one pass for min, + one for max) while the streaming engine folds both into one morsel pass; + measured on an M5 at 10M rows 1.87 vs 3.51 ms (0.53x), at 200M rows + 34.7 vs 63.4 ms, and 0.62x at 4 threads; below 4 threads the streaming + engine has no win, hence the skip. + + Check: the probe takes at most PROBE_TIME_RATIO of the in-memory select, + and its peak memory stays under PROBE_PEAK_MB (the engine must never copy + the column). + """ + rng = np.random.default_rng(0) + df = pl.DataFrame({"v": rng.standard_normal(PROBE_ROWS)}) + exprs = [pl.col("v").min().alias("__min_v__"), pl.col("v").max().alias("__max_v__")] + + # One sampler over both timings: it costs each engine the same, and a 1 ms + # poll is needed because a probe lasts ~2 ms. + with PeakSampler(interval=0.001) as sampler: + # The old code path, for reference. + reference_ms = _median_ms( + lambda: df.lazy().select(exprs).collect(engine="in-memory") + ) + # A fresh builder per call: the min/max memo is per builder. + probe_ms = _median_ms(lambda: LFQueryBuilder(df.lazy()).physical_minmax(["v"])) + + assert probe_ms <= PROBE_TIME_RATIO * reference_ms, ( + f"probe took {probe_ms:.2f} ms, over {PROBE_TIME_RATIO} x the " + f"{reference_ms:.2f} ms in-memory select" + ) + assert sampler.peak_mb <= PROBE_PEAK_MB, ( + f"the min/max select peaked {sampler.peak_mb:.1f} MB over its baseline, " + f"above the {PROBE_PEAK_MB} MB cap" + )