From 8af8d7dd7223a2c9548368e10c81c70e404dae6d Mon Sep 17 00:00:00 2001 From: Miguel Moncada Date: Mon, 10 Aug 2026 15:41:44 +0200 Subject: [PATCH 1/2] Replace searchsorted with a hash index in the reverse-pivot scatter The reverse pivot resolves each result row's dim-coord values to array positions before scatter-writing into the dense output. Irregular (non-uniformly-spaced) axes previously used np.argsort once plus np.searchsorted per batch: O(log n) per row. They now use a pd.Index whose hash table is built once per dimension and probed per batch with get_indexer: O(1) amortized per row. Measured 2.1-3.3x faster reconstruction on shuffled irregular-axis results from 24K to 24M cells (the speedup grows with axis cardinality); uniformly spaced axes keep the existing affine fast path and regular-grid workloads (e.g. ERA5 lat/lon/time) are unaffected. The three strategies now live in one place, _CoordLookup: * affine formula for uniformly spaced axes (unchanged); * hash index for irregular axes with unique values; * the previous argsort+searchsorted for axes with duplicate values, which a unique-key hash table cannot represent. A result value absent from the axis now raises a ValueError naming the dimension (previously a searchsorted misplacement surfaced as an opaque AssertionError or a silent wrong-cell write). Co-Authored-By: Claude Fable 5 --- tests/test_coord_lookup.py | 156 +++++++++++++++++++++++++++++++++++++ xarray_sql/ds.py | 122 ++++++++++++++++++++--------- 2 files changed, 242 insertions(+), 36 deletions(-) create mode 100644 tests/test_coord_lookup.py diff --git a/tests/test_coord_lookup.py b/tests/test_coord_lookup.py new file mode 100644 index 0000000..c534a81 --- /dev/null +++ b/tests/test_coord_lookup.py @@ -0,0 +1,156 @@ +"""Coordinate-value -> array-position lookup used by the reverse pivot. + +``_scatter_batches_to_ndarray`` places each result row into a dense N-D +array by resolving its dim-coord values to integer positions through +``_CoordLookup``: an affine formula for uniformly spaced axes, a hash +table (``pd.Index``, built once, probed per batch) for irregular unique +axes, and an ``argsort``/``searchsorted`` fallback for axes with +duplicate values. These tests pin the correctness of each strategy on +out-of-order (shuffled) input -- the row order a parallel engine +produces -- plus the error contract for values absent from the axis. +""" + +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 _CoordLookup, _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_batches( + time: np.ndarray, station: np.ndarray, values: np.ndarray, batch_size: int +) -> list[pa.RecordBatch]: + """The full (time, station) grid as shuffled row batches.""" + tt, ss = np.meshgrid(time, station, indexing="ij") + flat_t, flat_s, flat_v = tt.ravel(), ss.ravel(), values.ravel() + rng = np.random.default_rng(1) + order = rng.permutation(flat_v.shape[0]) + flat_t, flat_s, flat_v = flat_t[order], flat_s[order], flat_v[order] + schema = pa.schema( + [ + ("time", pa.from_numpy_dtype(time.dtype)), + ("station", pa.from_numpy_dtype(station.dtype)), + ("v", pa.from_numpy_dtype(values.dtype)), + ] + ) + return [ + pa.RecordBatch.from_arrays( + [ + pa.array(flat_t[i : i + batch_size]), + pa.array(flat_s[i : i + batch_size]), + pa.array(flat_v[i : i + batch_size]), + ], + schema=schema, + ) + for i in range(0, flat_v.shape[0], batch_size) + ] + + +def test_irregular_axis_scatter_matches_reference(): + """Shuffled rows over an irregular (hash-path) axis land correctly.""" + time = np.arange(6, dtype="int64") + station = _irregular_axis(50) + values = np.random.default_rng(2).standard_normal((6, 50)).astype("f4") + batches = _shuffled_batches(time, station, values, batch_size=37) + + out = _scatter_batches_to_ndarray( + batches=batches, + dimension_columns=["time", "station"], + requested={"time": time, "station": station}, + var_name="v", + out_shape=(6, 50), + dtype=np.dtype("float32"), + drop_axes=[], + ) + np.testing.assert_array_equal(out, values) + + +def test_descending_affine_axis_unchanged(): + """A descending uniformly spaced axis stays on the affine path.""" + lat = np.linspace(90.0, -90.0, 19) # descending, uniform + lookup = _CoordLookup(lat) + assert lookup._affine is not None + pos = lookup.positions_for(np.array([90.0, 0.0, -90.0]), dim="lat") + np.testing.assert_array_equal(pos, [0, 9, 18]) + + +def test_missing_value_raises_value_error(): + """A result value absent from the axis is a coordinate-discovery bug; + it must fail loudly instead of scattering to a wrong cell.""" + station = _irregular_axis(10) + lookup = _CoordLookup(station) + assert lookup._hash_index is not None + with pytest.raises(ValueError, match="dimension 'station'"): + lookup.positions_for(np.array([-1.0]), dim="station") + + +def test_duplicate_axis_values_fall_back_to_search(): + """An axis with duplicate values cannot key a unique hash table; the + sorted-search fallback keeps every value resolving to a position that + holds it (which of the duplicate positions is returned is + unspecified).""" + axis = np.array([3.0, 1.0, 2.0, 1.0]) # 1.0 appears twice + lookup = _CoordLookup(axis) + assert lookup._hash_index is None and lookup._sorted_req is not None + pos = lookup.positions_for(np.array([1.0, 2.0, 3.0]), dim="x") + assert axis[pos[0]] == 1.0 + assert axis[pos[1]] == 2.0 + assert axis[pos[2]] == 3.0 + + +def test_nan_in_irregular_axis_resolves(): + """A NaN dim value in the result resolves to the axis's NaN position + (pandas index lookups treat NaN as equal to NaN).""" + axis = np.array([2.0, np.nan, 5.0, 1.0]) # non-affine (NaN breaks it) + lookup = _CoordLookup(axis) + assert lookup._hash_index is not None + pos = lookup.positions_for(np.array([np.nan, 1.0]), dim="x") + np.testing.assert_array_equal(pos, [1, 3]) + + +def test_to_dataset_roundtrips_shuffled_irregular_result(): + """End to end through the engine-agnostic ``to_dataset``: a shuffled + Arrow result over an irregular axis reconstructs the exact Dataset.""" + time = np.arange(4, dtype="int64") + station = _irregular_axis(30, seed=3) + values = np.random.default_rng(4).standard_normal((4, 30)).astype("f8") + template = xr.Dataset( + {"v": (("time", "station"), values)}, + coords={"time": time, "station": station}, + ) + + tt, ss = np.meshgrid(time, station, indexing="ij") + rng = np.random.default_rng(5) + order = rng.permutation(values.size) + table = pa.table( + { + "time": tt.ravel()[order], + "station": ss.ravel()[order], + "v": values.ravel()[order], + } + ) + + out = to_dataset(table, dims=["time", "station"], template=template) + # Coordinate order follows first appearance in the (shuffled) result -- + # the documented behavior that lets an ORDER BY direction carry through + # -- so compare on a common sort. + xr.testing.assert_allclose(out.sortby(["time", "station"]), template) + + +def test_hash_index_reused_across_batches(): + """The pandas hash index is built once per reconstruction, not per + batch -- the property the speedup rests on.""" + station = _irregular_axis(100) + lookup = _CoordLookup(station) + first = lookup._hash_index + lookup.positions_for(station[:10], dim="station") + lookup.positions_for(station[50:60], dim="station") + assert lookup._hash_index is first diff --git a/xarray_sql/ds.py b/xarray_sql/ds.py index 99ef7b3..8543574 100644 --- a/xarray_sql/ds.py +++ b/xarray_sql/ds.py @@ -172,12 +172,85 @@ 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 with duplicate values, where a unique-key + hash table cannot represent the value-to-position mapping: + ``np.argsort`` once, then ``np.searchsorted`` per batch (O(log n) + per row). Duplicate dim values only reach this code from a + pathological result (``to_dataset`` raises on duplicate dim tuples + earlier on the main paths), but the fallback keeps the mapping + well-defined: each value resolves to one of the positions holding + it (which one is unspecified). + + 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(requested) + if 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 +265,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 +282,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 +294,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 +343,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 +512,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: From e1acfb7f18836e1b6e50dbbb9518904737a7e223 Mon Sep 17 00:00:00 2001 From: Miguel Moncada Date: Mon, 10 Aug 2026 16:22:36 +0200 Subject: [PATCH 2/2] Address review: public-contract tests, float16 fallback, honest duplicate semantics - Tests now exercise each lookup strategy through the public to_dataset contract (values, dims, coords) on shuffled rows, with dims inferred from the template. Only two behaviors stay at the _scatter_batches_to_ndarray seam, with the reason documented: the missing-value error and the duplicate-axis fallback, both unreachable through the eager public path because to_dataset derives each axis from the same rows it scatters. - pd.Index construction falls back to sorted search for dtypes pandas cannot index (float16 raises NotImplementedError on pandas 2.3.0), preserving the previous behavior for those axes; regression-tested through to_dataset. - The _CoordLookup docstring no longer claims to_dataset raises on duplicate dim tuples (no reconstruction path does); it now describes the actual behavior: sorted-search resolution to one of the holding positions plus the scatter's last-write-wins overwrite. Co-Authored-By: Claude Fable 5 --- tests/test_coord_lookup.py | 253 +++++++++++++++++++------------------ xarray_sql/ds.py | 25 ++-- 2 files changed, 148 insertions(+), 130 deletions(-) diff --git a/tests/test_coord_lookup.py b/tests/test_coord_lookup.py index c534a81..7258261 100644 --- a/tests/test_coord_lookup.py +++ b/tests/test_coord_lookup.py @@ -1,13 +1,20 @@ """Coordinate-value -> array-position lookup used by the reverse pivot. -``_scatter_batches_to_ndarray`` places each result row into a dense N-D -array by resolving its dim-coord values to integer positions through -``_CoordLookup``: an affine formula for uniformly spaced axes, a hash -table (``pd.Index``, built once, probed per batch) for irregular unique -axes, and an ``argsort``/``searchsorted`` fallback for axes with -duplicate values. These tests pin the correctness of each strategy on -out-of-order (shuffled) input -- the row order a parallel engine -produces -- plus the error contract for values absent from the axis. +``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 @@ -16,7 +23,7 @@ import xarray as xr from xarray_sql import to_dataset -from xarray_sql.ds import _CoordLookup, _scatter_batches_to_ndarray +from xarray_sql.ds import _scatter_batches_to_ndarray def _irregular_axis(n: int, seed: int = 0) -> np.ndarray: @@ -25,132 +32,136 @@ def _irregular_axis(n: int, seed: int = 0) -> np.ndarray: return np.cumsum(rng.exponential(scale=1.0, size=n) + 1e-6) -def _shuffled_batches( - time: np.ndarray, station: np.ndarray, values: np.ndarray, batch_size: int -) -> list[pa.RecordBatch]: - """The full (time, station) grid as shuffled row batches.""" - tt, ss = np.meshgrid(time, station, indexing="ij") - flat_t, flat_s, flat_v = tt.ravel(), ss.ravel(), values.ravel() - rng = np.random.default_rng(1) - order = rng.permutation(flat_v.shape[0]) - flat_t, flat_s, flat_v = flat_t[order], flat_s[order], flat_v[order] - schema = pa.schema( - [ - ("time", pa.from_numpy_dtype(time.dtype)), - ("station", pa.from_numpy_dtype(station.dtype)), - ("v", pa.from_numpy_dtype(values.dtype)), - ] - ) - return [ - pa.RecordBatch.from_arrays( - [ - pa.array(flat_t[i : i + batch_size]), - pa.array(flat_s[i : i + batch_size]), - pa.array(flat_v[i : i + batch_size]), - ], - schema=schema, - ) - for i in range(0, flat_v.shape[0], batch_size) - ] +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 test_irregular_axis_scatter_matches_reference(): - """Shuffled rows over an irregular (hash-path) axis land correctly.""" - time = np.arange(6, dtype="int64") - station = _irregular_axis(50) - values = np.random.default_rng(2).standard_normal((6, 50)).astype("f4") - batches = _shuffled_batches(time, station, values, batch_size=37) +def _roundtrip(template: xr.Dataset) -> xr.Dataset: + """Shuffle the template into rows, reconstruct, and sort back. - out = _scatter_batches_to_ndarray( - batches=batches, - dimension_columns=["time", "station"], - requested={"time": time, "station": station}, - var_name="v", - out_shape=(6, 50), - dtype=np.dtype("float32"), - drop_axes=[], - ) - np.testing.assert_array_equal(out, values) + 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_descending_affine_axis_unchanged(): - """A descending uniformly spaced axis stays on the affine path.""" - lat = np.linspace(90.0, -90.0, 19) # descending, uniform - lookup = _CoordLookup(lat) - assert lookup._affine is not None - pos = lookup.positions_for(np.array([90.0, 0.0, -90.0]), dim="lat") - np.testing.assert_array_equal(pos, [0, 9, 18]) +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_missing_value_raises_value_error(): - """A result value absent from the axis is a coordinate-discovery bug; - it must fail loudly instead of scattering to a wrong cell.""" - station = _irregular_axis(10) - lookup = _CoordLookup(station) - assert lookup._hash_index is not None - with pytest.raises(ValueError, match="dimension 'station'"): - lookup.positions_for(np.array([-1.0]), dim="station") +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_duplicate_axis_values_fall_back_to_search(): - """An axis with duplicate values cannot key a unique hash table; the - sorted-search fallback keeps every value resolving to a position that - holds it (which of the duplicate positions is returned is - unspecified).""" - axis = np.array([3.0, 1.0, 2.0, 1.0]) # 1.0 appears twice - lookup = _CoordLookup(axis) - assert lookup._hash_index is None and lookup._sorted_req is not None - pos = lookup.positions_for(np.array([1.0, 2.0, 3.0]), dim="x") - assert axis[pos[0]] == 1.0 - assert axis[pos[1]] == 2.0 - assert axis[pos[2]] == 3.0 - - -def test_nan_in_irregular_axis_resolves(): - """A NaN dim value in the result resolves to the axis's NaN position - (pandas index lookups treat NaN as equal to NaN).""" - axis = np.array([2.0, np.nan, 5.0, 1.0]) # non-affine (NaN breaks it) - lookup = _CoordLookup(axis) - assert lookup._hash_index is not None - pos = lookup.positions_for(np.array([np.nan, 1.0]), dim="x") - np.testing.assert_array_equal(pos, [1, 3]) - - -def test_to_dataset_roundtrips_shuffled_irregular_result(): - """End to end through the engine-agnostic ``to_dataset``: a shuffled - Arrow result over an irregular axis reconstructs the exact Dataset.""" - time = np.arange(4, dtype="int64") - station = _irregular_axis(30, seed=3) - values = np.random.default_rng(4).standard_normal((4, 30)).astype("f8") +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": (("time", "station"), values)}, - coords={"time": time, "station": station}, + {"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) + - tt, ss = np.meshgrid(time, station, indexing="ij") - rng = np.random.default_rng(5) - order = rng.permutation(values.size) +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( { - "time": tt.ravel()[order], - "station": ss.ravel()[order], - "v": values.ravel()[order], + "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=[], + ) - out = to_dataset(table, dims=["time", "station"], template=template) - # Coordinate order follows first appearance in the (shuffled) result -- - # the documented behavior that lets an ORDER BY direction carry through - # -- so compare on a common sort. - xr.testing.assert_allclose(out.sortby(["time", "station"]), template) - - -def test_hash_index_reused_across_batches(): - """The pandas hash index is built once per reconstruction, not per - batch -- the property the speedup rests on.""" - station = _irregular_axis(100) - lookup = _CoordLookup(station) - first = lookup._hash_index - lookup.positions_for(station[:10], dim="station") - lookup.positions_for(station[50:60], dim="station") - assert lookup._hash_index is first + +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 8543574..94339e0 100644 --- a/xarray_sql/ds.py +++ b/xarray_sql/ds.py @@ -198,14 +198,14 @@ class _CoordLookup: batches. The table holds roughly twice the transient memory of the sorted copy the search strategy needs; affine axes build neither. - * **Sorted search** — axes with duplicate values, where a unique-key - hash table cannot represent the value-to-position mapping: + * **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). Duplicate dim values only reach this code from a - pathological result (``to_dataset`` raises on duplicate dim tuples - earlier on the main paths), but the fallback keeps the mapping - well-defined: each value resolves to one of the positions holding - it (which one is unspecified). + 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 @@ -222,8 +222,15 @@ def __init__(self, requested: np.ndarray) -> None: self._sorted_idx: np.ndarray | None = None self._sorted_req: np.ndarray | None = None if self._affine is None: - index = pd.Index(requested) - if index.is_unique: + 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)