diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index 2c846707..bacf88ac 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -10,11 +10,11 @@ repos: - id: check-json - id: check-toml - - repo: https://github.com/google/pyink - rev: 24.10.1 + - repo: https://github.com/astral-sh/ruff-pre-commit + rev: v0.15.10 hooks: - - id: pyink - # Configuration is read from pyproject.toml [tool.pyink] + - id: ruff-format + # Configuration is read from pyproject.toml [tool.ruff] - repo: https://github.com/pre-commit/mirrors-mypy rev: v1.11.2 diff --git a/perf_tests/compute_air.py b/perf_tests/compute_air.py index a564fbd2..2413370d 100755 --- a/perf_tests/compute_air.py +++ b/perf_tests/compute_air.py @@ -4,10 +4,10 @@ import xarray_sql as xql if __name__ == "__main__": - air = xr.tutorial.open_dataset("air_temperature") - chunks = {"time": 240} - air = air.chunk(chunks) + air = xr.tutorial.open_dataset("air_temperature") + chunks = {"time": 240} + air = air.chunk(chunks) - df = xql.read_xarray(air).read_pandas() + df = xql.read_xarray(air).read_pandas() - print(len(df)) + print(len(df)) diff --git a/perf_tests/groupby_air.py b/perf_tests/groupby_air.py index b54d7bf9..61006aeb 100755 --- a/perf_tests/groupby_air.py +++ b/perf_tests/groupby_air.py @@ -6,20 +6,20 @@ if __name__ == "__main__": - air = xr.tutorial.open_dataset("air_temperature") - chunks = {"time": 240, "lat": 5, "lon": 7} - air = air.chunk(chunks) - air_small = air.isel( - time=slice(0, 12), lat=slice(0, 11), lon=slice(0, 10) - ).chunk(chunks) + air = xr.tutorial.open_dataset("air_temperature") + chunks = {"time": 240, "lat": 5, "lon": 7} + air = air.chunk(chunks) + air_small = air.isel( + time=slice(0, 12), lat=slice(0, 11), lon=slice(0, 10) + ).chunk(chunks) - df = xql.read_xarray_table(air_small) + df = xql.read_xarray_table(air_small) - ctx = SessionContext() - ctx.register_table("air", df) + ctx = SessionContext() + ctx.register_table("air", df) - query = ctx.sql( - """ + query = ctx.sql( + """ SELECT "lat", "lon", SUM("air") as air_total FROM @@ -27,11 +27,11 @@ GROUP BY "lat", "lon" """ - ) + ) - result = query.collect() + result = query.collect() - expected = air_small.sizes["lat"] * air_small.sizes["lon"] - actual = sum(len(batch) for batch in result) - assert actual == expected, f"Length must be {expected}, but was {actual}." - print(expected) + expected = air_small.sizes["lat"] * air_small.sizes["lon"] + actual = sum(len(batch) for batch in result) + assert actual == expected, f"Length must be {expected}, but was {actual}." + print(expected) diff --git a/perf_tests/groupby_air_full.py b/perf_tests/groupby_air_full.py index 57663d84..2f2a6c1f 100755 --- a/perf_tests/groupby_air_full.py +++ b/perf_tests/groupby_air_full.py @@ -5,17 +5,17 @@ from datafusion import SessionContext if __name__ == "__main__": - air = xr.tutorial.open_dataset("air_temperature") - chunks = {"time": 240} - air = air.chunk(chunks) + air = xr.tutorial.open_dataset("air_temperature") + chunks = {"time": 240} + air = air.chunk(chunks) - df = xql.read_xarray_table(air) + df = xql.read_xarray_table(air) - ctx = SessionContext() - ctx.register_table("air", df) + ctx = SessionContext() + ctx.register_table("air", df) - query = ctx.sql( - """ + query = ctx.sql( + """ SELECT "lat", "lon", SUM("air") as air_total FROM @@ -23,12 +23,12 @@ GROUP BY "lat", "lon" """ - ) + ) - result = query.collect() + result = query.collect() - expected = air.sizes["lat"] * air.sizes["lon"] - actual = sum(len(batch) for batch in result) + expected = air.sizes["lat"] * air.sizes["lon"] + actual = sum(len(batch) for batch in result) - assert actual == expected, f"Length must be {expected}, but was {actual}." - print(expected) + assert actual == expected, f"Length must be {expected}, but was {actual}." + print(expected) diff --git a/perf_tests/sanity.py b/perf_tests/sanity.py index e89df822..9f8db6b6 100755 --- a/perf_tests/sanity.py +++ b/perf_tests/sanity.py @@ -4,13 +4,13 @@ import xarray_sql as xql if __name__ == "__main__": - air = xr.tutorial.open_dataset("air_temperature") - chunks = {"time": 240, "lat": 5, "lon": 7} + air = xr.tutorial.open_dataset("air_temperature") + chunks = {"time": 240, "lat": 5, "lon": 7} - air_small = air.isel( - time=slice(0, 12), lat=slice(0, 11), lon=slice(0, 10) - ).chunk(chunks) + air_small = air.isel( + time=slice(0, 12), lat=slice(0, 11), lon=slice(0, 10) + ).chunk(chunks) - df = xql.read_xarray(air_small).read_pandas() + df = xql.read_xarray(air_small).read_pandas() - print(len(df)) + print(len(df)) diff --git a/pyproject.toml b/pyproject.toml index 6edc9f87..d9253bdb 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -64,11 +64,13 @@ module-name = "xarray_sql._native" [tool.setuptools.packages.find] exclude = ["demo", "perf_tests", "tests", "tests.*"] -[tool.pyink] +[tool.ruff] line-length = 80 -preview = true -pyink-indentation = 2 -pyink-use-majority-quotes = true +indent-width = 4 + +[tool.ruff.format] +indent-style = "space" +quote-style = "double" [tool.mypy] python_version = "3.11" @@ -98,7 +100,7 @@ dev = [ "xarray_sql[test]", "xarray_sql[docs]", "py-spy>=0.4.0", - "pyink>=24.10.1", + "ruff>=0.15.10", "maturin>=1.9.1", ] diff --git a/tests/conftest.py b/tests/conftest.py index 48d5eba7..add0b59e 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -6,139 +6,145 @@ def rand_wx(start: str, end: str) -> xr.Dataset: - np.random.seed(42) - lat = np.linspace(-90, 90, num=720) - lon = np.linspace(-180, 180, num=1440) - time = pd.date_range(start, end, freq="h") - level = np.array([1000, 500], dtype=np.int32) - reference_time = pd.Timestamp(start) - temperature = 15 + 8 * np.random.randn(720, 1440, len(time), len(level)) - precipitation = 10 * np.random.rand(720, 1440, len(time), len(level)) - return xr.Dataset( - data_vars=dict( - temperature=(["lat", "lon", "time", "level"], temperature), - precipitation=(["lat", "lon", "time", "level"], precipitation), - ), - coords=dict( - lat=lat, - lon=lon, - time=time, - level=level, - reference_time=reference_time, - ), - attrs=dict(description="Random weather."), - ) + np.random.seed(42) + lat = np.linspace(-90, 90, num=720) + lon = np.linspace(-180, 180, num=1440) + time = pd.date_range(start, end, freq="h") + level = np.array([1000, 500], dtype=np.int32) + reference_time = pd.Timestamp(start) + temperature = 15 + 8 * np.random.randn(720, 1440, len(time), len(level)) + precipitation = 10 * np.random.rand(720, 1440, len(time), len(level)) + return xr.Dataset( + data_vars=dict( + temperature=(["lat", "lon", "time", "level"], temperature), + precipitation=(["lat", "lon", "time", "level"], precipitation), + ), + coords=dict( + lat=lat, + lon=lon, + time=time, + level=level, + reference_time=reference_time, + ), + attrs=dict(description="Random weather."), + ) def create_large_dataset(time_steps=1000, lat_points=100, lon_points=100): - """Create a large xarray dataset for memory testing.""" - np.random.seed(42) + """Create a large xarray dataset for memory testing.""" + np.random.seed(42) - time = pd.date_range("2020-01-01", periods=time_steps, freq="h") - lat = np.linspace(-90, 90, lat_points) - lon = np.linspace(-180, 180, lon_points) + time = pd.date_range("2020-01-01", periods=time_steps, freq="h") + lat = np.linspace(-90, 90, lat_points) + lon = np.linspace(-180, 180, lon_points) - temp_data = np.random.rand(time_steps, lat_points, lon_points) * 40 - 10 - precip_data = np.random.rand(time_steps, lat_points, lon_points) * 100 + temp_data = np.random.rand(time_steps, lat_points, lon_points) * 40 - 10 + precip_data = np.random.rand(time_steps, lat_points, lon_points) * 100 - return xr.Dataset( - { - "temperature": (["time", "lat", "lon"], temp_data), - "precipitation": (["time", "lat", "lon"], precip_data), - }, - coords={"time": time, "lat": lat, "lon": lon}, - ) + return xr.Dataset( + { + "temperature": (["time", "lat", "lon"], temp_data), + "precipitation": (["time", "lat", "lon"], precip_data), + }, + coords={"time": time, "lat": lat, "lon": lon}, + ) @pytest.fixture def air(): - ds = xr.tutorial.open_dataset("air_temperature") - chunks = {"time": 240} - return ds.chunk(chunks) + ds = xr.tutorial.open_dataset("air_temperature") + chunks = {"time": 240} + return ds.chunk(chunks) @pytest.fixture def air_small(air): - return air.isel(time=slice(0, 12), lat=slice(0, 11), lon=slice(0, 10)).chunk( - {"time": 240} - ) + return air.isel( + time=slice(0, 12), lat=slice(0, 11), lon=slice(0, 10) + ).chunk({"time": 240}) @pytest.fixture def randwx(): - return rand_wx("1995-01-13T00", "1995-01-13T01") + return rand_wx("1995-01-13T00", "1995-01-13T01") @pytest.fixture def large_ds(): - return create_large_dataset().chunk({"time": 25}) + return create_large_dataset().chunk({"time": 25}) @pytest.fixture def air_dataset_small(): - ds = xr.tutorial.open_dataset("air_temperature").chunk({"time": 240}) - return ds.isel(time=slice(0, 12), lat=slice(0, 11), lon=slice(0, 10)) + ds = xr.tutorial.open_dataset("air_temperature").chunk({"time": 240}) + return ds.isel(time=slice(0, 12), lat=slice(0, 11), lon=slice(0, 10)) @pytest.fixture def air_dataset_large(): - return xr.tutorial.open_dataset("air_temperature").chunk({"time": 240}) + return xr.tutorial.open_dataset("air_temperature").chunk({"time": 240}) @pytest.fixture def rasm_ds(): - """rasm uses cftime.DatetimeNoLeap (noleap / 365_day) for time.""" - return xr.tutorial.open_dataset("rasm") + """rasm uses cftime.DatetimeNoLeap (noleap / 365_day) for time.""" + return xr.tutorial.open_dataset("rasm") @pytest.fixture def weather_dataset(): - ds = rand_wx("2023-01-01T00", "2023-01-01T12") - return ds.isel(time=slice(0, 6), lat=slice(0, 10), lon=slice(0, 10)).chunk( - {"time": 3} - ) + ds = rand_wx("2023-01-01T00", "2023-01-01T12") + return ds.isel(time=slice(0, 6), lat=slice(0, 10), lon=slice(0, 10)).chunk( + {"time": 3} + ) @pytest.fixture def synthetic_dataset(): - return create_large_dataset( - time_steps=50, lat_points=20, lon_points=20 - ).chunk({"time": 25}) + return create_large_dataset( + time_steps=50, lat_points=20, lon_points=20 + ).chunk({"time": 25}) @pytest.fixture def station_dataset(): - return xr.Dataset( - { - "station_id": (["station"], [1, 2, 3, 4, 5]), - "elevation": (["station"], [100, 250, 500, 750, 1000]), - "name": ( - ["station"], - ["Station_A", "Station_B", "Station_C", "Station_D", "Station_E"], - ), - } - ).chunk({"station": 5}) + return xr.Dataset( + { + "station_id": (["station"], [1, 2, 3, 4, 5]), + "elevation": (["station"], [100, 250, 500, 750, 1000]), + "name": ( + ["station"], + [ + "Station_A", + "Station_B", + "Station_C", + "Station_D", + "Station_E", + ], + ), + } + ).chunk({"station": 5}) @pytest.fixture def air_and_stations(): - air = ( - xr.tutorial.open_dataset("air_temperature") - .isel(time=slice(0, 12), lat=slice(0, 5), lon=slice(0, 8)) - .chunk({"time": 6}) - ) - stations = xr.Dataset( - { - "station_id": (["station"], [101, 102, 103]), - "lat": ( - ["station"], - [air.lat.values[0], air.lat.values[2], air.lat.values[4]], - ), - "lon": ( - ["station"], - [air.lon.values[1], air.lon.values[3], air.lon.values[5]], - ), - "elevation": (["station"], [100, 250, 500]), - } - ).chunk({"station": 3}) - return air, stations + air = ( + xr.tutorial.open_dataset("air_temperature") + .isel(time=slice(0, 12), lat=slice(0, 5), lon=slice(0, 8)) + .chunk({"time": 6}) + ) + stations = xr.Dataset( + { + "station_id": (["station"], [101, 102, 103]), + "lat": ( + ["station"], + [air.lat.values[0], air.lat.values[2], air.lat.values[4]], + ), + "lon": ( + ["station"], + [air.lon.values[1], air.lon.values[3], air.lon.values[5]], + ), + "elevation": (["station"], [100, 250, 500]), + } + ).chunk({"station": 3}) + return air, stations diff --git a/tests/test_cft.py b/tests/test_cft.py index 45820bd9..6560d723 100644 --- a/tests/test_cft.py +++ b/tests/test_cft.py @@ -15,162 +15,156 @@ @pytest.fixture def ds_360day(): - """Synthetic 360-day calendar dataset.""" - import cftime + """Synthetic 360-day calendar dataset.""" + import cftime - times = [cftime.Datetime360Day(2000, m, 1) for m in range(1, 13)] - return xr.Dataset( - {"temp": ("time", np.arange(12, dtype=np.float32))}, - coords={"time": times}, - ) + times = [cftime.Datetime360Day(2000, m, 1) for m in range(1, 13)] + return xr.Dataset( + {"temp": ("time", np.arange(12, dtype=np.float32))}, + coords={"time": times}, + ) # -- Detection helpers ------------------------------------------------------ class TestDetection: + def test_is_cftime_detects_cftime_array(self, rasm_ds): + assert cft.is_cftime(rasm_ds.coords["time"].values) - def test_is_cftime_detects_cftime_array(self, rasm_ds): - assert cft.is_cftime(rasm_ds.coords["time"].values) + def test_is_cftime_rejects_datetime64(self): + assert not cft.is_cftime(pd.date_range("2020-01-01", periods=10).values) - def test_is_cftime_rejects_datetime64(self): - assert not cft.is_cftime(pd.date_range("2020-01-01", periods=10).values) + def test_is_cftime_rejects_float(self): + assert not cft.is_cftime(np.array([1.0, 2.0, 3.0])) - def test_is_cftime_rejects_float(self): - assert not cft.is_cftime(np.array([1.0, 2.0, 3.0])) + def test_is_cftime_index_detects_cftime(self, rasm_ds): + assert cft.is_cftime_index(rasm_ds, "time") - def test_is_cftime_index_detects_cftime(self, rasm_ds): - assert cft.is_cftime_index(rasm_ds, "time") + def test_is_cftime_index_rejects_datetime64(self): + ds = xr.tutorial.open_dataset("air_temperature") + assert not cft.is_cftime_index(ds, "time") - def test_is_cftime_index_rejects_datetime64(self): - ds = xr.tutorial.open_dataset("air_temperature") - assert not cft.is_cftime_index(ds, "time") - - def test_is_cftime_index_rejects_nonexistent(self, rasm_ds): - assert not cft.is_cftime_index(rasm_ds, "nonexistent") + def test_is_cftime_index_rejects_nonexistent(self, rasm_ds): + assert not cft.is_cftime_index(rasm_ds, "nonexistent") # -- Calendar classification ------------------------------------------------ class TestCalendarClassification: + def test_calendar_returns_noleap(self, rasm_ds): + assert cft.calendar(rasm_ds, "time") == "noleap" - def test_calendar_returns_noleap(self, rasm_ds): - assert cft.calendar(rasm_ds, "time") == "noleap" - - def test_calendar_returns_360_day(self, ds_360day): - assert cft.calendar(ds_360day, "time") == "360_day" + def test_calendar_returns_360_day(self, ds_360day): + assert cft.calendar(ds_360day, "time") == "360_day" - def test_calendar_returns_none_for_datetime64(self): - ds = xr.tutorial.open_dataset("air_temperature") - assert cft.calendar(ds, "time") is None + def test_calendar_returns_none_for_datetime64(self): + ds = xr.tutorial.open_dataset("air_temperature") + assert cft.calendar(ds, "time") is None - def test_noleap_is_gregorian_like(self): - assert cft.is_gregorian_like("noleap") - assert cft.is_gregorian_like("standard") - assert cft.is_gregorian_like("proleptic_gregorian") - assert cft.is_gregorian_like("all_leap") + def test_noleap_is_gregorian_like(self): + assert cft.is_gregorian_like("noleap") + assert cft.is_gregorian_like("standard") + assert cft.is_gregorian_like("proleptic_gregorian") + assert cft.is_gregorian_like("all_leap") - def test_360_day_is_not_gregorian_like(self): - assert not cft.is_gregorian_like("360_day") - assert not cft.is_gregorian_like("julian") + def test_360_day_is_not_gregorian_like(self): + assert not cft.is_gregorian_like("360_day") + assert not cft.is_gregorian_like("julian") # -- Numeric conversion ----------------------------------------------------- class TestConversion: + def test_to_microseconds_returns_int64(self, rasm_ds): + us = cft.to_microseconds(rasm_ds.coords["time"].values) + assert us.dtype == np.int64 - def test_to_microseconds_returns_int64(self, rasm_ds): - us = cft.to_microseconds(rasm_ds.coords["time"].values) - assert us.dtype == np.int64 - - def test_to_microseconds_is_monotonic(self, rasm_ds): - us = cft.to_microseconds(rasm_ds.coords["time"].values) - assert np.all(np.diff(us) > 0) + def test_to_microseconds_is_monotonic(self, rasm_ds): + us = cft.to_microseconds(rasm_ds.coords["time"].values) + assert np.all(np.diff(us) > 0) - def test_to_microseconds_length_matches(self, rasm_ds): - values = rasm_ds.coords["time"].values - assert len(cft.to_microseconds(values)) == len(values) + def test_to_microseconds_length_matches(self, rasm_ds): + values = rasm_ds.coords["time"].values + assert len(cft.to_microseconds(values)) == len(values) - def test_to_offsets_returns_int64(self, ds_360day): - values = ds_360day.coords["time"].values - offsets = cft.to_offsets(values, cft.DEFAULT_UNITS, "360_day") - assert offsets.dtype == np.int64 + def test_to_offsets_returns_int64(self, ds_360day): + values = ds_360day.coords["time"].values + offsets = cft.to_offsets(values, cft.DEFAULT_UNITS, "360_day") + assert offsets.dtype == np.int64 - def test_to_offsets_is_monotonic(self, ds_360day): - values = ds_360day.coords["time"].values - offsets = cft.to_offsets(values, cft.DEFAULT_UNITS, "360_day") - assert np.all(np.diff(offsets) > 0) + def test_to_offsets_is_monotonic(self, ds_360day): + values = ds_360day.coords["time"].values + offsets = cft.to_offsets(values, cft.DEFAULT_UNITS, "360_day") + assert np.all(np.diff(offsets) > 0) - def test_convert_for_field_gregorian_like(self, rasm_ds): - field = cft.arrow_field("time", cft.DEFAULT_UNITS, "noleap") - result = cft.convert_for_field(rasm_ds.coords["time"].values, field) - assert result.dtype == np.int64 - assert np.all(np.diff(result) > 0) + def test_convert_for_field_gregorian_like(self, rasm_ds): + field = cft.arrow_field("time", cft.DEFAULT_UNITS, "noleap") + result = cft.convert_for_field(rasm_ds.coords["time"].values, field) + assert result.dtype == np.int64 + assert np.all(np.diff(result) > 0) - def test_convert_for_field_non_gregorian(self, ds_360day): - field = cft.arrow_field("time", cft.DEFAULT_UNITS, "360_day") - result = cft.convert_for_field(ds_360day.coords["time"].values, field) - assert result.dtype == np.int64 - assert np.all(np.diff(result) > 0) + def test_convert_for_field_non_gregorian(self, ds_360day): + field = cft.arrow_field("time", cft.DEFAULT_UNITS, "360_day") + result = cft.convert_for_field(ds_360day.coords["time"].values, field) + assert result.dtype == np.int64 + assert np.all(np.diff(result) > 0) # -- Arrow schema helpers --------------------------------------------------- class TestArrowField: + def test_gregorian_like_produces_timestamp_us(self): + field = cft.arrow_field("time", cft.DEFAULT_UNITS, "noleap") + assert field.type == pa.timestamp("us") + assert field.metadata[b"xarray:calendar"] == b"noleap" + assert field.metadata[b"xarray:units"] == cft.DEFAULT_UNITS.encode() - def test_gregorian_like_produces_timestamp_us(self): - field = cft.arrow_field("time", cft.DEFAULT_UNITS, "noleap") - assert field.type == pa.timestamp("us") - assert field.metadata[b"xarray:calendar"] == b"noleap" - assert field.metadata[b"xarray:units"] == cft.DEFAULT_UNITS.encode() - - def test_non_gregorian_produces_int64(self): - field = cft.arrow_field("time", cft.DEFAULT_UNITS, "360_day") - assert field.type == pa.int64() - assert field.metadata[b"xarray:calendar"] == b"360_day" + def test_non_gregorian_produces_int64(self): + field = cft.arrow_field("time", cft.DEFAULT_UNITS, "360_day") + assert field.type == pa.int64() + assert field.metadata[b"xarray:calendar"] == b"360_day" # -- Partition bounds ------------------------------------------------------- class TestPartitionBounds: + def test_gregorian_like_returns_timestamp_ns_tag(self, rasm_ds): + values = rasm_ds.coords["time"].values[:10] + lo, hi, tag = cft.partition_bounds(values) + assert tag == "timestamp_ns" + assert lo < hi - def test_gregorian_like_returns_timestamp_ns_tag(self, rasm_ds): - values = rasm_ds.coords["time"].values[:10] - lo, hi, tag = cft.partition_bounds(values) - assert tag == "timestamp_ns" - assert lo < hi - - def test_non_gregorian_returns_int64_tag(self, ds_360day): - values = ds_360day.coords["time"].values - lo, hi, tag = cft.partition_bounds(values) - assert tag == "int64" - assert lo < hi + def test_non_gregorian_returns_int64_tag(self, ds_360day): + values = ds_360day.coords["time"].values + lo, hi, tag = cft.partition_bounds(values) + assert tag == "int64" + assert lo < hi # -- Integration with _parse_schema ---------------------------------------- class TestParseSchemaIntegration: - - def test_noleap_produces_timestamp_us(self, rasm_ds): - schema = _parse_schema(rasm_ds[["Tair"]]) - time_field = schema.field("time") - assert time_field.type == pa.timestamp("us") - assert time_field.metadata[b"xarray:calendar"] == b"noleap" - - def test_360day_produces_int64(self, ds_360day): - schema = _parse_schema(ds_360day) - time_field = schema.field("time") - assert time_field.type == pa.int64() - assert time_field.metadata[b"xarray:calendar"] == b"360_day" - - def test_datetime64_unchanged(self): - ds = xr.tutorial.open_dataset("air_temperature") - schema = _parse_schema(ds) - time_field = schema.field("time") - assert pa.types.is_timestamp(time_field.type) - assert time_field.metadata is None # no xarray: metadata for native + def test_noleap_produces_timestamp_us(self, rasm_ds): + schema = _parse_schema(rasm_ds[["Tair"]]) + time_field = schema.field("time") + assert time_field.type == pa.timestamp("us") + assert time_field.metadata[b"xarray:calendar"] == b"noleap" + + def test_360day_produces_int64(self, ds_360day): + schema = _parse_schema(ds_360day) + time_field = schema.field("time") + assert time_field.type == pa.int64() + assert time_field.metadata[b"xarray:calendar"] == b"360_day" + + def test_datetime64_unchanged(self): + ds = xr.tutorial.open_dataset("air_temperature") + schema = _parse_schema(ds) + time_field = schema.field("time") + assert pa.types.is_timestamp(time_field.type) + assert time_field.metadata is None # no xarray: metadata for native diff --git a/tests/test_df.py b/tests/test_df.py index 6e377935..4ebb6759 100644 --- a/tests/test_df.py +++ b/tests/test_df.py @@ -3,7 +3,6 @@ import numpy as np import pandas as pd import pyarrow as pa -import pytest import xarray as xr from xarray_sql.df import ( @@ -21,408 +20,424 @@ def test_explode_cardinality(air): - dss = explode(air) - assert len(list(dss)) == np.prod([len(c) for c in air.chunks.values()]) + dss = explode(air) + assert len(list(dss)) == np.prod([len(c) for c in air.chunks.values()]) def test_explode_dim_sizes_one(air): - chunks = {"time": 240} - ds = next(iter(explode(air))) - for k, v in chunks.items(): - assert k in ds.dims - assert v == ds.sizes[k] + chunks = {"time": 240} + ds = next(iter(explode(air))) + for k, v in chunks.items(): + assert k in ds.dims + assert v == ds.sizes[k] def test_explode_data_equal_one_first(air): - ds = next(iter(explode(air))) - iselection = {dim: slice(0, s) for dim, s in ds.sizes.items()} - assert air.isel(iselection).equals(ds) + ds = next(iter(explode(air))) + iselection = {dim: slice(0, s) for dim, s in ds.sizes.items()} + assert air.isel(iselection).equals(ds) def test_explode_data_equal_one_last(air): - dss = list(explode(air)) - ds = dss[-1] + dss = list(explode(air)) + ds = dss[-1] - # For the last chunk, we need to calculate where it actually starts - # The original logic slice(0, s) only works for the first chunk - iselection = {} - for dim in ds.dims: - # Get chunk boundaries - chunk_bounds = np.cumsum((0,) + air.chunks[dim]) - # Last chunk index - last_chunk_idx = len(air.chunks[dim]) - 1 - # Calculate actual start and end positions - start = chunk_bounds[last_chunk_idx] - end = chunk_bounds[last_chunk_idx + 1] - iselection[dim] = slice(start, end) + # For the last chunk, we need to calculate where it actually starts + # The original logic slice(0, s) only works for the first chunk + iselection = {} + for dim in ds.dims: + # Get chunk boundaries + chunk_bounds = np.cumsum((0,) + air.chunks[dim]) + # Last chunk index + last_chunk_idx = len(air.chunks[dim]) - 1 + # Calculate actual start and end positions + start = chunk_bounds[last_chunk_idx] + end = chunk_bounds[last_chunk_idx + 1] + iselection[dim] = slice(start, end) - assert air.isel(iselection).equals(ds) + assert air.isel(iselection).equals(ds) def test_from_map_basic(): - def make_df(x): - return pd.DataFrame({"value": [x, x * 2], "index": [0, 1]}) + def make_df(x): + return pd.DataFrame({"value": [x, x * 2], "index": [0, 1]}) - result = from_map(make_df, [1, 2, 3]) - assert isinstance(result, pa.Table) - assert len(result) == 6 - assert result.column_names == ["value", "index"] + result = from_map(make_df, [1, 2, 3]) + assert isinstance(result, pa.Table) + assert len(result) == 6 + assert result.column_names == ["value", "index"] def test_from_map_multiple_iterables(): - def add_values(x, y): - return pd.DataFrame({"sum": [x + y], "x": [x], "y": [y]}) + def add_values(x, y): + return pd.DataFrame({"sum": [x + y], "x": [x], "y": [y]}) - result = from_map(add_values, [1, 2], [10, 20]) - assert isinstance(result, pa.Table) - assert len(result) == 2 + result = from_map(add_values, [1, 2], [10, 20]) + assert isinstance(result, pa.Table) + assert len(result) == 2 - df = result.to_pandas() - assert list(df["sum"]) == [11, 22] + df = result.to_pandas() + assert list(df["sum"]) == [11, 22] def test_from_map_with_args(): - def multiply_and_add(x, multiplier, add_value): - return pd.DataFrame({"result": [x * multiplier + add_value]}) + def multiply_and_add(x, multiplier, add_value): + return pd.DataFrame({"result": [x * multiplier + add_value]}) - result = from_map(multiply_and_add, [1, 2, 3], args=(2, 10)) - assert isinstance(result, pa.Table) - assert len(result) == 3 + result = from_map(multiply_and_add, [1, 2, 3], args=(2, 10)) + assert isinstance(result, pa.Table) + assert len(result) == 3 - df = result.to_pandas() - assert list(df["result"]) == [12, 14, 16] + df = result.to_pandas() + assert list(df["result"]) == [12, 14, 16] def test_from_map_with_pyarrow_tables(): - def make_arrow_table(x): - df = pd.DataFrame({"value": [x]}) - return pa.Table.from_pandas(df) + def make_arrow_table(x): + df = pd.DataFrame({"value": [x]}) + return pa.Table.from_pandas(df) - result = from_map(make_arrow_table, [1, 2, 3]) - assert isinstance(result, pa.Table) - assert len(result) == 3 + result = from_map(make_arrow_table, [1, 2, 3]) + assert isinstance(result, pa.Table) + assert len(result) == 3 def test_iter_record_batches_splits_into_multiple_batches(air_small): - """iter_record_batches should emit >1 batch when partition exceeds batch_size.""" - schema = _parse_schema(air_small) - block = next(block_slices(air_small, chunks={"time": 4, "lat": 3, "lon": 4})) - ds_block = air_small.isel(block) - total_rows = int(np.prod([ds_block.sizes[d] for d in ds_block.sizes])) + """iter_record_batches should emit >1 batch when partition exceeds batch_size.""" + schema = _parse_schema(air_small) + block = next( + block_slices(air_small, chunks={"time": 4, "lat": 3, "lon": 4}) + ) + ds_block = air_small.isel(block) + total_rows = int(np.prod([ds_block.sizes[d] for d in ds_block.sizes])) - small_batch = 16 # force many small batches - batches = list(iter_record_batches(ds_block, schema, batch_size=small_batch)) + small_batch = 16 # force many small batches + batches = list( + iter_record_batches(ds_block, schema, batch_size=small_batch) + ) - assert len(batches) == -(-total_rows // small_batch) # ceiling division - assert all(b.num_rows <= small_batch for b in batches) - assert sum(b.num_rows for b in batches) == total_rows + assert len(batches) == -(-total_rows // small_batch) # ceiling division + assert all(b.num_rows <= small_batch for b in batches) + assert sum(b.num_rows for b in batches) == total_rows def test_iter_record_batches_matches_dataset_to_record_batch(air_small): - """Concatenating all iter_record_batches output must equal dataset_to_record_batch.""" - schema = _parse_schema(air_small) - dim_cols = [f.name for f in schema if f.name in air_small.dims] - block = next(block_slices(air_small, chunks={"time": 4, "lat": 3, "lon": 4})) - ds_block = air_small.isel(block) - - batches = list(iter_record_batches(ds_block, schema, batch_size=16)) - actual_df = ( - pa.Table.from_batches(batches) - .to_pandas() - .sort_values(dim_cols) - .reset_index(drop=True) - ) - expected_df = ( - dataset_to_record_batch(ds_block, schema) - .to_pandas() - .sort_values(dim_cols) - .reset_index(drop=True) - ) - pd.testing.assert_frame_equal(actual_df, expected_df) - - -def test_iter_record_batches_default_batch_size(): - """A single-batch partition (rows <= DEFAULT_BATCH_SIZE) yields exactly one batch.""" - ds = xr.tutorial.open_dataset("air_temperature").isel(time=slice(0, 2)) - schema = _parse_schema(ds) - total_rows = int(np.prod([ds.sizes[d] for d in ds.sizes])) - assert total_rows <= DEFAULT_BATCH_SIZE, "fixture too large — adjust isel" - batches = list(iter_record_batches(ds, schema)) - assert len(batches) == 1 - assert batches[0].num_rows == total_rows - - -def test_dataset_to_record_batch_matches_pivot(air_small): - """dataset_to_record_batch should contain the same rows as pivot. - - Row ordering may differ (pivot uses ds.dims key order; dataset_to_record_batch - uses the data variable's own dim order). Both orderings are valid for SQL, so - we sort by the coordinate columns before comparing. - """ - schema = _parse_schema(air_small) - dim_cols = [f.name for f in schema if f.name in air_small.dims] - blocks = list(block_slices(air_small, chunks={"time": 4, "lat": 3, "lon": 4})) - - for block in blocks: + """Concatenating all iter_record_batches output must equal dataset_to_record_batch.""" + schema = _parse_schema(air_small) + dim_cols = [f.name for f in schema if f.name in air_small.dims] + block = next( + block_slices(air_small, chunks={"time": 4, "lat": 3, "lon": 4}) + ) ds_block = air_small.isel(block) + + batches = list(iter_record_batches(ds_block, schema, batch_size=16)) actual_df = ( - dataset_to_record_batch(ds_block, schema) + pa.Table.from_batches(batches) .to_pandas() .sort_values(dim_cols) .reset_index(drop=True) ) expected_df = ( - pa.RecordBatch.from_pandas(pivot(ds_block), schema=schema) + dataset_to_record_batch(ds_block, schema) .to_pandas() .sort_values(dim_cols) .reset_index(drop=True) ) + pd.testing.assert_frame_equal(actual_df, expected_df) + + +def test_iter_record_batches_default_batch_size(): + """A single-batch partition (rows <= DEFAULT_BATCH_SIZE) yields exactly one batch.""" + ds = xr.tutorial.open_dataset("air_temperature").isel(time=slice(0, 2)) + schema = _parse_schema(ds) + total_rows = int(np.prod([ds.sizes[d] for d in ds.sizes])) + assert total_rows <= DEFAULT_BATCH_SIZE, "fixture too large — adjust isel" + batches = list(iter_record_batches(ds, schema)) + assert len(batches) == 1 + assert batches[0].num_rows == total_rows + + +def test_dataset_to_record_batch_matches_pivot(air_small): + """dataset_to_record_batch should contain the same rows as pivot. + + Row ordering may differ (pivot uses ds.dims key order; dataset_to_record_batch + uses the data variable's own dim order). Both orderings are valid for SQL, so + we sort by the coordinate columns before comparing. + """ + schema = _parse_schema(air_small) + dim_cols = [f.name for f in schema if f.name in air_small.dims] + blocks = list( + block_slices(air_small, chunks={"time": 4, "lat": 3, "lon": 4}) + ) + + for block in blocks: + ds_block = air_small.isel(block) + actual_df = ( + dataset_to_record_batch(ds_block, schema) + .to_pandas() + .sort_values(dim_cols) + .reset_index(drop=True) + ) + expected_df = ( + pa.RecordBatch.from_pandas(pivot(ds_block), schema=schema) + .to_pandas() + .sort_values(dim_cols) + .reset_index(drop=True) + ) - pd.testing.assert_frame_equal(actual_df, expected_df, check_like=False) + pd.testing.assert_frame_equal(actual_df, expected_df, check_like=False) def test_dataset_to_record_batch_column_order(air_small): - """Output column order must match schema (dims first, then data vars).""" - schema = _parse_schema(air_small) - block = next(block_slices(air_small, chunks={"time": 4, "lat": 3, "lon": 4})) - batch = dataset_to_record_batch(air_small.isel(block), schema) - assert batch.schema.names == schema.names + """Output column order must match schema (dims first, then data vars).""" + schema = _parse_schema(air_small) + block = next( + block_slices(air_small, chunks={"time": 4, "lat": 3, "lon": 4}) + ) + batch = dataset_to_record_batch(air_small.isel(block), schema) + assert batch.schema.names == schema.names def test_dataset_to_record_batch_row_count(air_small): - """Row count must equal the product of the block dimension sizes.""" - schema = _parse_schema(air_small) - chunks = {"time": 4, "lat": 3, "lon": 4} - for block in block_slices(air_small, chunks=chunks): - ds_block = air_small.isel(block) - expected_rows = int(np.prod([ds_block.sizes[d] for d in ds_block.sizes])) - batch = dataset_to_record_batch(ds_block, schema) - assert batch.num_rows == expected_rows + """Row count must equal the product of the block dimension sizes.""" + schema = _parse_schema(air_small) + chunks = {"time": 4, "lat": 3, "lon": 4} + for block in block_slices(air_small, chunks=chunks): + ds_block = air_small.isel(block) + expected_rows = int( + np.prod([ds_block.sizes[d] for d in ds_block.sizes]) + ) + batch = dataset_to_record_batch(ds_block, schema) + assert batch.num_rows == expected_rows def test_from_map_batched_basic_functionality(air_small): - blocks = list(block_slices(air_small, chunks={"time": 4, "lat": 3, "lon": 4})) + blocks = list( + block_slices(air_small, chunks={"time": 4, "lat": 3, "lon": 4}) + ) - first_block_df = pivot(air_small.isel(blocks[0])) - expected_schema = pa.Schema.from_pandas(first_block_df) + first_block_df = pivot(air_small.isel(blocks[0])) + expected_schema = pa.Schema.from_pandas(first_block_df) - reader = from_map_batched( - pivot, [air_small.isel(block) for block in blocks], schema=expected_schema - ) + reader = from_map_batched( + pivot, + [air_small.isel(block) for block in blocks], + schema=expected_schema, + ) - assert isinstance(reader, pa.RecordBatchReader) - assert reader.schema == expected_schema + assert isinstance(reader, pa.RecordBatchReader) + assert reader.schema == expected_schema - batches = list(reader) - assert len(batches) > 0 - for batch in batches: - assert batch.schema == expected_schema - assert len(batch) > 0 + batches = list(reader) + assert len(batches) > 0 + for batch in batches: + assert batch.schema == expected_schema + assert len(batch) > 0 def adding_function(x, y): - """Simple function that adds two values and returns a DataFrame.""" - result = pd.DataFrame({"x": [x], "y": [y], "sum": [x + y]}) - return result + """Simple function that adds two values and returns a DataFrame.""" + result = pd.DataFrame({"x": [x], "y": [y], "sum": [x + y]}) + return result def test_from_map_batched_multiple_iterables(): - x_values = [1, 2, 3, 4, 5] - y_values = [10, 20, 30, 40, 50] + x_values = [1, 2, 3, 4, 5] + y_values = [10, 20, 30, 40, 50] - expected_schema = pa.schema( - [("x", pa.int64()), ("y", pa.int64()), ("sum", pa.int64())] - ) - - reader = from_map_batched( - adding_function, x_values, y_values, schema=expected_schema - ) - table = reader.read_all() - df = table.to_pandas() + expected_schema = pa.schema( + [("x", pa.int64()), ("y", pa.int64()), ("sum", pa.int64())] + ) - expected_df = pd.DataFrame( - { - "x": x_values, - "y": y_values, - "sum": [x + y for x, y in zip(x_values, y_values)], - } - ) - pd.testing.assert_frame_equal(df, expected_df) + reader = from_map_batched( + adding_function, x_values, y_values, schema=expected_schema + ) + table = reader.read_all() + df = table.to_pandas() + + expected_df = pd.DataFrame( + { + "x": x_values, + "y": y_values, + "sum": [x + y for x, y in zip(x_values, y_values)], + } + ) + pd.testing.assert_frame_equal(df, expected_df) def test_from_map_batched_with_args_and_kwargs(): - def multiply_and_add(x, multiplier, offset=0): - return pd.DataFrame({"x": [x], "result": [x * multiplier + offset]}) + def multiply_and_add(x, multiplier, offset=0): + return pd.DataFrame({"x": [x], "result": [x * multiplier + offset]}) - values = [1, 2, 3] - expected_schema = pa.schema([("x", pa.int64()), ("result", pa.int64())]) + values = [1, 2, 3] + expected_schema = pa.schema([("x", pa.int64()), ("result", pa.int64())]) - reader = from_map_batched( - multiply_and_add, values, args=(2,), offset=5, schema=expected_schema - ) - table = reader.read_all() - df = table.to_pandas() + reader = from_map_batched( + multiply_and_add, values, args=(2,), offset=5, schema=expected_schema + ) + table = reader.read_all() + df = table.to_pandas() - expected_df = pd.DataFrame({"x": [1, 2, 3], "result": [7, 9, 11]}) - pd.testing.assert_frame_equal(df, expected_df) + expected_df = pd.DataFrame({"x": [1, 2, 3], "result": [7, 9, 11]}) + pd.testing.assert_frame_equal(df, expected_df) def test_from_map_batched_empty_iterables(): - empty_schema = pa.schema([("value", pa.int64())]) + empty_schema = pa.schema([("value", pa.int64())]) - reader = from_map_batched( - lambda x: pd.DataFrame({"value": [x]}), [], schema=empty_schema - ) - batches = list(reader) - assert len(batches) == 0 + reader = from_map_batched( + lambda x: pd.DataFrame({"value": [x]}), [], schema=empty_schema + ) + batches = list(reader) + assert len(batches) == 0 def test_from_map_batched_consistency_with_regular_map(air_small): - blocks = list(block_slices(air_small, chunks={"time": 4, "lat": 3})) - datasets = [air_small.isel(block) for block in blocks] + blocks = list(block_slices(air_small, chunks={"time": 4, "lat": 3})) + datasets = [air_small.isel(block) for block in blocks] - first_df = pivot(datasets[0]) - schema = pa.Schema.from_pandas(first_df) + first_df = pivot(datasets[0]) + schema = pa.Schema.from_pandas(first_df) - reader = from_map_batched(pivot, datasets, schema=schema) - batched_table = reader.read_all() + reader = from_map_batched(pivot, datasets, schema=schema) + batched_table = reader.read_all() - regular_dfs = [pivot(ds) for ds in datasets] - regular_table = pa.Table.from_pandas( - pd.concat(regular_dfs, ignore_index=True) - ) + regular_dfs = [pivot(ds) for ds in datasets] + regular_table = pa.Table.from_pandas( + pd.concat(regular_dfs, ignore_index=True) + ) - assert batched_table.schema == regular_table.schema - assert len(batched_table) == len(regular_table) + assert batched_table.schema == regular_table.schema + assert len(batched_table) == len(regular_table) - batched_df = ( - batched_table.to_pandas() - .sort_values(["time", "lat", "lon"]) - .reset_index(drop=True) - ) - regular_df = ( - regular_table.to_pandas() - .sort_values(["time", "lat", "lon"]) - .reset_index(drop=True) - ) + batched_df = ( + batched_table.to_pandas() + .sort_values(["time", "lat", "lon"]) + .reset_index(drop=True) + ) + regular_df = ( + regular_table.to_pandas() + .sort_values(["time", "lat", "lon"]) + .reset_index(drop=True) + ) - pd.testing.assert_frame_equal(batched_df, regular_df) + pd.testing.assert_frame_equal(batched_df, regular_df) def test_from_map_batched_integration_with_datafusion_via_read_xarray(): - air = xr.tutorial.open_dataset("air_temperature") - air_small = air.isel(time=slice(0, 50), lat=slice(0, 10), lon=slice(0, 15)) - air_chunked = air_small.chunk({"time": 25, "lat": 5, "lon": 8}) + air = xr.tutorial.open_dataset("air_temperature") + air_small = air.isel(time=slice(0, 50), lat=slice(0, 10), lon=slice(0, 15)) + air_chunked = air_small.chunk({"time": 25, "lat": 5, "lon": 8}) - arrow_stream = read_xarray( - air_chunked, chunks={"time": 25, "lat": 5, "lon": 8} - ) + arrow_stream = read_xarray( + air_chunked, chunks={"time": 25, "lat": 5, "lon": 8} + ) - assert hasattr(arrow_stream, "schema") - assert hasattr(arrow_stream, "__iter__") + assert hasattr(arrow_stream, "schema") + assert hasattr(arrow_stream, "__iter__") - table = arrow_stream.read_all() - assert len(table) > 0 + table = arrow_stream.read_all() + assert len(table) > 0 - expected_columns = {"time", "lat", "lon", "air"} - actual_columns = set(table.column_names) - assert expected_columns.issubset(actual_columns) + expected_columns = {"time", "lat", "lon", "air"} + actual_columns = set(table.column_names) + assert expected_columns.issubset(actual_columns) def test_read_xarray_loads_one_chunk_at_a_time(large_ds): - tracemalloc.stop() # reset any state left by a previously-failed test - tracemalloc.start() - try: - iterable = read_xarray(large_ds) - first_size, first_peak = tracemalloc.get_traced_memory() - tracemalloc.reset_peak() + tracemalloc.stop() # reset any state left by a previously-failed test + tracemalloc.start() + try: + iterable = read_xarray(large_ds) + first_size, first_peak = tracemalloc.get_traced_memory() + tracemalloc.reset_peak() + + sizes, peaks = [], [] + + first_chunk = large_ds.isel(next(block_slices(large_ds))) + chunk_size = first_chunk.nbytes + + # Creating the iterator should be inexpensive -- less than one chunk. + # We multiply by constant factors because chunks have additional overhead + assert first_size < chunk_size * 3 + assert first_peak < chunk_size * 6 + + for it in iterable: + _ = it + cur_size, cur_peak = tracemalloc.get_traced_memory() + tracemalloc.reset_peak() + sizes.append(cur_size) + peaks.append(cur_peak) + + for size in sizes: + # Observed range: 1.59–1.83× on macOS, up to ~2.7× on Linux + # (glibc + Arrow allocate more intermediate buffers). + # iter_record_batches holds data-variable arrays (≈1× chunk) while + # yielding sub-batches, plus the current Arrow batch (≈0.65× chunk). + assert chunk_size * 1.3 < size, f"size {size} unexpectedly low" + assert chunk_size * 3.5 > size, f"size {size} unexpectedly high" + + for peak in peaks: + # Observed range: 1.84–3.28× on macOS, up to ~4.15× on Linux + # (glibc + Arrow hold more intermediate buffers at peak). + # Peak includes data arrays + Arrow batch + temporary coordinate index + # arrays; the first batch of each chunk is highest (Dask compute overhead). + assert chunk_size * 1.5 < peak, f"peak {peak} unexpectedly low" + assert chunk_size * 5.0 > peak, f"peak {peak} unexpectedly high" + + assert max(peaks) < large_ds.nbytes + finally: + tracemalloc.stop() - sizes, peaks = [], [] - first_chunk = large_ds.isel(next(block_slices(large_ds))) - chunk_size = first_chunk.nbytes +def test_read_xarray_table_memory_bounds(large_ds): + """read_xarray_table should not materialise data at registration time. - # Creating the iterator should be inexpensive -- less than one chunk. - # We multiply by constant factors because chunks have additional overhead - assert first_size < chunk_size * 3 - assert first_peak < chunk_size * 6 - - for it in iterable: - _ = it - cur_size, cur_peak = tracemalloc.get_traced_memory() - tracemalloc.reset_peak() - sizes.append(cur_size) - peaks.append(cur_peak) - - for size in sizes: - # Observed range: 1.59–1.83× on macOS, up to ~2.7× on Linux - # (glibc + Arrow allocate more intermediate buffers). - # iter_record_batches holds data-variable arrays (≈1× chunk) while - # yielding sub-batches, plus the current Arrow batch (≈0.65× chunk). - assert chunk_size * 1.3 < size, f"size {size} unexpectedly low" - assert chunk_size * 3.5 > size, f"size {size} unexpectedly high" - - for peak in peaks: - # Observed range: 1.84–3.28× on macOS, up to ~4.15× on Linux - # (glibc + Arrow hold more intermediate buffers at peak). - # Peak includes data arrays + Arrow batch + temporary coordinate index - # arrays; the first batch of each chunk is highest (Dask compute overhead). - assert chunk_size * 1.5 < peak, f"peak {peak} unexpectedly low" - assert chunk_size * 5.0 > peak, f"peak {peak} unexpectedly high" - - assert max(peaks) < large_ds.nbytes - finally: - tracemalloc.stop() + Registration should only hold coordinate arrays and Rust partition metadata + (no data variables). Peak memory during a full-table query should be a + small fraction of the whole dataset — i.e. partitions are processed without + loading all of them simultaneously. + """ + from datafusion import SessionContext + first_chunk = large_ds.isel(next(block_slices(large_ds))) + chunk_size = first_chunk.nbytes -def test_read_xarray_table_memory_bounds(large_ds): - """read_xarray_table should not materialise data at registration time. - - Registration should only hold coordinate arrays and Rust partition metadata - (no data variables). Peak memory during a full-table query should be a - small fraction of the whole dataset — i.e. partitions are processed without - loading all of them simultaneously. - """ - from datafusion import SessionContext - - first_chunk = large_ds.isel(next(block_slices(large_ds))) - chunk_size = first_chunk.nbytes - - tracemalloc.stop() # reset any state left by a previously-failed test - # --- Registration phase --- - tracemalloc.start() - try: - table = read_xarray_table(large_ds) - reg_size, reg_peak = tracemalloc.get_traced_memory() - tracemalloc.reset_peak() - - # The lazy generator only materialises coord arrays (~O(dim sizes)) and - # factory closure objects — no data arrays. Both metrics should be well - # below one chunk of data. - assert reg_size < chunk_size, ( - f"Registration held {reg_size} bytes >= chunk_size {chunk_size}: " - "data may have been loaded eagerly" - ) - assert ( - reg_peak < chunk_size * 2 - ), f"Registration peak {reg_peak} too high (expected < 2× chunk_size {chunk_size})" - - # --- Query phase --- - ctx = SessionContext() - ctx.register_table("weather", table) - ctx.sql( - "SELECT AVG(temperature), AVG(precipitation) FROM weather" - ).collect() - _, query_peak = tracemalloc.get_traced_memory() - - # tracemalloc measures Python-heap allocations, which include Arrow - # buffer copies and object overhead on top of the raw data. The - # observed peak is typically 1.1–1.5× the raw dataset size; we use - # 2× as a generous bound that would still catch catastrophic regressions - # (e.g. loading all partitions twice simultaneously). - assert query_peak < large_ds.nbytes * 2, ( - f"Query peak {query_peak} >= 2× dataset {large_ds.nbytes}: " - "may be holding excessive data in memory" - ) - finally: - tracemalloc.stop() + tracemalloc.stop() # reset any state left by a previously-failed test + # --- Registration phase --- + tracemalloc.start() + try: + table = read_xarray_table(large_ds) + reg_size, reg_peak = tracemalloc.get_traced_memory() + tracemalloc.reset_peak() + + # The lazy generator only materialises coord arrays (~O(dim sizes)) and + # factory closure objects — no data arrays. Both metrics should be well + # below one chunk of data. + assert reg_size < chunk_size, ( + f"Registration held {reg_size} bytes >= chunk_size {chunk_size}: " + "data may have been loaded eagerly" + ) + assert reg_peak < chunk_size * 2, ( + f"Registration peak {reg_peak} too high (expected < 2× chunk_size {chunk_size})" + ) + + # --- Query phase --- + ctx = SessionContext() + ctx.register_table("weather", table) + ctx.sql( + "SELECT AVG(temperature), AVG(precipitation) FROM weather" + ).collect() + _, query_peak = tracemalloc.get_traced_memory() + + # tracemalloc measures Python-heap allocations, which include Arrow + # buffer copies and object overhead on top of the raw data. The + # observed peak is typically 1.1–1.5× the raw dataset size; we use + # 2× as a generous bound that would still catch catastrophic regressions + # (e.g. loading all partitions twice simultaneously). + assert query_peak < large_ds.nbytes * 2, ( + f"Query peak {query_peak} >= 2× dataset {large_ds.nbytes}: " + "may be holding excessive data in memory" + ) + finally: + tracemalloc.stop() diff --git a/tests/test_reader.py b/tests/test_reader.py index d970b4c4..a8e7e33f 100644 --- a/tests/test_reader.py +++ b/tests/test_reader.py @@ -25,1348 +25,1353 @@ from xarray_sql._native import LazyArrowStreamTable from xarray_sql.reader import XarrayRecordBatchReader, read_xarray_table -from xarray_sql.df import _parse_schema @pytest.fixture def small_ds(): - """Create a small dataset for testing.""" - np.random.seed(42) - time = pd.date_range("2020-01-01", periods=100, freq="h") - lat = np.linspace(-90, 90, 10) - lon = np.linspace(-180, 180, 10) + """Create a small dataset for testing.""" + np.random.seed(42) + time = pd.date_range("2020-01-01", periods=100, freq="h") + lat = np.linspace(-90, 90, 10) + lon = np.linspace(-180, 180, 10) - data = np.random.rand(100, 10, 10).astype(np.float32) + data = np.random.rand(100, 10, 10).astype(np.float32) - return xr.Dataset( - {"temperature": (["time", "lat", "lon"], data)}, - coords={"time": time, "lat": lat, "lon": lon}, - ) + return xr.Dataset( + {"temperature": (["time", "lat", "lon"], data)}, + coords={"time": time, "lat": lat, "lon": lon}, + ) class IterationTracker: - """Tracks when iteration occurs for testing lazy evaluation. + """Tracks when iteration occurs for testing lazy evaluation. - The callback signature is ``(block, projection_names)`` where - ``projection_names`` is the list of column names requested by the query - (``None`` when no projection pushdown occurred, e.g. for - ``XarrayRecordBatchReader`` or a ``SELECT *`` query). - """ + The callback signature is ``(block, projection_names)`` where + ``projection_names`` is the list of column names requested by the query + (``None`` when no projection pushdown occurred, e.g. for + ``XarrayRecordBatchReader`` or a ``SELECT *`` query). + """ - def __init__(self): - self.iteration_count = 0 - self.blocks_seen = [] - self.projections_seen = [] + def __init__(self): + self.iteration_count = 0 + self.blocks_seen = [] + self.projections_seen = [] - def __call__(self, block, projection_names=None): - self.iteration_count += 1 - self.blocks_seen.append(block) - self.projections_seen.append(projection_names) + def __call__(self, block, projection_names=None): + self.iteration_count += 1 + self.blocks_seen.append(block) + self.projections_seen.append(projection_names) - def reset(self): - self.iteration_count = 0 - self.blocks_seen = [] - self.projections_seen = [] + def reset(self): + self.iteration_count = 0 + self.blocks_seen = [] + self.projections_seen = [] class TestXarrayRecordBatchReaderCreation: - """Tests that reader creation does NOT trigger data iteration.""" + """Tests that reader creation does NOT trigger data iteration.""" - def test_reader_creation_does_not_iterate(self, small_ds): - """Creating a reader should NOT iterate through any data.""" - tracker = IterationTracker() + def test_reader_creation_does_not_iterate(self, small_ds): + """Creating a reader should NOT iterate through any data.""" + tracker = IterationTracker() - reader = XarrayRecordBatchReader( - small_ds, - chunks={"time": 25}, - _iteration_callback=tracker, - ) + XarrayRecordBatchReader( + small_ds, + chunks={"time": 25}, + _iteration_callback=tracker, + ) - assert tracker.iteration_count == 0, ( - f"Expected 0 iterations during reader creation, " - f"but got {tracker.iteration_count}" - ) + assert tracker.iteration_count == 0, ( + f"Expected 0 iterations during reader creation, " + f"but got {tracker.iteration_count}" + ) - def test_schema_access_does_not_iterate(self, small_ds): - """Accessing the schema should NOT trigger iteration.""" - tracker = IterationTracker() + def test_schema_access_does_not_iterate(self, small_ds): + """Accessing the schema should NOT trigger iteration.""" + tracker = IterationTracker() - reader = XarrayRecordBatchReader( - small_ds, - chunks={"time": 25}, - _iteration_callback=tracker, - ) + reader = XarrayRecordBatchReader( + small_ds, + chunks={"time": 25}, + _iteration_callback=tracker, + ) - # Access schema - _ = reader.schema - _ = reader.__arrow_c_schema__() + # Access schema + _ = reader.schema + _ = reader.__arrow_c_schema__() - assert tracker.iteration_count == 0, ( - f"Expected 0 iterations when accessing schema, " - f"but got {tracker.iteration_count}" - ) + assert tracker.iteration_count == 0, ( + f"Expected 0 iterations when accessing schema, " + f"but got {tracker.iteration_count}" + ) class TestDataFusionRegistration: - """Tests that DataFusion table registration does NOT trigger iteration. - - These tests use read_xarray_table with register_table() - to achieve true lazy evaluation. - """ + """Tests that DataFusion table registration does NOT trigger iteration. - def test_register_table_does_not_iterate(self, small_ds): - """Registering a LazyArrowStreamTable should NOT iterate data. - - This is the KEY test for lazy evaluation. LazyArrowStreamTable wraps - a factory and implements __datafusion_table_provider__ with StreamingTable, - ensuring data is only read during query execution. + These tests use read_xarray_table with register_table() + to achieve true lazy evaluation. """ - tracker = IterationTracker() - - # Use read_xarray_table which creates a factory-based table - table = read_xarray_table( - small_ds, - chunks={"time": 25}, - _iteration_callback=tracker, - ) - - ctx = SessionContext() - ctx.register_table("test_table", table) - - assert tracker.iteration_count == 0, ( - f"LAZY EVALUATION FAILED: Expected 0 iterations during " - f"register_table(), but got {tracker.iteration_count}." - ) - def test_sql_planning_does_not_iterate(self, small_ds): - """Creating a SQL query plan should NOT iterate data.""" - tracker = IterationTracker() - - table = read_xarray_table( - small_ds, - chunks={"time": 25}, - _iteration_callback=tracker, - ) - - ctx = SessionContext() - ctx.register_table("test_table", table) - - # Create a query but don't execute it - query = ctx.sql("SELECT AVG(temperature) FROM test_table") - - # Just creating the query shouldn't iterate - assert tracker.iteration_count == 0, ( - f"Expected 0 iterations during SQL planning, " - f"but got {tracker.iteration_count}. " - f"DataFusion may be scanning data during query planning." - ) + def test_register_table_does_not_iterate(self, small_ds): + """Registering a LazyArrowStreamTable should NOT iterate data. + This is the KEY test for lazy evaluation. LazyArrowStreamTable wraps + a factory and implements __datafusion_table_provider__ with StreamingTable, + ensuring data is only read during query execution. + """ + tracker = IterationTracker() -class TestDataFusionCollect: - """Tests that data iteration ONLY occurs during collect(). - - These tests use read_xarray_table to verify lazy evaluation. - """ - - def test_collect_triggers_iteration(self, small_ds): - """collect() should trigger data iteration.""" - tracker = IterationTracker() - - table = read_xarray_table( - small_ds, - chunks={"time": 25}, - _iteration_callback=tracker, - ) - - ctx = SessionContext() - ctx.register_table("test_table", table) + # Use read_xarray_table which creates a factory-based table + table = read_xarray_table( + small_ds, + chunks={"time": 25}, + _iteration_callback=tracker, + ) - # Verify no iteration yet (lazy registration) - iteration_before_collect = tracker.iteration_count - assert ( - iteration_before_collect == 0 - ), "Should have 0 iterations before collect" + ctx = SessionContext() + ctx.register_table("test_table", table) - # Now collect - this SHOULD iterate - result = ctx.sql("SELECT * FROM test_table LIMIT 10").collect() + assert tracker.iteration_count == 0, ( + f"LAZY EVALUATION FAILED: Expected 0 iterations during " + f"register_table(), but got {tracker.iteration_count}." + ) - assert tracker.iteration_count > 0, ( - f"Expected iterations during collect(), but got 0. " - f"Data was never read!" - ) - assert ( - tracker.iteration_count > iteration_before_collect - ), f"Expected more iterations after collect()" - - def test_full_query_iterates_all_blocks(self, small_ds): - """A query that reads all data should iterate all blocks.""" - tracker = IterationTracker() - - chunks = {"time": 25} - table = read_xarray_table( - small_ds, - chunks=chunks, - _iteration_callback=tracker, - ) + def test_sql_planning_does_not_iterate(self, small_ds): + """Creating a SQL query plan should NOT iterate data.""" + tracker = IterationTracker() - ctx = SessionContext() - ctx.register_table("test_table", table) + table = read_xarray_table( + small_ds, + chunks={"time": 25}, + _iteration_callback=tracker, + ) - # Run a query that needs to scan all data - result = ctx.sql("SELECT COUNT(*) FROM test_table").collect() + ctx = SessionContext() + ctx.register_table("test_table", table) - # With time=100 and chunks=25, we expect 4 blocks - expected_blocks = 100 // 25 - assert tracker.iteration_count == expected_blocks, ( - f"Expected {expected_blocks} block iterations, " - f"but got {tracker.iteration_count}" - ) + # Create a query but don't execute it + ctx.sql("SELECT AVG(temperature) FROM test_table") - def test_aggregation_query_iterates_correctly(self, small_ds): - """Aggregation queries should iterate all necessary blocks.""" - tracker = IterationTracker() + # Just creating the query shouldn't iterate + assert tracker.iteration_count == 0, ( + f"Expected 0 iterations during SQL planning, " + f"but got {tracker.iteration_count}. " + f"DataFusion may be scanning data during query planning." + ) - table = read_xarray_table( - small_ds, - chunks={"time": 25}, - _iteration_callback=tracker, - ) - ctx = SessionContext() - ctx.register_table("test_table", table) +class TestDataFusionCollect: + """Tests that data iteration ONLY occurs during collect(). - # Run aggregation - result = ctx.sql( - "SELECT lat, AVG(temperature) as avg_temp " - "FROM test_table GROUP BY lat" - ).collect() + These tests use read_xarray_table to verify lazy evaluation. + """ - # Should have iterated some blocks - assert tracker.iteration_count > 0 - assert len(result) > 0 + def test_collect_triggers_iteration(self, small_ds): + """collect() should trigger data iteration.""" + tracker = IterationTracker() + + table = read_xarray_table( + small_ds, + chunks={"time": 25}, + _iteration_callback=tracker, + ) + + ctx = SessionContext() + ctx.register_table("test_table", table) + + # Verify no iteration yet (lazy registration) + iteration_before_collect = tracker.iteration_count + assert iteration_before_collect == 0, ( + "Should have 0 iterations before collect" + ) + + # Now collect - this SHOULD iterate + ctx.sql("SELECT * FROM test_table LIMIT 10").collect() + + assert tracker.iteration_count > 0, ( + "Expected iterations during collect(), but got 0. " + "Data was never read!" + ) + assert tracker.iteration_count > iteration_before_collect, ( + "Expected more iterations after collect()" + ) + + def test_full_query_iterates_all_blocks(self, small_ds): + """A query that reads all data should iterate all blocks.""" + tracker = IterationTracker() + + chunks = {"time": 25} + table = read_xarray_table( + small_ds, + chunks=chunks, + _iteration_callback=tracker, + ) + + ctx = SessionContext() + ctx.register_table("test_table", table) + + # Run a query that needs to scan all data + ctx.sql("SELECT COUNT(*) FROM test_table").collect() + + # With time=100 and chunks=25, we expect 4 blocks + expected_blocks = 100 // 25 + assert tracker.iteration_count == expected_blocks, ( + f"Expected {expected_blocks} block iterations, " + f"but got {tracker.iteration_count}" + ) + + def test_aggregation_query_iterates_correctly(self, small_ds): + """Aggregation queries should iterate all necessary blocks.""" + tracker = IterationTracker() + + table = read_xarray_table( + small_ds, + chunks={"time": 25}, + _iteration_callback=tracker, + ) + + ctx = SessionContext() + ctx.register_table("test_table", table) + + # Run aggregation + result = ctx.sql( + "SELECT lat, AVG(temperature) as avg_temp " + "FROM test_table GROUP BY lat" + ).collect() + + # Should have iterated some blocks + assert tracker.iteration_count > 0 + assert len(result) > 0 class TestLazyEvaluationEndToEnd: - """End-to-end tests verifying lazy evaluation through the full pipeline. - - These tests use read_xarray_table to achieve true lazy evaluation. - """ - - def test_lazy_evaluation_sequence(self, small_ds): - """Verify the exact sequence of lazy evaluation stages. + """End-to-end tests verifying lazy evaluation through the full pipeline. - This is the comprehensive test that proves true lazy evaluation: - 1. Table creation: 0 iterations - 2. Table registration: 0 iterations - 3. Query planning: 0 iterations - 4. collect(): N iterations (where N = number of blocks) + These tests use read_xarray_table to achieve true lazy evaluation. """ - tracker = IterationTracker() - - # Stage 1: Table creation (with factory) - table = read_xarray_table( - small_ds, - chunks={"time": 25}, - _iteration_callback=tracker, - ) - iterations_after_table = tracker.iteration_count - assert iterations_after_table == 0, ( - f"Stage 1 FAILED: Table creation triggered " - f"{iterations_after_table} iterations" - ) - - # Stage 2: Table registration - ctx = SessionContext() - ctx.register_table("test_table", table) - iterations_after_registration = tracker.iteration_count - assert iterations_after_registration == 0, ( - f"Stage 2 FAILED: Table registration triggered " - f"{iterations_after_registration} iterations" - ) - - # Stage 3: Query planning - query = ctx.sql("SELECT * FROM test_table") - iterations_after_planning = tracker.iteration_count - assert iterations_after_planning == 0, ( - f"Stage 3 FAILED: Query planning triggered " - f"{iterations_after_planning} iterations" - ) - - # Stage 4: collect() - NOW iteration should happen - result = query.collect() - iterations_after_collect = tracker.iteration_count - assert ( - iterations_after_collect > 0 - ), f"Stage 4 FAILED: collect() triggered 0 iterations - no data was read!" - - # Verify we got the expected number of blocks (100 time steps / 25 = 4) - expected_blocks = 4 - assert ( - iterations_after_collect == expected_blocks - ), f"Expected {expected_blocks} iterations, got {iterations_after_collect}" - - def test_multiple_queries_on_same_table(self, small_ds): - """Same table can be queried multiple times with fresh iteration each time.""" - tracker = IterationTracker() - - table = read_xarray_table( - small_ds, - chunks={"time": 50}, - _iteration_callback=tracker, - ) - ctx = SessionContext() - ctx.register_table("test_table", table) + def test_lazy_evaluation_sequence(self, small_ds): + """Verify the exact sequence of lazy evaluation stages. - # First query - ctx.sql("SELECT COUNT(*) FROM test_table").collect() - first_query_iterations = tracker.iteration_count - assert first_query_iterations > 0, "First query should iterate" - - # Second query on same table - should iterate again - ctx.sql("SELECT AVG(temperature) FROM test_table").collect() - second_query_iterations = tracker.iteration_count - assert ( - second_query_iterations > first_query_iterations - ), "Second query should trigger additional iterations" - - def test_stream_consumed_error(self, small_ds): - """Once consumed, a single XarrayRecordBatchReader should not be reusable.""" - reader = XarrayRecordBatchReader(small_ds, chunks={"time": 25}) - - # Consume the reader by converting to a PyArrow reader and reading - import pyarrow as pa - - pa_reader = pa.RecordBatchReader.from_stream(reader) - _ = pa_reader.read_all() - - # Reader is now consumed, calling __arrow_c_stream__ again should fail - with pytest.raises(RuntimeError, match="already consumed"): - reader.__arrow_c_stream__() + This is the comprehensive test that proves true lazy evaluation: + 1. Table creation: 0 iterations + 2. Table registration: 0 iterations + 3. Query planning: 0 iterations + 4. collect(): N iterations (where N = number of blocks) + """ + tracker = IterationTracker() + + # Stage 1: Table creation (with factory) + table = read_xarray_table( + small_ds, + chunks={"time": 25}, + _iteration_callback=tracker, + ) + iterations_after_table = tracker.iteration_count + assert iterations_after_table == 0, ( + f"Stage 1 FAILED: Table creation triggered " + f"{iterations_after_table} iterations" + ) + + # Stage 2: Table registration + ctx = SessionContext() + ctx.register_table("test_table", table) + iterations_after_registration = tracker.iteration_count + assert iterations_after_registration == 0, ( + f"Stage 2 FAILED: Table registration triggered " + f"{iterations_after_registration} iterations" + ) + + # Stage 3: Query planning + query = ctx.sql("SELECT * FROM test_table") + iterations_after_planning = tracker.iteration_count + assert iterations_after_planning == 0, ( + f"Stage 3 FAILED: Query planning triggered " + f"{iterations_after_planning} iterations" + ) + + # Stage 4: collect() - NOW iteration should happen + query.collect() + iterations_after_collect = tracker.iteration_count + assert iterations_after_collect > 0, ( + "Stage 4 FAILED: collect() triggered 0 iterations - no data was read!" + ) + + # Verify we got the expected number of blocks (100 time steps / 25 = 4) + expected_blocks = 4 + assert iterations_after_collect == expected_blocks, ( + f"Expected {expected_blocks} iterations, got {iterations_after_collect}" + ) + + def test_multiple_queries_on_same_table(self, small_ds): + """Same table can be queried multiple times with fresh iteration each time.""" + tracker = IterationTracker() + + table = read_xarray_table( + small_ds, + chunks={"time": 50}, + _iteration_callback=tracker, + ) + + ctx = SessionContext() + ctx.register_table("test_table", table) + + # First query + ctx.sql("SELECT COUNT(*) FROM test_table").collect() + first_query_iterations = tracker.iteration_count + assert first_query_iterations > 0, "First query should iterate" + + # Second query on same table - should iterate again + ctx.sql("SELECT AVG(temperature) FROM test_table").collect() + second_query_iterations = tracker.iteration_count + assert second_query_iterations > first_query_iterations, ( + "Second query should trigger additional iterations" + ) + + def test_stream_consumed_error(self, small_ds): + """Once consumed, a single XarrayRecordBatchReader should not be reusable.""" + reader = XarrayRecordBatchReader(small_ds, chunks={"time": 25}) + + # Consume the reader by converting to a PyArrow reader and reading + import pyarrow as pa + + pa_reader = pa.RecordBatchReader.from_stream(reader) + _ = pa_reader.read_all() + + # Reader is now consumed, calling __arrow_c_stream__ again should fail + with pytest.raises(RuntimeError, match="already consumed"): + reader.__arrow_c_stream__() class TestDataIntegrity: - """Tests that verify data correctness alongside lazy evaluation. + """Tests that verify data correctness alongside lazy evaluation. - These tests use read_xarray_table for lazy streaming. - """ + These tests use read_xarray_table for lazy streaming. + """ - def test_query_results_are_correct(self, small_ds): - """Verify that lazy evaluation produces correct results.""" - table = read_xarray_table(small_ds, chunks={"time": 25}) + def test_query_results_are_correct(self, small_ds): + """Verify that lazy evaluation produces correct results.""" + table = read_xarray_table(small_ds, chunks={"time": 25}) - ctx = SessionContext() - ctx.register_table("test_table", table) + ctx = SessionContext() + ctx.register_table("test_table", table) - # Get count - result = ctx.sql("SELECT COUNT(*) as cnt FROM test_table").collect() - count = result[0].to_pandas()["cnt"].iloc[0] + # Get count + result = ctx.sql("SELECT COUNT(*) as cnt FROM test_table").collect() + count = result[0].to_pandas()["cnt"].iloc[0] - # Expected: 100 time steps * 10 lat * 10 lon = 10,000 rows - expected_count = 100 * 10 * 10 - assert ( - count == expected_count - ), f"Expected {expected_count} rows, got {count}" + # Expected: 100 time steps * 10 lat * 10 lon = 10,000 rows + expected_count = 100 * 10 * 10 + assert count == expected_count, ( + f"Expected {expected_count} rows, got {count}" + ) - def test_aggregation_results_are_correct(self, small_ds): - """Verify aggregation produces correct results.""" - table = read_xarray_table(small_ds, chunks={"time": 25}) + def test_aggregation_results_are_correct(self, small_ds): + """Verify aggregation produces correct results.""" + table = read_xarray_table(small_ds, chunks={"time": 25}) - ctx = SessionContext() - ctx.register_table("test_table", table) + ctx = SessionContext() + ctx.register_table("test_table", table) - # Get average temperature - result = ctx.sql( - "SELECT AVG(temperature) as avg_temp FROM test_table" - ).collect() - avg_temp = result[0].to_pandas()["avg_temp"].iloc[0] + # Get average temperature + result = ctx.sql( + "SELECT AVG(temperature) as avg_temp FROM test_table" + ).collect() + avg_temp = result[0].to_pandas()["avg_temp"].iloc[0] - # With seed 42 and random data in [0, 1), average should be ~0.5 - assert ( - 0.4 < avg_temp < 0.6 - ), f"Expected average temperature ~0.5, got {avg_temp}" + # With seed 42 and random data in [0, 1), average should be ~0.5 + assert 0.4 < avg_temp < 0.6, ( + f"Expected average temperature ~0.5, got {avg_temp}" + ) class TestPyArrowInterop: - """Tests for PyArrow interoperability.""" + """Tests for PyArrow interoperability.""" - def test_from_stream_does_not_iterate(self, small_ds): - """pa.RecordBatchReader.from_stream() should not iterate.""" - tracker = IterationTracker() + def test_from_stream_does_not_iterate(self, small_ds): + """pa.RecordBatchReader.from_stream() should not iterate.""" + tracker = IterationTracker() - reader = XarrayRecordBatchReader( - small_ds, - chunks={"time": 25}, - _iteration_callback=tracker, - ) + reader = XarrayRecordBatchReader( + small_ds, + chunks={"time": 25}, + _iteration_callback=tracker, + ) - # Create PyArrow reader from our stream - pa_reader = pa.RecordBatchReader.from_stream(reader) + # Create PyArrow reader from our stream + pa.RecordBatchReader.from_stream(reader) - assert tracker.iteration_count == 0, ( - f"Expected 0 iterations when creating PyArrow reader, " - f"but got {tracker.iteration_count}" - ) + assert tracker.iteration_count == 0, ( + f"Expected 0 iterations when creating PyArrow reader, " + f"but got {tracker.iteration_count}" + ) - def test_pyarrow_iteration_triggers_callbacks(self, small_ds): - """Iterating via PyArrow should trigger our callbacks.""" - tracker = IterationTracker() + def test_pyarrow_iteration_triggers_callbacks(self, small_ds): + """Iterating via PyArrow should trigger our callbacks.""" + tracker = IterationTracker() - reader = XarrayRecordBatchReader( - small_ds, - chunks={"time": 25}, - _iteration_callback=tracker, - ) + reader = XarrayRecordBatchReader( + small_ds, + chunks={"time": 25}, + _iteration_callback=tracker, + ) - pa_reader = pa.RecordBatchReader.from_stream(reader) + pa_reader = pa.RecordBatchReader.from_stream(reader) - # Now iterate - for batch in pa_reader: - pass + # Now iterate + for batch in pa_reader: + pass - assert ( - tracker.iteration_count == 4 - ), f"Expected 4 iterations, got {tracker.iteration_count}" + assert tracker.iteration_count == 4, ( + f"Expected 4 iterations, got {tracker.iteration_count}" + ) - def test_read_all_iterates_all(self, small_ds): - """read_all() should iterate through all blocks.""" - tracker = IterationTracker() + def test_read_all_iterates_all(self, small_ds): + """read_all() should iterate through all blocks.""" + tracker = IterationTracker() - reader = XarrayRecordBatchReader( - small_ds, - chunks={"time": 25}, - _iteration_callback=tracker, - ) + reader = XarrayRecordBatchReader( + small_ds, + chunks={"time": 25}, + _iteration_callback=tracker, + ) - pa_reader = pa.RecordBatchReader.from_stream(reader) - table = pa_reader.read_all() + pa_reader = pa.RecordBatchReader.from_stream(reader) + table = pa_reader.read_all() - assert tracker.iteration_count == 4 - assert len(table) == 100 * 10 * 10 + assert tracker.iteration_count == 4 + assert len(table) == 100 * 10 * 10 class StreamingTracker: - """Tracks timing of batch iterations to verify streaming behavior. + """Tracks timing of batch iterations to verify streaming behavior. - This tracker records when each batch is processed, allowing us to verify - that batches are streamed incrementally rather than all loaded at once. - """ + This tracker records when each batch is processed, allowing us to verify + that batches are streamed incrementally rather than all loaded at once. + """ - def __init__(self): - self.batch_times = [] - self.batch_count = 0 - self._lock = threading.Lock() + def __init__(self): + self.batch_times = [] + self.batch_count = 0 + self._lock = threading.Lock() - def __call__(self, block, projection_names=None): - with self._lock: - self.batch_times.append(time.monotonic()) - self.batch_count += 1 + def __call__(self, block, projection_names=None): + with self._lock: + self.batch_times.append(time.monotonic()) + self.batch_count += 1 - def reset(self): - with self._lock: - self.batch_times = [] - self.batch_count = 0 + def reset(self): + with self._lock: + self.batch_times = [] + self.batch_count = 0 - @property - def max_concurrent_batches_estimate(self): - """Estimate max batches that could have been in memory simultaneously. + @property + def max_concurrent_batches_estimate(self): + """Estimate max batches that could have been in memory simultaneously. - If all batches are loaded at once, all batch_times will be very close. - If streaming works correctly, batch_times should be spread out. - """ - if len(self.batch_times) < 2: - return len(self.batch_times) + If all batches are loaded at once, all batch_times will be very close. + If streaming works correctly, batch_times should be spread out. + """ + if len(self.batch_times) < 2: + return len(self.batch_times) - # Sort times and look at gaps - sorted_times = sorted(self.batch_times) - # If times are spread out, streaming is working - # If all times are within a tiny window, all batches loaded at once - total_duration = sorted_times[-1] - sorted_times[0] + # Sort times and look at gaps + sorted_times = sorted(self.batch_times) + # If times are spread out, streaming is working + # If all times are within a tiny window, all batches loaded at once + sorted_times[-1] - sorted_times[0] - # If the spread is very small compared to number of batches, - # batches were likely all loaded at once - return len(self.batch_times) + # If the spread is very small compared to number of batches, + # batches were likely all loaded at once + return len(self.batch_times) class TestStreamingBehavior: - """Tests that verify true streaming with bounded memory. - - These tests ensure that the Rust implementation streams batches through - a bounded channel rather than loading all data into memory at once. - """ - - def test_batches_processed_incrementally(self, small_ds): - """Verify batches are processed one at a time, not all at once. + """Tests that verify true streaming with bounded memory. - This test uses a callback that tracks when each batch is processed. - With true streaming, batches should be processed incrementally. + These tests ensure that the Rust implementation streams batches through + a bounded channel rather than loading all data into memory at once. """ - tracker = StreamingTracker() - - table = read_xarray_table( - small_ds, - chunks={"time": 25}, - _iteration_callback=tracker, - ) - - ctx = SessionContext() - ctx.register_table("test_table", table) - - # Run query that scans all data - ctx.sql("SELECT COUNT(*) FROM test_table").collect() - - # All 4 batches should have been processed - assert ( - tracker.batch_count == 4 - ), f"Expected 4 batches, got {tracker.batch_count}" - - def test_all_partitions_processed(self, small_ds): - """Verify that all partitions are processed (order may vary with parallelism).""" - blocks_seen = [] - - def track_order(block, projection_names=None): - # Record the time slice for ordering verification - blocks_seen.append(block.get("time", None)) - - table = read_xarray_table( - small_ds, - chunks={"time": 25}, - _iteration_callback=track_order, - ) - ctx = SessionContext() - ctx.register_table("test_table", table) - ctx.sql("SELECT * FROM test_table").collect() + def test_batches_processed_incrementally(self, small_ds): + """Verify batches are processed one at a time, not all at once. - # Should have 4 blocks/partitions - assert len(blocks_seen) == 4 - - # All blocks should be present (though order may vary due to parallelism) - # Extract start positions and verify they cover all expected ranges - starts = sorted([b.start for b in blocks_seen]) - expected_starts = [0, 25, 50, 75] - assert ( - starts == expected_starts - ), f"Expected partition starts {expected_starts}, got {starts}" - - def test_large_dataset_streams_correctly(self): - """Test streaming with a larger dataset to verify memory behavior. - - This test creates a dataset with many blocks to verify that - streaming works correctly at scale. - """ - # Create a dataset with 20 blocks - np.random.seed(42) - time = pd.date_range("2020-01-01", periods=200, freq="h") - lat = np.linspace(-90, 90, 10) - lon = np.linspace(-180, 180, 10) + This test uses a callback that tracks when each batch is processed. + With true streaming, batches should be processed incrementally. + """ + tracker = StreamingTracker() + + table = read_xarray_table( + small_ds, + chunks={"time": 25}, + _iteration_callback=tracker, + ) + + ctx = SessionContext() + ctx.register_table("test_table", table) + + # Run query that scans all data + ctx.sql("SELECT COUNT(*) FROM test_table").collect() + + # All 4 batches should have been processed + assert tracker.batch_count == 4, ( + f"Expected 4 batches, got {tracker.batch_count}" + ) + + def test_all_partitions_processed(self, small_ds): + """Verify that all partitions are processed (order may vary with parallelism).""" + blocks_seen = [] + + def track_order(block, projection_names=None): + # Record the time slice for ordering verification + blocks_seen.append(block.get("time", None)) + + table = read_xarray_table( + small_ds, + chunks={"time": 25}, + _iteration_callback=track_order, + ) + + ctx = SessionContext() + ctx.register_table("test_table", table) + ctx.sql("SELECT * FROM test_table").collect() + + # Should have 4 blocks/partitions + assert len(blocks_seen) == 4 + + # All blocks should be present (though order may vary due to parallelism) + # Extract start positions and verify they cover all expected ranges + starts = sorted([b.start for b in blocks_seen]) + expected_starts = [0, 25, 50, 75] + assert starts == expected_starts, ( + f"Expected partition starts {expected_starts}, got {starts}" + ) + + def test_large_dataset_streams_correctly(self): + """Test streaming with a larger dataset to verify memory behavior. + + This test creates a dataset with many blocks to verify that + streaming works correctly at scale. + """ + # Create a dataset with 20 blocks + np.random.seed(42) + time = pd.date_range("2020-01-01", periods=200, freq="h") + lat = np.linspace(-90, 90, 10) + lon = np.linspace(-180, 180, 10) - data = np.random.rand(200, 10, 10).astype(np.float32) + data = np.random.rand(200, 10, 10).astype(np.float32) - large_ds = xr.Dataset( - {"temperature": (["time", "lat", "lon"], data)}, - coords={"time": time, "lat": lat, "lon": lon}, - ) + large_ds = xr.Dataset( + {"temperature": (["time", "lat", "lon"], data)}, + coords={"time": time, "lat": lat, "lon": lon}, + ) - tracker = StreamingTracker() + tracker = StreamingTracker() - # Use small chunks to create many blocks - table = read_xarray_table( - large_ds, - chunks={"time": 10}, # 200 / 10 = 20 blocks - _iteration_callback=tracker, - ) + # Use small chunks to create many blocks + table = read_xarray_table( + large_ds, + chunks={"time": 10}, # 200 / 10 = 20 blocks + _iteration_callback=tracker, + ) - ctx = SessionContext() - ctx.register_table("test_table", table) + ctx = SessionContext() + ctx.register_table("test_table", table) - # Run a query that needs all data - result = ctx.sql("SELECT COUNT(*) as cnt FROM test_table").collect() - count = result[0].to_pandas()["cnt"].iloc[0] + # Run a query that needs all data + result = ctx.sql("SELECT COUNT(*) as cnt FROM test_table").collect() + count = result[0].to_pandas()["cnt"].iloc[0] - # Verify all blocks were processed - assert ( - tracker.batch_count == 20 - ), f"Expected 20 batches for large dataset, got {tracker.batch_count}" + # Verify all blocks were processed + assert tracker.batch_count == 20, ( + f"Expected 20 batches for large dataset, got {tracker.batch_count}" + ) - # Verify data integrity - expected_count = 200 * 10 * 10 - assert ( - count == expected_count - ), f"Expected {expected_count} rows, got {count}" + # Verify data integrity + expected_count = 200 * 10 * 10 + assert count == expected_count, ( + f"Expected {expected_count} rows, got {count}" + ) class TestBoundedMemoryBehavior: - """Tests that verify memory usage remains bounded during streaming. - - The key property we're testing: only a small number of batches should - be in memory at once (the channel buffer size, which is 4), not the - entire dataset. - - These tests verify that: - 1. Many batches can be processed without loading all into memory - 2. Production times are spread out (indicating back-pressure) - 3. Large datasets complete successfully (memory doesn't explode) - """ - - def test_many_batches_stream_successfully(self): - """Verify streaming works with many more batches than buffer size. - - With buffer size = 4, if we have 16 batches and streaming works, - the query should complete successfully. If all batches were loaded - at once (no streaming), this would use 4x more memory. - """ - # Create dataset with 16 batches (4x buffer size) - np.random.seed(42) - time_coord = pd.date_range("2020-01-01", periods=160, freq="h") - lat = np.linspace(-90, 90, 5) - lon = np.linspace(-180, 180, 5) - data = np.random.rand(160, 5, 5).astype(np.float32) - - ds = xr.Dataset( - {"temperature": (["time", "lat", "lon"], data)}, - coords={"time": time_coord, "lat": lat, "lon": lon}, - ) - - tracker = StreamingTracker() - - # 16 batches (160 / 10 = 16) - table = read_xarray_table( - ds, - chunks={"time": 10}, - _iteration_callback=tracker, - ) - - ctx = SessionContext() - ctx.register_table("test_table", table) - - result = ctx.sql("SELECT COUNT(*) as cnt FROM test_table").collect() - count = result[0].to_pandas()["cnt"].iloc[0] - - # All 16 batches should have been processed - assert ( - tracker.batch_count == 16 - ), f"Expected 16 batches, got {tracker.batch_count}" + """Tests that verify memory usage remains bounded during streaming. - # Verify data integrity - expected = 160 * 5 * 5 - assert count == expected, f"Expected {expected} rows, got {count}" + The key property we're testing: only a small number of batches should + be in memory at once (the channel buffer size, which is 4), not the + entire dataset. - def test_production_times_spread_out(self): - """Verify batch production is spread over time, not instant. - - If back-pressure works, later batches can only be produced after - earlier batches have been consumed. Production times should span - a non-zero duration. + These tests verify that: + 1. Many batches can be processed without loading all into memory + 2. Production times are spread out (indicating back-pressure) + 3. Large datasets complete successfully (memory doesn't explode) """ - np.random.seed(123) - time_coord = pd.date_range("2020-01-01", periods=100, freq="h") - lat = np.linspace(-90, 90, 5) - lon = np.linspace(-180, 180, 5) - data = np.random.rand(100, 5, 5).astype(np.float32) - - ds = xr.Dataset( - {"temperature": (["time", "lat", "lon"], data)}, - coords={"time": time_coord, "lat": lat, "lon": lon}, - ) - - tracker = StreamingTracker() - - # 10 batches, more than buffer size of 4 - table = read_xarray_table( - ds, - chunks={"time": 10}, - _iteration_callback=tracker, - ) - - ctx = SessionContext() - ctx.register_table("test_table", table) - ctx.sql("SELECT AVG(temperature) FROM test_table").collect() - - # All 10 batches should be produced - assert tracker.batch_count == 10 - # Production should span some time (not all instant) - sorted_times = sorted(tracker.batch_times) - production_span = sorted_times[-1] - sorted_times[0] + def test_many_batches_stream_successfully(self): + """Verify streaming works with many more batches than buffer size. - # With streaming and back-pressure, production_span should be > 0 - # (If all batches were produced simultaneously, span would be ~0) - assert production_span >= 0, "Production span should be non-negative" - - def test_large_batch_count_completes(self): - """Verify that processing many batches completes successfully. + With buffer size = 4, if we have 16 batches and streaming works, + the query should complete successfully. If all batches were loaded + at once (no streaming), this would use 4x more memory. + """ + # Create dataset with 16 batches (4x buffer size) + np.random.seed(42) + time_coord = pd.date_range("2020-01-01", periods=160, freq="h") + lat = np.linspace(-90, 90, 5) + lon = np.linspace(-180, 180, 5) + data = np.random.rand(160, 5, 5).astype(np.float32) + + ds = xr.Dataset( + {"temperature": (["time", "lat", "lon"], data)}, + coords={"time": time_coord, "lat": lat, "lon": lon}, + ) + + tracker = StreamingTracker() + + # 16 batches (160 / 10 = 16) + table = read_xarray_table( + ds, + chunks={"time": 10}, + _iteration_callback=tracker, + ) + + ctx = SessionContext() + ctx.register_table("test_table", table) + + result = ctx.sql("SELECT COUNT(*) as cnt FROM test_table").collect() + count = result[0].to_pandas()["cnt"].iloc[0] + + # All 16 batches should have been processed + assert tracker.batch_count == 16, ( + f"Expected 16 batches, got {tracker.batch_count}" + ) + + # Verify data integrity + expected = 160 * 5 * 5 + assert count == expected, f"Expected {expected} rows, got {count}" + + def test_production_times_spread_out(self): + """Verify batch production is spread over time, not instant. + + If back-pressure works, later batches can only be produced after + earlier batches have been consumed. Production times should span + a non-zero duration. + """ + np.random.seed(123) + time_coord = pd.date_range("2020-01-01", periods=100, freq="h") + lat = np.linspace(-90, 90, 5) + lon = np.linspace(-180, 180, 5) + data = np.random.rand(100, 5, 5).astype(np.float32) + + ds = xr.Dataset( + {"temperature": (["time", "lat", "lon"], data)}, + coords={"time": time_coord, "lat": lat, "lon": lon}, + ) + + tracker = StreamingTracker() + + # 10 batches, more than buffer size of 4 + table = read_xarray_table( + ds, + chunks={"time": 10}, + _iteration_callback=tracker, + ) + + ctx = SessionContext() + ctx.register_table("test_table", table) + ctx.sql("SELECT AVG(temperature) FROM test_table").collect() + + # All 10 batches should be produced + assert tracker.batch_count == 10 + + # Production should span some time (not all instant) + sorted_times = sorted(tracker.batch_times) + production_span = sorted_times[-1] - sorted_times[0] + + # With streaming and back-pressure, production_span should be > 0 + # (If all batches were produced simultaneously, span would be ~0) + assert production_span >= 0, "Production span should be non-negative" + + def test_large_batch_count_completes(self): + """Verify that processing many batches completes successfully. + + This is a stress test: 50 batches is well above the buffer size of 4. + If streaming works correctly, this should complete without memory issues. + """ + np.random.seed(456) + time_coord = pd.date_range("2020-01-01", periods=500, freq="h") + lat = np.linspace(-90, 90, 10) + lon = np.linspace(-180, 180, 10) + data = np.random.rand(500, 10, 10).astype(np.float32) + + ds = xr.Dataset( + {"temperature": (["time", "lat", "lon"], data)}, + coords={"time": time_coord, "lat": lat, "lon": lon}, + ) + + tracker = StreamingTracker() + + # 50 batches (500 / 10 = 50) + table = read_xarray_table( + ds, + chunks={"time": 10}, + _iteration_callback=tracker, + ) + + ctx = SessionContext() + ctx.register_table("test_table", table) + + result = ctx.sql("SELECT COUNT(*) as cnt FROM test_table").collect() + count = result[0].to_pandas()["cnt"].iloc[0] + + # All 50 batches processed + assert tracker.batch_count == 50, ( + f"Expected 50 batches, got {tracker.batch_count}" + ) + + # Data integrity + expected = 500 * 10 * 10 + assert count == expected, f"Expected {expected} rows, got {count}" + + def test_aggregation_with_many_batches(self): + """Verify aggregation queries work correctly with many batches. + + GROUP BY queries require processing all data, making them a good + test for streaming behavior. Uses collect() to verify that parallel + aggregation returns complete results (fixed in DataFusion 52+). + """ + np.random.seed(789) + time_coord = pd.date_range("2020-01-01", periods=120, freq="h") + # Use integer lat/lon to avoid floating point grouping issues + lat = np.array([0, 1, 2, 3, 4], dtype=np.float64) + lon = np.array([0, 1, 2, 3, 4], dtype=np.float64) + data = np.random.rand(120, 5, 5).astype(np.float32) - This is a stress test: 50 batches is well above the buffer size of 4. - If streaming works correctly, this should complete without memory issues. - """ - np.random.seed(456) - time_coord = pd.date_range("2020-01-01", periods=500, freq="h") - lat = np.linspace(-90, 90, 10) - lon = np.linspace(-180, 180, 10) - data = np.random.rand(500, 10, 10).astype(np.float32) + ds = xr.Dataset( + {"temperature": (["time", "lat", "lon"], data)}, + coords={"time": time_coord, "lat": lat, "lon": lon}, + ) - ds = xr.Dataset( - {"temperature": (["time", "lat", "lon"], data)}, - coords={"time": time_coord, "lat": lat, "lon": lon}, - ) + tracker = StreamingTracker() - tracker = StreamingTracker() + # 12 partitions (one per chunk) + table = read_xarray_table( + ds, + chunks={"time": 10}, + _iteration_callback=tracker, + ) - # 50 batches (500 / 10 = 50) - table = read_xarray_table( - ds, - chunks={"time": 10}, - _iteration_callback=tracker, - ) + ctx = SessionContext() + ctx.register_table("test_table", table) - ctx = SessionContext() - ctx.register_table("test_table", table) + # GROUP BY requires scanning all data; collect() must return complete results + df = ctx.sql( + "SELECT lat, AVG(temperature) as avg_temp FROM test_table GROUP BY lat" + ).to_pandas() - result = ctx.sql("SELECT COUNT(*) as cnt FROM test_table").collect() - count = result[0].to_pandas()["cnt"].iloc[0] + assert len(df) == 5, f"Expected 5 lat groups, got {len(df)}" - # All 50 batches processed - assert ( - tracker.batch_count == 50 - ), f"Expected 50 batches, got {tracker.batch_count}" + # All partitions processed + assert tracker.batch_count == 12, ( + f"Expected 12 partitions, got {tracker.batch_count}" + ) - # Data integrity - expected = 500 * 10 * 10 - assert count == expected, f"Expected {expected} rows, got {count}" - def test_aggregation_with_many_batches(self): - """Verify aggregation queries work correctly with many batches. +class TestErrorPropagation: + """Tests that verify errors are properly propagated through the stream. - GROUP BY queries require processing all data, making them a good - test for streaming behavior. Uses collect() to verify that parallel - aggregation returns complete results (fixed in DataFusion 52+). + These tests ensure that errors during batch reading surface to the user + rather than being silently swallowed. """ - np.random.seed(789) - time_coord = pd.date_range("2020-01-01", periods=120, freq="h") - # Use integer lat/lon to avoid floating point grouping issues - lat = np.array([0, 1, 2, 3, 4], dtype=np.float64) - lon = np.array([0, 1, 2, 3, 4], dtype=np.float64) - data = np.random.rand(120, 5, 5).astype(np.float32) - - ds = xr.Dataset( - {"temperature": (["time", "lat", "lon"], data)}, - coords={"time": time_coord, "lat": lat, "lon": lon}, - ) - - tracker = StreamingTracker() - - # 12 partitions (one per chunk) - table = read_xarray_table( - ds, - chunks={"time": 10}, - _iteration_callback=tracker, - ) - - ctx = SessionContext() - ctx.register_table("test_table", table) - # GROUP BY requires scanning all data; collect() must return complete results - df = ctx.sql( - "SELECT lat, AVG(temperature) as avg_temp FROM test_table GROUP BY lat" - ).to_pandas() + def test_factory_error_propagates(self): + """Errors from the factory function should propagate to the user.""" - assert len(df) == 5, f"Expected 5 lat groups, got {len(df)}" + def failing_factory(): + raise ValueError("Factory intentionally failed") - # All partitions processed - assert ( - tracker.batch_count == 12 - ), f"Expected 12 partitions, got {tracker.batch_count}" + schema = pa.schema([("value", pa.int64())]) + # partitions is an iterable of (factory, metadata_dict) pairs + table = LazyArrowStreamTable([(failing_factory, {})], schema) + ctx = SessionContext() + ctx.register_table("test_table", table) -class TestErrorPropagation: - """Tests that verify errors are properly propagated through the stream. - - These tests ensure that errors during batch reading surface to the user - rather than being silently swallowed. - """ - - def test_factory_error_propagates(self): - """Errors from the factory function should propagate to the user.""" + # The error should surface when we try to collect + with pytest.raises(Exception) as exc_info: + ctx.sql("SELECT * FROM test_table").collect() - def failing_factory(): - raise ValueError("Factory intentionally failed") + # Verify the error message mentions the factory failure + error_message = str(exc_info.value).lower() + assert "factory" in error_message or "failed" in error_message, ( + f"Expected error about factory failure, got: {exc_info.value}" + ) - schema = pa.schema([("value", pa.int64())]) - # partitions is an iterable of (factory, metadata_dict) pairs - table = LazyArrowStreamTable([(failing_factory, {})], schema) + def test_iteration_error_propagates(self, small_ds): + """Errors during batch iteration should propagate to the user.""" + error_on_batch = 2 # Fail on the third batch - ctx = SessionContext() - ctx.register_table("test_table", table) + def failing_callback(block, projection_names=None): + # Track which batch we're on using a mutable default + if not hasattr(failing_callback, "count"): + failing_callback.count = 0 + failing_callback.count += 1 - # The error should surface when we try to collect - with pytest.raises(Exception) as exc_info: - ctx.sql("SELECT * FROM test_table").collect() + if failing_callback.count == error_on_batch: + raise RuntimeError("Intentional batch processing error") - # Verify the error message mentions the factory failure - error_message = str(exc_info.value).lower() - assert ( - "factory" in error_message or "failed" in error_message - ), f"Expected error about factory failure, got: {exc_info.value}" - - def test_iteration_error_propagates(self, small_ds): - """Errors during batch iteration should propagate to the user.""" - error_on_batch = 2 # Fail on the third batch - - def failing_callback(block, projection_names=None): - # Track which batch we're on using a mutable default - if not hasattr(failing_callback, "count"): + # Reset the counter failing_callback.count = 0 - failing_callback.count += 1 - if failing_callback.count == error_on_batch: - raise RuntimeError("Intentional batch processing error") + table = read_xarray_table( + small_ds, + chunks={"time": 25}, + _iteration_callback=failing_callback, + ) - # Reset the counter - failing_callback.count = 0 + ctx = SessionContext() + ctx.register_table("test_table", table) - table = read_xarray_table( - small_ds, - chunks={"time": 25}, - _iteration_callback=failing_callback, - ) + # The error should surface when we try to collect + with pytest.raises(Exception): + ctx.sql("SELECT * FROM test_table").collect() - ctx = SessionContext() - ctx.register_table("test_table", table) - - # The error should surface when we try to collect - with pytest.raises(Exception): - ctx.sql("SELECT * FROM test_table").collect() - - def test_empty_dataset_handled_gracefully(self): - """Empty datasets should work without errors.""" - # Create an empty dataset with the right structure - empty_ds = xr.Dataset( - { - "temperature": ( - ["time", "lat", "lon"], - np.array([]).reshape(0, 0, 0), - ) - }, - coords={ - "time": pd.DatetimeIndex([]), - "lat": np.array([]), - "lon": np.array([]), - }, - ) + def test_empty_dataset_handled_gracefully(self): + """Empty datasets should work without errors.""" + # Create an empty dataset with the right structure + empty_ds = xr.Dataset( + { + "temperature": ( + ["time", "lat", "lon"], + np.array([]).reshape(0, 0, 0), + ) + }, + coords={ + "time": pd.DatetimeIndex([]), + "lat": np.array([]), + "lon": np.array([]), + }, + ) - # This should work without crashing - table = read_xarray_table(empty_ds, chunks={"time": 10}) + # This should work without crashing + table = read_xarray_table(empty_ds, chunks={"time": 10}) - ctx = SessionContext() - ctx.register_table("test_table", table) + ctx = SessionContext() + ctx.register_table("test_table", table) - result = ctx.sql("SELECT COUNT(*) as cnt FROM test_table").collect() - count = result[0].to_pandas()["cnt"].iloc[0] + result = ctx.sql("SELECT COUNT(*) as cnt FROM test_table").collect() + count = result[0].to_pandas()["cnt"].iloc[0] - assert count == 0, f"Expected 0 rows for empty dataset, got {count}" + assert count == 0, f"Expected 0 rows for empty dataset, got {count}" class TestMultiplePartitions: - """Tests for scenarios with multiple queries and table reuse.""" - - def test_fresh_stream_per_query(self, small_ds): - """Each query should get a fresh stream from the factory.""" - call_count = {"value": 0} - original_callback = None - - def counting_callback(block, projection_names=None): - call_count["value"] += 1 - if original_callback: - original_callback(block) - - table = read_xarray_table( - small_ds, - chunks={"time": 50}, # 2 blocks per query - _iteration_callback=counting_callback, - ) - - ctx = SessionContext() - ctx.register_table("test_table", table) - - # First query - ctx.sql("SELECT COUNT(*) FROM test_table").collect() - first_query_count = call_count["value"] - assert ( - first_query_count == 2 - ), f"First query: expected 2, got {first_query_count}" - - # Second query should trigger fresh iteration - ctx.sql("SELECT AVG(temperature) FROM test_table").collect() - second_query_count = call_count["value"] - assert ( - second_query_count == 4 - ), f"After second query: expected 4 total, got {second_query_count}" - - # Third query - ctx.sql("SELECT MAX(temperature) FROM test_table").collect() - third_query_count = call_count["value"] - assert ( - third_query_count == 6 - ), f"After third query: expected 6 total, got {third_query_count}" - - def test_parallel_queries_independent(self, small_ds): - """Multiple contexts with the same table should work independently.""" - tracker1 = IterationTracker() - tracker2 = IterationTracker() - - table1 = read_xarray_table( - small_ds, - chunks={"time": 25}, - _iteration_callback=tracker1, - ) - - table2 = read_xarray_table( - small_ds, - chunks={"time": 50}, - _iteration_callback=tracker2, - ) - - ctx1 = SessionContext() - ctx2 = SessionContext() - - ctx1.register_table("test_table", table1) - ctx2.register_table("test_table", table2) - - # Execute queries - ctx1.sql("SELECT COUNT(*) FROM test_table").collect() - ctx2.sql("SELECT COUNT(*) FROM test_table").collect() - - # Each should have its own iteration count - assert ( - tracker1.iteration_count == 4 - ), f"Table1: expected 4 blocks, got {tracker1.iteration_count}" - assert ( - tracker2.iteration_count == 2 - ), f"Table2: expected 2 blocks, got {tracker2.iteration_count}" + """Tests for scenarios with multiple queries and table reuse.""" + + def test_fresh_stream_per_query(self, small_ds): + """Each query should get a fresh stream from the factory.""" + call_count = {"value": 0} + original_callback = None + + def counting_callback(block, projection_names=None): + call_count["value"] += 1 + if original_callback: + original_callback(block) + + table = read_xarray_table( + small_ds, + chunks={"time": 50}, # 2 blocks per query + _iteration_callback=counting_callback, + ) + + ctx = SessionContext() + ctx.register_table("test_table", table) + + # First query + ctx.sql("SELECT COUNT(*) FROM test_table").collect() + first_query_count = call_count["value"] + assert first_query_count == 2, ( + f"First query: expected 2, got {first_query_count}" + ) + + # Second query should trigger fresh iteration + ctx.sql("SELECT AVG(temperature) FROM test_table").collect() + second_query_count = call_count["value"] + assert second_query_count == 4, ( + f"After second query: expected 4 total, got {second_query_count}" + ) + + # Third query + ctx.sql("SELECT MAX(temperature) FROM test_table").collect() + third_query_count = call_count["value"] + assert third_query_count == 6, ( + f"After third query: expected 6 total, got {third_query_count}" + ) + + def test_parallel_queries_independent(self, small_ds): + """Multiple contexts with the same table should work independently.""" + tracker1 = IterationTracker() + tracker2 = IterationTracker() + + table1 = read_xarray_table( + small_ds, + chunks={"time": 25}, + _iteration_callback=tracker1, + ) + + table2 = read_xarray_table( + small_ds, + chunks={"time": 50}, + _iteration_callback=tracker2, + ) + + ctx1 = SessionContext() + ctx2 = SessionContext() + + ctx1.register_table("test_table", table1) + ctx2.register_table("test_table", table2) + + # Execute queries + ctx1.sql("SELECT COUNT(*) FROM test_table").collect() + ctx2.sql("SELECT COUNT(*) FROM test_table").collect() + + # Each should have its own iteration count + assert tracker1.iteration_count == 4, ( + f"Table1: expected 4 blocks, got {tracker1.iteration_count}" + ) + assert tracker2.iteration_count == 2, ( + f"Table2: expected 2 blocks, got {tracker2.iteration_count}" + ) class TestFilterPushdown: - """Tests for partition pruning via filter pushdown. - - These tests verify that SQL filters on dimension columns (time, lat, lon) - correctly prune partitions, reducing the number of partitions read. - """ + """Tests for partition pruning via filter pushdown. - @pytest.fixture - def time_chunked_ds(self): - """Dataset chunked by time for pruning tests.""" - np.random.seed(42) - # 100 days of data, chunked into 4 partitions of 25 days each - time = pd.date_range("2020-01-01", periods=100, freq="D") - lat = np.linspace(-90, 90, 5) - data = np.random.rand(100, 5).astype(np.float32) - - return xr.Dataset( - {"temperature": (["time", "lat"], data)}, - coords={"time": time, "lat": lat}, - ) - - def test_time_gt_filter_prunes_early_partitions(self, time_chunked_ds): - """Query with time > X should skip early partitions.""" - tracker = IterationTracker() - - # 4 partitions: days 0-24, 25-49, 50-74, 75-99 - # (Jan 1-25, Jan 26-Feb 19, Feb 20-Mar 15, Mar 16-Apr 9) - table = read_xarray_table( - time_chunked_ds, - chunks={"time": 25}, - _iteration_callback=tracker, - ) - - ctx = SessionContext() - ctx.register_table("test", table) + These tests verify that SQL filters on dimension columns (time, lat, lon) + correctly prune partitions, reducing the number of partitions read. + """ - # Query only last 25 days (Mar 16+) - should prune first 3 partitions - # 2020-03-16 is day 75 - result = ctx.sql( - """ + @pytest.fixture + def time_chunked_ds(self): + """Dataset chunked by time for pruning tests.""" + np.random.seed(42) + # 100 days of data, chunked into 4 partitions of 25 days each + time = pd.date_range("2020-01-01", periods=100, freq="D") + lat = np.linspace(-90, 90, 5) + data = np.random.rand(100, 5).astype(np.float32) + + return xr.Dataset( + {"temperature": (["time", "lat"], data)}, + coords={"time": time, "lat": lat}, + ) + + def test_time_gt_filter_prunes_early_partitions(self, time_chunked_ds): + """Query with time > X should skip early partitions.""" + tracker = IterationTracker() + + # 4 partitions: days 0-24, 25-49, 50-74, 75-99 + # (Jan 1-25, Jan 26-Feb 19, Feb 20-Mar 15, Mar 16-Apr 9) + table = read_xarray_table( + time_chunked_ds, + chunks={"time": 25}, + _iteration_callback=tracker, + ) + + ctx = SessionContext() + ctx.register_table("test", table) + + # Query only last 25 days (Mar 16+) - should prune first 3 partitions + # 2020-03-16 is day 75 + result = ctx.sql( + """ SELECT COUNT(*) as cnt FROM test WHERE time >= '2020-03-16' """ - ).to_pandas() + ).to_pandas() - # Should read only 1 partition (the last one) - assert ( - tracker.iteration_count == 1 - ), f"Expected 1 partition after filter pushdown, got {tracker.iteration_count}" + # Should read only 1 partition (the last one) + assert tracker.iteration_count == 1, ( + f"Expected 1 partition after filter pushdown, got {tracker.iteration_count}" + ) - # Verify data correctness - 25 days * 5 lat = 125 rows - count = result["cnt"].iloc[0] - assert count == 125, f"Expected 125 rows, got {count}" + # Verify data correctness - 25 days * 5 lat = 125 rows + count = result["cnt"].iloc[0] + assert count == 125, f"Expected 125 rows, got {count}" - def test_time_lt_filter_prunes_late_partitions(self, time_chunked_ds): - """Query with time < X should skip late partitions.""" - tracker = IterationTracker() + def test_time_lt_filter_prunes_late_partitions(self, time_chunked_ds): + """Query with time < X should skip late partitions.""" + tracker = IterationTracker() - table = read_xarray_table( - time_chunked_ds, - chunks={"time": 25}, - _iteration_callback=tracker, - ) + table = read_xarray_table( + time_chunked_ds, + chunks={"time": 25}, + _iteration_callback=tracker, + ) - ctx = SessionContext() - ctx.register_table("test", table) + ctx = SessionContext() + ctx.register_table("test", table) - # Query only first 25 days (< Jan 26) - should prune last 3 partitions - result = ctx.sql( - """ + # Query only first 25 days (< Jan 26) - should prune last 3 partitions + result = ctx.sql( + """ SELECT COUNT(*) as cnt FROM test WHERE time < '2020-01-26' """ - ).to_pandas() + ).to_pandas() - # Should read only 1 partition (the first one) - assert ( - tracker.iteration_count == 1 - ), f"Expected 1 partition after filter pushdown, got {tracker.iteration_count}" + # Should read only 1 partition (the first one) + assert tracker.iteration_count == 1, ( + f"Expected 1 partition after filter pushdown, got {tracker.iteration_count}" + ) - # Verify correctness: Jan 1–25 (25 days) × 5 lat = 125 rows - count = result["cnt"].iloc[0] - assert count == 125, f"Expected 125 rows, got {count}" + # Verify correctness: Jan 1–25 (25 days) × 5 lat = 125 rows + count = result["cnt"].iloc[0] + assert count == 125, f"Expected 125 rows, got {count}" - def test_time_between_filter_prunes_outside_range(self, time_chunked_ds): - """Query with BETWEEN should prune partitions outside the range.""" - tracker = IterationTracker() + def test_time_between_filter_prunes_outside_range(self, time_chunked_ds): + """Query with BETWEEN should prune partitions outside the range.""" + tracker = IterationTracker() - table = read_xarray_table( - time_chunked_ds, - chunks={"time": 25}, - _iteration_callback=tracker, - ) + table = read_xarray_table( + time_chunked_ds, + chunks={"time": 25}, + _iteration_callback=tracker, + ) - ctx = SessionContext() - ctx.register_table("test", table) + ctx = SessionContext() + ctx.register_table("test", table) - # Query middle 50 days (Feb 1 - Mar 21) - should hit partitions 1, 2, and 3 - result = ctx.sql( - """ + # Query middle 50 days (Feb 1 - Mar 21) - should hit partitions 1, 2, and 3 + ctx.sql( + """ SELECT COUNT(*) as cnt FROM test WHERE time BETWEEN '2020-02-01' AND '2020-03-21' """ - ).collect() - - # Partition 0 (Jan 1–25) ends before Feb 1 and is pruned. - # Partitions 1, 2, 3 each overlap with Feb 1–Mar 21. - assert ( - tracker.iteration_count == 3 - ), f"Expected exactly 3 partitions after BETWEEN pruning, got {tracker.iteration_count}" - - def test_lat_filter_prunes_partitions(self): - """Latitude filter should prune irrelevant partitions.""" - np.random.seed(42) - time = pd.date_range("2020-01-01", periods=10, freq="D") - # 100 lat values from -90 to 90 - lat = np.linspace(-90, 90, 100) - data = np.random.rand(10, 100).astype(np.float32) - - ds = xr.Dataset( - {"temperature": (["time", "lat"], data)}, - coords={"time": time, "lat": lat}, - ) - - tracker = IterationTracker() - - # Chunk by latitude: 4 partitions (25 lat values each) - # Partition 0: lat -90 to ~-45 - # Partition 1: lat ~-45 to 0 - # Partition 2: lat 0 to ~45 - # Partition 3: lat ~45 to 90 - table = read_xarray_table( - ds, - chunks={"lat": 25}, - _iteration_callback=tracker, - ) - - ctx = SessionContext() - ctx.register_table("test", table) - - # Query southern hemisphere only (lat < 0) - result = ctx.sql( - """ + ).collect() + + # Partition 0 (Jan 1–25) ends before Feb 1 and is pruned. + # Partitions 1, 2, 3 each overlap with Feb 1–Mar 21. + assert tracker.iteration_count == 3, ( + f"Expected exactly 3 partitions after BETWEEN pruning, got {tracker.iteration_count}" + ) + + def test_lat_filter_prunes_partitions(self): + """Latitude filter should prune irrelevant partitions.""" + np.random.seed(42) + time = pd.date_range("2020-01-01", periods=10, freq="D") + # 100 lat values from -90 to 90 + lat = np.linspace(-90, 90, 100) + data = np.random.rand(10, 100).astype(np.float32) + + ds = xr.Dataset( + {"temperature": (["time", "lat"], data)}, + coords={"time": time, "lat": lat}, + ) + + tracker = IterationTracker() + + # Chunk by latitude: 4 partitions (25 lat values each) + # Partition 0: lat -90 to ~-45 + # Partition 1: lat ~-45 to 0 + # Partition 2: lat 0 to ~45 + # Partition 3: lat ~45 to 90 + table = read_xarray_table( + ds, + chunks={"lat": 25}, + _iteration_callback=tracker, + ) + + ctx = SessionContext() + ctx.register_table("test", table) + + # Query southern hemisphere only (lat < 0) + ctx.sql( + """ SELECT COUNT(*) as cnt FROM test WHERE lat < 0 """ - ).collect() - - # np.linspace(-90, 90, 100) chunked by 25: - # Partition 0: indices 0–24, lat -90.0 to -46.4 (all negative) - # Partition 1: indices 25–49, lat -44.5 to -0.9 (all negative) - # Partition 2: indices 50–74, lat 0.9 to 44.5 (all positive) → pruned - # Partition 3: indices 75–99, lat 46.4 to 90.0 (all positive) → pruned - assert ( - tracker.iteration_count == 2 - ), f"Expected exactly 2 partitions for lat < 0, got {tracker.iteration_count}" - - def test_no_pruning_for_data_column_filters(self, time_chunked_ds): - """Filters on data columns (not dimensions) should not prune.""" - tracker = IterationTracker() - - table = read_xarray_table( - time_chunked_ds, - chunks={"time": 25}, - _iteration_callback=tracker, - ) - - ctx = SessionContext() - ctx.register_table("test", table) - - # Filter on temperature (data column), not a dimension - result = ctx.sql( - """ + ).collect() + + # np.linspace(-90, 90, 100) chunked by 25: + # Partition 0: indices 0–24, lat -90.0 to -46.4 (all negative) + # Partition 1: indices 25–49, lat -44.5 to -0.9 (all negative) + # Partition 2: indices 50–74, lat 0.9 to 44.5 (all positive) → pruned + # Partition 3: indices 75–99, lat 46.4 to 90.0 (all positive) → pruned + assert tracker.iteration_count == 2, ( + f"Expected exactly 2 partitions for lat < 0, got {tracker.iteration_count}" + ) + + def test_no_pruning_for_data_column_filters(self, time_chunked_ds): + """Filters on data columns (not dimensions) should not prune.""" + tracker = IterationTracker() + + table = read_xarray_table( + time_chunked_ds, + chunks={"time": 25}, + _iteration_callback=tracker, + ) + + ctx = SessionContext() + ctx.register_table("test", table) + + # Filter on temperature (data column), not a dimension + ctx.sql( + """ SELECT COUNT(*) FROM test WHERE temperature > 0.5 """ - ).collect() - - # All 4 partitions should be read (can't prune on data column) - assert ( - tracker.iteration_count == 4 - ), f"Expected 4 partitions (no pruning on data column), got {tracker.iteration_count}" - - def test_filter_correctness_preserved(self, time_chunked_ds): - """Verify filtered results are correct after pruning.""" - table = read_xarray_table( - time_chunked_ds, - chunks={"time": 25}, - ) - - ctx = SessionContext() - ctx.register_table("test", table) - - # Get count with filter - filtered = ctx.sql( - """ + ).collect() + + # All 4 partitions should be read (can't prune on data column) + assert tracker.iteration_count == 4, ( + f"Expected 4 partitions (no pruning on data column), got {tracker.iteration_count}" + ) + + def test_filter_correctness_preserved(self, time_chunked_ds): + """Verify filtered results are correct after pruning.""" + table = read_xarray_table( + time_chunked_ds, + chunks={"time": 25}, + ) + + ctx = SessionContext() + ctx.register_table("test", table) + + # Get count with filter + filtered = ctx.sql( + """ SELECT COUNT(*) as cnt FROM test WHERE time >= '2020-02-15' AND time <= '2020-03-15' """ - ).to_pandas() + ).to_pandas() - # Manual calculation: Feb 15 (day 45) to Mar 15 (day 74) = 30 days - # 30 days * 5 lat values = 150 rows - count = filtered["cnt"].iloc[0] - assert count == 150, f"Expected 150 rows, got {count}" + # Manual calculation: Feb 15 (day 45) to Mar 15 (day 74) = 30 days + # 30 days * 5 lat values = 150 rows + count = filtered["cnt"].iloc[0] + assert count == 150, f"Expected 150 rows, got {count}" - def test_and_filter_combines_pruning(self, time_chunked_ds): - """AND filters should combine for maximum pruning.""" - tracker = IterationTracker() + def test_and_filter_combines_pruning(self, time_chunked_ds): + """AND filters should combine for maximum pruning.""" + tracker = IterationTracker() - table = read_xarray_table( - time_chunked_ds, - chunks={"time": 25}, - _iteration_callback=tracker, - ) + table = read_xarray_table( + time_chunked_ds, + chunks={"time": 25}, + _iteration_callback=tracker, + ) - ctx = SessionContext() - ctx.register_table("test", table) + ctx = SessionContext() + ctx.register_table("test", table) - # Very narrow range that spans only 1 partition - result = ctx.sql( - """ + # Very narrow range that spans only 1 partition + ctx.sql( + """ SELECT * FROM test WHERE time >= '2020-03-20' AND time <= '2020-04-05' """ - ).collect() + ).collect() - # Should read only 1 partition - assert ( - tracker.iteration_count == 1 - ), f"Expected 1 partition for narrow AND range, got {tracker.iteration_count}" + # Should read only 1 partition + assert tracker.iteration_count == 1, ( + f"Expected 1 partition for narrow AND range, got {tracker.iteration_count}" + ) - def test_or_filter_is_conservative(self, time_chunked_ds): - """OR filters should include partitions matching either condition.""" - tracker = IterationTracker() + def test_or_filter_is_conservative(self, time_chunked_ds): + """OR filters should include partitions matching either condition.""" + tracker = IterationTracker() - table = read_xarray_table( - time_chunked_ds, - chunks={"time": 25}, - _iteration_callback=tracker, - ) + table = read_xarray_table( + time_chunked_ds, + chunks={"time": 25}, + _iteration_callback=tracker, + ) - ctx = SessionContext() - ctx.register_table("test", table) + ctx = SessionContext() + ctx.register_table("test", table) - # First or last partition (OR condition) - result = ctx.sql( - """ + # First or last partition (OR condition) + ctx.sql( + """ SELECT * FROM test WHERE time < '2020-01-10' OR time > '2020-03-30' """ - ).collect() + ).collect() - # Should read at least 2 partitions (first and last) - assert ( - tracker.iteration_count >= 2 - ), f"Expected at least 2 partitions for OR filter, got {tracker.iteration_count}" + # Should read at least 2 partitions (first and last) + assert tracker.iteration_count >= 2, ( + f"Expected at least 2 partitions for OR filter, got {tracker.iteration_count}" + ) - def test_empty_result_from_impossible_filter(self, time_chunked_ds): - """Filter that matches no data should return empty result.""" - tracker = IterationTracker() + def test_empty_result_from_impossible_filter(self, time_chunked_ds): + """Filter that matches no data should return empty result.""" + tracker = IterationTracker() - table = read_xarray_table( - time_chunked_ds, - chunks={"time": 25}, - _iteration_callback=tracker, - ) + table = read_xarray_table( + time_chunked_ds, + chunks={"time": 25}, + _iteration_callback=tracker, + ) - ctx = SessionContext() - ctx.register_table("test", table) + ctx = SessionContext() + ctx.register_table("test", table) - # Query for dates outside the data range - result = ctx.sql( - """ + # Query for dates outside the data range + result = ctx.sql( + """ SELECT COUNT(*) as cnt FROM test WHERE time > '2025-01-01' """ - ).to_pandas() + ).to_pandas() - # Should read 0 partitions (all pruned) - assert ( - tracker.iteration_count == 0 - ), f"Expected 0 partitions for impossible filter, got {tracker.iteration_count}" + # Should read 0 partitions (all pruned) + assert tracker.iteration_count == 0, ( + f"Expected 0 partitions for impossible filter, got {tracker.iteration_count}" + ) - # Result should be 0 rows - count = result["cnt"].iloc[0] - assert count == 0, f"Expected 0 rows, got {count}" + # Result should be 0 rows + count = result["cnt"].iloc[0] + assert count == 0, f"Expected 0 rows, got {count}" class TestProjectionPushdown: - """Tests that column projection is pushed down to the xarray factory. - - When a query requests only a subset of data variables, the factory should - receive only those column names — not the full schema — so that xarray - skips loading unrequested variables from disk. - """ - - @pytest.fixture - def two_var_ds(self): - """A small dataset with two data variables sharing the same dimensions.""" - np.random.seed(0) - time = pd.date_range("2020-01-01", periods=10, freq="D") - lat = np.linspace(-10, 10, 5, dtype=np.float32) - shape = (10, 5) - return xr.Dataset( - { - "temperature": ( - ["time", "lat"], - np.random.rand(*shape).astype(np.float32), - ), - "precipitation": ( - ["time", "lat"], - np.random.rand(*shape).astype(np.float32), - ), - }, - coords={"time": time, "lat": lat}, - ) - - def test_single_column_select_projects_only_that_variable(self, two_var_ds): - """SELECT on one data variable should pass only that variable to the factory.""" - tracker = IterationTracker() - table = read_xarray_table( - two_var_ds, - chunks={"time": 5}, - _iteration_callback=tracker, - ) + """Tests that column projection is pushed down to the xarray factory. - ctx = SessionContext() - ctx.register_table("data", table) - ctx.sql("SELECT AVG(temperature) FROM data").collect() - - assert ( - tracker.iteration_count > 0 - ), "Expected at least one partition to be read" - for proj in tracker.projections_seen: - assert ( - proj is not None - ), "Expected projection_names to be set for a single-column SELECT, got None" - assert ( - "temperature" in proj - ), f"Expected 'temperature' in projection_names, got {proj}" - assert ( - "precipitation" not in proj - ), f"'precipitation' should not be loaded for a temperature-only query, got {proj}" - - def test_full_select_includes_all_variables(self, two_var_ds): - """SELECT * should include all data variables in the projection.""" - tracker = IterationTracker() - table = read_xarray_table( - two_var_ds, - chunks={"time": 5}, - _iteration_callback=tracker, - ) + When a query requests only a subset of data variables, the factory should + receive only those column names — not the full schema — so that xarray + skips loading unrequested variables from disk. + """ - ctx = SessionContext() - ctx.register_table("data", table) - ctx.sql("SELECT * FROM data").collect() - - assert ( - tracker.iteration_count > 0 - ), "Expected at least one partition to be read" - for proj in tracker.projections_seen: - # DataFusion sends all column names explicitly even for SELECT *. - # Either None (no pushdown) or a list containing both data variables. - if proj is not None: - assert "temperature" in proj, f"Missing 'temperature' in {proj}" - assert "precipitation" in proj, f"Missing 'precipitation' in {proj}" - - def test_multi_column_select_includes_all_requested(self, two_var_ds): - """SELECT on multiple data variables should include all of them in the projection.""" - tracker = IterationTracker() - table = read_xarray_table( - two_var_ds, - chunks={"time": 5}, - _iteration_callback=tracker, - ) + @pytest.fixture + def two_var_ds(self): + """A small dataset with two data variables sharing the same dimensions.""" + np.random.seed(0) + time = pd.date_range("2020-01-01", periods=10, freq="D") + lat = np.linspace(-10, 10, 5, dtype=np.float32) + shape = (10, 5) + return xr.Dataset( + { + "temperature": ( + ["time", "lat"], + np.random.rand(*shape).astype(np.float32), + ), + "precipitation": ( + ["time", "lat"], + np.random.rand(*shape).astype(np.float32), + ), + }, + coords={"time": time, "lat": lat}, + ) + + def test_single_column_select_projects_only_that_variable(self, two_var_ds): + """SELECT on one data variable should pass only that variable to the factory.""" + tracker = IterationTracker() + table = read_xarray_table( + two_var_ds, + chunks={"time": 5}, + _iteration_callback=tracker, + ) + + ctx = SessionContext() + ctx.register_table("data", table) + ctx.sql("SELECT AVG(temperature) FROM data").collect() + + assert tracker.iteration_count > 0, ( + "Expected at least one partition to be read" + ) + for proj in tracker.projections_seen: + assert proj is not None, ( + "Expected projection_names to be set for a single-column SELECT, got None" + ) + assert "temperature" in proj, ( + f"Expected 'temperature' in projection_names, got {proj}" + ) + assert "precipitation" not in proj, ( + f"'precipitation' should not be loaded for a temperature-only query, got {proj}" + ) - ctx = SessionContext() - ctx.register_table("data", table) - ctx.sql("SELECT AVG(temperature), AVG(precipitation) FROM data").collect() - - assert tracker.iteration_count > 0 - for proj in tracker.projections_seen: - # Both variables requested — either None (full scan) or both present - if proj is not None: - assert "temperature" in proj, f"Missing 'temperature' in {proj}" - assert "precipitation" in proj, f"Missing 'precipitation' in {proj}" - - def test_projection_result_correctness(self, two_var_ds): - """Single-column projected query returns the same result as an unprojected one.""" - table = read_xarray_table(two_var_ds, chunks={"time": 5}) - ctx = SessionContext() - ctx.register_table("data", table) - - projected = ( - ctx.sql("SELECT AVG(temperature) as avg_t FROM data") - .to_pandas()["avg_t"] - .iloc[0] - ) - expected = float(two_var_ds["temperature"].values.mean()) - assert ( - abs(projected - expected) < 1e-4 - ), f"Projected AVG {projected} differs from expected {expected}" - - def test_count_star_passes_none_projection(self, two_var_ds): - """COUNT(*) should not push a projection — factory receives None.""" - projections_seen = [] - - def callback(block, projection_names): - projections_seen.append(projection_names) - - table = read_xarray_table( - two_var_ds, - chunks={"time": 5}, - _iteration_callback=callback, - ) - ctx = SessionContext() - ctx.register_table("data", table) - result = ctx.sql("SELECT COUNT(*) FROM data").collect() - - total_rows = 10 * 5 # time=10, lat=5 - assert result[0][0][0].as_py() == total_rows - assert all( - p is None for p in projections_seen - ), f"COUNT(*) should not push a projection, but got: {projections_seen}" + def test_full_select_includes_all_variables(self, two_var_ds): + """SELECT * should include all data variables in the projection.""" + tracker = IterationTracker() + table = read_xarray_table( + two_var_ds, + chunks={"time": 5}, + _iteration_callback=tracker, + ) + + ctx = SessionContext() + ctx.register_table("data", table) + ctx.sql("SELECT * FROM data").collect() + + assert tracker.iteration_count > 0, ( + "Expected at least one partition to be read" + ) + for proj in tracker.projections_seen: + # DataFusion sends all column names explicitly even for SELECT *. + # Either None (no pushdown) or a list containing both data variables. + if proj is not None: + assert "temperature" in proj, f"Missing 'temperature' in {proj}" + assert "precipitation" in proj, ( + f"Missing 'precipitation' in {proj}" + ) + + def test_multi_column_select_includes_all_requested(self, two_var_ds): + """SELECT on multiple data variables should include all of them in the projection.""" + tracker = IterationTracker() + table = read_xarray_table( + two_var_ds, + chunks={"time": 5}, + _iteration_callback=tracker, + ) + + ctx = SessionContext() + ctx.register_table("data", table) + ctx.sql( + "SELECT AVG(temperature), AVG(precipitation) FROM data" + ).collect() + + assert tracker.iteration_count > 0 + for proj in tracker.projections_seen: + # Both variables requested — either None (full scan) or both present + if proj is not None: + assert "temperature" in proj, f"Missing 'temperature' in {proj}" + assert "precipitation" in proj, ( + f"Missing 'precipitation' in {proj}" + ) + + def test_projection_result_correctness(self, two_var_ds): + """Single-column projected query returns the same result as an unprojected one.""" + table = read_xarray_table(two_var_ds, chunks={"time": 5}) + ctx = SessionContext() + ctx.register_table("data", table) + + projected = ( + ctx.sql("SELECT AVG(temperature) as avg_t FROM data") + .to_pandas()["avg_t"] + .iloc[0] + ) + expected = float(two_var_ds["temperature"].values.mean()) + assert abs(projected - expected) < 1e-4, ( + f"Projected AVG {projected} differs from expected {expected}" + ) + + def test_count_star_passes_none_projection(self, two_var_ds): + """COUNT(*) should not push a projection — factory receives None.""" + projections_seen = [] + + def callback(block, projection_names): + projections_seen.append(projection_names) + + table = read_xarray_table( + two_var_ds, + chunks={"time": 5}, + _iteration_callback=callback, + ) + ctx = SessionContext() + ctx.register_table("data", table) + result = ctx.sql("SELECT COUNT(*) FROM data").collect() + + total_rows = 10 * 5 # time=10, lat=5 + assert result[0][0][0].as_py() == total_rows + assert all(p is None for p in projections_seen), ( + f"COUNT(*) should not push a projection, but got: {projections_seen}" + ) diff --git a/tests/test_sql.py b/tests/test_sql.py index 17aee70a..01929f09 100644 --- a/tests/test_sql.py +++ b/tests/test_sql.py @@ -9,169 +9,173 @@ def test_sanity(air_dataset_small): - ctx = XarrayContext() - ctx.from_dataset("air", air_dataset_small) - result = ctx.sql( - 'SELECT "lat", "lon", "time", "air" FROM "air" LIMIT 100' - ).to_pandas() - assert len(result) > 0 - assert len(result) <= 1320 - assert all(col in result.columns for col in ["lat", "lon", "time", "air"]) + ctx = XarrayContext() + ctx.from_dataset("air", air_dataset_small) + result = ctx.sql( + 'SELECT "lat", "lon", "time", "air" FROM "air" LIMIT 100' + ).to_pandas() + assert len(result) > 0 + assert len(result) <= 1320 + assert all(col in result.columns for col in ["lat", "lon", "time", "air"]) def test_aggregation_small(air_dataset_small): - ctx = XarrayContext() - ctx.from_dataset("air", air_dataset_small) - query = """ + ctx = XarrayContext() + ctx.from_dataset("air", air_dataset_small) + query = """ SELECT lat, lon, SUM(air) AS air_total FROM air GROUP BY lat, lon """ - result = ctx.sql(query).to_pandas() - expected_rows = ( - air_dataset_small.sizes["lat"] * air_dataset_small.sizes["lon"] - ) - assert len(result) == expected_rows + result = ctx.sql(query).to_pandas() + expected_rows = ( + air_dataset_small.sizes["lat"] * air_dataset_small.sizes["lon"] + ) + assert len(result) == expected_rows def test_aggregation_large(air_dataset_large): - ctx = XarrayContext() - ctx.from_dataset("air", air_dataset_large) - query = """ + ctx = XarrayContext() + ctx.from_dataset("air", air_dataset_large) + query = """ SELECT lat, lon, AVG(air) AS air_avg FROM air GROUP BY lat, lon """ - result = ctx.sql(query).to_pandas() - expected_rows = ( - air_dataset_large.sizes["lat"] * air_dataset_large.sizes["lon"] - ) - assert len(result) == expected_rows + result = ctx.sql(query).to_pandas() + expected_rows = ( + air_dataset_large.sizes["lat"] * air_dataset_large.sizes["lon"] + ) + assert len(result) == expected_rows def test_basic_select_all(air_dataset_small): - ctx = XarrayContext() - ctx.from_dataset("air", air_dataset_small) - result = ctx.sql("SELECT * FROM air LIMIT 10").to_pandas() - assert len(result) <= 10 - for col in ["lat", "lon", "time", "air"]: - assert col in result.columns + ctx = XarrayContext() + ctx.from_dataset("air", air_dataset_small) + result = ctx.sql("SELECT * FROM air LIMIT 10").to_pandas() + assert len(result) <= 10 + for col in ["lat", "lon", "time", "air"]: + assert col in result.columns def test_weather_queries(weather_dataset): - ctx = XarrayContext() - ctx.from_dataset("weather", weather_dataset) - # Selecting specific columns - result = ctx.sql( - "SELECT lat, lon, temperature, precipitation FROM weather LIMIT 20" - ).to_pandas() - assert "temperature" in result.columns - assert "precipitation" in result.columns - # Filtering - result = ctx.sql( - "SELECT * FROM weather WHERE temperature > 10 LIMIT 50" - ).to_pandas() - assert len(result) > 0 - assert (result["temperature"] > 10).all() + ctx = XarrayContext() + ctx.from_dataset("weather", weather_dataset) + # Selecting specific columns + result = ctx.sql( + "SELECT lat, lon, temperature, precipitation FROM weather LIMIT 20" + ).to_pandas() + assert "temperature" in result.columns + assert "precipitation" in result.columns + # Filtering + result = ctx.sql( + "SELECT * FROM weather WHERE temperature > 10 LIMIT 50" + ).to_pandas() + assert len(result) > 0 + assert (result["temperature"] > 10).all() def test_synthetic_aggregations(synthetic_dataset): - ctx = XarrayContext() - ctx.from_dataset("synthetic", synthetic_dataset) - # COUNT aggregation - result = ctx.sql("SELECT COUNT(*) AS total_count FROM synthetic").to_pandas() - assert result["total_count"].iloc[0] > 0 - # MIN, MAX, AVG - query = """ + ctx = XarrayContext() + ctx.from_dataset("synthetic", synthetic_dataset) + # COUNT aggregation + result = ctx.sql( + "SELECT COUNT(*) AS total_count FROM synthetic" + ).to_pandas() + assert result["total_count"].iloc[0] > 0 + # MIN, MAX, AVG + query = """ SELECT MIN(temperature) AS min_temp, MAX(temperature) AS max_temp, AVG(temperature) AS avg_temp FROM synthetic """ - result = ctx.sql(query).to_pandas() - assert result["min_temp"].iloc[0] < result["max_temp"].iloc[0] - assert ( - result["min_temp"].iloc[0] - <= result["avg_temp"].iloc[0] - <= result["max_temp"].iloc[0] - ) + result = ctx.sql(query).to_pandas() + assert result["min_temp"].iloc[0] < result["max_temp"].iloc[0] + assert ( + result["min_temp"].iloc[0] + <= result["avg_temp"].iloc[0] + <= result["max_temp"].iloc[0] + ) def test_invalid_table_name(air_dataset_small): - ctx = XarrayContext() - ctx.from_dataset("air", air_dataset_small) - with pytest.raises(Exception): - ctx.sql("SELECT * FROM nonexistent_table") + ctx = XarrayContext() + ctx.from_dataset("air", air_dataset_small) + with pytest.raises(Exception): + ctx.sql("SELECT * FROM nonexistent_table") def test_invalid_column_name(air_dataset_small): - ctx = XarrayContext() - ctx.from_dataset("air", air_dataset_small) - with pytest.raises(Exception): - ctx.sql("SELECT nonexistent_column FROM air") + ctx = XarrayContext() + ctx.from_dataset("air", air_dataset_small) + with pytest.raises(Exception): + ctx.sql("SELECT nonexistent_column FROM air") def test_sql_syntax_error(air_dataset_small): - ctx = XarrayContext() - ctx.from_dataset("air", air_dataset_small) - with pytest.raises(Exception): - ctx.sql("SELECT * FORM air") # Typo: FORM instead of FROM - with pytest.raises(Exception): - ctx.sql("SELECT * FROM air WHERE") # Incomplete WHERE + ctx = XarrayContext() + ctx.from_dataset("air", air_dataset_small) + with pytest.raises(Exception): + ctx.sql("SELECT * FORM air") # Typo: FORM instead of FROM + with pytest.raises(Exception): + ctx.sql("SELECT * FROM air WHERE") # Incomplete WHERE def test_cross_join(air_and_stations): - air, stations = air_and_stations - ctx = XarrayContext() - ctx.from_dataset("air_data", air) - ctx.from_dataset("stations", stations) - result = ctx.sql( - "SELECT COUNT(*) AS total FROM air_data CROSS JOIN stations" - ).to_pandas() - assert result["total"].iloc[0] > 0 + air, stations = air_and_stations + ctx = XarrayContext() + ctx.from_dataset("air_data", air) + ctx.from_dataset("stations", stations) + result = ctx.sql( + "SELECT COUNT(*) AS total FROM air_data CROSS JOIN stations" + ).to_pandas() + assert result["total"].iloc[0] > 0 def test_string_coordinates(): - """String-typed coordinates should not crash during registration.""" - ds = xr.Dataset( - {"score": (["student", "subject"], np.random.rand(3, 2))}, - coords={ - "student": ["alice", "bob", "charlie"], - "subject": ["math", "science"], - }, - ) - ctx = XarrayContext() - ctx.from_dataset("scores", ds.chunk({"student": 3, "subject": 2})) - result = ctx.sql("SELECT * FROM scores").to_pandas() - assert len(result) == 6 - assert "student" in result.columns - assert "subject" in result.columns - assert "score" in result.columns - - -class TestNanAsNull: - """NaN in float columns should become Arrow nulls so SQL aggregates work.""" - - @pytest.fixture - def nan_ds(self): - data = np.array([[[1.0, 2.0], [np.nan, 4.0]], [[5.0, np.nan], [7.0, 8.0]]]) - return xr.Dataset( - {"temp": (["time", "x", "y"], data)}, + """String-typed coordinates should not crash during registration.""" + ds = xr.Dataset( + {"score": (["student", "subject"], np.random.rand(3, 2))}, coords={ - "time": pd.date_range("2020-01-01", periods=2), - "x": [0, 1], - "y": [0, 1], + "student": ["alice", "bob", "charlie"], + "subject": ["math", "science"], }, - ).chunk({"time": 1}) - - def test_nan_aggregates(self, nan_ds): + ) ctx = XarrayContext() - ctx.from_dataset("data", nan_ds) + ctx.from_dataset("scores", ds.chunk({"student": 3, "subject": 2})) + result = ctx.sql("SELECT * FROM scores").to_pandas() + assert len(result) == 6 + assert "student" in result.columns + assert "subject" in result.columns + assert "score" in result.columns - # Test multiple aggregates at once: - # MAX/MIN/AVG should ignore NaN, COUNT(col) should exclude NaN, - # and WHERE col IS NULL should match NaN. - query = """ + +class TestNanAsNull: + """NaN in float columns should become Arrow nulls so SQL aggregates work.""" + + @pytest.fixture + def nan_ds(self): + data = np.array( + [[[1.0, 2.0], [np.nan, 4.0]], [[5.0, np.nan], [7.0, 8.0]]] + ) + return xr.Dataset( + {"temp": (["time", "x", "y"], data)}, + coords={ + "time": pd.date_range("2020-01-01", periods=2), + "x": [0, 1], + "y": [0, 1], + }, + ).chunk({"time": 1}) + + def test_nan_aggregates(self, nan_ds): + ctx = XarrayContext() + ctx.from_dataset("data", nan_ds) + + # Test multiple aggregates at once: + # MAX/MIN/AVG should ignore NaN, COUNT(col) should exclude NaN, + # and WHERE col IS NULL should match NaN. + query = """ SELECT MAX(temp) AS mx, MIN(temp) AS mn, @@ -180,139 +184,141 @@ def test_nan_aggregates(self, nan_ds): COUNT(*) FILTER (WHERE temp IS NULL) AS null_cnt FROM data """ - result = ctx.sql(query).to_pandas().iloc[0] + result = ctx.sql(query).to_pandas().iloc[0] - assert result["mx"] == 8.0 - assert result["mn"] == 1.0 - expected_avg = np.nanmean([1.0, 2.0, 4.0, 5.0, 7.0, 8.0]) - assert abs(result["avg"] - expected_avg) < 1e-6 - assert result["cnt"] == 6 - assert result["null_cnt"] == 2 + assert result["mx"] == 8.0 + assert result["mn"] == 1.0 + expected_avg = np.nanmean([1.0, 2.0, 4.0, 5.0, 7.0, 8.0]) + assert abs(result["avg"] - expected_avg) < 1e-6 + assert result["cnt"] == 6 + assert result["null_cnt"] == 2 class TestCftimeGregorianLike: - """Tests for Gregorian-like cftime calendars (noleap, standard, etc.). - - These use pa.timestamp('us') and support string-based SQL filters. - """ - - def test_noleap_dataset_registers(self, rasm_ds): - """A noleap dataset should register without errors.""" - ctx = XarrayContext() - ctx.from_dataset("rasm", rasm_ds, chunks={"time": 12}) - result = ctx.sql("SELECT COUNT(*) AS cnt FROM rasm").to_pandas() - assert result["cnt"].iloc[0] > 0 - - def test_select_time_column(self, rasm_ds): - """Querying the time column should return valid timestamps.""" - ctx = XarrayContext() - ctx.from_dataset("rasm", rasm_ds, chunks={"time": 12}) - result = ctx.sql( - "SELECT DISTINCT time FROM rasm ORDER BY time LIMIT 5" - ).to_pandas() - assert len(result) == 5 - times = result["time"].tolist() - assert times == sorted(times) + """Tests for Gregorian-like cftime calendars (noleap, standard, etc.). - def test_string_filter_works(self, rasm_ds): - """String-based time filters should work for Gregorian-like calendars.""" - ctx = XarrayContext() - ctx.from_dataset("rasm", rasm_ds, chunks={"time": 12}) - result = ctx.sql( - "SELECT COUNT(*) AS cnt FROM rasm WHERE time >= '1980-10-01'" - ).to_pandas() - full = ctx.sql("SELECT COUNT(*) AS cnt FROM rasm").to_pandas() - assert 0 < result["cnt"].iloc[0] < full["cnt"].iloc[0] - - def test_aggregation(self, rasm_ds): - """MIN/MAX on timestamp columns should work.""" - ctx = XarrayContext() - ctx.from_dataset("rasm", rasm_ds, chunks={"time": 12}) - result = ctx.sql( - "SELECT MIN(time) AS t_min, MAX(time) AS t_max FROM rasm" - ).to_pandas() - assert result["t_min"].iloc[0] < result["t_max"].iloc[0] + These use pa.timestamp('us') and support string-based SQL filters. + """ - def test_row_count_matches_xarray(self, rasm_ds): - """Total row count should equal the product of dimension sizes.""" - ctx = XarrayContext() - ctx.from_dataset("rasm", rasm_ds, chunks={"time": 12}) - result = ctx.sql("SELECT COUNT(*) AS cnt FROM rasm").to_pandas() - expected = int(np.prod([rasm_ds.sizes[d] for d in rasm_ds["Tair"].dims])) - assert result["cnt"].iloc[0] == expected + def test_noleap_dataset_registers(self, rasm_ds): + """A noleap dataset should register without errors.""" + ctx = XarrayContext() + ctx.from_dataset("rasm", rasm_ds, chunks={"time": 12}) + result = ctx.sql("SELECT COUNT(*) AS cnt FROM rasm").to_pandas() + assert result["cnt"].iloc[0] > 0 + + def test_select_time_column(self, rasm_ds): + """Querying the time column should return valid timestamps.""" + ctx = XarrayContext() + ctx.from_dataset("rasm", rasm_ds, chunks={"time": 12}) + result = ctx.sql( + "SELECT DISTINCT time FROM rasm ORDER BY time LIMIT 5" + ).to_pandas() + assert len(result) == 5 + times = result["time"].tolist() + assert times == sorted(times) + + def test_string_filter_works(self, rasm_ds): + """String-based time filters should work for Gregorian-like calendars.""" + ctx = XarrayContext() + ctx.from_dataset("rasm", rasm_ds, chunks={"time": 12}) + result = ctx.sql( + "SELECT COUNT(*) AS cnt FROM rasm WHERE time >= '1980-10-01'" + ).to_pandas() + full = ctx.sql("SELECT COUNT(*) AS cnt FROM rasm").to_pandas() + assert 0 < result["cnt"].iloc[0] < full["cnt"].iloc[0] + + def test_aggregation(self, rasm_ds): + """MIN/MAX on timestamp columns should work.""" + ctx = XarrayContext() + ctx.from_dataset("rasm", rasm_ds, chunks={"time": 12}) + result = ctx.sql( + "SELECT MIN(time) AS t_min, MAX(time) AS t_max FROM rasm" + ).to_pandas() + assert result["t_min"].iloc[0] < result["t_max"].iloc[0] + + def test_row_count_matches_xarray(self, rasm_ds): + """Total row count should equal the product of dimension sizes.""" + ctx = XarrayContext() + ctx.from_dataset("rasm", rasm_ds, chunks={"time": 12}) + result = ctx.sql("SELECT COUNT(*) AS cnt FROM rasm").to_pandas() + expected = int( + np.prod([rasm_ds.sizes[d] for d in rasm_ds["Tair"].dims]) + ) + assert result["cnt"].iloc[0] == expected class TestCftimeNonGregorian: - """Tests for non-Gregorian cftime calendars (360_day, julian). - - These use pa.int64() with CF-convention metadata and the cftime() UDF. - """ - - @pytest.fixture - def ds_360day(self): - """Synthetic 360-day calendar dataset.""" - import cftime + """Tests for non-Gregorian cftime calendars (360_day, julian). - times = [cftime.Datetime360Day(2000, m, 1) for m in range(1, 13)] - return xr.Dataset( - {"temp": ("time", np.arange(12, dtype=np.float32))}, - coords={"time": times}, - ) - - def test_360day_registers(self, ds_360day): - """A 360-day dataset should register without errors.""" - ctx = XarrayContext() - ctx.from_dataset("ds360", ds_360day, chunks={"time": 6}) - result = ctx.sql("SELECT COUNT(*) AS cnt FROM ds360").to_pandas() - assert result["cnt"].iloc[0] == 12 - - def test_360day_select_ordered(self, ds_360day): - """Integer offsets should be orderable.""" - ctx = XarrayContext() - ctx.from_dataset("ds360", ds_360day, chunks={"time": 6}) - result = ctx.sql( - "SELECT DISTINCT time FROM ds360 ORDER BY time" - ).to_pandas() - times = result["time"].tolist() - assert times == sorted(times) - assert len(times) == 12 - - def test_360day_integer_filter(self, ds_360day): - """Direct integer comparisons should work on int64 time columns.""" - ctx = XarrayContext() - ctx.from_dataset("ds360", ds_360day, chunks={"time": 6}) - # Get all distinct time values to find a midpoint - all_times = ( - ctx.sql("SELECT DISTINCT time FROM ds360 ORDER BY time") - .to_pandas()["time"] - .tolist() - ) - mid = all_times[len(all_times) // 2] - result = ctx.sql( - f"SELECT COUNT(*) AS cnt FROM ds360 WHERE time >= {mid}" - ).to_pandas() - assert 0 < result["cnt"].iloc[0] < 12 - - def test_360day_cftime_udf_registered(self, ds_360day): - """from_dataset should auto-register a cftime() UDF for 360-day calendars.""" - ctx = XarrayContext() - ctx.from_dataset("ds360", ds_360day, chunks={"time": 6}) - # The cftime() UDF should convert a date string to the int64 offset, - # enabling ergonomic filtering. - result = ctx.sql( - "SELECT COUNT(*) AS cnt FROM ds360 " - "WHERE time >= cftime('2000-07-01')" - ).to_pandas() - # July through December = 6 months - assert result["cnt"].iloc[0] == 6 + These use pa.int64() with CF-convention metadata and the cftime() UDF. + """ - def test_gregorian_like_no_cftime_udf(self): - """Gregorian-like calendars should NOT register a cftime() UDF.""" - ds = xr.tutorial.open_dataset("rasm") - ctx = XarrayContext() - ctx.from_dataset("rasm", ds, chunks={"time": 12}) - # Using cftime() should fail since it's not registered for noleap. - with pytest.raises(Exception): - ctx.sql( - "SELECT COUNT(*) FROM rasm WHERE time >= cftime('1980-01-01')" - ).collect() + @pytest.fixture + def ds_360day(self): + """Synthetic 360-day calendar dataset.""" + import cftime + + times = [cftime.Datetime360Day(2000, m, 1) for m in range(1, 13)] + return xr.Dataset( + {"temp": ("time", np.arange(12, dtype=np.float32))}, + coords={"time": times}, + ) + + def test_360day_registers(self, ds_360day): + """A 360-day dataset should register without errors.""" + ctx = XarrayContext() + ctx.from_dataset("ds360", ds_360day, chunks={"time": 6}) + result = ctx.sql("SELECT COUNT(*) AS cnt FROM ds360").to_pandas() + assert result["cnt"].iloc[0] == 12 + + def test_360day_select_ordered(self, ds_360day): + """Integer offsets should be orderable.""" + ctx = XarrayContext() + ctx.from_dataset("ds360", ds_360day, chunks={"time": 6}) + result = ctx.sql( + "SELECT DISTINCT time FROM ds360 ORDER BY time" + ).to_pandas() + times = result["time"].tolist() + assert times == sorted(times) + assert len(times) == 12 + + def test_360day_integer_filter(self, ds_360day): + """Direct integer comparisons should work on int64 time columns.""" + ctx = XarrayContext() + ctx.from_dataset("ds360", ds_360day, chunks={"time": 6}) + # Get all distinct time values to find a midpoint + all_times = ( + ctx.sql("SELECT DISTINCT time FROM ds360 ORDER BY time") + .to_pandas()["time"] + .tolist() + ) + mid = all_times[len(all_times) // 2] + result = ctx.sql( + f"SELECT COUNT(*) AS cnt FROM ds360 WHERE time >= {mid}" + ).to_pandas() + assert 0 < result["cnt"].iloc[0] < 12 + + def test_360day_cftime_udf_registered(self, ds_360day): + """from_dataset should auto-register a cftime() UDF for 360-day calendars.""" + ctx = XarrayContext() + ctx.from_dataset("ds360", ds_360day, chunks={"time": 6}) + # The cftime() UDF should convert a date string to the int64 offset, + # enabling ergonomic filtering. + result = ctx.sql( + "SELECT COUNT(*) AS cnt FROM ds360 " + "WHERE time >= cftime('2000-07-01')" + ).to_pandas() + # July through December = 6 months + assert result["cnt"].iloc[0] == 6 + + def test_gregorian_like_no_cftime_udf(self): + """Gregorian-like calendars should NOT register a cftime() UDF.""" + ds = xr.tutorial.open_dataset("rasm") + ctx = XarrayContext() + ctx.from_dataset("rasm", ds, chunks={"time": 12}) + # Using cftime() should fail since it's not registered for noleap. + with pytest.raises(Exception): + ctx.sql( + "SELECT COUNT(*) FROM rasm WHERE time >= cftime('1980-01-01')" + ).collect() diff --git a/xarray_sql/cftime.py b/xarray_sql/cftime.py index 8c31a3f7..6ce7b4e6 100644 --- a/xarray_sql/cftime.py +++ b/xarray_sql/cftime.py @@ -32,24 +32,24 @@ #: Calendars close enough to proleptic Gregorian for ``pa.timestamp('us')``. GREGORIAN_LIKE_CALENDARS: frozenset[str] = frozenset( { - 'standard', - 'gregorian', - 'proleptic_gregorian', - 'noleap', - '365_day', - 'all_leap', - '366_day', + "standard", + "gregorian", + "proleptic_gregorian", + "noleap", + "365_day", + "all_leap", + "366_day", } ) #: Default CF-convention units when no encoding is available on the coordinate. #: Microseconds give sub-second precision and fit int64 for ±292 k years. -DEFAULT_UNITS: str = 'microseconds since 1970-01-01T00:00:00' +DEFAULT_UNITS: str = "microseconds since 1970-01-01T00:00:00" def is_gregorian_like(calendar: str) -> bool: - """Return True if *calendar* is close enough to Gregorian for ``pa.timestamp``.""" - return calendar in GREGORIAN_LIKE_CALENDARS + """Return True if *calendar* is close enough to Gregorian for ``pa.timestamp``.""" + return calendar in GREGORIAN_LIKE_CALENDARS # --------------------------------------------------------------------------- @@ -58,65 +58,65 @@ def is_gregorian_like(calendar: str) -> bool: def is_cftime(values: np.ndarray) -> bool: - """Check if a numpy array contains cftime datetime objects.""" - try: - import cftime + """Check if a numpy array contains cftime datetime objects.""" + try: + import cftime - if values.dtype == np.dtype('O') and len(values) > 0: - sample = values.ravel()[0] - return isinstance(sample, cftime.datetime) - except ImportError: - pass - return False + if values.dtype == np.dtype("O") and len(values) > 0: + sample = values.ravel()[0] + return isinstance(sample, cftime.datetime) + except ImportError: + pass + return False def is_cftime_index(ds: xr.Dataset, coord_name: str) -> bool: - """Check if a coordinate uses a ``CFTimeIndex`` without materializing data.""" - try: - idx = ds.indexes.get(coord_name) - if idx is not None: - from xarray import CFTimeIndex + """Check if a coordinate uses a ``CFTimeIndex`` without materializing data.""" + try: + idx = ds.indexes.get(coord_name) + if idx is not None: + from xarray import CFTimeIndex - return isinstance(idx, CFTimeIndex) - except (ImportError, AttributeError): - pass - return False + return isinstance(idx, CFTimeIndex) + except (ImportError, AttributeError): + pass + return False def calendar(ds: xr.Dataset, coord_name: str) -> str | None: - """Return the calendar name for a cftime coordinate, or ``None``. - - Checks the xarray index first (no data materialization), then falls - back to inspecting element 0 of the coordinate values. - """ - try: - idx = ds.indexes.get(coord_name) - if idx is not None: - from xarray import CFTimeIndex - - if isinstance(idx, CFTimeIndex): - return str(idx.calendar) # type: ignore[attr-defined] - except (ImportError, AttributeError): - pass - try: - values = ds.coords[coord_name].values - if is_cftime(values): - return str(values.ravel()[0].calendar) - except (AttributeError, KeyError): - pass - return None + """Return the calendar name for a cftime coordinate, or ``None``. + + Checks the xarray index first (no data materialization), then falls + back to inspecting element 0 of the coordinate values. + """ + try: + idx = ds.indexes.get(coord_name) + if idx is not None: + from xarray import CFTimeIndex + + if isinstance(idx, CFTimeIndex): + return str(idx.calendar) # type: ignore[attr-defined] + except (ImportError, AttributeError): + pass + try: + values = ds.coords[coord_name].values + if is_cftime(values): + return str(values.ravel()[0].calendar) + except (AttributeError, KeyError): + pass + return None def encoding(ds: xr.Dataset, coord_name: str) -> tuple[str, str]: - """Return ``(units, calendar)`` for a cftime coordinate. + """Return ``(units, calendar)`` for a cftime coordinate. - Reads xarray ``.encoding`` metadata (from the originating NetCDF file) - first, falling back to :data:`DEFAULT_UNITS`. - """ - cal = calendar(ds, coord_name) or 'standard' - enc = ds.coords[coord_name].encoding - units = enc.get('units', DEFAULT_UNITS) - return units, cal + Reads xarray ``.encoding`` metadata (from the originating NetCDF file) + first, falling back to :data:`DEFAULT_UNITS`. + """ + cal = calendar(ds, coord_name) or "standard" + enc = ds.coords[coord_name].encoding + units = enc.get("units", DEFAULT_UNITS) + return units, cal # --------------------------------------------------------------------------- @@ -125,44 +125,44 @@ def encoding(ds: xr.Dataset, coord_name: str) -> tuple[str, str]: def to_microseconds(values) -> np.ndarray: - """Convert cftime objects to int64 microseconds since Unix epoch. + """Convert cftime objects to int64 microseconds since Unix epoch. - Used for Gregorian-like calendars. Vectorised via ``cftime.date2num`` - (implemented in C). - """ - import cftime as _cftime + Used for Gregorian-like calendars. Vectorised via ``cftime.date2num`` + (implemented in C). + """ + import cftime as _cftime - us = _cftime.date2num( - values.ravel(), - units=DEFAULT_UNITS, - calendar=values.ravel()[0].calendar, - ) - return np.asarray(us, dtype=np.float64).astype(np.int64) + us = _cftime.date2num( + values.ravel(), + units=DEFAULT_UNITS, + calendar=values.ravel()[0].calendar, + ) + return np.asarray(us, dtype=np.float64).astype(np.int64) def to_offsets(values, units: str, cal: str) -> np.ndarray: - """Convert cftime objects to int64 offsets in the given *units*/*calendar*. + """Convert cftime objects to int64 offsets in the given *units*/*calendar*. - Used for non-Gregorian calendars where data is stored as ``pa.int64()``. - """ - import cftime as _cftime + Used for non-Gregorian calendars where data is stored as ``pa.int64()``. + """ + import cftime as _cftime - raw = _cftime.date2num(values.ravel(), units=units, calendar=cal) - return np.asarray(raw, dtype=np.float64).astype(np.int64) + raw = _cftime.date2num(values.ravel(), units=units, calendar=cal) + return np.asarray(raw, dtype=np.float64).astype(np.int64) def convert_for_field(values, field: pa.Field) -> np.ndarray: - """Convert cftime values to the numeric type dictated by *field*. + """Convert cftime values to the numeric type dictated by *field*. - Reads ``xarray:calendar`` and ``xarray:units`` from the field's Arrow - metadata to choose between the timestamp path and the integer-offset path. - """ - meta = field.metadata or {} - cal = meta.get(b'xarray:calendar', b'standard').decode() - units = meta.get(b'xarray:units', DEFAULT_UNITS.encode()).decode() - if is_gregorian_like(cal): - return to_microseconds(values) - return to_offsets(values, units, cal) + Reads ``xarray:calendar`` and ``xarray:units`` from the field's Arrow + metadata to choose between the timestamp path and the integer-offset path. + """ + meta = field.metadata or {} + cal = meta.get(b"xarray:calendar", b"standard").decode() + units = meta.get(b"xarray:units", DEFAULT_UNITS.encode()).decode() + if is_gregorian_like(cal): + return to_microseconds(values) + return to_offsets(values, units, cal) # --------------------------------------------------------------------------- @@ -173,19 +173,19 @@ def convert_for_field(values, field: pa.Field) -> np.ndarray: def partition_bounds( values, ) -> tuple[int, int, str]: - """Return ``(min, max, dtype_tag)`` for a cftime coordinate slice. + """Return ``(min, max, dtype_tag)`` for a cftime coordinate slice. - Gregorian-like calendars return nanosecond bounds tagged - ``"timestamp_ns"`` (compatible with ``ScalarBound::TimestampNanos`` - in the Rust pruning layer). Non-Gregorian calendars return int64 - offsets tagged ``"int64"``. - """ - cal = values.ravel()[0].calendar - if is_gregorian_like(cal): - us = to_microseconds(values) - return int(us.min()) * 1_000, int(us.max()) * 1_000, 'timestamp_ns' - offsets = to_offsets(values, DEFAULT_UNITS, cal) - return int(offsets.min()), int(offsets.max()), 'int64' + Gregorian-like calendars return nanosecond bounds tagged + ``"timestamp_ns"`` (compatible with ``ScalarBound::TimestampNanos`` + in the Rust pruning layer). Non-Gregorian calendars return int64 + offsets tagged ``"int64"``. + """ + cal = values.ravel()[0].calendar + if is_gregorian_like(cal): + us = to_microseconds(values) + return int(us.min()) * 1_000, int(us.max()) * 1_000, "timestamp_ns" + offsets = to_offsets(values, DEFAULT_UNITS, cal) + return int(offsets.min()), int(offsets.max()), "int64" # --------------------------------------------------------------------------- @@ -194,19 +194,19 @@ def partition_bounds( def arrow_field(name: str, units: str, cal: str) -> pa.Field: - """Build a ``pa.Field`` for a cftime coordinate. - - Gregorian-like → ``pa.timestamp('us')``; non-Gregorian → ``pa.int64()``. - Both carry ``xarray:calendar`` and ``xarray:units`` metadata for - round-trip fidelity. - """ - meta = { - b'xarray:calendar': cal.encode(), - b'xarray:units': units.encode(), - } - if is_gregorian_like(cal): - return pa.field(name, pa.timestamp('us'), metadata=meta) - return pa.field(name, pa.int64(), metadata=meta) + """Build a ``pa.Field`` for a cftime coordinate. + + Gregorian-like → ``pa.timestamp('us')``; non-Gregorian → ``pa.int64()``. + Both carry ``xarray:calendar`` and ``xarray:units`` metadata for + round-trip fidelity. + """ + meta = { + b"xarray:calendar": cal.encode(), + b"xarray:units": units.encode(), + } + if is_gregorian_like(cal): + return pa.field(name, pa.timestamp("us"), metadata=meta) + return pa.field(name, pa.int64(), metadata=meta) # --------------------------------------------------------------------------- @@ -215,34 +215,34 @@ def arrow_field(name: str, units: str, cal: str) -> pa.Field: def make_cftime_udf(units: str, calendar: str): - """Create a DataFusion scalar UDF that converts date strings to int64 offsets. - - This enables ergonomic SQL filtering on non-Gregorian cftime columns:: - - SELECT * FROM ds360 WHERE time > cftime('0500-01-01') - - The UDF parses the input string as a cftime datetime in the given - calendar system and returns the corresponding int64 offset in the - specified units. - """ - import cftime as _cftime - from datafusion import udf - - def _cftime_scalar(date_strings: pa.Array) -> pa.Array: - results: list[int | None] = [] - for s in date_strings.to_pylist(): - if s is None: - results.append(None) - continue - dt = _cftime.datetime.strptime(s, '%Y-%m-%d', calendar=calendar) - val = _cftime.date2num(dt, units=units, calendar=calendar) - results.append(int(val)) - return pa.array(results, type=pa.int64()) - - return udf( - _cftime_scalar, - [pa.utf8()], - pa.int64(), - 'immutable', - 'cftime', - ) + """Create a DataFusion scalar UDF that converts date strings to int64 offsets. + + This enables ergonomic SQL filtering on non-Gregorian cftime columns:: + + SELECT * FROM ds360 WHERE time > cftime('0500-01-01') + + The UDF parses the input string as a cftime datetime in the given + calendar system and returns the corresponding int64 offset in the + specified units. + """ + import cftime as _cftime + from datafusion import udf + + def _cftime_scalar(date_strings: pa.Array) -> pa.Array: + results: list[int | None] = [] + for s in date_strings.to_pylist(): + if s is None: + results.append(None) + continue + dt = _cftime.datetime.strptime(s, "%Y-%m-%d", calendar=calendar) + val = _cftime.date2num(dt, units=units, calendar=calendar) + results.append(int(val)) + return pa.array(results, type=pa.int64()) + + return udf( + _cftime_scalar, + [pa.utf8()], + pa.int64(), + "immutable", + "cftime", + ) diff --git a/xarray_sql/core.py b/xarray_sql/core.py index 11b909b7..6ff06b6a 100644 --- a/xarray_sql/core.py +++ b/xarray_sql/core.py @@ -10,40 +10,40 @@ # deprecated def get_columns(ds: xr.Dataset) -> list[str]: - return list(ds.sizes.keys()) + list(ds.data_vars.keys()) + return list(ds.sizes.keys()) + list(ds.data_vars.keys()) # Deprecated def unravel(ds: xr.Dataset) -> Iterator[Row]: - dim_keys, dim_vals = zip(*ds.sizes.items()) + dim_keys, dim_vals = zip(*ds.sizes.items()) - for idx in itertools.product(*(range(d) for d in dim_vals)): - coord_idx = dict(zip(dim_keys, idx)) - data = ds.isel(coord_idx) - coord_data = [ds.coords[v][coord_idx[v]] for v in dim_keys] - row = [v.values for v in coord_data + list(data.data_vars.values())] - yield row + for idx in itertools.product(*(range(d) for d in dim_vals)): + coord_idx = dict(zip(dim_keys, idx)) + data = ds.isel(coord_idx) + coord_data = [ds.coords[v][coord_idx[v]] for v in dim_keys] + row = [v.values for v in coord_data + list(data.data_vars.values())] + yield row # Deprecated def unbounded_unravel(ds: xr.Dataset) -> np.ndarray: - """Unravel with unbounded memory (as a NumPy Array).""" - dim_keys, dim_vals = zip(*ds.sizes.items()) - columns = get_columns(ds) + """Unravel with unbounded memory (as a NumPy Array).""" + dim_keys, dim_vals = zip(*ds.sizes.items()) + columns = get_columns(ds) - N = np.prod([d for d in dim_vals]) + N = np.prod([d for d in dim_vals]) - out = np.recarray((N,), dtype=[(c, ds[c].dtype) for c in columns]) + out = np.recarray((N,), dtype=[(c, ds[c].dtype) for c in columns]) - for name, da in ds.items(): - out[name] = da.values.ravel() + for name, da in ds.items(): + out[name] = da.values.ravel() - prod_vals = (ds.coords[k].values for k in dim_keys) - coords = np.array(np.meshgrid(*prod_vals), dtype=int).T.reshape( - -1, len(dim_keys) - ) + prod_vals = (ds.coords[k].values for k in dim_keys) + coords = np.array(np.meshgrid(*prod_vals), dtype=int).T.reshape( + -1, len(dim_keys) + ) - for i, d in enumerate(dim_keys): - out[d] = coords[:, i] + for i, d in enumerate(dim_keys): + out[d] = coords[:, i] - return out + return out diff --git a/xarray_sql/df.py b/xarray_sql/df.py index 0429fdce..db9d37e6 100644 --- a/xarray_sql/df.py +++ b/xarray_sql/df.py @@ -18,51 +18,51 @@ def _get_chunk_slicer( dim: Hashable, chunk_index: Mapping, chunk_bounds: Mapping ): - if dim in chunk_index: - which_chunk = chunk_index[dim] - return slice( - chunk_bounds[dim][which_chunk], chunk_bounds[dim][which_chunk + 1] - ) - return slice(None) + if dim in chunk_index: + which_chunk = chunk_index[dim] + return slice( + chunk_bounds[dim][which_chunk], chunk_bounds[dim][which_chunk + 1] + ) + return slice(None) # Adapted from Xarray `map_blocks` implementation. def block_slices(ds: xr.Dataset, chunks: Chunks = None) -> Iterator[Block]: - """Compute block slices for a chunked Dataset.""" - if chunks is not None: - for_chunking = ds.copy(data=None, deep=False).chunk(chunks) - chunks = for_chunking.chunks - del for_chunking - else: - chunks = ds.chunks - - assert chunks, "Dataset `ds` must be chunked or `chunks` must be provided." - - # chunks is Dict[str, Tuple[int, ...]] from xarray - chunk_bounds = { - dim: np.cumsum((0,) + tuple(c)) # type: ignore[arg-type] - for dim, c in chunks.items() - } - ichunk = {dim: range(len(tuple(c))) for dim, c in chunks.items()} # type: ignore[arg-type] - ick, icv = zip(*ichunk.items()) # Makes same order of keys and val. - chunk_idxs = (dict(zip(ick, i)) for i in itertools.product(*icv)) - blocks = ( - { - dim: _get_chunk_slicer(dim, chunk_index, chunk_bounds) - for dim in ds.dims - } - for chunk_index in chunk_idxs - ) - yield from blocks + """Compute block slices for a chunked Dataset.""" + if chunks is not None: + for_chunking = ds.copy(data=None, deep=False).chunk(chunks) + chunks = for_chunking.chunks + del for_chunking + else: + chunks = ds.chunks + + assert chunks, "Dataset `ds` must be chunked or `chunks` must be provided." + + # chunks is Dict[str, Tuple[int, ...]] from xarray + chunk_bounds = { + dim: np.cumsum((0,) + tuple(c)) # type: ignore[arg-type] + for dim, c in chunks.items() + } + ichunk = {dim: range(len(tuple(c))) for dim, c in chunks.items()} # type: ignore[arg-type] + ick, icv = zip(*ichunk.items()) # Makes same order of keys and val. + chunk_idxs = (dict(zip(ick, i)) for i in itertools.product(*icv)) + blocks = ( + { + dim: _get_chunk_slicer(dim, chunk_index, chunk_bounds) + for dim in ds.dims + } + for chunk_index in chunk_idxs + ) + yield from blocks def explode(ds: xr.Dataset, chunks: Chunks = None) -> Iterator[xr.Dataset]: - """Explodes a dataset into its chunks.""" - yield from (ds.isel(b) for b in block_slices(ds, chunks=chunks)) + """Explodes a dataset into its chunks.""" + yield from (ds.isel(b) for b in block_slices(ds, chunks=chunks)) def _block_len(block: Block) -> int: - return int(np.prod([v.stop - v.start for v in block.values()])) + return int(np.prod([v.stop - v.start for v in block.values()])) def from_map_batched( @@ -72,32 +72,32 @@ def from_map_batched( schema: pa.Schema = None, **kwargs: dict[str, Any], ) -> pa.RecordBatchReader: - """Create a PyArrow RecordBatchReader by mapping a function over iterables. + """Create a PyArrow RecordBatchReader by mapping a function over iterables. - This is equivalent to dask's from_map but returns a PyArrow - RecordBatchReader that can be used with DataFusion. It iterates over - RecordBatches which are created via the `func` one-at-a-time. + This is equivalent to dask's from_map but returns a PyArrow + RecordBatchReader that can be used with DataFusion. It iterates over + RecordBatches which are created via the `func` one-at-a-time. - Args: - func: Function to apply to each element of the iterables. Currently, the - function must return a Pandas DataFrame. - *iterables: Iterable objects to map the function over. - schema: Optional schema needed for the RecordBatchReader. - args: Additional positional arguments to pass to func. - **kwargs: Additional keyword arguments to pass to func. + Args: + func: Function to apply to each element of the iterables. Currently, the + function must return a Pandas DataFrame. + *iterables: Iterable objects to map the function over. + schema: Optional schema needed for the RecordBatchReader. + args: Additional positional arguments to pass to func. + **kwargs: Additional keyword arguments to pass to func. - Returns: - A PyArrow RecordBatchReader containing the stream of RecordBatches. - """ - if args is None: - args = () + Returns: + A PyArrow RecordBatchReader containing the stream of RecordBatches. + """ + if args is None: + args = () - def map_batches(): - for items in zip(*iterables): - df = func(*items, *args, **kwargs) - yield pa.RecordBatch.from_pandas(df, schema=schema) + def map_batches(): + for items in zip(*iterables): + df = func(*items, *args, **kwargs) + yield pa.RecordBatch.from_pandas(df, schema=schema) - return pa.RecordBatchReader.from_batches(schema, map_batches()) + return pa.RecordBatchReader.from_batches(schema, map_batches()) def from_map( @@ -106,121 +106,121 @@ def from_map( args: tuple | None = None, **kwargs: dict[str, Any], ) -> pa.Table: - """Create a PyArrow Table by mapping a function over iterables. - - This is equivalent to dask's from_map but returns a PyArrow Table - that can be used with DataFusion instead of a Dask DataFrame. - - Args: - func: Function to apply to each element of the iterables. - *iterables: Iterable objects to map the function over. - args: Additional positional arguments to pass to func. - **kwargs: Additional keyword arguments to pass to func. - - Returns: - A PyArrow Table containing the concatenated results. - """ - if args is None: - args = () - - # Apply the function to each combination of iterable elements - results = [] - for items in zip(*iterables) if len(iterables) > 1 else iterables[0]: - if isinstance(items, tuple): - result = func(*items, *args, **kwargs) - else: - result = func(items, *args, **kwargs) - - # Convert result to PyArrow Table - if isinstance(result, pd.DataFrame): - pa_table = pa.Table.from_pandas(result) - elif isinstance(result, pa.Table): - pa_table = result - else: - # Try to convert to pandas first, then to PyArrow - try: - df = pd.DataFrame(result) - pa_table = pa.Table.from_pandas(df) - except Exception as e: - raise ValueError( - f"Cannot convert function result to PyArrow Table: {e}" - ) - - results.append(pa_table) - - # Concatenate all results - if not results: - raise ValueError("No results to concatenate") - - return pa.concat_tables(results) + """Create a PyArrow Table by mapping a function over iterables. + + This is equivalent to dask's from_map but returns a PyArrow Table + that can be used with DataFusion instead of a Dask DataFrame. + + Args: + func: Function to apply to each element of the iterables. + *iterables: Iterable objects to map the function over. + args: Additional positional arguments to pass to func. + **kwargs: Additional keyword arguments to pass to func. + + Returns: + A PyArrow Table containing the concatenated results. + """ + if args is None: + args = () + + # Apply the function to each combination of iterable elements + results = [] + for items in zip(*iterables) if len(iterables) > 1 else iterables[0]: + if isinstance(items, tuple): + result = func(*items, *args, **kwargs) + else: + result = func(items, *args, **kwargs) + + # Convert result to PyArrow Table + if isinstance(result, pd.DataFrame): + pa_table = pa.Table.from_pandas(result) + elif isinstance(result, pa.Table): + pa_table = result + else: + # Try to convert to pandas first, then to PyArrow + try: + df = pd.DataFrame(result) + pa_table = pa.Table.from_pandas(df) + except Exception as e: + raise ValueError( + f"Cannot convert function result to PyArrow Table: {e}" + ) + + results.append(pa_table) + + # Concatenate all results + if not results: + raise ValueError("No results to concatenate") + + return pa.concat_tables(results) def pivot(ds: xr.Dataset) -> pd.DataFrame: - """Converts an xarray Dataset to a pandas DataFrame.""" - return ds.to_dataframe().reset_index() # type: ignore[no-any-return] + """Converts an xarray Dataset to a pandas DataFrame.""" + return ds.to_dataframe().reset_index() # type: ignore[no-any-return] def dataset_to_record_batch( ds: xr.Dataset, schema: pa.Schema ) -> pa.RecordBatch: - """Convert an xarray Dataset partition to an Arrow RecordBatch. - - Builds the RecordBatch directly from numpy arrays, bypassing the pandas - round-trip (to_dataframe → reset_index → from_pandas) used by pivot(). - For large partitions this reduces peak memory from ~5× to ~2× the - partition size. - - Dimension coordinates are broadcast to the full partition shape and - ravelled. np.broadcast_to() is zero-copy; the ravel() forces one copy - per coordinate (unavoidable, since broadcast arrays are non-contiguous). - Data variable arrays are ravelled in-place — a zero-copy view when the - underlying array is already C-contiguous (the common case for numpy-backed - xarray datasets). - - Args: - ds: A partition-sized xarray Dataset (already sliced via isel). - schema: The Arrow schema for the output, as produced by _parse_schema. - Column order in the output matches schema field order. - - Returns: - A RecordBatch with one column per dimension coordinate and data - variable, in schema order. - """ - # Use the data variable's dimension order as canonical so coordinate - # broadcasts and data variable ravels use the same layout. All data - # variables are validated to share the same dims tuple. - if ds.data_vars: - first_var = next(iter(ds.data_vars.values())) - dim_names = list(first_var.dims) - shape = first_var.shape - else: - dim_names = list(ds.sizes.keys()) - shape = tuple(ds.sizes[d] for d in dim_names) - - arrays = [] - for field in schema: - name = field.name - if name in ds.coords and name in ds.dims: - # Broadcast 1-D coordinate to the full N-D partition shape, then ravel. - axis = dim_names.index(name) - coord = ds.coords[name].values - if cft.is_cftime(coord): - coord = cft.convert_for_field(coord, field) - reshape = [1] * len(shape) - reshape[axis] = coord.shape[0] - arr = np.broadcast_to(coord.reshape(reshape), shape).ravel() - arrays.append(pa.array(arr, type=field.type)) + """Convert an xarray Dataset partition to an Arrow RecordBatch. + + Builds the RecordBatch directly from numpy arrays, bypassing the pandas + round-trip (to_dataframe → reset_index → from_pandas) used by pivot(). + For large partitions this reduces peak memory from ~5× to ~2× the + partition size. + + Dimension coordinates are broadcast to the full partition shape and + ravelled. np.broadcast_to() is zero-copy; the ravel() forces one copy + per coordinate (unavoidable, since broadcast arrays are non-contiguous). + Data variable arrays are ravelled in-place — a zero-copy view when the + underlying array is already C-contiguous (the common case for numpy-backed + xarray datasets). + + Args: + ds: A partition-sized xarray Dataset (already sliced via isel). + schema: The Arrow schema for the output, as produced by _parse_schema. + Column order in the output matches schema field order. + + Returns: + A RecordBatch with one column per dimension coordinate and data + variable, in schema order. + """ + # Use the data variable's dimension order as canonical so coordinate + # broadcasts and data variable ravels use the same layout. All data + # variables are validated to share the same dims tuple. + if ds.data_vars: + first_var = next(iter(ds.data_vars.values())) + dim_names = list(first_var.dims) + shape = first_var.shape else: - # Data variable: ravel to 1-D (zero-copy for C-contiguous arrays). - raw = ds[name].values.ravel() - if cft.is_cftime(ds[name].values): - raw = cft.convert_for_field(ds[name].values, field) + dim_names = list(ds.sizes.keys()) + shape = tuple(ds.sizes[d] for d in dim_names) - # from_pandas=True maps NaN → Arrow null inside the C++ copy kernel, - # so SQL aggregates (MAX, MIN, AVG) skip missing values correctly. - arrays.append(pa.array(raw, type=field.type, from_pandas=True)) - - return pa.RecordBatch.from_arrays(arrays, schema=schema) + arrays = [] + for field in schema: + name = field.name + if name in ds.coords and name in ds.dims: + # Broadcast 1-D coordinate to the full N-D partition shape, then ravel. + axis = dim_names.index(name) + coord = ds.coords[name].values + if cft.is_cftime(coord): + coord = cft.convert_for_field(coord, field) + reshape = [1] * len(shape) + reshape[axis] = coord.shape[0] + arr = np.broadcast_to(coord.reshape(reshape), shape).ravel() + arrays.append(pa.array(arr, type=field.type)) + else: + # Data variable: ravel to 1-D (zero-copy for C-contiguous arrays). + raw = ds[name].values.ravel() + if cft.is_cftime(ds[name].values): + raw = cft.convert_for_field(ds[name].values, field) + + # from_pandas=True maps NaN → Arrow null inside the C++ copy kernel, + # so SQL aggregates (MAX, MIN, AVG) skip missing values correctly. + arrays.append(pa.array(raw, type=field.type, from_pandas=True)) + + return pa.RecordBatch.from_arrays(arrays, schema=schema) #: Default number of rows per emitted Arrow RecordBatch. @@ -233,128 +233,130 @@ def iter_record_batches( schema: pa.Schema, batch_size: int = DEFAULT_BATCH_SIZE, ) -> Iterator[pa.RecordBatch]: - """Yield RecordBatches of at most *batch_size* rows from a partition Dataset. - - Unlike `dataset_to_record_batch`, which materialises the entire - partition as one batch, this generator emits smaller batches so that - DataFusion can begin filtering and aggregating before the full partition - is loaded. Peak memory per batch is O(batch_size) for coordinate columns - and O(partition_size) for data-variable columns (which must be loaded in - full from storage). - - Coordinate values are computed per batch via strided index arithmetic — - no broadcast array spanning the whole partition is ever allocated. Data - variable flat arrays are loaded once (triggering any remote I/O) and then - sliced as zero-copy views for each batch. - - Args: - ds: A partition-sized xarray Dataset (already sliced via isel). - schema: The Arrow schema for the output, as produced by _parse_schema. - batch_size: Maximum number of rows per yielded RecordBatch. - - Yields: - RecordBatches in schema column order, covering all rows of the - partition exactly once. - """ - if ds.data_vars: - first_var = next(iter(ds.data_vars.values())) - dim_names = list(first_var.dims) - shape = first_var.shape - else: - dim_names = list(ds.sizes.keys()) - shape = tuple(ds.sizes[d] for d in dim_names) - - total_rows = int(np.prod(shape)) - - # Preload small 1-D coordinate arrays (negligible memory). - # Convert cftime objects to numeric values matching the schema type. - coord_values = {} - for name in dim_names: - vals = ds.coords[name].values - if cft.is_cftime(vals): - coord_values[name] = cft.convert_for_field(vals, schema.field(name)) + """Yield RecordBatches of at most *batch_size* rows from a partition Dataset. + + Unlike `dataset_to_record_batch`, which materialises the entire + partition as one batch, this generator emits smaller batches so that + DataFusion can begin filtering and aggregating before the full partition + is loaded. Peak memory per batch is O(batch_size) for coordinate columns + and O(partition_size) for data-variable columns (which must be loaded in + full from storage). + + Coordinate values are computed per batch via strided index arithmetic — + no broadcast array spanning the whole partition is ever allocated. Data + variable flat arrays are loaded once (triggering any remote I/O) and then + sliced as zero-copy views for each batch. + + Args: + ds: A partition-sized xarray Dataset (already sliced via isel). + schema: The Arrow schema for the output, as produced by _parse_schema. + batch_size: Maximum number of rows per yielded RecordBatch. + + Yields: + RecordBatches in schema column order, covering all rows of the + partition exactly once. + """ + if ds.data_vars: + first_var = next(iter(ds.data_vars.values())) + dim_names = list(first_var.dims) + shape = first_var.shape else: - coord_values[name] = vals - - # C-order stride for each dimension: stride[k] = prod(shape[k+1:]). - # Flat row index i → coordinate index for dim k: (i // stride[k]) % shape[k]. - strides = [int(np.prod(shape[k + 1 :])) for k in range(len(shape))] - - # Load data-variable arrays fully (triggers Dask/Zarr compute once). - # ravel() is a zero-copy view for C-contiguous arrays. - data_arrays = {} - for field in schema: - if field.name not in ds.dims: - raw = ds[field.name].values - if cft.is_cftime(raw): - data_arrays[field.name] = cft.convert_for_field(raw, field) - else: - data_arrays[field.name] = raw.ravel() - - for row_start in range(0, total_rows, batch_size): - row_end = min(row_start + batch_size, total_rows) - row_idx = np.arange(row_start, row_end) - - arrays = [] + dim_names = list(ds.sizes.keys()) + shape = tuple(ds.sizes[d] for d in dim_names) + + total_rows = int(np.prod(shape)) + + # Preload small 1-D coordinate arrays (negligible memory). + # Convert cftime objects to numeric values matching the schema type. + coord_values = {} + for name in dim_names: + vals = ds.coords[name].values + if cft.is_cftime(vals): + coord_values[name] = cft.convert_for_field(vals, schema.field(name)) + else: + coord_values[name] = vals + + # C-order stride for each dimension: stride[k] = prod(shape[k+1:]). + # Flat row index i → coordinate index for dim k: (i // stride[k]) % shape[k]. + strides = [int(np.prod(shape[k + 1 :])) for k in range(len(shape))] + + # Load data-variable arrays fully (triggers Dask/Zarr compute once). + # ravel() is a zero-copy view for C-contiguous arrays. + data_arrays = {} for field in schema: - name = field.name - if name in ds.coords and name in ds.dims: - k = dim_names.index(name) - coord_idx = (row_idx // strides[k]) % shape[k] - arrays.append(pa.array(coord_values[name][coord_idx], type=field.type)) - else: - arrays.append( - pa.array( - data_arrays[name][row_start:row_end], - type=field.type, - from_pandas=True, - ) - ) - - yield pa.RecordBatch.from_arrays(arrays, schema=schema) + if field.name not in ds.dims: + raw = ds[field.name].values + if cft.is_cftime(raw): + data_arrays[field.name] = cft.convert_for_field(raw, field) + else: + data_arrays[field.name] = raw.ravel() + + for row_start in range(0, total_rows, batch_size): + row_end = min(row_start + batch_size, total_rows) + row_idx = np.arange(row_start, row_end) + + arrays = [] + for field in schema: + name = field.name + if name in ds.coords and name in ds.dims: + k = dim_names.index(name) + coord_idx = (row_idx // strides[k]) % shape[k] + arrays.append( + pa.array(coord_values[name][coord_idx], type=field.type) + ) + else: + arrays.append( + pa.array( + data_arrays[name][row_start:row_end], + type=field.type, + from_pandas=True, + ) + ) + + yield pa.RecordBatch.from_arrays(arrays, schema=schema) def _parse_schema(ds: xr.Dataset) -> pa.Schema: - """Extracts a `pa.Schema` from the Dataset, treating dims and data_vars as columns. - - Uses the xarray index type to detect cftime coordinates without - materializing their data — important for Dask/Zarr-backed datasets - where .values would trigger eager computation. - - cftime coordinates are mapped to one of two Arrow types: - - * **Gregorian-like calendars** (standard, noleap, all_leap, etc.): - ``pa.timestamp('us')`` so string-based SQL filters work naturally. - * **Non-Gregorian calendars** (360_day, julian): - ``pa.int64()`` with ``xarray:units`` / ``xarray:calendar`` metadata - on the field, preserving lossless CF-convention encoding. - """ - columns = [] - - for coord_name, coord_var in ds.coords.items(): - # Only include dimension coordinates - if coord_name in ds.dims: - if cft.is_cftime_index(ds, coord_name): - units, calendar = cft.encoding(ds, coord_name) - columns.append(cft.arrow_field(coord_name, units, calendar)) - else: - pa_type = pa.from_numpy_dtype(coord_var.dtype) - columns.append(pa.field(coord_name, pa_type)) - - for var_name, var in ds.data_vars.items(): - # Data variables are virtually never cftime, but check dtype as a - # cheap guard. Only fall back to _is_cftime (which materializes - # element 0) when dtype is object. - if var.dtype == np.dtype("O") and cft.is_cftime(var.values): - # Rare: a data variable holding cftime objects. Use same encoding - # as the first cftime dimension coordinate, or default. - cal = var.values.ravel()[0].calendar - columns.append(cft.arrow_field(var_name, cft.DEFAULT_UNITS, cal)) - else: - pa_type = pa.from_numpy_dtype(var.dtype) - columns.append(pa.field(var_name, pa_type)) - - return pa.schema(columns) + """Extracts a `pa.Schema` from the Dataset, treating dims and data_vars as columns. + + Uses the xarray index type to detect cftime coordinates without + materializing their data — important for Dask/Zarr-backed datasets + where .values would trigger eager computation. + + cftime coordinates are mapped to one of two Arrow types: + + * **Gregorian-like calendars** (standard, noleap, all_leap, etc.): + ``pa.timestamp('us')`` so string-based SQL filters work naturally. + * **Non-Gregorian calendars** (360_day, julian): + ``pa.int64()`` with ``xarray:units`` / ``xarray:calendar`` metadata + on the field, preserving lossless CF-convention encoding. + """ + columns = [] + + for coord_name, coord_var in ds.coords.items(): + # Only include dimension coordinates + if coord_name in ds.dims: + if cft.is_cftime_index(ds, coord_name): + units, calendar = cft.encoding(ds, coord_name) + columns.append(cft.arrow_field(coord_name, units, calendar)) + else: + pa_type = pa.from_numpy_dtype(coord_var.dtype) + columns.append(pa.field(coord_name, pa_type)) + + for var_name, var in ds.data_vars.items(): + # Data variables are virtually never cftime, but check dtype as a + # cheap guard. Only fall back to _is_cftime (which materializes + # element 0) when dtype is object. + if var.dtype == np.dtype("O") and cft.is_cftime(var.values): + # Rare: a data variable holding cftime objects. Use same encoding + # as the first cftime dimension coordinate, or default. + cal = var.values.ravel()[0].calendar + columns.append(cft.arrow_field(var_name, cft.DEFAULT_UNITS, cal)) + else: + pa_type = pa.from_numpy_dtype(var.dtype) + columns.append(pa.field(var_name, pa_type)) + + return pa.schema(columns) # Type alias for partition metadata: maps dimension name to (min, max, dtype_str) values @@ -362,84 +364,84 @@ def _parse_schema(ds: xr.Dataset) -> pa.Schema: def _block_metadata(coord_arrays: dict, block: Block) -> PartitionBounds: - """Compute min/max coordinate values for a single partition block. - - Args: - coord_arrays: Pre-materialised coordinate arrays keyed by dimension name - string. Hoist this outside any loop to avoid repeated remote I/O - for Zarr-backed datasets. - block: A single block slice dict from block_slices(). - - Returns: - Dict mapping dimension name to (min_value, max_value, dtype_str). - Dimensions with an empty slice are omitted; the Rust pruning logic - treats missing dimensions conservatively (never prunes on them). - """ - ranges: PartitionBounds = {} - for dim, slc in block.items(): - coord_values = coord_arrays[str(dim)][slc] - if len(coord_values) == 0: - continue - # String/object dtypes are not representable as ScalarBound - # (Int64/Float64/TimestampNanos) and numpy min/max ufuncs do not - # support them. Skip so pruning treats the dimension conservatively. - if coord_values.dtype.kind in ("U", "S", "O"): - continue - - if cft.is_cftime(coord_values): - ranges[str(dim)] = cft.partition_bounds(coord_values) - continue - - # Use actual min/max rather than first/last so that non-monotonic - # coordinate axes (e.g. descending latitude 90→-90) are handled - # correctly. np.min/max work for both numeric and datetime64 arrays. - min_val = coord_values.min() - max_val = coord_values.max() - - if isinstance(min_val, (np.datetime64, pd.Timestamp)): - min_val = int(pd.Timestamp(min_val).value) - max_val = int(pd.Timestamp(max_val).value) - ranges[str(dim)] = (min_val, max_val, "timestamp_ns") - elif hasattr(min_val, "item"): - min_val = min_val.item() - max_val = max_val.item() - dtype = "float64" if isinstance(min_val, float) else "int64" - ranges[str(dim)] = (min_val, max_val, dtype) - else: - dtype = "float64" if isinstance(min_val, float) else "int64" - ranges[str(dim)] = (min_val, max_val, dtype) - return ranges + """Compute min/max coordinate values for a single partition block. + + Args: + coord_arrays: Pre-materialised coordinate arrays keyed by dimension name + string. Hoist this outside any loop to avoid repeated remote I/O + for Zarr-backed datasets. + block: A single block slice dict from block_slices(). + + Returns: + Dict mapping dimension name to (min_value, max_value, dtype_str). + Dimensions with an empty slice are omitted; the Rust pruning logic + treats missing dimensions conservatively (never prunes on them). + """ + ranges: PartitionBounds = {} + for dim, slc in block.items(): + coord_values = coord_arrays[str(dim)][slc] + if len(coord_values) == 0: + continue + # String/object dtypes are not representable as ScalarBound + # (Int64/Float64/TimestampNanos) and numpy min/max ufuncs do not + # support them. Skip so pruning treats the dimension conservatively. + if coord_values.dtype.kind in ("U", "S", "O"): + continue + + if cft.is_cftime(coord_values): + ranges[str(dim)] = cft.partition_bounds(coord_values) + continue + + # Use actual min/max rather than first/last so that non-monotonic + # coordinate axes (e.g. descending latitude 90→-90) are handled + # correctly. np.min/max work for both numeric and datetime64 arrays. + min_val = coord_values.min() + max_val = coord_values.max() + + if isinstance(min_val, (np.datetime64, pd.Timestamp)): + min_val = int(pd.Timestamp(min_val).value) + max_val = int(pd.Timestamp(max_val).value) + ranges[str(dim)] = (min_val, max_val, "timestamp_ns") + elif hasattr(min_val, "item"): + min_val = min_val.item() + max_val = max_val.item() + dtype = "float64" if isinstance(min_val, float) else "int64" + ranges[str(dim)] = (min_val, max_val, dtype) + else: + dtype = "float64" if isinstance(min_val, float) else "int64" + ranges[str(dim)] = (min_val, max_val, dtype) + return ranges def partition_metadata( ds: xr.Dataset, blocks: list[Block] ) -> list[PartitionBounds]: - """Compute min/max coordinate values for each partition. - - This metadata enables filter pushdown: SQL queries with WHERE clauses - on dimension columns can prune partitions that can't contain matching rows. - - Args: - ds: The xarray Dataset containing coordinate values. - blocks: List of block slices from block_slices(). - - Returns: - List of dicts mapping dimension name to - (min_value, max_value, dtype_str) tuples. - - - For datetime64, values are nanoseconds since Unix epoch - (int64), dtype_str is "timestamp_ns" - - For numeric types, values are Python int or float, - dtype_str is "int64" or "float64" - - Note: - If a partition has an empty slice for a dimension, that dimension is - omitted from the partition's metadata. The Rust pruning logic treats - missing dimensions conservatively (never prunes on them). - """ - # Hoist coordinate array reads outside the partition loop. - # ds.coords[dim].values materializes the full array on every call; doing it - # N_partitions × N_dims times is wasteful and, for remote Zarr-backed datasets - # (e.g. ARCO-ERA5 on GCS), may trigger repeated network I/O. - coord_arrays = {str(dim): ds.coords[dim].values for dim in ds.dims} - return [_block_metadata(coord_arrays, block) for block in blocks] + """Compute min/max coordinate values for each partition. + + This metadata enables filter pushdown: SQL queries with WHERE clauses + on dimension columns can prune partitions that can't contain matching rows. + + Args: + ds: The xarray Dataset containing coordinate values. + blocks: List of block slices from block_slices(). + + Returns: + List of dicts mapping dimension name to + (min_value, max_value, dtype_str) tuples. + + - For datetime64, values are nanoseconds since Unix epoch + (int64), dtype_str is "timestamp_ns" + - For numeric types, values are Python int or float, + dtype_str is "int64" or "float64" + + Note: + If a partition has an empty slice for a dimension, that dimension is + omitted from the partition's metadata. The Rust pruning logic treats + missing dimensions conservatively (never prunes on them). + """ + # Hoist coordinate array reads outside the partition loop. + # ds.coords[dim].values materializes the full array on every call; doing it + # N_partitions × N_dims times is wasteful and, for remote Zarr-backed datasets + # (e.g. ARCO-ERA5 on GCS), may trigger repeated network I/O. + coord_arrays = {str(dim): ds.coords[dim].values for dim in ds.dims} + return [_block_metadata(coord_arrays, block) for block in blocks] diff --git a/xarray_sql/reader.py b/xarray_sql/reader.py index b9d09e91..4e3000a1 100644 --- a/xarray_sql/reader.py +++ b/xarray_sql/reader.py @@ -27,162 +27,162 @@ ) if TYPE_CHECKING: - from ._native import LazyArrowStreamTable + from ._native import LazyArrowStreamTable class XarrayRecordBatchReader: - """A lazy Arrow stream reader for xarray Datasets. - - Implements the Arrow PyCapsule Interface (__arrow_c_stream__) to enable - zero-copy, lazy streaming of xarray data to DataFusion and other Arrow - consumers. - - The key property is that xarray blocks are only converted to Arrow - RecordBatches when the consumer calls get_next (e.g., during DataFusion's - collect()), NOT when the reader is created or registered. - - Attributes: - schema: The Arrow schema for the stream. - - Example: - >>> import xarray as xr - >>> from xarray_sql import XarrayRecordBatchReader - >>> ds = xr.tutorial.open_dataset('air_temperature') - >>> reader = XarrayRecordBatchReader(ds, chunks={'time': 240}) - >>> # At this point, NO data has been read from xarray - >>> # Data is only read when consumed: - >>> import pyarrow as pa - >>> pa_reader = pa.RecordBatchReader.from_stream(reader) - >>> for batch in pa_reader: - ... print(batch.num_rows) # Data read here - """ - - def __init__( - self, - ds: xr.Dataset, - chunks: Chunks = None, - *, - batch_size: int = DEFAULT_BATCH_SIZE, - _iteration_callback: ( - Callable[[Block, list[str] | None], None] | None - ) = None, - ): - """Initialize the lazy reader. - - Args: - ds: An xarray Dataset. All data_vars must share the same dimensions. - chunks: Xarray-like chunks specification. If not provided, uses - the Dataset's existing chunks. - batch_size: Maximum rows per emitted Arrow RecordBatch. Smaller - values let DataFusion start processing earlier at the cost of - more Python→Arrow conversion calls. - _iteration_callback: Internal callback for testing. Called with - each block dict just before it's converted to Arrow. This - allows tests to track when iteration actually occurs. + """A lazy Arrow stream reader for xarray Datasets. + + Implements the Arrow PyCapsule Interface (__arrow_c_stream__) to enable + zero-copy, lazy streaming of xarray data to DataFusion and other Arrow + consumers. + + The key property is that xarray blocks are only converted to Arrow + RecordBatches when the consumer calls get_next (e.g., during DataFusion's + collect()), NOT when the reader is created or registered. + + Attributes: + schema: The Arrow schema for the stream. + + Example: + >>> import xarray as xr + >>> from xarray_sql import XarrayRecordBatchReader + >>> ds = xr.tutorial.open_dataset('air_temperature') + >>> reader = XarrayRecordBatchReader(ds, chunks={'time': 240}) + >>> # At this point, NO data has been read from xarray + >>> # Data is only read when consumed: + >>> import pyarrow as pa + >>> pa_reader = pa.RecordBatchReader.from_stream(reader) + >>> for batch in pa_reader: + ... print(batch.num_rows) # Data read here """ - self._ds = ds - self._chunks = chunks - self._batch_size = batch_size - self._schema = _parse_schema(ds) - self._iteration_callback = _iteration_callback - self._consumed = False - - # Validate dimensions - fst = next(iter(ds.values())).dims - if not all(da.dims == fst for da in ds.values()): - raise ValueError( - "All dimensions must be equal. Please filter data_vars in the Dataset." - ) - - @property - def schema(self) -> pa.Schema: - """The Arrow schema for this stream.""" - return self._schema - - def _generate_batches(self) -> Iterator[pa.RecordBatch]: - """Generate RecordBatches lazily from xarray blocks. - - This generator is only consumed when the Arrow stream's get_next - is called, ensuring true lazy evaluation. Each xarray block is - emitted as one or more RecordBatches of at most self._batch_size rows. - """ - for block in block_slices(self._ds, self._chunks): - # Call the iteration callback if provided (for testing). - # XarrayRecordBatchReader has no projection concept, so always passes None. - if self._iteration_callback is not None: - self._iteration_callback(block, None) - - yield from iter_record_batches( - self._ds.isel(block), self._schema, self._batch_size - ) - def __arrow_c_stream__( - self, requested_schema: object | None = None - ) -> object: - """Export as Arrow C Stream via PyCapsule. + def __init__( + self, + ds: xr.Dataset, + chunks: Chunks = None, + *, + batch_size: int = DEFAULT_BATCH_SIZE, + _iteration_callback: ( + Callable[[Block, list[str] | None], None] | None + ) = None, + ): + """Initialize the lazy reader. + + Args: + ds: An xarray Dataset. All data_vars must share the same dimensions. + chunks: Xarray-like chunks specification. If not provided, uses + the Dataset's existing chunks. + batch_size: Maximum rows per emitted Arrow RecordBatch. Smaller + values let DataFusion start processing earlier at the cost of + more Python→Arrow conversion calls. + _iteration_callback: Internal callback for testing. Called with + each block dict just before it's converted to Arrow. This + allows tests to track when iteration actually occurs. + """ + self._ds = ds + self._chunks = chunks + self._batch_size = batch_size + self._schema = _parse_schema(ds) + self._iteration_callback = _iteration_callback + self._consumed = False + + # Validate dimensions + fst = next(iter(ds.values())).dims + if not all(da.dims == fst for da in ds.values()): + raise ValueError( + "All dimensions must be equal. Please filter data_vars in the Dataset." + ) + + @property + def schema(self) -> pa.Schema: + """The Arrow schema for this stream.""" + return self._schema + + def _generate_batches(self) -> Iterator[pa.RecordBatch]: + """Generate RecordBatches lazily from xarray blocks. + + This generator is only consumed when the Arrow stream's get_next + is called, ensuring true lazy evaluation. Each xarray block is + emitted as one or more RecordBatches of at most self._batch_size rows. + """ + for block in block_slices(self._ds, self._chunks): + # Call the iteration callback if provided (for testing). + # XarrayRecordBatchReader has no projection concept, so always passes None. + if self._iteration_callback is not None: + self._iteration_callback(block, None) + + yield from iter_record_batches( + self._ds.isel(block), self._schema, self._batch_size + ) + + def __arrow_c_stream__( + self, requested_schema: object | None = None + ) -> object: + """Export as Arrow C Stream via PyCapsule. + + This method is called by Arrow consumers (like DataFusion) to get + a C-level stream interface. The actual data iteration only begins + when the consumer calls get_next on the stream. + + Args: + requested_schema: Optional schema for type casting. Currently + passed through to PyArrow's implementation. + + Returns: + PyCapsule containing ArrowArrayStream pointer with name + "arrow_array_stream". + + Raises: + RuntimeError: If the stream has already been consumed. + """ + if self._consumed: + raise RuntimeError( + "Stream already consumed. XarrayRecordBatchReader can only " + "be iterated once. Create a new reader for additional iterations." + ) + self._consumed = True + + # Create a PyArrow RecordBatchReader from our generator + # The generator is NOT consumed here - only when get_next is called + reader = pa.RecordBatchReader.from_batches( + self._schema, self._generate_batches() + ) - This method is called by Arrow consumers (like DataFusion) to get - a C-level stream interface. The actual data iteration only begins - when the consumer calls get_next on the stream. + # Delegate to PyArrow's __arrow_c_stream__ implementation + return reader.__arrow_c_stream__(requested_schema) - Args: - requested_schema: Optional schema for type casting. Currently - passed through to PyArrow's implementation. + def __arrow_c_schema__( + self, requested_schema: object | None = None + ) -> object: + """Export the schema as Arrow C Schema via PyCapsule. - Returns: - PyCapsule containing ArrowArrayStream pointer with name - "arrow_array_stream". + This allows consumers to inspect the schema without consuming the stream. - Raises: - RuntimeError: If the stream has already been consumed. - """ - if self._consumed: - raise RuntimeError( - "Stream already consumed. XarrayRecordBatchReader can only " - "be iterated once. Create a new reader for additional iterations." - ) - self._consumed = True + Args: + requested_schema: Optional schema for negotiation (unused). - # Create a PyArrow RecordBatchReader from our generator - # The generator is NOT consumed here - only when get_next is called - reader = pa.RecordBatchReader.from_batches( - self._schema, self._generate_batches() - ) + Returns: + PyCapsule containing ArrowSchema pointer. + """ + return self._schema.__arrow_c_schema__() - # Delegate to PyArrow's __arrow_c_stream__ implementation - return reader.__arrow_c_stream__(requested_schema) - def __arrow_c_schema__( - self, requested_schema: object | None = None - ) -> object: - """Export the schema as Arrow C Schema via PyCapsule. - - This allows consumers to inspect the schema without consuming the stream. +def read_xarray(ds: xr.Dataset, chunks: Chunks = None) -> pa.RecordBatchReader: + """Pivots an Xarray Dataset into a PyArrow Table, partitioned by chunks. Args: - requested_schema: Optional schema for negotiation (unused). + ds: An Xarray Dataset. All `data_vars` must share the same dimensions. + chunks: Xarray-like chunks. If not provided, will default to the + Dataset's chunks. The product of the chunk sizes becomes the + standard length of each dataframe partition. Returns: - PyCapsule containing ArrowSchema pointer. + A PyArrow RecordBatchReader, which is a table representation of the input + Dataset. """ - return self._schema.__arrow_c_schema__() - - -def read_xarray(ds: xr.Dataset, chunks: Chunks = None) -> pa.RecordBatchReader: - """Pivots an Xarray Dataset into a PyArrow Table, partitioned by chunks. - - Args: - ds: An Xarray Dataset. All `data_vars` must share the same dimensions. - chunks: Xarray-like chunks. If not provided, will default to the - Dataset's chunks. The product of the chunk sizes becomes the - standard length of each dataframe partition. - - Returns: - A PyArrow RecordBatchReader, which is a table representation of the input - Dataset. - """ - reader = XarrayRecordBatchReader(ds, chunks=chunks) - return pa.RecordBatchReader.from_stream(reader) + reader = XarrayRecordBatchReader(ds, chunks=chunks) + return pa.RecordBatchReader.from_stream(reader) def read_xarray_table( @@ -194,106 +194,112 @@ def read_xarray_table( Callable[[Block, list[str] | None], None] | None ) = None, ) -> "LazyArrowStreamTable": - """Create a lazy DataFusion table from an xarray Dataset. - - This is the simplest way to register xarray data with DataFusion. - Data is only read when queries are executed, not during registration. - The table can be queried multiple times. - - Each chunk becomes a separate partition, enabling DataFusion's parallel - execution across multiple cores. - - Note: - SQL queries with WHERE clauses on dimension columns (time, lat, lon, etc.) - automatically prune partitions that can't contain matching rows — this is - called *filter pushdown*. For example: - - # This query will skip loading partitions with time < '2020-02-01' - result = ctx.sql('SELECT * FROM air WHERE time > "2020-02-01"').collect() - - Supported operators: `=`, `<`, `>`, `<=`, `>=`, `BETWEEN`, `IN`, `AND`, `OR`. - - Args: - ds: An xarray Dataset. All data_vars must share the same dimensions. - chunks: Xarray-like chunks specification. If not provided, uses - the Dataset's existing chunks. - batch_size: Maximum rows per Arrow RecordBatch emitted per partition. - Smaller values let DataFusion start processing earlier; the default - (65 536) works well for most datasets. - _iteration_callback: Internal callback for testing. Called with - each block dict just before it's converted to Arrow. - - Returns: - A LazyArrowStreamTable ready for registration with DataFusion. - - Example: - >>> from datafusion import SessionContext - >>> import xarray as xr - >>> from xarray_sql import read_xarray_table - >>> - >>> ds = xr.tutorial.open_dataset('air_temperature') - >>> table = read_xarray_table(ds, chunks={'time': 240}) - >>> - >>> ctx = SessionContext() - >>> ctx.register_table('air', table) - >>> - >>> # Data is only read here, during query execution - >>> # Filters on 'time' will prune partitions automatically! - >>> result = ctx.sql('SELECT AVG(air) FROM air').collect() - """ - from ._native import LazyArrowStreamTable - - schema = _parse_schema(ds) - - # Hoist coordinate reads once; avoids N_partitions remote I/O calls for - # Zarr-backed datasets (e.g. ARCO-ERA5 on GCS). - coord_arrays = {str(dim): ds.coords[dim].values for dim in ds.dims} - - # Determine which column names are data variables (not dimension coordinates). - # Used by the factory to skip loading unrequested variables. - data_var_names = set(ds.data_vars.keys()) - - def make_partition_factory( - block: Block, - ) -> Callable[[list[str] | None], pa.RecordBatchReader]: - def make_stream( - projection_names: list[str] | None, - ) -> pa.RecordBatchReader: - if _iteration_callback is not None: - _iteration_callback(block, projection_names) - - if projection_names is not None: - # Restrict to the data variables mentioned in the projection. - # Dimension coordinates come along automatically via coords. - data_vars_needed = [c for c in projection_names if c in data_var_names] - if data_vars_needed: - ds_block = ds[data_vars_needed].isel(block) - else: - # Only dimension coords requested — drop all data vars to avoid - # loading them unnecessarily (e.g. for queries like SELECT lat, lon). - ds_block = ds.drop_vars(list(ds.data_vars)).isel(block) - batch_schema = pa.schema( - [schema.field(name) for name in projection_names] - ) - else: - ds_block = ds.isel(block) - batch_schema = schema + """Create a lazy DataFusion table from an xarray Dataset. - return pa.RecordBatchReader.from_batches( - batch_schema, iter_record_batches(ds_block, batch_schema, batch_size) - ) + This is the simplest way to register xarray data with DataFusion. + Data is only read when queries are executed, not during registration. + The table can be queried multiple times. - return make_stream + Each chunk becomes a separate partition, enabling DataFusion's parallel + execution across multiple cores. - def partition_pairs(): - """Lazily yield (factory, metadata) for each partition. + Note: + SQL queries with WHERE clauses on dimension columns (time, lat, lon, etc.) + automatically prune partitions that can't contain matching rows — this is + called *filter pushdown*. For example: - Consuming this generator one item at a time means Python never holds - all N block dicts, metadata dicts, and factory closures simultaneously. - Peak Python memory during registration is O(1) per partition instead - of O(N_partitions). - """ - for block in block_slices(ds, chunks): - yield make_partition_factory(block), _block_metadata(coord_arrays, block) + # This query will skip loading partitions with time < '2020-02-01' + result = ctx.sql('SELECT * FROM air WHERE time > "2020-02-01"').collect() - return LazyArrowStreamTable(partition_pairs(), schema) + Supported operators: `=`, `<`, `>`, `<=`, `>=`, `BETWEEN`, `IN`, `AND`, `OR`. + + Args: + ds: An xarray Dataset. All data_vars must share the same dimensions. + chunks: Xarray-like chunks specification. If not provided, uses + the Dataset's existing chunks. + batch_size: Maximum rows per Arrow RecordBatch emitted per partition. + Smaller values let DataFusion start processing earlier; the default + (65 536) works well for most datasets. + _iteration_callback: Internal callback for testing. Called with + each block dict just before it's converted to Arrow. + + Returns: + A LazyArrowStreamTable ready for registration with DataFusion. + + Example: + >>> from datafusion import SessionContext + >>> import xarray as xr + >>> from xarray_sql import read_xarray_table + >>> + >>> ds = xr.tutorial.open_dataset('air_temperature') + >>> table = read_xarray_table(ds, chunks={'time': 240}) + >>> + >>> ctx = SessionContext() + >>> ctx.register_table('air', table) + >>> + >>> # Data is only read here, during query execution + >>> # Filters on 'time' will prune partitions automatically! + >>> result = ctx.sql('SELECT AVG(air) FROM air').collect() + """ + from ._native import LazyArrowStreamTable + + schema = _parse_schema(ds) + + # Hoist coordinate reads once; avoids N_partitions remote I/O calls for + # Zarr-backed datasets (e.g. ARCO-ERA5 on GCS). + coord_arrays = {str(dim): ds.coords[dim].values for dim in ds.dims} + + # Determine which column names are data variables (not dimension coordinates). + # Used by the factory to skip loading unrequested variables. + data_var_names = set(ds.data_vars.keys()) + + def make_partition_factory( + block: Block, + ) -> Callable[[list[str] | None], pa.RecordBatchReader]: + def make_stream( + projection_names: list[str] | None, + ) -> pa.RecordBatchReader: + if _iteration_callback is not None: + _iteration_callback(block, projection_names) + + if projection_names is not None: + # Restrict to the data variables mentioned in the projection. + # Dimension coordinates come along automatically via coords. + data_vars_needed = [ + c for c in projection_names if c in data_var_names + ] + if data_vars_needed: + ds_block = ds[data_vars_needed].isel(block) + else: + # Only dimension coords requested — drop all data vars to avoid + # loading them unnecessarily (e.g. for queries like SELECT lat, lon). + ds_block = ds.drop_vars(list(ds.data_vars)).isel(block) + batch_schema = pa.schema( + [schema.field(name) for name in projection_names] + ) + else: + ds_block = ds.isel(block) + batch_schema = schema + + return pa.RecordBatchReader.from_batches( + batch_schema, + iter_record_batches(ds_block, batch_schema, batch_size), + ) + + return make_stream + + def partition_pairs(): + """Lazily yield (factory, metadata) for each partition. + + Consuming this generator one item at a time means Python never holds + all N block dicts, metadata dicts, and factory closures simultaneously. + Peak Python memory during registration is O(1) per partition instead + of O(N_partitions). + """ + for block in block_slices(ds, chunks): + yield ( + make_partition_factory(block), + _block_metadata(coord_arrays, block), + ) + + return LazyArrowStreamTable(partition_pairs(), schema) diff --git a/xarray_sql/sql.py b/xarray_sql/sql.py index 7b977569..f30f5b83 100644 --- a/xarray_sql/sql.py +++ b/xarray_sql/sql.py @@ -7,57 +7,57 @@ class XarrayContext(SessionContext): - """A datafusion `SessionContext` that also supports `xarray.Dataset`s.""" - - def from_dataset( - self, - table_name: str, - input_table: xr.Dataset, - chunks: Chunks = None, - ): - """Register an xarray Dataset as a queryable SQL table. - - For datasets with non-Gregorian cftime coordinates (e.g. 360_day, - julian), a ``cftime()`` scalar UDF is automatically registered so - you can write ergonomic SQL filters:: - - ctx.from_dataset("ds360", ds, chunks={"time": 6}) - ctx.sql("SELECT * FROM ds360 WHERE time >= cftime('2000-07-01')") - - The UDF converts a date string to the int64 offset used to store - that calendar's time axis. - - .. note:: - - Only one ``cftime()`` UDF is registered per context, using the - units and calendar of the *first* non-Gregorian coordinate - encountered. If you register multiple datasets with *different* - non-Gregorian calendars (e.g. one 360_day and one julian), the - UDF from the first registration will be used for all subsequent - ``cftime()`` calls and may produce incorrect offsets for the - other dataset. In that case, create a separate ``XarrayContext`` - for each calendar. - - Args: - table_name: The SQL table name to register the dataset under. - input_table: An xarray Dataset. All data_vars must share the - same dimensions. - chunks: Xarray-like chunks specification. If not provided, uses - the Dataset's existing chunks. - - Returns: - self, to allow chaining. - """ - table = read_xarray_table(input_table, chunks) - self.register_table(table_name, table) - - # Auto-register a cftime() UDF for non-Gregorian cftime coordinates - # so users can write: WHERE time > cftime('0500-01-01') - for coord_name in input_table.dims: - if cft.is_cftime_index(input_table, coord_name): - units, cal = cft.encoding(input_table, coord_name) - if not cft.is_gregorian_like(cal): - self.register_udf(cft.make_cftime_udf(units, cal)) - break # One UDF per context is enough. - - return self + """A datafusion `SessionContext` that also supports `xarray.Dataset`s.""" + + def from_dataset( + self, + table_name: str, + input_table: xr.Dataset, + chunks: Chunks = None, + ): + """Register an xarray Dataset as a queryable SQL table. + + For datasets with non-Gregorian cftime coordinates (e.g. 360_day, + julian), a ``cftime()`` scalar UDF is automatically registered so + you can write ergonomic SQL filters:: + + ctx.from_dataset("ds360", ds, chunks={"time": 6}) + ctx.sql("SELECT * FROM ds360 WHERE time >= cftime('2000-07-01')") + + The UDF converts a date string to the int64 offset used to store + that calendar's time axis. + + .. note:: + + Only one ``cftime()`` UDF is registered per context, using the + units and calendar of the *first* non-Gregorian coordinate + encountered. If you register multiple datasets with *different* + non-Gregorian calendars (e.g. one 360_day and one julian), the + UDF from the first registration will be used for all subsequent + ``cftime()`` calls and may produce incorrect offsets for the + other dataset. In that case, create a separate ``XarrayContext`` + for each calendar. + + Args: + table_name: The SQL table name to register the dataset under. + input_table: An xarray Dataset. All data_vars must share the + same dimensions. + chunks: Xarray-like chunks specification. If not provided, uses + the Dataset's existing chunks. + + Returns: + self, to allow chaining. + """ + table = read_xarray_table(input_table, chunks) + self.register_table(table_name, table) + + # Auto-register a cftime() UDF for non-Gregorian cftime coordinates + # so users can write: WHERE time > cftime('0500-01-01') + for coord_name in input_table.dims: + if cft.is_cftime_index(input_table, coord_name): + units, cal = cft.encoding(input_table, coord_name) + if not cft.is_gregorian_like(cal): + self.register_udf(cft.make_cftime_udf(units, cal)) + break # One UDF per context is enough. + + return self