diff --git a/tests/test_coord_lookup.py b/tests/test_coord_lookup.py new file mode 100644 index 0000000..7258261 --- /dev/null +++ b/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]) diff --git a/xarray_sql/ds.py b/xarray_sql/ds.py index 99ef7b3..94339e0 100644 --- a/xarray_sql/ds.py +++ b/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: