diff --git a/README.md b/README.md index ddbd497..60cf7c6 100644 --- a/README.md +++ b/README.md @@ -20,22 +20,21 @@ import xarray as xr import xarray_sql as xql -# Open a year of ARCO-ERA5 — all 273 variables. Selecting a year up front -# keeps Dask's partition setup cheap before any chunks are read from GCS. -ds = ( - xr.open_zarr('gs://gcp-public-data-arco-era5/ar/full_37-1h-0p25deg-chunk-1.zarr-v3', - chunks=dict(time=1), - storage_options={'token': 'anon'}) # Anonymous read from the public GCS bucket — no auth required. - .sel(time='2020') +# Open ARCO-ERA5 — a weather dataset with 273 variables since 1940. +# Turning off dask means we don't have to wait to construct a task graph. +ds = xr.open_zarr( + 'gs://gcp-public-data-arco-era5/ar/full_37-1h-0p25deg-chunk-1.zarr-v3', + chunks=None, # Turn dask off + storage_options={'token': 'anon'} # Anonymous read from the public GCS bucket — no auth required. ) ctx = xql.XarrayContext() -ctx.from_dataset('era5', ds, table_names={ +# Make sure to pass `chunks`! +ctx.from_dataset('era5', ds, chunks=dict(time=6), table_names={ ('time', 'latitude', 'longitude'): 'surface', ('time', 'level', 'latitude', 'longitude'): 'atmosphere', }) -# Registration: ~0.5s for a full year of hourly ERA5, all variables. - +# Registration takes ~10s on my machine. # Heads up: ARCO-ERA5 has 262 surface + 11 atmospheric variables. The library # pushes column projection down to Zarr, so SELECT only fetches what you ask @@ -81,13 +80,26 @@ result = ctx.sql(''' # | 775 | -2.3064649711534457 | # +-------+----------------------+ -avg_temp_ds = result.to_dataset(dims=["level"]) -# Size: 592B -# Dimensions: (level: 37) +ctx.sql(''' + SELECT latitude, longitude, AVG("2m_temperature") - 273.15 AS avg_c + FROM era5.surface + WHERE time BETWEEN TIMESTAMP '2020-01-01' + AND TIMESTAMP '2020-01-01 05:00:00' + GROUP BY latitude, longitude + ORDER BY latitude DESC, longitude +''').to_dataset(dims=['latitude', 'longitude'], template=ds) +# Size: 8MB +# Dimensions: (latitude: 721, longitude: 1440) # Coordinates: -# * level (level) int64 296B 1000 975 950 925 900 875 850 ... 20 10 7 5 3 2 1 +# * latitude (latitude) float32 3kB 90.0 89.75 89.5 ... -89.5 -89.75 -90.0 +# * longitude (longitude) float32 6kB 0.0 0.25 0.5 0.75 ... 359.2 359.5 359.8 # Data variables: -# avg_c (level) float64 296B 6.621 5.186 4.028 ... -21.51 -13.36 -9.021 +# avg_c (latitude, longitude) float64 8MB -26.84 -26.84 ... -27.38 -27.38 +# Attributes: +# last_updated: 2026-06-20 02:33:34.265980+00:00 +# valid_time_start: 1940-01-01 +# valid_time_stop: 2025-12-31 +# valid_time_stop_era5t: 2026-06-14 ``` _(A runnable version of this example lives at diff --git a/perf_tests/era5_temp_profile.py b/perf_tests/era5_temp_profile.py index 28275a9..55cb8be 100644 --- a/perf_tests/era5_temp_profile.py +++ b/perf_tests/era5_temp_profile.py @@ -1,7 +1,7 @@ #!/usr/bin/env python3 """Surface and global-atmospheric temperatures on 2020-01-01, in SQL. -Two queries against ARCO-ERA5 on the morning of January 1, 2020: +Three queries against ARCO-ERA5 on the morning of January 1, 2020: * **Surface (local).** Average 2m-temperature over a small grid covering the New York City area for the first six hours. @@ -9,16 +9,19 @@ over the entire planet for the same six hours — a classic atmospheric temperature profile (surface around 1000 hPa is warmest, tropopause near 100 hPa is coldest). + * **Surface (global, gridded).** Average 2m-temperature per (lat, lon) cell + for the same six hours, returned as an xarray Dataset. -Both queries express their filters entirely in SQL: ``xr.open_zarr`` is given -a single calendar year and no spatial slicing. The library's table provider -prunes time partitions for ``WHERE time …`` filters, and pushes ``WHERE +All filters live in SQL: the dataset is opened with no time or spatial +slicing on the xarray side. The library's table provider prunes time +partitions for ``WHERE time …`` filters, and pushes ``WHERE latitude/longitude …`` down to dimension columns. ARCO-ERA5's atmospheric variables are stored in native Zarr chunks of shape -``(1, 37, 721, 1440)`` — about 150 MB per hour. We align Dask chunks to that -shape with ``chunks={'time': 1}`` so chunks fetch from GCS concurrently. The -global atmospheric query scans ~230M rows after pruning. +``(1, 37, 721, 1440)`` — about 150 MB per hour. ``chunks=dict(time=6)`` groups +six native chunks per DataFusion partition: large enough to keep partition +count (and registration time) low, small enough that a 6-hour WHERE clause +hits a single partition with no wasted I/O. The Zarr is read anonymously from the public GCS bucket — no auth required. """ @@ -34,27 +37,29 @@ def main() -> None: - full = xr.open_zarr(URL, chunks=None, storage_options={"token": "anon"}) - - # Open a full calendar year — all 273 variables. No spatial slicing on - # the xarray side; SQL WHERE clauses below express the filters. - # - # Heads up: the library pushes column projection down to Zarr, so SELECT - # only fetches what you ask for — but `SELECT * FROM era5.surface` would - # try to read every variable across the year (terabytes from GCS). - # Always SELECT specific columns. - ds = full.sel(time="2020").chunk({"time": 1}) + # Open the full ARCO-ERA5 archive — all 273 variables since 1940. No + # time or spatial slicing on the xarray side; SQL WHERE clauses below + # express the filters. Turning dask off (chunks=None) skips task-graph + # construction at open time. + ds = xr.open_zarr(URL, chunks=None, storage_options={"token": "anon"}) print( - "ARCO-ERA5 opened: year 2020, " + "ARCO-ERA5 opened: " f"{ds.sizes['time']:,} hourly time steps, " - f"{len(ds.data_vars)} variables (no spatial pre-slicing)." + f"{len(ds.data_vars)} variables (no pre-slicing)." ) + # Heads up: ARCO-ERA5 has 262 surface + 11 atmospheric variables. The + # library pushes column projection down to Zarr, so SELECT only fetches + # what you ask for — but `SELECT * FROM era5.surface` would try to pull + # every variable across the archive (terabytes from GCS). + # ---> Always SELECT specific columns. <--- ctx = xql.XarrayContext() t0 = time.perf_counter() + # Make sure to pass `chunks`! ctx.from_dataset( "era5", ds, + chunks=dict(time=6), table_names={ ("time", "latitude", "longitude"): "surface", ("time", "level", "latitude", "longitude"): "atmosphere", @@ -94,7 +99,25 @@ def main() -> None: """ ).to_pandas() print(profile.to_string(index=False)) - print(f" ({time.perf_counter() - t0:.2f}s, ~230M rows scanned)") + print(f" ({time.perf_counter() - t0:.2f}s)") + + print( + "\nAverage 2m-temperature per (lat, lon) cell, globally, " + "2020-01-01 00:00-05:00 UTC (°C):" + ) + t0 = time.perf_counter() + gridded = ctx.sql( + """ + SELECT latitude, longitude, AVG("2m_temperature") - 273.15 AS avg_c + FROM era5.surface + WHERE time BETWEEN TIMESTAMP '2020-01-01' + AND TIMESTAMP '2020-01-01 05:00:00' + GROUP BY latitude, longitude + ORDER BY latitude DESC, longitude + """ + ).to_dataset(dims=["latitude", "longitude"], template=ds) + print(gridded) + print(f" ({time.perf_counter() - t0:.2f}s)") if __name__ == "__main__": diff --git a/tests/test_df.py b/tests/test_df.py index a2b443f..2c81144 100644 --- a/tests/test_df.py +++ b/tests/test_df.py @@ -10,6 +10,7 @@ DEFAULT_BATCH_SIZE, _parse_schema, block_slices, + compute_chunks, dataset_to_record_batch, explode, from_map, @@ -473,3 +474,61 @@ def test_read_xarray_table_memory_bounds(large_ds): ) finally: tracemalloc.stop() + + +# --------------------------------------------------------------------------- +# compute_chunks: arithmetic replacement for ds.chunk(...).chunks. +# Dask serves as the source of truth. +# --------------------------------------------------------------------------- + + +def _dask_chunks(ds: xr.Dataset, chunks: dict) -> dict: + rechunked = ds.copy(data=None, deep=False).chunk(chunks) + return {str(k): tuple(v) for k, v in rechunked.chunks.items()} + + +def _normalise(result: dict) -> dict: + return {str(k): tuple(v) for k, v in result.items()} + + +def _simple_ds(shape: tuple[int, ...], dims: tuple[str, ...]) -> xr.Dataset: + return xr.Dataset( + {"v": (dims, np.zeros(shape))}, + coords={d: np.arange(s) for d, s in zip(dims, shape)}, + ) + + +@pytest.mark.parametrize( + "ds,chunks", + [ + # Even divide on a single dim. + (_simple_ds((10,), ("x",)), {"x": 5}), + # Uneven divide: trailing remainder chunk. + (_simple_ds((10,), ("x",)), {"x": 3}), + # Requested chunk size larger than the dim → single chunk. + (_simple_ds((5,), ("x",)), {"x": 100}), + # Multi-dim spec with a dim left unspecified (kept as one chunk). + (_simple_ds((4, 6), ("x", "y")), {"x": 2}), + # Multi-dim spec rechunking every dim. + (_simple_ds((7, 11, 13), ("a", "b", "c")), {"a": 3, "b": 4, "c": 5}), + ], +) +def test_compute_chunks_matches_dask(ds, chunks): + assert _normalise(compute_chunks(ds, chunks)) == _dask_chunks(ds, chunks) + + +def test_compute_chunks_preserves_existing_dask_chunking(): + # When the dataset is already dask-backed, rechunking one dim must + # leave other dims' existing chunk tuples alone. + ds = _simple_ds((4, 5), ("x", "y")).chunk({"x": 1, "y": 2}) + chunks = {"x": 2} + assert _normalise(compute_chunks(ds, chunks)) == _dask_chunks(ds, chunks) + + +def test_compute_chunks_tuples_sum_to_dim_size(): + # Dask-independent invariant: every per-dim chunk tuple must fully + # cover its dimension. + ds = _simple_ds((7, 11, 13), ("a", "b", "c")) + result = compute_chunks(ds, {"a": 3, "b": 4, "c": 5}) + for dim, tup in result.items(): + assert sum(tup) == ds.sizes[dim] diff --git a/tests/test_reader.py b/tests/test_reader.py index a8e7e33..173fd0f 100644 --- a/tests/test_reader.py +++ b/tests/test_reader.py @@ -1102,6 +1102,44 @@ def test_lat_filter_prunes_partitions(self): f"Expected exactly 2 partitions for lat < 0, got {tracker.iteration_count}" ) + def test_unchunked_dim_filter_still_prunes(self): + """Filters on an *unchunked* dim still prune via static bounds. + + ``read_xarray_table`` precomputes bounds for unchunked dims once + rather than re-scanning their full coord array on every partition. + Regression guard: if the static-range merge ever stops attaching + those bounds to each partition, the Rust pruner falls back to + "never prune" for the unchunked dim and reads every partition. + """ + np.random.seed(42) + time = pd.date_range("2020-01-01", periods=100, freq="D") + lat = np.linspace(-90, 90, 50) # unchunked + data = np.random.rand(100, 50).astype(np.float32) + + ds = xr.Dataset( + {"temperature": (["time", "lat"], data)}, + coords={"time": time, "lat": lat}, + ) + + tracker = IterationTracker() + # Chunk only on time → lat is "static" (one chunk spanning the axis). + # Every partition still spans lat -90 to +90. + table = read_xarray_table( + ds, chunks={"time": 25}, _iteration_callback=tracker + ) + + ctx = SessionContext() + ctx.register_table("test", table) + + # WHERE lat > 100 matches no rows. With static bounds (lat ∈ [-90, 90]) + # attached to every partition, the pruner drops *all* partitions and + # the table is never iterated. Without them, all 4 are read. + ctx.sql("SELECT COUNT(*) FROM test WHERE lat > 100").collect() + assert tracker.iteration_count == 0, ( + "Static lat bounds should let the pruner skip every partition; " + f"got {tracker.iteration_count} partitions read." + ) + def test_no_pruning_for_data_column_filters(self, time_chunked_ds): """Filters on data columns (not dimensions) should not prune.""" tracker = IterationTracker() diff --git a/xarray_sql/df.py b/xarray_sql/df.py index f55b6ad..99c8408 100644 --- a/xarray_sql/df.py +++ b/xarray_sql/df.py @@ -1,5 +1,5 @@ import itertools -from collections.abc import Callable, Hashable, Iterator, Mapping +from collections.abc import Callable, Hashable, Iterable, Iterator, Mapping from typing import Any import numpy as np @@ -26,22 +26,58 @@ def _get_chunk_slicer( 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.""" +def compute_chunks( + ds: xr.Dataset, chunks: dict[str, int] +) -> dict[Hashable, tuple[int, ...]]: + """Per-dim chunk-size tuples matching ``ds.chunk(chunks).chunks``. + + Pure arithmetic replacement for the dask rechunk round-trip; dask's + ``.chunk()`` eagerly builds a task graph, which dominates + ``block_slices()`` cost on large datasets. + """ + existing = dict(ds.chunks) if ds.chunks else {} + result: dict[Hashable, tuple[int, ...]] = {} + for dim in ds.dims: + size = ds.sizes[dim] + if dim in chunks: + cs = chunks[dim] + if cs <= 0 or cs >= size: + result[dim] = (size,) + else: + n_full, rem = divmod(size, cs) + result[dim] = (cs,) * n_full + ((rem,) if rem else ()) + elif dim in existing: + result[dim] = tuple(existing[dim]) + else: + result[dim] = (size,) + return result + + +def resolve_chunks( + ds: xr.Dataset, chunks: Chunks +) -> Mapping[Hashable, tuple[int, ...]]: + """Normalise the user's ``chunks`` argument to per-dim size tuples. + + Filters out keys for dims this dataset doesn't have (sub-datasets in a + heterogeneous group need not contain every dimension named in the + spec), then either rechunks arithmetically via ``compute_chunks`` or + falls back to the dataset's existing dask chunks. + + Returns an empty mapping for scalar datasets; callers should treat that + as "one block covering everything". + """ if chunks is not None: - # Only chunk dimensions this dataset actually has. A sub-dataset - # (e.g. one dimension group of a heterogeneous Dataset) need not - # contain every dimension named in the chunks spec. chunks = {dim: size for dim, size in chunks.items() if dim in ds.sizes} if chunks: - for_chunking = ds.copy(data=None, deep=False).chunk(chunks) - chunks = for_chunking.chunks - del for_chunking - else: - chunks = ds.chunks + return compute_chunks(ds, chunks) + return {d: tuple(c) for d, c in ds.chunks.items()} - if not chunks: + +def _block_slices_from_resolved( + ds: xr.Dataset, resolved: Mapping[Hashable, tuple[int, ...]] +) -> Iterator[Block]: + """Emit blocks given pre-resolved per-dim chunk tuples.""" + if not resolved: # No chunkable dimensions. A dimensionless dataset (e.g. scalar # metadata variables) is a single block; a dataset that has # dimensions but no chunking is a user error. @@ -51,22 +87,25 @@ def block_slices(ds: xr.Dataset, chunks: Chunks = None) -> Iterator[Block]: yield {} return - # 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() + dim: np.cumsum((0,) + tuple(c)) for dim, c in resolved.items() } - ichunk = {dim: range(len(tuple(c))) for dim, c in chunks.items()} # type: ignore[arg-type] + ichunk = {dim: range(len(tuple(c))) for dim, c in resolved.items()} 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 = ( + yield from ( { dim: _get_chunk_slicer(dim, chunk_index, chunk_bounds) for dim in ds.dims } for chunk_index in chunk_idxs ) - yield from blocks + + +# Adapted from Xarray `map_blocks` implementation. +def block_slices(ds: xr.Dataset, chunks: Chunks = None) -> Iterator[Block]: + """Compute block slices for a chunked Dataset.""" + yield from _block_slices_from_resolved(ds, resolve_chunks(ds, chunks)) def explode(ds: xr.Dataset, chunks: Chunks = None) -> Iterator[xr.Dataset]: @@ -376,7 +415,11 @@ def _parse_schema(ds: xr.Dataset) -> pa.Schema: PartitionBounds = dict[str, tuple[Any, Any, str]] -def _block_metadata(coord_arrays: dict, block: Block) -> PartitionBounds: +def _block_metadata( + coord_arrays: dict, + block: Block, + dims: Iterable[Hashable] | None = None, +) -> PartitionBounds: """Compute min/max coordinate values for a single partition block. Args: @@ -384,14 +427,19 @@ def _block_metadata(coord_arrays: dict, block: Block) -> PartitionBounds: 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(). + dims: Optional restriction to a subset of dims to compute. Used by + ``read_xarray_table`` to skip unchunked dims whose bounds are + constant across all partitions and have been precomputed once. + Defaults to all dims present in ``block``. 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). """ + items = ((d, block[d]) for d in dims) if dims is not None else block.items() ranges: PartitionBounds = {} - for dim, slc in block.items(): + for dim, slc in items: coord_values = coord_arrays[str(dim)][slc] if len(coord_values) == 0: continue diff --git a/xarray_sql/reader.py b/xarray_sql/reader.py index 4e3000a..f8c5975 100644 --- a/xarray_sql/reader.py +++ b/xarray_sql/reader.py @@ -13,6 +13,7 @@ from collections.abc import Callable, Iterator from typing import TYPE_CHECKING +import numpy as np import pyarrow as pa import xarray as xr @@ -21,9 +22,11 @@ Chunks, DEFAULT_BATCH_SIZE, _block_metadata, + _block_slices_from_resolved, _parse_schema, block_slices, iter_record_batches, + resolve_chunks, ) if TYPE_CHECKING: @@ -190,6 +193,7 @@ def read_xarray_table( chunks: Chunks = None, *, batch_size: int = DEFAULT_BATCH_SIZE, + coord_arrays: dict[str, np.ndarray] | None = None, _iteration_callback: ( Callable[[Block, list[str] | None], None] | None ) = None, @@ -220,6 +224,12 @@ def read_xarray_table( 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. + coord_arrays: Pre-materialised coordinate arrays keyed by dim-name + string. Hand in to share a single read across multiple tables + built from the same parent Dataset (e.g. surface + atmosphere + from ARCO-ERA5); the dim coords are otherwise read once per + ``read_xarray_table`` call, which is a network round-trip for + Zarr-backed datasets. _iteration_callback: Internal callback for testing. Called with each block dict just before it's converted to Arrow. @@ -246,8 +256,11 @@ def read_xarray_table( 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} + # Zarr-backed datasets (e.g. ARCO-ERA5 on GCS). When the caller supplies + # pre-materialised arrays (e.g. shared across surface + atmosphere + # tables), reuse them and skip the extra read. + if coord_arrays is None: + 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. @@ -288,6 +301,19 @@ def make_stream( return make_stream + # Separate dims whose chunk bounds vary across partitions from those + # whose bounds are constant (one chunk spanning the whole axis). For the + # latter we compute min/max once instead of re-scanning the full coord + # array on every partition — dominant cost when registering hundreds of + # thousands of single-time-step partitions on a 4-D dataset like ERA5. + resolved = resolve_chunks(ds, chunks) + varying_dims = [d for d, tup in resolved.items() if len(tup) > 1] + static_dims = [d for d in ds.dims if d not in varying_dims] + static_block: Block = {d: slice(None) for d in static_dims} + static_ranges = _block_metadata( + coord_arrays, static_block, dims=static_dims + ) + def partition_pairs(): """Lazily yield (factory, metadata) for each partition. @@ -296,10 +322,11 @@ def partition_pairs(): Peak Python memory during registration is O(1) per partition instead of O(N_partitions). """ - for block in block_slices(ds, chunks): + for block in _block_slices_from_resolved(ds, resolved): + dynamic = _block_metadata(coord_arrays, block, dims=varying_dims) yield ( make_partition_factory(block), - _block_metadata(coord_arrays, block), + {**static_ranges, **dynamic}, ) return LazyArrowStreamTable(partition_pairs(), schema) diff --git a/xarray_sql/sql.py b/xarray_sql/sql.py index 99efea0..0ec60ad 100644 --- a/xarray_sql/sql.py +++ b/xarray_sql/sql.py @@ -89,9 +89,18 @@ def from_dataset( """ groups = _group_vars_by_dims(input_table) + # Materialise dim coordinates once and share across every sub-table. + # For Zarr-backed parents (e.g. ARCO-ERA5 on GCS) this saves one + # network round-trip per dim per dim-group. + coord_arrays = { + str(dim): input_table.coords[dim].values for dim in input_table.dims + } + if len(groups) <= 1: self._registered_datasets[name] = input_table - return self._from_dataset(name, input_table, chunks) + return self._from_dataset( + name, input_table, chunks, coord_arrays=coord_arrays + ) table_names = table_names or {} schema = Schema.memory_schema(self) @@ -102,7 +111,13 @@ def from_dataset( # the empty string; fall back to a valid default table name. sub_name = table_names.get(dims, "_".join(dims) or "scalar") sub_ds = input_table[var_names] - self._from_dataset(sub_name, sub_ds, chunks, schema=schema) + self._from_dataset( + sub_name, + sub_ds, + chunks, + schema=schema, + coord_arrays=coord_arrays, + ) # Track the fully-qualified name so XarrayDataFrame metadata # recovery can find this Dataset on round-trip. self._registered_datasets[f"{name}.{sub_name}"] = sub_ds @@ -115,6 +130,7 @@ def _from_dataset( input_table: xr.Dataset, chunks: Chunks = None, schema: Schema | None = None, + coord_arrays: dict | None = None, ): """Register a Dataset as a single SQL table. @@ -124,7 +140,10 @@ def _from_dataset( register = ( self.register_table if schema is None else schema.register_table ) - register(table_name, read_xarray_table(input_table, chunks)) + register( + table_name, + read_xarray_table(input_table, chunks, coord_arrays=coord_arrays), + ) self._maybe_register_cftime_udf(input_table) return self