From 52e37c79236ed6b1a896ec57b734809a463b711b Mon Sep 17 00:00:00 2001 From: ghostiee-11 Date: Mon, 29 Jun 2026 23:32:29 +0530 Subject: [PATCH 1/2] test: peak-memory integration tests for the lazy round-trip Adds two opt-in perf tests behind an 'integration' pytest marker: - A single .sel(time=t0).load() peaks far below the eager whole-grid materialization on NCEP air_temperature (lazy < 10 MB; eager > 50 MB; eager / lazy > 10x). - A GROUP BY reduction along the long axis stays under 4x the source's nominal size, with values matching the xarray-native mean. Registers the marker in pyproject.toml and defaults pytest to '-m "not integration"' so the default test run skips them; opt in with 'pytest -m integration'. Matches the CI workflow's existing 'pytest -v . -m "not integration"' invocation. --- pyproject.toml | 4 ++ tests/test_to_dataset_perf.py | 118 ++++++++++++++++++++++++++++++++++ 2 files changed, 122 insertions(+) create mode 100644 tests/test_to_dataset_perf.py diff --git a/pyproject.toml b/pyproject.toml index 409da03..2711840 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -111,3 +111,7 @@ cache-keys = [{file = "pyproject.toml"}, {file = "rust/Cargo.toml"}, {file = "** [tool.pytest.ini_options] testpaths = ["tests"] +addopts = "-m 'not integration'" +markers = [ + "integration: heavier perf / memory tests, opt-in via -m integration", +] diff --git a/tests/test_to_dataset_perf.py b/tests/test_to_dataset_perf.py new file mode 100644 index 0000000..82883cc --- /dev/null +++ b/tests/test_to_dataset_perf.py @@ -0,0 +1,118 @@ +"""Opt-in peak-memory tests for the lazy SQL -> xarray round-trip. + +Asserts the lazy backend honors its contract: a single-slab access +peaks far below an eager whole-grid materialization, and a streaming +aggregation does not balloon past the source size. Gated by the +``integration`` marker; run with ``pytest -m integration``. +""" + +import gc +import tracemalloc + +import numpy as np +import pytest +import xarray as xr + +from xarray_sql import XarrayContext + +pytestmark = pytest.mark.integration + + +def _peak_mb(fn): + """Return ``(result, peak_memory_mb)`` for a single call to ``fn``.""" + gc.collect() + tracemalloc.start() + tracemalloc.reset_peak() + out = fn() + _, peak = tracemalloc.get_traced_memory() + tracemalloc.stop() + return out, peak / 1e6 + + +@pytest.fixture(scope="module") +def air_source(): + """NCEP ``air_temperature`` chunked along time, ~31 MB dense. + + Module-scoped so the pooch download (~3.5 MB) is amortized across + tests in this file. Skips when the tutorial dataset is unreachable. + """ + try: + return xr.tutorial.open_dataset("air_temperature").chunk({"time": 24}) + except (OSError, ValueError, ImportError) as e: + pytest.skip(f"air_temperature tutorial dataset unavailable: {e}") + + +def test_lazy_slab_peak_memory_is_bounded(air_source): + """``.sel(time=t0).load()`` materializes only the slab, not the cube. + + Reference observation: lazy slab peak is ~1.8 MB on a 31 MB + dense source. A regression that quietly buffers the whole result + would push past 10 MB. Eager `to_dataset(chunks=None)` is measured + too as a sanity floor for the gap: the eager path should be at + least an order of magnitude heavier than the lazy slab, otherwise + the lazy path isn't actually lazy. + """ + ctx = XarrayContext() + ctx.from_dataset("air", air_source, chunks={"time": 24}) + t0 = air_source["time"].values[0] + + out = ctx.sql('SELECT * FROM "air"').to_dataset() + slab, slab_peak = _peak_mb(lambda: out["air"].sel(time=t0).load()) + assert slab.size == air_source.sizes["lat"] * air_source.sizes["lon"] + assert slab_peak < 10.0, ( + f"lazy single-slab access should stay under 10 MB on " + f"air_temperature, got {slab_peak:.2f} MB" + ) + + _, eager_peak = _peak_mb( + lambda: ctx.sql('SELECT * FROM "air"').to_dataset(chunks=None) + ) + assert eager_peak > 50.0, ( + f"eager whole-grid materialization should peak above 50 MB; " + f"if it doesn't, the eager path may have silently gone lazy " + f"and the lazy assertion above no longer means anything. " + f"Got {eager_peak:.1f} MB" + ) + assert eager_peak / max(slab_peak, 0.1) > 10.0, ( + f"eager peak should be at least 10x the lazy slab peak; " + f"if it isn't, the lazy path isn't actually lazy. " + f"Got eager={eager_peak:.1f} MB, lazy={slab_peak:.2f} MB" + ) + + +def test_streaming_aggregation_does_not_explode(air_source): + """A ``GROUP BY`` reducing the long axis streams in a single pass. + + Reduces 3.86M rows of ``air_temperature`` to ~1.3K group cells. The + aggregation must stream the source once and emit a tiny result, not + buffer the entire row set into memory. Reference observation: peak + ~54 MB on a 31 MB dense source (within 2x, well below a "buffer + the whole thing twice" failure mode). Threshold is 4x source size + so transient pandas / DataFusion buffers fit. + """ + ctx = XarrayContext() + ctx.from_dataset("air", air_source, chunks={"time": 24}) + source_mb = air_source.nbytes / 1e6 + + agg, agg_peak = _peak_mb( + lambda: ctx.sql( + 'SELECT lat, lon, AVG(air) AS air_avg FROM "air" GROUP BY lat, lon' + ).to_dataset(dims=["lat", "lon"]) + ) + assert agg.sizes["lat"] * agg.sizes["lon"] == ( + air_source.sizes["lat"] * air_source.sizes["lon"] + ) + assert agg_peak < 4 * source_mb, ( + f"GROUP BY reduction should not balloon past 4x source size; " + f"got peak={agg_peak:.1f} MB on a {source_mb:.1f} MB source" + ) + + # Sanity: values agree with the xarray-native reduction. + ref = ( + air_source.compute() + .mean(dim="time")["air"] + .sortby(["lat", "lon"]) + .values + ) + got = agg.sortby(["lat", "lon"])["air_avg"].values + np.testing.assert_allclose(got, ref, rtol=1e-5) From 4dbb148613b54efdfc5b70f2a3664d0ed16473e0 Mon Sep 17 00:00:00 2001 From: ghostiee-11 Date: Tue, 30 Jun 2026 01:46:25 +0530 Subject: [PATCH 2/2] review: rename slab to chunk, drop integration marker Per review on PR #199: - 'Let's not use the term slab' -- renamed slab -> chunk throughout - 'These could be unit tests' -- dropped the pytestmark integration gate; tests now run by default since air_temperature is small - Reverted pyproject.toml marker registration + addopts since no test uses it any more --- pyproject.toml | 4 ---- tests/test_to_dataset_perf.py | 39 ++++++++++++++++------------------- 2 files changed, 18 insertions(+), 25 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index 2711840..409da03 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -111,7 +111,3 @@ cache-keys = [{file = "pyproject.toml"}, {file = "rust/Cargo.toml"}, {file = "** [tool.pytest.ini_options] testpaths = ["tests"] -addopts = "-m 'not integration'" -markers = [ - "integration: heavier perf / memory tests, opt-in via -m integration", -] diff --git a/tests/test_to_dataset_perf.py b/tests/test_to_dataset_perf.py index 82883cc..4e46d2e 100644 --- a/tests/test_to_dataset_perf.py +++ b/tests/test_to_dataset_perf.py @@ -1,9 +1,8 @@ -"""Opt-in peak-memory tests for the lazy SQL -> xarray round-trip. +"""Peak-memory tests for the lazy SQL -> xarray round-trip. -Asserts the lazy backend honors its contract: a single-slab access +Asserts the lazy backend honors its contract: a single-chunk access peaks far below an eager whole-grid materialization, and a streaming -aggregation does not balloon past the source size. Gated by the -``integration`` marker; run with ``pytest -m integration``. +aggregation does not balloon past the source size. """ import gc @@ -15,8 +14,6 @@ from xarray_sql import XarrayContext -pytestmark = pytest.mark.integration - def _peak_mb(fn): """Return ``(result, peak_memory_mb)`` for a single call to ``fn``.""" @@ -42,26 +39,26 @@ def air_source(): pytest.skip(f"air_temperature tutorial dataset unavailable: {e}") -def test_lazy_slab_peak_memory_is_bounded(air_source): - """``.sel(time=t0).load()`` materializes only the slab, not the cube. +def test_lazy_chunk_peak_memory_is_bounded(air_source): + """``.sel(time=t0).load()`` materializes only one chunk, not the cube. - Reference observation: lazy slab peak is ~1.8 MB on a 31 MB + Reference observation: lazy chunk peak is ~1.8 MB on a 31 MB dense source. A regression that quietly buffers the whole result - would push past 10 MB. Eager `to_dataset(chunks=None)` is measured - too as a sanity floor for the gap: the eager path should be at - least an order of magnitude heavier than the lazy slab, otherwise - the lazy path isn't actually lazy. + would push past 10 MB. Eager ``to_dataset(chunks=None)`` is + measured too as a sanity floor for the gap: the eager path should + be at least an order of magnitude heavier than the lazy chunk, + otherwise the lazy path isn't actually lazy. """ ctx = XarrayContext() ctx.from_dataset("air", air_source, chunks={"time": 24}) t0 = air_source["time"].values[0] out = ctx.sql('SELECT * FROM "air"').to_dataset() - slab, slab_peak = _peak_mb(lambda: out["air"].sel(time=t0).load()) - assert slab.size == air_source.sizes["lat"] * air_source.sizes["lon"] - assert slab_peak < 10.0, ( - f"lazy single-slab access should stay under 10 MB on " - f"air_temperature, got {slab_peak:.2f} MB" + chunk, chunk_peak = _peak_mb(lambda: out["air"].sel(time=t0).load()) + assert chunk.size == air_source.sizes["lat"] * air_source.sizes["lon"] + assert chunk_peak < 10.0, ( + f"lazy single-chunk access should stay under 10 MB on " + f"air_temperature, got {chunk_peak:.2f} MB" ) _, eager_peak = _peak_mb( @@ -73,10 +70,10 @@ def test_lazy_slab_peak_memory_is_bounded(air_source): f"and the lazy assertion above no longer means anything. " f"Got {eager_peak:.1f} MB" ) - assert eager_peak / max(slab_peak, 0.1) > 10.0, ( - f"eager peak should be at least 10x the lazy slab peak; " + assert eager_peak / max(chunk_peak, 0.1) > 10.0, ( + f"eager peak should be at least 10x the lazy chunk peak; " f"if it isn't, the lazy path isn't actually lazy. " - f"Got eager={eager_peak:.1f} MB, lazy={slab_peak:.2f} MB" + f"Got eager={eager_peak:.1f} MB, lazy={chunk_peak:.2f} MB" )