Replace searchsorted with a hash index in the reverse-pivot scatter

tests/test_coord_lookup.py

@@ -0,0 +1,167 @@
+"""Coordinate-value -> array-position lookup used by the reverse pivot.
+
+``to_dataset`` places each result row into a dense N-D array by
+resolving its dim-coord values to integer positions: an affine formula
+for uniformly spaced axes, a hash table for irregular unique axes, and
+an ``argsort``/``searchsorted`` fallback for axes a hash table cannot
+represent. These tests exercise each strategy through the public
+``to_dataset`` contract on out-of-order (shuffled) rows — the arrival
+order a parallel engine produces.
+
+Two behaviors are unreachable through the eager public path, because
+``to_dataset`` derives each axis from the same rows it scatters
+(``pd.unique``), so the axis can neither miss a row's value nor carry
+duplicates. Those two are covered at the ``_scatter_batches_to_ndarray``
+seam directly: the error raised for a value absent from the axis (which
+arises when an engine-backed lazy read's pre-computed coords go stale),
+and the duplicate-axis fallback.
+"""
+
+import numpy as np
+import pyarrow as pa
+import pytest
+import xarray as xr
+
+from xarray_sql import to_dataset
+from xarray_sql.ds import _scatter_batches_to_ndarray
+
+
+def _irregular_axis(n: int, seed: int = 0) -> np.ndarray:
+    """n unique, strictly increasing, non-uniformly spaced float64 values."""
+    rng = np.random.default_rng(seed)
+    return np.cumsum(rng.exponential(scale=1.0, size=n) + 1e-6)
+
+
+def _shuffled_table(template: xr.Dataset, seed: int = 1) -> pa.Table:
+    """The template's full grid as one Arrow table with rows shuffled."""
+    dims = list(next(iter(template.data_vars.values())).dims)
+    grids = np.meshgrid(*(template[d].values for d in dims), indexing="ij")
+    columns = {d: g.ravel() for d, g in zip(dims, grids)}
+    for name, var in template.data_vars.items():
+        columns[name] = var.values.ravel()
+    rng = np.random.default_rng(seed)
+    order = rng.permutation(len(columns[dims[0]]))
+    return pa.table({name: col[order] for name, col in columns.items()})
+
+
+def _roundtrip(template: xr.Dataset) -> xr.Dataset:
+    """Shuffle the template into rows, reconstruct, and sort back.
+
+    Output coordinate order follows first appearance in the (shuffled)
+    result — the behavior that lets an ORDER BY direction carry through
+    — so the reconstruction is sorted before comparing.
+    """
+    out = to_dataset(_shuffled_table(template), template=template)
+    dims = list(next(iter(template.data_vars.values())).dims)
+    return out.sortby(dims)
+
+
+def test_irregular_axis_roundtrip():
+    """Shuffled rows over an irregular (hash-strategy) axis reconstruct
+    the exact Dataset."""
+    template = xr.Dataset(
+        {
+            "v": (
+                ("time", "station"),
+                np.random.default_rng(2).standard_normal((6, 50)),
+            )
+        },
+        coords={
+            "time": np.arange(6, dtype="int64"),
+            "station": _irregular_axis(50),
+        },
+    )
+    xr.testing.assert_allclose(_roundtrip(template), template)
+
+
+def test_descending_uniform_axis_roundtrip():
+    """A descending uniformly spaced axis (affine strategy) reconstructs;
+    the descending order itself survives via first-appearance coords."""
+    template = xr.Dataset(
+        {
+            "v": (
+                ("lat",),
+                np.random.default_rng(3).standard_normal(19),
+            )
+        },
+        coords={"lat": np.linspace(90.0, -90.0, 19)},
+    )
+    # Unshuffled rows: the descending source order carries through as-is.
+    out = to_dataset(_shuffled_table(template, seed=0), template=template)
+    xr.testing.assert_allclose(out.sortby("lat"), template.sortby("lat"))
+
+
+def test_nan_dim_value_roundtrip():
+    """A NaN dim value in the result resolves to its own cell (the hash
+    strategy matches NaN by value equality)."""
+    station = np.array([2.0, np.nan, 5.0, 1.0])
+    template = xr.Dataset(
+        {"v": (("station",), np.array([10.0, 20.0, 30.0, 40.0]))},
+        coords={"station": station},
+    )
+    table = pa.table({"station": station, "v": template["v"].values})
+    out = to_dataset(table, template=template)
+    np.testing.assert_array_equal(out["v"].values, template["v"].values)
+
+
+def test_float16_axis_roundtrip():
+    """A float16 coordinate axis reconstructs through the sorted-search
+    strategy (pandas indexes do not support float16)."""
+    station = np.array([0.5, 1.5, 4.0, 9.0], dtype="float16")
+    template = xr.Dataset(
+        {"v": (("station",), np.array([1.0, 2.0, 3.0, 4.0], dtype="f4"))},
+        coords={"station": station},
+    )
+    table = pa.table(
+        {
+            "station": pa.array(station, type=pa.float16()),
+            "v": template["v"].values,
+        }
+    )
+    out = to_dataset(table, template=template)
+    np.testing.assert_array_equal(out["v"].values, template["v"].values)
+
+
+def _one_batch(station: np.ndarray, v: np.ndarray) -> list[pa.RecordBatch]:
+    return [
+        pa.RecordBatch.from_arrays(
+            [pa.array(station), pa.array(v)], names=["station", "v"]
+        )
+    ]
+
+
+def test_missing_value_raises_value_error():
+    """A row value absent from a unique irregular axis fails loudly
+    instead of scattering to a wrong cell."""
+    axis = _irregular_axis(10)
+    with pytest.raises(ValueError, match="dimension 'station'"):
+        _scatter_batches_to_ndarray(
+            batches=_one_batch(np.array([-1.0]), np.array([0.0], dtype="f4")),
+            dimension_columns=["station"],
+            requested={"station": axis},
+            var_name="v",
+            out_shape=(10,),
+            dtype=np.dtype("float32"),
+            drop_axes=[],
+        )
+
+
+def test_duplicate_axis_values_scatter_to_a_holding_position():
+    """An axis with duplicate values takes the sorted-search fallback:
+    each value lands on a position that holds it in the axis."""
+    axis = np.array([3.0, 1.0, 2.0, 1.0])  # 1.0 appears twice
+    out = _scatter_batches_to_ndarray(
+        batches=_one_batch(
+            np.array([1.0, 2.0, 3.0]), np.array([10.0, 20.0, 30.0], dtype="f4")
+        ),
+        dimension_columns=["station"],
+        requested={"station": axis},
+        var_name="v",
+        out_shape=(4,),
+        dtype=np.dtype("float32"),
+        drop_axes=[],
+    )
+    # 2.0 and 3.0 have unique positions; 10.0 landed on one of the two
+    # cells whose coordinate is 1.0.
+    assert out[2] == 20.0 and out[0] == 30.0
+    assert 10.0 in (out[1], out[3])

xarray_sql/ds.py

@@ -172,12 +172,92 @@ def _affine_axis(requested: np.ndarray) -> tuple[float, float] | None:
     predicted = numeric[0] + step * np.arange(len(numeric))
     # Written as a <= comparison so a NaN anywhere in the axis (e.g. a
     # NULL dim value in the result) fails the check and falls back to
-    # the searchsorted path, which handles it positionally.
+    # the non-affine lookup, whose hash strategy matches NaN by value
+    # equality (pandas index lookups treat NaN as equal to NaN).
     if not (np.abs(numeric - predicted) <= 0.25 * abs(step)).all():
         return None
     return float(numeric[0]), float(step)
 
 
+class _CoordLookup:
+    """Maps one dimension's coordinate values to their array positions.
+
+    Built once per dimension per reconstruction, then probed with every
+    batch's coordinate column. Three strategies, fastest applicable wins:
+
+    * **Affine** — the axis is uniformly spaced (the norm for rasters and
+      regular time steps, ascending or descending): the position is
+      ``rint((value - origin) / step)``, a fused vector op with no per-row
+      lookup at all.
+    * **Hash** — irregular axes (station networks, arbitrary point sets)
+      with unique values: a ``pd.Index`` built once; ``get_indexer``
+      probes its persistent hash table per batch, O(1) amortized per
+      row instead of a per-row binary search. Building the table once
+      here, outside the batch loop, is what makes this fast, and it is
+      only correct to reuse because the axis does not change between
+      batches. The table holds roughly twice the transient memory of
+      the sorted copy the search strategy needs; affine axes build
+      neither.
+    * **Sorted search** — axes the hash strategy cannot represent:
+      duplicate values (a unique-key table has no single position per
+      value) or dtypes ``pd.Index`` does not support (e.g. float16):
+      ``np.argsort`` once, then ``np.searchsorted`` per batch (O(log n)
+      per row). With duplicates, each value resolves to one of the
+      positions holding it (which one is unspecified), and rows
+      targeting the same cell overwrite in batch order — the scatter's
+      last-write-wins semantics.
+
+    Only the hash strategy can detect a probe value absent from the
+    axis; ``positions_for`` then raises ``ValueError`` — a symptom of a
+    filtered query whose coordinate discovery missed a value, which
+    would otherwise scatter to a wrong cell. The affine strategy rounds
+    such a value to the nearest grid position and the sorted strategy
+    resolves it to a neighbor, both preserving their historical
+    semantics.
+    """
+
+    def __init__(self, requested: np.ndarray) -> None:
+        self._affine: tuple[float, float] | None = _affine_axis(requested)
+        self._hash_index: pd.Index | None = None
+        self._sorted_idx: np.ndarray | None = None
+        self._sorted_req: np.ndarray | None = None
+        if self._affine is None:
+            index: pd.Index | None
+            try:
+                index = pd.Index(requested)
+            except (NotImplementedError, TypeError):
+                # Dtypes pandas cannot index (e.g. float16) take the
+                # sorted-search strategy, which only needs numpy
+                # comparisons.
+                index = None
+            if index is not None and index.is_unique:
+                self._hash_index = index
+            else:
+                self._sorted_idx = np.argsort(requested)
+                self._sorted_req = requested[self._sorted_idx]
+
+    def positions_for(self, vals: np.ndarray, dim: str) -> np.ndarray:
+        """Positions of ``vals`` within the axis, in axis order."""
+        if self._affine is not None:
+            origin, step = self._affine
+            pos = np.rint((_axis_numeric(vals) - origin) / step).astype(np.intp)
+            return cast(np.ndarray, pos)
+        if self._hash_index is not None:
+            pos = np.asarray(self._hash_index.get_indexer(vals))
+            if (pos < 0).any():
+                missing = vals[pos < 0]
+                raise ValueError(
+                    f"result contains {len(missing)} value(s) for dimension "
+                    f"{dim!r} not present in its coordinate array "
+                    f"(first: {missing[0]!r}); the query result does not "
+                    "match the reconstruction's coordinates."
+                )
+            return cast(np.ndarray, pos)
+        assert self._sorted_req is not None and self._sorted_idx is not None
+        pos_in_sorted = np.searchsorted(self._sorted_req, vals)
+        return cast(np.ndarray, self._sorted_idx[pos_in_sorted])
+
+
 def _scatter_batches_to_ndarray(
     batches: list[pa.RecordBatch],
     dimension_columns: list[str],
@@ -192,9 +272,9 @@ def _scatter_batches_to_ndarray(
     SQL query results arrive as flat rows; xarray expects N-D arrays.
     This bridges the two: each row carries the dim-coord values that
     identify its cell in the output cube plus the value to write there.
-    We look up the row's N-D position by binary-searching its coord
-    values within the caller's requested coord arrays
-    (``np.searchsorted``), then scatter-write the value at that index.
+    We look up each row's N-D position within the caller's requested
+    coord arrays (see ``_CoordLookup``), then scatter-write the value
+    at that index.
 
     Missing combinations (sparse results from filtered queries) stay as
     ``NaN`` for floating-point outputs by pre-filling the buffer; integer
@@ -209,21 +289,7 @@ def _scatter_batches_to_ndarray(
         else np.empty(out_shape, dtype=dtype)
     )
 
-    # ``requested[d]`` may be in any order (callers can iselect arbitrary
-    # positions, and template coords like air_temperature.lat are descending).
-    # ``np.searchsorted`` requires ascending input, so we sort each requested
-    # array once, search there, and remap back to the original positions.
-    # Uniformly spaced axes (the norm for rasters and regular time steps,
-    # ascending or descending) skip the search entirely: the position is
-    # ``rint((value - origin) / step)``, a fused vector op several times
-    # faster than a per-row binary search.
-    affine = {d: _affine_axis(requested[d]) for d in dimension_columns}
-    sorted_idx = {
-        d: np.argsort(requested[d])
-        for d in dimension_columns
-        if affine[d] is None
-    }
-    sorted_req = {d: requested[d][sorted_idx[d]] for d in sorted_idx}
+    lookups = {d: _CoordLookup(requested[d]) for d in dimension_columns}
 
     for batch in batches:
         if batch.num_rows == 0:
@@ -235,16 +301,7 @@ def _scatter_batches_to_ndarray(
         for d in dimension_columns:
             col_arr = batch.column(schema_names.index(d))
             vals = col_arr.to_numpy(zero_copy_only=False)
-            pair = affine[d]
-            if pair is not None:
-                origin, step = pair
-                pos = np.rint((_axis_numeric(vals) - origin) / step).astype(
-                    np.intp
-                )
-                positions.append(pos)
-            else:
-                pos_in_sorted = np.searchsorted(sorted_req[d], vals)
-                positions.append(sorted_idx[d][pos_in_sorted])
+            positions.append(lookups[d].positions_for(vals, dim=d))
         value_arr = batch.column(schema_names.index(var_name)).to_numpy(
             zero_copy_only=False
         )
@@ -293,11 +350,11 @@ class SQLBackendArray(xr.backends.BackendArray):
             filter/project/execute chain if a predicate refers to a
             missing column, the dtype of a literal is incompatible, or
             the execution itself fails.
-        AssertionError: from ``np.searchsorted`` mis-alignment, which
-            indicates the result contains coordinate values not present
-            in the wrapper's pre-computed coord arrays -- usually a
-            symptom of a filtered query whose coord discovery missed a
-            value.
+        ValueError: from the coordinate lookup (``_CoordLookup``) when
+            the result contains values for an irregular unique axis
+            that are not present in the wrapper's pre-computed coord
+            arrays -- usually a symptom of a filtered query whose coord
+            discovery missed a value.
 
     Constructed by ``_build_lazy_scan``; users should not instantiate
     this class directly.
@@ -462,8 +519,8 @@ def _c_order_grid(
     dimension column is its coordinates repeated/tiled in C order — the
     shape any unfiltered or bbox-windowed scan produces. When it holds,
     data variables are dense row-major arrays already and can be
-    reshaped instead of scatter-written (one memcpy versus a
-    ``searchsorted`` per dimension per row).
+    reshaped instead of scatter-written (one memcpy versus a per-row
+    position lookup and write; see ``_CoordLookup``).
     """
     shape = tuple(len(coord_arrays[d]) for d in dimension_columns)
     if total_rows != int(np.prod(shape)) or total_rows == 0: