Integer grid-position columns for exact grid joins

[alxmrs]

Add opt-in int32 <dim>_idx columns so grid joins (regridding, forecast alignment) can key on exact integer positions instead of floating-point coordinate values.

Why

Regridding is a sparse matmul expressed as JOIN source grid TO weight table ON the source coordinate. Keying that join on the float coordinate value is fragile: any sub-ULP drift — e.g. a reproject/interp UDF that computes in float32 (cases 07/09) — makes the equality join silently drop rows and return a wrong answer with no error.

A prototype at realistic scale (900×900 source, 3.24M weight rows) makes this concrete:

=== correctness when weight coords pass through float32 (a reproject UDF) ===
  float-key join :    1,077 / 810,000 dst cells matched
  int-key  join  :  810,000 / 810,000 dst cells matched
  -> float join silently DROPPED 99.87% of cells; int join is exact

Speed is a secondary ~1.1× (the GROUP BY/SUM dominates) and the keys are 2× smaller (int32 vs float64).

What

from_dataset(..., index_columns=True) (and read_xarray_table(..., index_columns=True)) emit, for every dimension, an int32 <dim>_idx column carrying each row's absolute integer position on that axis:

ctx.from_dataset("src", src, chunks={"time": 24}, index_columns=True)
# join the weight table on integer grid position, not float coords:
ctx.sql('''
  SELECT w.dst_id, SUM(s.value * w.weight) AS out
  FROM weights w JOIN src s
    ON s.lat_idx = w.src_lat_idx AND s.lon_idx = w.src_lon_idx
  GROUP BY w.dst_id
''')

• Plain Int32 columns — not dictionary-encoded — so none of DataFusion's join/aggregate/scalar-function paths are stressed (this is the safe alternative to the shelved dictionary-encoding approach in Dictionary-encode coordinate columns #217).
• Global, not per-partition: the reader adds each block's start offset, so the index lines up across chunks. A local index would restart at 0 in every partition and mis-join — this is the key correctness property, and it's tested directly.
• Coordinate columns stay dense and available for value predicates (WHERE lat > 45, date_part) and display. Index columns are opt-in, off by default, so nothing changes for existing users.

Implementation

• df.py: _parse_schema(index_columns=) appends the <dim>_idx fields (with a collision guard); iter_record_batches / dataset_to_record_batch emit them from the strided position plus a block offset.
• reader.py / sql.py: thread index_columns through read_xarray_table and from_dataset, computing per-block offsets.

Tests

tests/test_grid_index.py: schema/dtype, indices global across chunks, an exact index-keyed regrid matching a numpy gather, and the float32-drift case where the float-equality join drops cells but the index join stays exact. Full suite green (190), plus ruff and mypy. No Rust changes.

🤖 Generated with Claude Code

https://claude.ai/code/session_019VuSeCio99NcME5eubcN3N

---

Generated by Claude Code

[alxmrs]

🤖 Samudra's canonical loader would use this as WHERE time_idx IN (...), so index-column partition pruning is important in addition to exact joins.

The current diff emits <dim>_idx fields and values, but _block_metadata still appears to report bounds only for the original dimension coordinates. Could this PR (or a linked follow-up) add (block.start, block.stop - 1, "int64"/"int32") bounds for each index column to the native provider's partition metadata?

A regression test using _iteration_callback could register many one-time-step partitions, query a small time_idx IN (...) set, and assert that only the matching factories are instantiated. Without index bounds, the row filter remains correct but may scan every partition.

[alxmrs]

Good call — done in 1636614, in this PR.

_block_metadata now emits (block.start, block.start + n - 1, "int64") bounds for each <dim>_idx column when index columns are enabled, threaded through both the static and per-partition (dynamic) metadata passes in read_xarray_table. Blocks are contiguous slices, so a partition's positions are exactly [start, start+n-1] — and this holds for every coordinate dtype, so index pruning works even on axes whose coordinate values aren't prunable (strings, out-of-ns-range datetimes).

No native change was needed: PrunableStreamingTable already derives dimension_columns from meta.ranges.keys(), and ScalarBound::Int64 already compares against Int32 literals — so the new <dim>_idx keys become prunable (=, IN, <, BETWEEN, …) automatically, and time_idx IN (...) prunes rather than full-scans.

Added the regression test you described: one partition per time step, WHERE time_idx IN (2, 5, 7), asserting via _iteration_callback that only those three partition factories are instantiated. The time coordinate values are deliberately unrelated to the 0..n−1 index, so pruning can only succeed via time_idx, not the time coord.

Verified locally (native build + tests): the 5 grid-index tests pass and the 21 test_reader.py pruning/iteration tests still pass. The remaining suite failures in my sandbox are all xr.tutorial.open_dataset → HTTP 403 (no egress to pydata/xarray-data here), unrelated to the change; CI has the data.

---

Generated by Claude Code

[alxmrs]

Rebased onto main and force-pushed as a single clean commit (283e9b2); all checks are green.

What was failing: the branch had fallen ~7 commits behind main, which rewrote iter_record_batches (#227, #230). CI tests the PR merge commit, so it combined this PR's _parse_schema change (which adds <dim>_idx to the schema) with main's new emit loop — which had no index-column handling and routed those columns to data_arrays[name], raising KeyError: '<dim>_idx'. It only reproduced on the merge, never on the branch head, which is why it looked like an environment issue for a while.

Fix: re-implemented the feature on current main, integrating index-column emission into both of the new code paths (the full-pivot fast path and the per-batch path) and the _as_single_array helper.

Your pruning review is included in this commit: _block_metadata now emits (start, start+n-1, "int64") bounds for each <dim>_idx, so WHERE <dim>_idx IN (...) prunes whole partitions — no native change needed, since the provider derives its prunable set from the metadata keys and ScalarBound::Int64 already compares against Int32 literals. There's a regression test (test_index_column_prunes_partitions) asserting, via the reader's iteration callback, that time_idx IN (2, 5, 7) instantiates only those three partition factories.

---

Generated by Claude Code