fix: four adversarial-review findings

- NaN/NaT in a chunk's dimension coordinate poisoned the shadow
  guarantee into (dim >= NaN), which Arrow simplifies every predicate
  against as false — chunks with matching rows were silently pruned.
  Such spans now carry an always-true guarantee (unprunable).
- _affine_axis accepted NaN-poisoned axes (NaN > tol is false), letting
  the affine scatter place values in wrong cells when a result's dim
  column contains NULLs. The acceptance test is now a <= .all() so NaN
  rejects and the positional searchsorted path handles it.
- Projected scans omit dimension columns from the scan schema; the
  cftime coordinate-conversion loop assumed every dim had a schema
  field and raised KeyError on e.g. SELECT SUM(v) over a 360_day
  dataset. Dims absent from the schema are now skipped.
- pyramid() rebinned float cell origins level-over-level, occasionally
  aliasing boundary points into a neighboring cell. Cells are now
  tracked as integer indices that halve exactly per level, with float
  origins kept as query labels.

Also: docs notes from stress testing (Polars is_in float-literal
pushdown caveat, per-thread DuckDB cursor re-registration pattern).
Regression tests for all four fixes; battle-test matrix passed on
duckdb 1.4.5 LTS and 1.5.4, 4-thread concurrency, dtype zoo, and a
600-query memory soak.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

Author: 2 people

docs/engines.md

@@ -131,6 +131,12 @@ Polars pushes its predicate and column selection into the dataset scan
 (verified: a filtered group-by read 1 of 20 chunks and 3 of 5 columns),
 and its results round-trip through `xql.to_dataset` unchanged.
 
+One upstream caveat: Polars' translation of `is_in` **float** literals
+into pyarrow expressions loses precision and can drop matching rows
+(reproducible without xarray-sql). Prefer range predicates
+(`is_between`) for float coordinates; integer and timestamp value sets
+are unaffected.
+
 ### Relation to duckdb-zarr
 
 [duckdb-zarr](https://github.com/xqlsystems/duckdb-zarr) reads Zarr

docs/performance.md

@@ -89,6 +89,24 @@ only the variables a query references are read. Corollaries:
 - A query with no `WHERE` on dimension columns is a full scan on any
   engine; that's physics, not a missing optimization.
 
+## Threads and DuckDB connections
+
+Registered Python objects are connection-local in DuckDB: `con.cursor()`
+does not inherit them, and one connection's result slot is not
+thread-safe. For multithreaded querying, give each thread its own
+cursor and register the *same* dataset object on it:
+
+```python
+dataset = xql.arrow_dataset(ds)
+def worker():
+    cur = con.cursor()
+    cur.register("t", dataset)   # cheap; shares the pruning index
+    ...
+```
+
+The dataset object itself is safe to share across threads (verified
+under concurrent query load).
+
 ## Stop re-scanning: materialize and pyramid
 
 Repeated statistics should pay the scan once:

tests/test_duckdb_backend.py

@@ -344,3 +344,44 @@ def test_register_dispatches_to_datafusion():
 def test_register_rejects_unknown_connection(ds):
     with pytest.raises(TypeError, match="No xarray-sql engine adapter"):
         xql.register(object(), "weather", ds)
+
+
+def test_nan_coordinate_chunk_is_not_pruned():
+    # NaN in a chunk's coordinate must disable pruning for that span,
+    # never poison the range guarantee (which would silently drop rows).
+    ds = xr.Dataset(
+        {"v": (["lat"], np.arange(6.0))},
+        coords={"lat": [np.nan, 5.0, 10.0, 20.0, 30.0, 40.0]},
+    ).chunk({"lat": 2})
+    con = duckdb.connect()
+    xql.register(con, "t", ds)
+    assert con.sql("SELECT v FROM t WHERE lat = 5.0").fetchall() == [(1.0,)]
+
+
+def test_cftime_dataset_aggregates_under_projection():
+    cftime = pytest.importorskip("cftime")
+
+    times = xr.date_range(
+        "2000-01-01", periods=6, calendar="360_day", use_cftime=True
+    )
+    ds = xr.Dataset(
+        {"v": (["time"], np.arange(6.0))}, coords={"time": times}
+    ).chunk({"time": 3})
+    con = duckdb.connect()
+    xql.register(con, "t", ds)
+    # The scan projects only `v`; the cftime dim column is absent from
+    # the scan schema but still shapes the iteration.
+    assert con.sql("SELECT SUM(v) FROM t").fetchone()[0] == 15.0
+
+
+def test_null_dimension_value_round_trips_positionally():
+    # A NULL in a result's dim column must reject the affine fast path
+    # and fall back to positional scatter.
+    table = pa.table(
+        {
+            "lat": pa.array([0.0, None, 1.0, 3.0], type=pa.float64()),
+            "v": [10.0, 99.0, 11.0, 13.0],
+        }
+    )
+    out = xql.to_dataset(table, dims=["lat"])
+    np.testing.assert_allclose(out["v"].values, [10.0, 99.0, 11.0, 13.0])

tests/test_materialize.py

@@ -140,3 +140,23 @@ def test_pyramid_filter_prunes_source_scan(con):
     n = _rows(con, "SELECT SUM(n) FROM pyr")[0][0]
     direct = _rows(con, "SELECT COUNT(*) FROM grid WHERE y > -32")[0][0]
     assert n == direct
+
+
+def test_pyramid_cell_membership_is_exact_across_levels():
+    # Rebinning float origins level-over-level can alias points across
+    # cell boundaries; integer indices must halve exactly instead.
+    import math
+
+    con = duckdb.connect()
+    con.execute("CREATE TABLE pts AS SELECT -130.943 AS x, 0.0005 AS y")
+    xql.pyramid(
+        con,
+        "pyr",
+        "pts",
+        aggs={"n": ("count", "*")},
+        base_cell=0.001,
+        levels=2,
+    )
+    rows = con.sql("SELECT level, x_idx FROM pyr ORDER BY level").fetchall()
+    assert rows[0][1] == math.floor(-130.943 / 0.001)
+    assert rows[1][1] == math.floor(-130.943 / 0.002)

xarray_sql/backends/pyarrow.py

@@ -122,13 +122,22 @@ def _build(self, spans: list[tuple[int, int]]) -> pads.FileSystemDataset:
         fragments = []
         for i, (lo_chunk, hi_chunk) in enumerate(spans):
             vals = self._coord[self._bounds[lo_chunk] : self._bounds[hi_chunk]]
-            # min/max (not first/last) so descending axes like latitude
-            # 90→-90 carry correct ranges.
-            lo = pa.scalar(vals.min(), type=self._field_type)
-            hi = pa.scalar(vals.max(), type=self._field_type)
-            guarantee = (pc.field(self._name) >= lo) & (
-                pc.field(self._name) <= hi
-            )
+            if (vals.dtype.kind == "f" and np.isnan(vals).any()) or (
+                vals.dtype.kind == "M" and np.isnat(vals).any()
+            ):
+                # NaN/NaT poisons min/max into a (dim >= NaN) guarantee
+                # that Arrow simplifies every predicate against as false,
+                # silently pruning rows. An always-true guarantee keeps
+                # the span unprunable instead.
+                guarantee = pc.scalar(True)
+            else:
+                # min/max (not first/last) so descending axes like
+                # latitude 90→-90 carry correct ranges.
+                lo = pa.scalar(vals.min(), type=self._field_type)
+                hi = pa.scalar(vals.max(), type=self._field_type)
+                guarantee = (pc.field(self._name) >= lo) & (
+                    pc.field(self._name) <= hi
+                )
             fragments.append(
                 fmt.make_fragment(str(i), fs, partition_expression=guarantee)
             )

xarray_sql/df.py

@@ -349,8 +349,12 @@ def iter_record_batches(
 
     # Preload small 1-D coordinate arrays (negligible memory).
     # Convert cftime objects to numeric values matching the schema type.
+    # Projected scans may omit dimension columns from the schema; those
+    # dims still shape the iteration but never emit a column.
     coord_values = {}
     for name in dim_names:
+        if name not in schema.names:
+            continue
         vals = ds.coords[name].values
         if cft.is_cftime(vals):
             coord_values[name] = cft.convert_for_field(vals, schema.field(name))

xarray_sql/ds.py

@@ -169,7 +169,10 @@ def _affine_axis(requested: np.ndarray) -> tuple[float, float] | None:
     if step == 0 or not np.isfinite(step):
         return None
     predicted = numeric[0] + step * np.arange(len(numeric))
-    if np.abs(numeric - predicted).max() > 0.25 * abs(step):
+    # 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.
+    if not (np.abs(numeric - predicted) <= 0.25 * abs(step)).all():
         return None
     return float(numeric[0]), float(step)
 

xarray_sql/materialize.py

@@ -110,11 +110,14 @@ def pyramid(
     the total cost beyond the single scan is negligible. The result is
     one long table::
 
-        (level INTEGER, x_bin DOUBLE, y_bin DOUBLE, <agg columns...>)
+        (level, x_idx BIGINT, y_idx BIGINT, x_bin, y_bin, <aggs...>)
 
-    where ``x_bin``/``y_bin`` are the cell origins at that level's cell
-    size (``base_cell * 2**level``). Query the level whose cells match
-    the resolution you need::
+    ``x_idx``/``y_idx`` are exact integer cell indices at that level
+    (they halve from level to level, so cell membership never drifts
+    across levels); ``x_bin``/``y_bin`` are the float cell origins
+    (``idx * base_cell * 2**level``) for human-readable querying —
+    filter them with ranges rather than equality. Query the level whose
+    cells match the resolution you need::
 
         SELECT x_bin, y_bin, class4_n / n AS share
         FROM grid_pyramid
@@ -160,19 +163,32 @@ def pyramid(
 
     run_sql = get_adapter(con).run_sql
 
+    # Cells are tracked as integer indices (x_idx, y_idx) and only
+    # labeled with float origins (x_bin, y_bin) for querying: rebinning
+    # float origins at each level occasionally lands boundary points in
+    # a different cell than direct binning would (float aliasing);
+    # integer indices halve exactly at every level.
     # Level 0: the single scan of the source.
     run_sql(
         con,
         f"""
         CREATE OR REPLACE TABLE {_ident(name)} AS
         SELECT
             0 AS level,
-            FLOOR({_ident(x)} / {base_cell}) * {base_cell} AS x_bin,
-            FLOOR({_ident(y)} / {base_cell}) * {base_cell} AS y_bin,
+            x_idx,
+            y_idx,
+            x_idx * {base_cell} AS x_bin,
+            y_idx * {base_cell} AS y_bin,
             {base_aggs}
-        FROM {_ident(table)}
-        {where}
-        GROUP BY 2, 3
+        FROM (
+            SELECT
+                CAST(FLOOR({_ident(x)} / {base_cell}) AS BIGINT) AS x_idx,
+                CAST(FLOOR({_ident(y)} / {base_cell}) AS BIGINT) AS y_idx,
+                *
+            FROM {_ident(table)}
+            {where}
+        )
+        GROUP BY x_idx, y_idx
         """,
     )
 
@@ -183,18 +199,27 @@ def pyramid(
             f"{_ROLLUP[kind]}({_ident(n)}) AS {_ident(n)}"
             for n, (kind, _) in aggs.items()
         )
+        pass_cols = ", ".join(_ident(n) for n in aggs)
         run_sql(
             con,
             f"""
             INSERT INTO {_ident(name)}
             SELECT
                 {level} AS level,
-                FLOOR(x_bin / {cell}) * {cell} AS x_bin,
-                FLOOR(y_bin / {cell}) * {cell} AS y_bin,
+                x_idx,
+                y_idx,
+                x_idx * {cell} AS x_bin,
+                y_idx * {cell} AS y_bin,
                 {rollup_aggs}
-            FROM {_ident(name)}
-            WHERE level = {level - 1}
-            GROUP BY 2, 3
+            FROM (
+                SELECT
+                    CAST(FLOOR(x_idx / 2.0) AS BIGINT) AS x_idx,
+                    CAST(FLOOR(y_idx / 2.0) AS BIGINT) AS y_idx,
+                    {pass_cols}
+                FROM {_ident(name)}
+                WHERE level = {level - 1}
+            )
+            GROUP BY x_idx, y_idx
             """,
         )
     return con