feat(vortex-spatial): support GeoArrow Geometry unions

[HarukiMoriarty]

Rationale for this change

GeoArrow Geometry columns use dense unions to mix geometry kinds and dimensions. Vortex needs a logical Geometry extension over its canonical Union dtype while retaining Arrow compact children through the external DenseUnion physical encoding.

What changes are included in this PR?

• Add the vortex.st.geometry extension mapped to geoarrow.geometry.
• Validate the standardized GeoArrow kind and dimension type IDs.
• Import Arrow dense unions as DenseUnion and export both DenseUnion and canonical sparse Union storage.
• Preserve null rows through the selected Arrow child.
• Decode Geometry values for existing spatial execution, AABB aggregation, and envelope computation.
• Cover all six supported geometry kinds, slicing, nulls, canonical sparse export, and GeometryCollection rejection.

What APIs are changed? Are there any user-facing changes?

Adds the public Geometry and GeometryData spatial extension types and registers GeoArrow Geometry import/export in vortex_spatial::initialize. GeometryCollection fields are accepted as part of the standard schema, but selected GeometryCollection values remain unsupported.

[connortsui20]

Three things need fixing before this lands, and one structural question about where the decode belongs. Everything after that is style.

to_geometry panics on POINT EMPTY

decode_mixed_geometries (geometry.rs:366) calls ToGeoGeometry::to_geometry, which geo-traits documents as panicking:

/// This will panic on an empty point or a MultiPoint containing empty points.
fn to_geometry(&self) -> Geometry<T> {
    self.try_to_geometry().expect(
        "geo-types does not support empty point or a MultiPoint containing empty points.",
    )
}

GeoArrow stores an empty point as NaN coordinates and geoarrow's PointTrait::coord returns None for those, so a POINT EMPTY in the Point child panics. aabb.rs already knows this. Its comment three lines below the new branch says the raw-buffer path "avoids to_geometry's panic on empty points (which decoding would hit)", and the new branch at aabb.rs:211 routes Geometry straight into the decode. ST_Envelope inherits it through mixed_geometry_boxes, and so do ST_Area, ST_Intersects, ST_Contains and ST_Distance through decode_geometries.

Use try_to_geometry and give None the same treatment the homogeneous path gives an empty geometry: a null box, not a panic. A POINT EMPTY row in supported_geometries() would have caught it.

Export rejects nulls that ordinary compute produces

geometry.rs:494 requires every null row to also be null in its selected child:

vortex_ensure!(
    union.child(type_id).is_null(offset),
    "GeoArrow Geometry null at row {row} is not represented by a null selected child"
);

Nothing maintains that. Extension::mask pushes the mask into the storage array, and DenseUnion::mask nulls the type IDs only:

// encodings/dense-union/src/compute/mask.rs
DenseUnion::try_new(
    array.type_ids().clone().mask(mask.clone())?,
    array.offsets().clone(),
    array.variants().clone(),
    array.iter_children().cloned(),
)

So from_arrow_array then any null-introducing compute then execute_arrow fails. Only arrays fresh off the importer satisfy the double representation, because the importer is what writes it.

The export owns the Arrow-level null, so it should write it: build each child's null buffer from the row validity instead of asserting the child already carries it. That also drops the second pass over the union.

geo::BoundingRect seeds from the first coordinate

get_bounding_rect starts the range at the first vertex and narrows with </>, so a NaN in position zero leaves the whole rect NaN. The path this replaces for other geometry types, box_corners, seeds at infinity and folds f64::min/f64::max, which skip NaN and keep the finite ordinates.

AabbPartial::merge also folds with f64::min/f64::max, so the NaN rect contributes nothing and the row's real coordinates never reach the zone box. SpatialIntersectsPrune can then drop a chunk that holds a matching row. Wrong answers, not just a wider box.

The decode is dispatched at the wrong layer

decode_geometries has an established shape: take the storage array, pick a <type>_geometries decoder.

if ext.is::<Geometry>() {
    return decode_mixed_geometries(array, ctx);
}
let storage = array.clone().execute::<ExtensionArray>(ctx)?.storage_array().clone();
if ext.is::<Point>() {
    point_geometries(&storage, ctx)
} else if ...

The early return exists because decode_mixed_geometries needs the extension dtype to build the target Arrow field, not the storage. That is one decoder's implementation detail leaking into the shared dispatch, and the same leak repeats in GeometryAabb::accumulate and envelope_array. Adding a geometry kind now means editing is_native_geometry, decode_geometries, aabb.rs and envelope.rs.

Passing the whole ArrayRef to every *_geometries decoder would let Geometry join the chain, and the two other branches would follow from that.

Two consequences worth calling out:

• is_native_geometry now returns true for a type whose storage is a DType::Union, which makes flatten_coordinates's guard and its doc ("Flatten a native geometry column into a single coordinate Struct<x, y, ...>") false for a type that guard accepts. Any new caller gets an opaque canonicalization error instead of the clear one.
• Going Vortex to Arrow to geoarrow to geo_types per row costs a lot more than the paths it sits in. Each union child is already a native geometry array that the columnar AABB and envelope code handles, so recursing into the children and merging would keep both off the row-oriented fallback. The aabb.rs comment justifying the current cost ("this is a write-time zone stat, and the common non-nullable case already costs nothing") no longer holds for this type. Fine as a follow-up, but note it.

vortex-duckdb was not updated

GeometryData and to_wkb mirror PointData, PolygonData and the rest exactly, but none of the three DuckDB dispatch sites learned about Geometry:

• convert/dtype.rs:258 bails Unsupported extension type, so a geoarrow.geometry column cannot be read through DuckDB at all.
• exporter/extension.rs:66 bails no non-temporal extension exporter.
• convert/expr.rs:129 silently declines ST_* pushdown.

So GeometryData::to_wkb has no caller anywhere in the repo. Either wire the three sites or drop the type until something needs it.

Structure and style

• geometry.rs is 633 lines against 306 to 426 for its siblings, and it is the only extension file with no mod tests. validate_dtype and the 58-line validate_variant_dtype have no direct coverage, while every sibling rstests its dimension and invalid-storage cases. The two Arrow vtables and GeometryUnionParts are separable.
• GeoArrowGeometryKind and its variants, GeometryUnionParts and its four fields, and all eight free functions are undocumented. offsets in particular means dense-union offsets on one path and synthesized row indices on the other, and that invariant currently lives in an in-function comment next to the producer.
• The type_id / 10 and type_id % 10 decoding rests on one unsourced comment. Link the GeoArrow spec section that fixes the scheme. Both bails also report the same message for two different failures.
• Error messages state a requirement without the reality: "GeoArrow geometry variant {type_id} must be an extension", "type ID {type_id} must contain Point values" and six others have the value in scope. geometry_variants gets it right with "...must be a union, got {dtype}".
• The GeometryCollection deferral appears four times as a bare string with no tag, reason, or issue link.
• Geometry::unpack_native clones the whole ScalarValue into a Scalar per call and validates nothing beyond what Scalar construction already does. Rect returns &'a ScalarValue for the same job.
• native_child_dtype calls validate_variant_dtype, which recomputes the geoarrow_type_id_parts it just computed, and validate_dtype runs the same check again on import.
• decode_mixed_geometries is point_geometries with a different array constructor, down to both error strings.
• Seven tests repeat the same three-line prelude, and aabb in tests/geometry.rs reimplements the one in aggregate_fn/aabb.rs instead of reusing rect_from_storage. None of the tests has a doc comment, unlike the sibling test files.
• Every fixture is XY and non-empty. No XYZ/XYM/XYZM, no empty geometry, no zero-length column, no non-nullable field, no interleaved-coordinate rejection, and no input union carrying a subset of the canonical type IDs. A geometry_column builder in test_harness.rs would let aabb.rs's every_native_column reach these too, which it currently cannot.
• vortex_dense_union::initialize(session) sits above the "Register the spatial extension types." comment with no note that Geometry storage needs the encoding registered.

---

Generated by Claude Code

[connortsui20]

A second pass turned up two more correctness problems and three structural ones, same commit.

The unknown-type-ID fallback keeps the original offset

execute_arrow substitutes the first variant's type ID for a null row whose type ID is not a known variant:

(false, false) => fallback_type_id,

but the offsets are copied straight across with no matching substitution:

let arrow_offsets = parts.offsets.as_slice::<i32>().to_vec();

So the row leaves with a type ID naming one child and an offset that indexed a different one. Arrow validates dense offsets against the child the type ID selects, so UnionArray::try_new rejects the entire export with Offsets must be non-negative and within the length of the Array whenever the fallback child is shorter than the original. A Geometry column holding no XY points has an empty Point XY child, and that child is the fallback.

Null slots are allowed to hold garbage: DenseUnion::canonicalize and scalar_at both skip null rows without reading their offsets. So this is a legal array that cannot be exported. Force the offset to 0 alongside the type ID.

exports_constant_nulls cannot catch it. A constant null takes the sparse path, where offsets are 0..len and every child is exactly len rows, so the fallback child always covers the offset by construction.

to_arrow_field ignores the dtype's variants

let (_, nullability) = geometry_variants(ext_dtype.storage_dtype())?;
Ok(Some(
    geoarrow_geometry_type(metadata).to_field(name, nullability.is_nullable()),
))

The variants are bound to _. from_arrow_field records whatever subset of children the input field carried, but schema inference always reports the full canonical union and execute_arrow pads the remainder with new_empty_array. A producer that hands Vortex a Point plus LineString union gets the canonical union back, so a round trip through Vortex rewrites the consumer's type.

Two consequences: the target GeoArrow geometry union is missing type ID {source_id} check can never fire against a self-inferred target, and every decode builds the full child set per batch. infers_canonical_geoarrow_field only feeds the canonical type in, so it cannot see the widening.

Export converts the whole column, not the window

DenseUnion::slice and filter both slice type_ids and offsets and pass iter_children().cloned() through untouched, by design. execute_arrow then exports each child in full, so the data converted to Arrow is proportional to the original column rather than to the rows requested.

That compounds with the decode. ST_Area, ST_Distance, SpatialEnvelope and GeometryAabb all reach decode_mixed_geometries through array.filter(valid), so one null row in a 10M-row column makes every kernel call convert all 10M rows of child data. A one-row slice costs the same.

rejects_selected_geometry_collection cannot tell which bail fired

The assertion is error.to_string().contains("GeometryCollection"), and that string appears at geometry.rs:171, :207 and :578. Only the third is the per-row rejection the test is named for. If from_arrow_field ever stops skipping GeometryCollection child fields, dtype construction bails with a message that also contains the word, the test stays green, and the row-level check goes untested. Match the specific message, or assert on the type ID.

Row validity on import

union.child(*type_id) followed by child.is_valid(offset) per row is a per-element accessor in a hot loop over every imported row. Hoist the child null buffers once, or skip the scan entirely when every child has null_count() == 0, which is the common case for a GeoArrow import.

---

Generated by Claude Code

[codspeed-hq]

Merging this PR will improve performance by 13.85%

⚠️ Unknown Walltime execution environment detected

Using the Walltime instrument on standard Hosted Runners will lead to inconsistent data.

For the most accurate results, we recommend using CodSpeed Macro Runners: bare-metal machines fine-tuned for performance measurement consistency.

⚡ 1 improved benchmark
✅ 2001 untouched benchmarks
⏩ 89 skipped benchmarks¹

Performance Changes

	Mode	Benchmark	BASE	HEAD	Efficiency
⚡	WallTime	words_gather_scalar[65536]	9.4 µs	8.3 µs	+13.85%

Tip

Curious why this is faster? Comment @codspeedbot explain why this is faster on this PR, or directly use the CodSpeed MCP with your agent.

---

_{Comparing nemo/geo-union (bcac46e) with nemo/dense-union (7459934)}

Footnotes

1. 89 benchmarks were skipped, so the baseline results were used instead. If they were deleted from the codebase, click here and archive them to remove them from the performance reports. ↩