feat: GeoArrow point-geometry columns at registration (geometry=)

register(..., geometry=(x_dim, y_dim)) appends a geometry point column
derived from the coordinate dims, synthesized per batch at scan time —
the pivot's coordinate columns already carry the values, so the column
costs an annotation plus (for WKB) a vectorized 21-byte encode of the
rows actually scanned.

Two encodings, chosen per destination:
- "wkb" (default): geoarrow.wkb extension metadata + CRS. DuckDB >=1.2
  with spatial loaded ingests the column as GEOMETRY('OGC:CRS84'), so
  ST_Within(geometry, ...) works with no ST_Point construction in SQL.
- "point": GeoArrow-native separated coordinates; the struct children
  ARE the coordinate arrays (no copy, no parse) for consumers that
  execute on native layouts — verified with GeoPandas 1.x
  GeoDataFrame.from_arrow, CRS carried through.

Documented sharp edge, measured: engines do not push ST_* functions
into the scan, so a geometry-only predicate defeats chunk pruning and
encodes every row (~29x slower than the paired form on a 10M-row
grid: 101ms bbox vs 2.9s ST_Within-only vs 118ms bbox+ST_Within).
The geospatial docs now state the idiom: bbox conjuncts for pruning,
geometry for exactness.

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

Author: 2 people

docs/geospatial.md

@@ -472,3 +472,29 @@ scale. The point of this suite is not to crown a winner but to show that the lin
 between the two is exactly where the operation is dense versus where it is
 relational, and that for a surprising share of geoscience, the operation is
 relational.
+
+## GeoArrow point-geometry columns
+
+`register(..., geometry=("x", "y"))` derives a `geometry` point column
+from two coordinate dims. With the default `"wkb"` encoding DuckDB
+(spatial loaded) ingests it as a native `GEOMETRY` with the CRS
+attached, so geometry predicates need no `ST_Point(x, y)` construction:
+
+```sql
+SELECT avg(risk) FROM eri
+WHERE y BETWEEN -29 AND -28 AND x BETWEEN -58 AND -57   -- prunes chunks
+  AND ST_Within(geometry, ST_GeomFromText('POLYGON (...)'))  -- refines
+```
+
+**Always pair geometry predicates with bbox conjuncts on the coordinate
+columns.** Engines do not push functions like `ST_Within` into the
+scan, so a geometry-only predicate scans (and encodes) every chunk —
+measured ~29x slower than the paired form on a 10M-row grid, where the
+bbox prunes first and the exact polygon test is nearly free.
+
+`geometry_encoding="point"` emits GeoArrow-native separated coordinates
+instead (the struct children *are* the coordinate arrays): zero-parse
+for GeoPandas 1.x (`GeoDataFrame.from_arrow`), lonboard, geoarrow-rs
+and SedonaDB. DuckDB does not consume this encoding — pick per
+destination. The CRS tag defaults to `OGC:CRS84`; pass
+`geometry_crs=...` for anything else.

tests/test_geometry.py

@@ -0,0 +1,123 @@
+"""GeoArrow point-geometry columns derived at registration."""
+
+import json
+
+import numpy as np
+import pandas as pd
+import pyarrow as pa
+import pytest
+import xarray as xr
+
+import xarray_sql as xql
+from xarray_sql.backends.pyarrow import XarrayPushdownDataset
+
+
+@pytest.fixture
+def grid() -> xr.Dataset:
+    return xr.Dataset(
+        {"risk": (["y", "x"], np.arange(8.0 * 6).reshape(8, 6))},
+        coords={
+            "y": np.linspace(-28.0, -29.4, 8),  # descending, like rasters
+            "x": np.linspace(-58.0, -57.0, 6),
+        },
+    )
+
+
+def test_geometry_field_annotation(grid):
+    dataset = xql.arrow_dataset(grid, {"y": 4}, geometry=("x", "y"))
+    field = dataset.schema.field("geometry")
+    assert field.type == pa.binary()
+    assert field.metadata[b"ARROW:extension:name"] == b"geoarrow.wkb"
+    meta = json.loads(field.metadata[b"ARROW:extension:metadata"])
+    assert meta == {"crs": "OGC:CRS84"}
+
+
+def test_wkb_points_decode_exactly(grid):
+    dataset = xql.arrow_dataset(grid, {"y": 4}, geometry=("x", "y"))
+    table = dataset.to_table(columns=["geometry", "x", "y"])
+    blob = table["geometry"][0].as_py()
+    assert len(blob) == 21 and blob[0] == 1
+    x = np.frombuffer(blob, "<f8", count=1, offset=5)[0]
+    y = np.frombuffer(blob, "<f8", count=1, offset=13)[0]
+    assert (x, y) == (table["x"][0].as_py(), table["y"][0].as_py())
+
+
+def test_native_points_are_the_coordinate_arrays(grid):
+    dataset = xql.arrow_dataset(
+        grid, {"y": 4}, geometry=("x", "y"), geometry_encoding="point"
+    )
+    table = dataset.to_table(columns=["x", "y", "geometry"])
+    geom = table["geometry"].combine_chunks()
+    np.testing.assert_array_equal(
+        geom.field("x").to_numpy(), table["x"].to_numpy()
+    )
+    np.testing.assert_array_equal(
+        geom.field("y").to_numpy(), table["y"].to_numpy()
+    )
+
+
+def test_geometry_without_projecting_coords(grid):
+    # geometry can be requested alone; the dims it derives from are
+    # scanned but projected away.
+    dataset = xql.arrow_dataset(grid, {"y": 4}, geometry=("x", "y"))
+    table = dataset.to_table(columns=["geometry", "risk"])
+    assert table.column_names == ["geometry", "risk"]
+    assert table.num_rows == 48
+
+
+def test_duckdb_st_within_on_geometry_column(grid):
+    duckdb = pytest.importorskip("duckdb")
+
+    reads: list = []
+    dataset = XarrayPushdownDataset(
+        grid,
+        {"y": 4},
+        geometry=("x", "y"),
+        _iteration_callback=lambda b, n: reads.append(b),
+    )
+    con = duckdb.connect()
+    try:
+        con.execute("INSTALL spatial; LOAD spatial;")
+    except duckdb.Error:
+        pytest.skip("duckdb spatial extension unavailable")
+    con.register("t", dataset)
+
+    described = dict(
+        (row[0], row[1]) for row in con.execute("DESCRIBE t").fetchall()
+    )
+    assert described["geometry"].startswith("GEOMETRY")
+
+    # bbox conjuncts prune chunks; the polygon refines row-exactly on
+    # the ingested GEOMETRY column — no ST_Point construction needed.
+    reads.clear()
+    got = con.execute(
+        "SELECT count(*), round(avg(risk), 3) FROM t "
+        "WHERE y BETWEEN -28.7 AND -28.0 "
+        "AND ST_Within(geometry, ST_GeomFromText("
+        "'POLYGON ((-58.1 -28.75, -56.9 -28.75, -56.9 -27.9, "
+        "-58.1 -27.9, -58.1 -28.75))'))"
+    ).fetchone()
+    assert len(reads) == 1  # y-range kept only the first chunk
+    inside = (grid.y >= -28.7) & (grid.y <= -28.0)
+    expected = grid.risk.values[inside.values, :]
+    assert got == (expected.size, round(float(expected.mean()), 3))
+
+
+def test_geopandas_consumes_native_points(grid):
+    gpd = pytest.importorskip("geopandas")
+
+    dataset = xql.arrow_dataset(
+        grid, {"y": 4}, geometry=("x", "y"), geometry_encoding="point"
+    )
+    gdf = gpd.GeoDataFrame.from_arrow(dataset.to_table())
+    assert gdf.geometry.iloc[0].x == float(grid.x[0])
+    assert str(gdf.crs).endswith("CRS84")
+
+
+def test_geometry_name_collision_raises():
+    clash = xr.Dataset(
+        {"geometry": (["x"], np.arange(3.0))},
+        coords={"x": np.arange(3.0)},
+    )
+    with pytest.raises(ValueError, match="shadow"):
+        xql.arrow_dataset(clash, {"x": 3}, geometry=("x", "x"))

xarray_sql/backends/pyarrow.py

@@ -48,6 +48,7 @@
     iter_record_batches,
     resolve_chunks,
 )
+from ..geometry import GEOMETRY_COLUMN, build_geometry, geometry_field
 from ..reader import XarrayRecordBatchReader
 
 DEFAULT_PREFETCH = 4
@@ -271,6 +272,9 @@ def __init__(
         batch_size: int = DEFAULT_BATCH_SIZE,
         prefetch: int = DEFAULT_PREFETCH,
         coalesce_rows: int | None = None,
+        geometry: tuple[str, str] | None = None,
+        geometry_encoding: str = "wkb",
+        geometry_crs: str | None = "OGC:CRS84",
         coord_arrays: dict[str, np.ndarray] | None = None,
         _iteration_callback: (
             Callable[[Block, list[str] | None], None] | None
@@ -287,6 +291,25 @@ def __init__(
                 )
         self._ds = ds
         self._schema = _parse_schema(ds)
+        self._geometry = tuple(geometry) if geometry else None
+        self._geometry_encoding = geometry_encoding
+        if self._geometry:
+            if GEOMETRY_COLUMN in self._schema.names:
+                raise ValueError(
+                    f"geometry= would shadow an existing column named "
+                    f"{GEOMETRY_COLUMN!r}."
+                )
+            missing = [
+                d for d in self._geometry if d not in self._schema.names
+            ]
+            if missing:
+                raise ValueError(
+                    f"geometry dims {missing} are not columns of the "
+                    f"table; available: {self._schema.names}"
+                )
+            self._schema = self._schema.append(
+                geometry_field(geometry_encoding, geometry_crs)
+            )
         self._resolved = resolve_chunks(ds, chunks)
         if not self._resolved and ds.sizes:
             raise ValueError(
@@ -372,13 +395,50 @@ def _scanner_for_blocks(
         proj = list(self._schema.names) if columns is None else list(columns)
         scan_names = self._scan_columns(proj, filter)
         scan_schema = pa.schema([self._schema.field(n) for n in scan_names])
-        batches = self._batch_generator(
-            scan_schema, blocks, batch_size or self._batch_size
-        )
+        size = batch_size or self._batch_size
+        if self._geometry and GEOMETRY_COLUMN in scan_names:
+            batches = self._with_geometry(scan_schema, blocks, size)
+        else:
+            batches = self._batch_generator(scan_schema, blocks, size)
         return pads.Scanner.from_batches(
             batches, schema=scan_schema, columns=proj, filter=filter
         )
 
+    def _with_geometry(
+        self,
+        scan_schema: pa.Schema,
+        blocks: Iterator[Block] | list[Block],
+        batch_size: int,
+    ) -> Iterator[pa.RecordBatch]:
+        """Emit ``scan_schema`` batches, synthesizing the geometry column.
+
+        The pivot never materializes geometry: batches are produced with
+        the coordinate dims the geometry derives from, and the geometry
+        column is built per batch from those columns (for the native
+        encoding the point struct's children *are* the coordinate
+        arrays — a schema annotation, not a copy).
+        """
+        assert self._geometry is not None
+        x_dim, y_dim = self._geometry
+        base_names = [n for n in scan_schema.names if n != GEOMETRY_COLUMN]
+        for d in (x_dim, y_dim):
+            if d not in base_names:
+                base_names.append(d)
+        base_schema = pa.schema([self._schema.field(n) for n in base_names])
+        for batch in self._batch_generator(base_schema, blocks, batch_size):
+            geom = build_geometry(
+                self._geometry_encoding,
+                batch.column(base_names.index(x_dim)),
+                batch.column(base_names.index(y_dim)),
+            )
+            arrays = [
+                geom
+                if name == GEOMETRY_COLUMN
+                else batch.column(base_names.index(name))
+                for name in scan_schema.names
+            ]
+            yield pa.RecordBatch.from_arrays(arrays, schema=scan_schema)
+
     def get_fragments(
         self, filter: pc.Expression | None = None
     ) -> list["_XarrayFragment"]:
@@ -857,6 +917,9 @@ def arrow_dataset(
     batch_size: int = DEFAULT_BATCH_SIZE,
     prefetch: int = DEFAULT_PREFETCH,
     coalesce_rows: int | None = None,
+    geometry: tuple[str, str] | None = None,
+    geometry_encoding: str = "wkb",
+    geometry_crs: str | None = "OGC:CRS84",
 ) -> XarrayPushdownDataset:
     """A pushdown-capable ``pyarrow.dataset.Dataset`` view of ``ds``.
 
@@ -886,6 +949,18 @@ def arrow_dataset(
             fetches its member chunks through the store's own concurrent
             batching. Memory scales with ``prefetch`` x the *merged*
             block size, so size accordingly (e.g. ``8_000_000``).
+        geometry: ``(x_dim, y_dim)`` coordinate dims to derive a
+            ``geometry`` point column from (see
+            :mod:`xarray_sql.geometry`). With the default ``"wkb"``
+            encoding, DuckDB (spatial loaded) sees a native ``GEOMETRY``
+            column, so ``ST_Within(geometry, ...)`` works directly.
+        geometry_encoding: ``"wkb"`` (default; DuckDB-consumable) or
+            ``"point"`` (GeoArrow native separated coordinates — the
+            struct children are the coordinate arrays; for GeoPandas,
+            lonboard, geoarrow-rs consumers).
+        geometry_crs: CRS tag carried in the extension metadata.
+            Defaults to ``OGC:CRS84`` (plain longitude/latitude); pass
+            ``None`` to omit, or an authority code / PROJJSON string.
 
     Returns:
         An :class:`XarrayPushdownDataset`.
@@ -896,4 +971,7 @@ def arrow_dataset(
         batch_size=batch_size,
         prefetch=prefetch,
         coalesce_rows=coalesce_rows,
+        geometry=geometry,
+        geometry_encoding=geometry_encoding,
+        geometry_crs=geometry_crs,
     )

xarray_sql/geometry.py

@@ -0,0 +1,83 @@
+"""GeoArrow point-geometry columns derived from coordinate dimensions.
+
+A regular grid's pivot already materializes per-row x/y coordinate
+columns; a point-geometry column is those same values under a GeoArrow
+extension annotation. Two encodings:
+
+* ``"wkb"`` (default) — 21-byte WKB points under the ``geoarrow.wkb``
+  extension name. DuckDB (>= 1.2, spatial loaded) ingests the column as
+  a native ``GEOMETRY`` with the CRS attached, so ``ST_Within(geometry,
+  ...)`` works with no ``ST_Point(x, y)`` construction in user SQL.
+* ``"point"`` — GeoArrow native points with *separated* coordinates
+  (``struct<x: double, y: double>`` under ``geoarrow.point``): the
+  child arrays are the coordinate columns themselves, no per-row
+  parsing for consumers that execute on native layouts (GeoPandas 1.x,
+  geoarrow-rs, lonboard, SedonaDB). DuckDB does not consume this
+  encoding.
+
+The CRS rides in the extension metadata (GeoArrow 0.2 allows
+authority:code strings alongside PROJJSON). ``OGC:CRS84`` is the
+correct tag for plain longitude/latitude grids.
+"""
+
+from __future__ import annotations
+
+import json
+
+import numpy as np
+import pyarrow as pa
+
+GEOMETRY_COLUMN = "geometry"
+
+_ENCODINGS = ("wkb", "point")
+
+
+def geometry_field(encoding: str, crs: str | None) -> pa.Field:
+    """The schema field for the derived geometry column."""
+    if encoding not in _ENCODINGS:
+        raise ValueError(
+            f"geometry_encoding must be one of {_ENCODINGS}, got {encoding!r}"
+        )
+    metadata = {
+        b"ARROW:extension:name": f"geoarrow.{encoding}".encode(),
+    }
+    if crs is not None:
+        metadata[b"ARROW:extension:metadata"] = json.dumps(
+            {"crs": crs}
+        ).encode()
+    storage = (
+        pa.binary()
+        if encoding == "wkb"
+        else pa.struct([("x", pa.float64()), ("y", pa.float64())])
+    )
+    return pa.field(GEOMETRY_COLUMN, storage, metadata=metadata)
+
+
+def build_geometry(
+    encoding: str, x: pa.Array, y: pa.Array
+) -> pa.Array:
+    """Point geometries for one batch's x/y coordinate columns."""
+    if encoding == "point":
+        return pa.StructArray.from_arrays(
+            [x.cast(pa.float64()), y.cast(pa.float64())], ["x", "y"]
+        )
+    return _wkb_points(
+        np.ascontiguousarray(x.to_numpy(zero_copy_only=False), "<f8"),
+        np.ascontiguousarray(y.to_numpy(zero_copy_only=False), "<f8"),
+    )
+
+
+def _wkb_points(x: np.ndarray, y: np.ndarray) -> pa.Array:
+    """Vectorized 21-byte little-endian WKB point encoding."""
+    n = len(x)
+    buf = np.empty((n, 21), dtype=np.uint8)
+    buf[:, 0] = 1  # little-endian byte order mark
+    buf[:, 1:5] = np.array([1, 0, 0, 0], dtype=np.uint8)  # type: Point
+    buf[:, 5:13] = x.view(np.uint8).reshape(n, 8)
+    buf[:, 13:21] = y.view(np.uint8).reshape(n, 8)
+    offsets = pa.py_buffer(
+        np.arange(0, (n + 1) * 21, 21, dtype=np.int32).tobytes()
+    )
+    return pa.Array.from_buffers(
+        pa.binary(), n, [None, offsets, pa.py_buffer(buf.tobytes())]
+    )