Optimizations to make registration faster (95s -> 11s).

Author: alxmrs

README.md

@@ -30,7 +30,7 @@ ds = xr.open_zarr(
 
 ctx = xql.XarrayContext()
 # Make sure to pass `chunks`!
-ctx.from_dataset('era5', ds, chunks=dict(time=1), table_names={
+ctx.from_dataset('era5', ds, chunks=dict(time=6), table_names={
     ('time', 'latitude', 'longitude'): 'surface',
     ('time', 'level', 'latitude', 'longitude'): 'atmosphere',
 })

perf_tests/era5_temp_profile.py

@@ -18,8 +18,10 @@
 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=dict(time=1)`` so chunks fetch from GCS concurrently.
+``(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.
 """
@@ -57,7 +59,7 @@ def main() -> None:
     ctx.from_dataset(
         "era5",
         ds,
-        chunks=dict(time=1),
+        chunks=dict(time=6),
         table_names={
             ("time", "latitude", "longitude"): "surface",
             ("time", "level", "latitude", "longitude"): "atmosphere",

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
@@ -402,22 +402,31 @@ 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:
         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().
+        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

xarray_sql/reader.py

@@ -23,6 +23,7 @@
     _block_metadata,
     _parse_schema,
     block_slices,
+    compute_chunks,
     iter_record_batches,
 )
 
@@ -288,6 +289,24 @@ 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.
+    if chunks:
+        resolved = compute_chunks(ds, chunks)
+    elif ds.chunks:
+        resolved = {d: tuple(c) for d, c in ds.chunks.items()}
+    else:
+        resolved = {}
+    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.
 
@@ -297,9 +316,10 @@ def partition_pairs():
         of O(N_partitions).
         """
         for block in block_slices(ds, chunks):
+            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)