Integer grid-position columns for exact grid joins

tests/test_grid_index.py

@@ -0,0 +1,169 @@
+"""Integer grid-position columns (``<dim>_idx``) for exact grid joins.
+
+Regridding and forecast alignment join a source grid to a table on the source
+*coordinate*. Joining on the floating-point coordinate value is fragile — any
+sub-ULP drift (e.g. a reproject/interp computed in float32) makes the equality
+join silently drop rows. The opt-in ``<dim>_idx`` columns give exact integer
+grid keys instead. These tests pin that the indices are global (partition
+independent) and that an index-keyed regrid stays exact where a float join does
+not.
+"""
+
+import numpy as np
+import pyarrow as pa
+import xarray as xr
+from datafusion import SessionContext
+
+from xarray_sql import XarrayContext
+from xarray_sql.df import _ensure_default_indexes, _parse_schema
+from xarray_sql.reader import read_xarray_table
+
+
+# Deliberately not float32-exact, so a float32 round-trip actually perturbs the
+# bits (a float32-exact axis like 10.0/20.0 would round-trip unchanged).
+_X = np.array([10.1, 20.2, 30.3])
+_Y = np.array([1.1, 2.2, 3.3, 4.4])
+
+
+def _src() -> xr.Dataset:
+    return xr.Dataset(
+        {"v": (("x", "y"), np.arange(12.0).reshape(3, 4))},
+        coords={"x": _X, "y": _Y},
+    )
+
+
+def test_index_columns_are_int32_in_schema():
+    schema = _parse_schema(_ensure_default_indexes(_src()), index_columns=True)
+    assert schema.field("x_idx").type == pa.int32()
+    assert schema.field("y_idx").type == pa.int32()
+    # Not present unless requested.
+    plain = _parse_schema(_ensure_default_indexes(_src()))
+    assert "x_idx" not in plain.names
+
+
+def test_index_columns_are_global_across_chunks():
+    """Indices are absolute axis positions, not per-partition local ones."""
+    ctx = XarrayContext()
+    # chunk x into 3 single-row partitions: a local index would restart at 0
+    # in every partition; the global index must run 0,1,2.
+    ctx.from_dataset("g", _src(), chunks={"x": 1}, index_columns=True)
+    df = ctx.sql("SELECT x, y, x_idx, y_idx FROM g").to_pandas()
+
+    assert df["x_idx"].dtype == np.int32
+    xpos = {v: i for i, v in enumerate(_X)}
+    ypos = {v: i for i, v in enumerate(_Y)}
+    assert (df["x_idx"] == df["x"].map(xpos)).all()
+    assert (df["y_idx"] == df["y"].map(ypos)).all()
+
+
+def _weights(perturb_f32: bool = False) -> xr.Dataset:
+    """A tiny gather 'regrid': dst cell k reads one src cell, weight 1.0."""
+    # dst 0..4 map to source cells (x_idx, y_idx):
+    sx = np.array([0, 2, 1, 0, 2], dtype=np.int32)
+    sy = np.array([0, 3, 1, 3, 0], dtype=np.int32)
+    src_x = _X[sx]
+    src_y = _Y[sy]
+    if perturb_f32:  # coords computed in single precision, as a reproject UDF
+        src_x = src_x.astype(np.float32).astype(np.float64)
+        src_y = src_y.astype(np.float32).astype(np.float64)
+    return xr.Dataset(
+        {
+            "dst_id": (("pair",), np.arange(5, dtype=np.int32)),
+            "src_x_idx": (("pair",), sx),
+            "src_y_idx": (("pair",), sy),
+            "src_x": (("pair",), src_x),
+            "src_y": (("pair",), src_y),
+            "weight": (("pair",), np.ones(5)),
+        }
+    )
+
+
+INDEX_JOIN = """
+    SELECT w.dst_id, SUM(s.v * w.weight) AS out
+    FROM weights w JOIN src s
+      ON s.x_idx = w.src_x_idx AND s.y_idx = w.src_y_idx
+    GROUP BY w.dst_id ORDER BY w.dst_id
+"""
+FLOAT_JOIN = """
+    SELECT w.dst_id, SUM(s.v * w.weight) AS out
+    FROM weights w JOIN src s
+      ON s.x = w.src_x AND s.y = w.src_y
+    GROUP BY w.dst_id ORDER BY w.dst_id
+"""
+
+
+def _ctx(weights: xr.Dataset) -> XarrayContext:
+    ctx = XarrayContext()
+    ctx.from_dataset("src", _src(), chunks={"x": 1}, index_columns=True)
+    ctx.from_dataset("weights", weights, chunks={"pair": 5})
+    return ctx
+
+
+def test_index_join_regrids_exactly():
+    """Index-keyed regrid matches a direct numpy gather, across chunks."""
+    ctx = _ctx(_weights())
+    got = ctx.sql(INDEX_JOIN).to_pandas().sort_values("dst_id")
+
+    v = np.arange(12.0).reshape(3, 4)
+    sx = np.array([0, 2, 1, 0, 2])
+    sy = np.array([0, 3, 1, 3, 0])
+    expected = v[sx, sy]
+    np.testing.assert_allclose(got["out"].to_numpy(), expected)
+
+
+def test_index_join_survives_float32_drift_where_float_join_does_not():
+    ctx = _ctx(_weights(perturb_f32=True))
+    idx = ctx.sql(INDEX_JOIN).to_pandas()
+    flt = ctx.sql(FLOAT_JOIN).to_pandas()
+    # Index join keeps every destination cell; the float-equality join drops the
+    # cells whose float32-roundtripped coord no longer bit-matches the source.
+    assert len(idx) == 5
+    assert len(flt) < 5
+
+
+def test_index_column_prunes_partitions():
+    """`WHERE <dim>_idx IN (...)` prunes partitions, not just filters rows.
+
+    Register one partition per time step, then query a small `time_idx` set and
+    assert (via the reader's iteration callback) that only the matching
+    partitions are ever instantiated — the row filter alone would be correct but
+    would scan every partition.
+    """
+    n_time = 20
+    ds = xr.Dataset(
+        {
+            "v": (
+                ("time", "x"),
+                np.arange(n_time * 3, dtype=float).reshape(n_time, 3),
+            )
+        },
+        # time coordinate values are deliberately unrelated to the 0..n-1 index,
+        # so pruning can only succeed via time_idx (not via the time coord).
+        coords={
+            "time": 100 + np.arange(n_time) * 10,
+            "x": np.array([0.0, 1.0, 2.0]),
+        },
+    )
+
+    scanned: list = []
+    table = read_xarray_table(
+        ds,
+        chunks={"time": 1},  # one partition per time step
+        index_columns=True,
+        _iteration_callback=lambda block, projection: scanned.append(block),
+    )
+    ctx = SessionContext()
+    ctx.register_table("t", table)
+
+    wanted = {2, 5, 7}
+    result = ctx.sql(
+        f"SELECT time_idx, v FROM t WHERE time_idx IN ({', '.join(map(str, sorted(wanted)))})"
+    ).to_arrow_table()
+
+    # Correct rows...
+    assert set(result.column("time_idx").to_pylist()) == wanted
+    # ...and only the matching partitions were scanned (pruning, not full scan).
+    scanned_starts = {b["time"].start for b in scanned}
+    assert scanned_starts == wanted, (
+        f"expected only partitions {wanted} scanned, got {scanned_starts}"
+    )

xarray_sql/df.py

@@ -249,7 +249,10 @@ def pivot(ds: xr.Dataset) -> pd.DataFrame:
 
 
 def dataset_to_record_batch(
-    ds: xr.Dataset, schema: pa.Schema
+    ds: xr.Dataset,
+    schema: pa.Schema,
+    *,
+    index_offsets: Mapping[str, int] | None = None,
 ) -> pa.RecordBatch:
     """Convert an xarray Dataset partition to an Arrow RecordBatch.
 
@@ -285,10 +288,24 @@ def dataset_to_record_batch(
         dim_names = list(ds.sizes.keys())
         shape = tuple(ds.sizes[d] for d in dim_names)
 
+    offsets = index_offsets or {}
+    index_fields = {
+        index_column_name(str(d)): (k, int(offsets.get(d, 0)))
+        for k, d in enumerate(dim_names)
+    }
+
     arrays = []
     for field in schema:
         name = field.name
-        if name in ds.coords and name in ds.dims:
+        if name in index_fields:
+            # Absolute integer grid position for this dim, broadcast+ravelled.
+            axis, offset = index_fields[name]
+            idx = offset + np.arange(shape[axis], dtype=np.int64)
+            reshape = [1] * len(shape)
+            reshape[axis] = shape[axis]
+            arr = np.broadcast_to(idx.reshape(reshape), shape).ravel()
+            arrays.append(pa.array(arr, type=field.type))
+        elif 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
@@ -350,6 +367,8 @@ def iter_record_batches(
     ds: xr.Dataset,
     schema: pa.Schema,
     batch_size: int = DEFAULT_BATCH_SIZE,
+    *,
+    index_offsets: Mapping[str, int] | None = None,
 ) -> Iterator[pa.RecordBatch]:
     """Yield RecordBatches of at most *batch_size* rows from a partition Dataset.
 
@@ -407,11 +426,22 @@ def iter_record_batches(
     # 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))]
 
+    # Integer grid-position (`<dim>_idx`) columns, if the schema requests them:
+    # map the column name to its dim position and absolute block offset so the
+    # emitted index is global (partition-independent), which is what a join key
+    # across partitions requires. The index is computed from position, never
+    # loaded, so it is not a data-variable column below.
+    offsets = index_offsets or {}
+    index_fields = {
+        index_column_name(str(d)): (k, int(offsets.get(d, 0)))
+        for k, d in enumerate(dim_names)
+    }
+
     # 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:
+        if field.name not in ds.dims and field.name not in index_fields:
             raw = ds[field.name].values
             if cft.is_cftime(raw):
                 data_arrays[field.name] = cft.convert_for_field(raw, field)
@@ -428,7 +458,14 @@ def iter_record_batches(
         full_arrays = []
         for field in schema:
             name = field.name
-            if name in ds.coords and name in ds.dims:
+            if name in index_fields:
+                k, offset = index_fields[name]
+                outer = int(np.prod(shape[:k]))
+                col = np.repeat(np.arange(shape[k]), strides[k])
+                if outer > 1:
+                    col = np.tile(col, outer)
+                full_arrays.append(_as_single_array(offset + col, field.type))
+            elif name in ds.coords and name in ds.dims:
                 k = dim_names.index(name)
                 outer = int(np.prod(shape[:k]))
                 col = np.repeat(coord_values[name], strides[k])
@@ -455,7 +492,11 @@ def iter_record_batches(
         arrays = []
         for field in schema:
             name = field.name
-            if name in ds.coords and name in ds.dims:
+            if name in index_fields:
+                k, offset = index_fields[name]
+                idx = offset + (row_idx // strides[k]) % shape[k]
+                arrays.append(_as_single_array(idx, field.type))
+            elif name in ds.coords and name in ds.dims:
                 k = dim_names.index(name)
                 coord_idx = (row_idx // strides[k]) % shape[k]
                 arrays.append(
@@ -487,14 +528,30 @@ def _arrow_type_for_object(values: np.ndarray) -> pa.DataType:
     return pa.array(np.asarray(values).ravel()).type
 
 
-def _parse_schema(ds: xr.Dataset) -> pa.Schema:
+#: Suffix for the integer grid-position column of a dimension.
+INDEX_COLUMN_SUFFIX = "_idx"
+
+
+def index_column_name(dim: str) -> str:
+    """Name of the integer grid-position column for dimension ``dim``."""
+    return f"{dim}{INDEX_COLUMN_SUFFIX}"
+
+
+def _parse_schema(ds: xr.Dataset, *, index_columns: bool = False) -> pa.Schema:
     """Extracts a `pa.Schema` from the Dataset, treating dims and data_vars as columns.
 
     Only *dimension coordinates* become dimension columns, so a dimension
     without a coordinate would be dropped. Callers must run the Dataset through
     ``_ensure_default_indexes`` first (the readers do) so every dimension
     has a coordinate and appears as a column.
 
+    When ``index_columns`` is set, an ``int32`` ``<dim>_idx`` column is appended
+    for every dimension, carrying each row's absolute integer position along that
+    axis. These are exact integer keys for grid joins (regridding weight tables,
+    forecast alignment) — faster than, and free of the float-equality fragility
+    of, joining on the floating-point coordinate values, while the coordinate
+    columns remain available for value predicates and display.
+
     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.
@@ -542,6 +599,17 @@ def _parse_schema(ds: xr.Dataset) -> pa.Schema:
             pa_type = pa.from_numpy_dtype(var.dtype)
             columns.append(pa.field(var_name, pa_type))
 
+    if index_columns:
+        existing = {f.name for f in columns}
+        for dim in ds.dims:
+            name = index_column_name(str(dim))
+            if name in existing:
+                raise ValueError(
+                    f"cannot add index column {name!r}: a column with that "
+                    f"name already exists (dimension {dim!r})"
+                )
+            columns.append(pa.field(name, pa.int32()))
+
     return pa.schema(columns)
 
 
@@ -553,6 +621,8 @@ def _block_metadata(
     coord_arrays: dict,
     block: Block,
     dims: Iterable[Hashable] | None = None,
+    *,
+    index_columns: bool = False,
 ) -> PartitionBounds:
     """Compute min/max coordinate values for a single partition block.
 
@@ -577,6 +647,18 @@ def _block_metadata(
         coord_values = coord_arrays[str(dim)][slc]
         if len(coord_values) == 0:
             continue
+        if index_columns:
+            # Blocks are contiguous slices, so this partition's positions on
+            # `dim` are exactly [start, start + n - 1]. Emitting these int64
+            # bounds lets `WHERE <dim>_idx IN (...)` prune whole partitions,
+            # and it works for every coordinate dtype (including strings and
+            # out-of-range datetimes, whose values are not otherwise prunable).
+            start = slc.start or 0
+            ranges[index_column_name(str(dim))] = (
+                start,
+                start + len(coord_values) - 1,
+                "int64",
+            )
         # cftime coordinates are object dtype but carry their own bound
         # encoding, so they must be handled before the string/object skip
         # below (otherwise pruning is silently disabled for them).