feat(vortex-spatial): support GeoArrow Geometry unions

vortex-spatial/Cargo.toml

@@ -15,6 +15,7 @@ version.workspace = true
 
 [dependencies]
 arrow-array = { workspace = true }
+arrow-buffer = { workspace = true }
 arrow-schema = { workspace = true }
 geo = { workspace = true }
 geo-traits = { workspace = true }

vortex-spatial/src/aggregate_fn/aabb.rs

@@ -3,6 +3,7 @@
 
 //! The 2D axis-aligned bounding-box (AABB) aggregate for native geometry columns.
 
+use geo::BoundingRect;
 use geo::Rect as SpatialRect;
 use vortex_array::ArrayRef;
 use vortex_array::Columnar;
@@ -29,7 +30,9 @@ use crate::extension::box_storage_dtype;
 use crate::extension::coordinate::Dimension;
 use crate::extension::coordinate::box_corners;
 use crate::extension::coordinate::ordinates;
+use crate::extension::decode_mixed_geometries;
 use crate::extension::flatten_coordinates;
+use crate::extension::is_mixed_geometry;
 use crate::extension::is_native_geometry;
 
 /// Aggregates a native geometry column's 2D axis-aligned bounding box (AABB) as a native
@@ -205,6 +208,14 @@ impl AggregateFnVTable for GeometryAabb {
         // non-nullable case already costs nothing (the all-true mask makes `filter` a no-op).
         let valid = array.validity()?.execute_mask(array.len(), ctx)?;
         let array = array.filter(valid)?;
+        if is_mixed_geometry(array.dtype()) {
+            for geometry in decode_mixed_geometries(&array, ctx)? {
+                if let Some(rect) = geometry.bounding_rect() {
+                    partial.merge(rect);
+                }
+            }
+            return Ok(());
+        }
         // Null rows are gone, so every coordinate below belongs to a present geometry — the
         // `unmasked_field_by_name` reads are therefore safe. Min/max the raw x/y buffers directly:
         // cheap, and avoids `to_geometry`'s panic on empty points (which decoding would hit).

vortex-spatial/src/dense_union/canonical.rs

@@ -16,18 +16,17 @@ use vortex_buffer::BufferMut;
 use vortex_error::VortexResult;
 use vortex_error::vortex_ensure;
 use vortex_error::vortex_err;
-use vortex_mask::AllOr;
 
 use super::array::DenseUnion;
 use super::array::DenseUnionArrayExt;
 use super::array::DenseUnionArraySlotsExt;
+use super::tag_lookup;
 
 /// Converts a dense union to its canonical sparse representation.
 ///
-/// Each child is a dictionary over the original compact child, so values are not copied. Codes for
-/// other variants stay zero because their type IDs make them unreachable. Unused variants use a
-/// constant zero code array, and empty children use a one-value constant because dictionaries
-/// require non-empty values.
+/// A selected variant becomes a dictionary over the original compact child, so values are not
+/// copied. Codes for rows that select another variant stay zero, which their type IDs make
+/// unreachable. A variant no row selects becomes a constant of its default value.
 ///
 /// # Errors
 ///
@@ -45,62 +44,44 @@ pub(crate) fn canonicalize(
     let valid_rows = type_ids.validity()?.execute_mask(len, ctx)?;
     let child_lengths = array.iter_children().map(ArrayRef::len).collect::<Vec<_>>();
 
-    let mut child_indices = [None; 256];
-    for (child_index, type_id) in variants.type_ids().iter().copied().enumerate() {
-        child_indices[usize::from(type_id)] = Some(child_index);
-    }
+    let child_indices = tag_lookup(&variants);
     let mut codes_by_child: Vec<Option<BufferMut<u32>>> = vec![None; variants.len()];
 
-    let mut assign_row = |row: usize| -> VortexResult<()> {
-        let type_id = type_id_values[row];
-        let child_index = child_indices[usize::from(type_id)]
+    for (row, ((type_id, offset), valid)) in type_id_values
+        .iter()
+        .zip(offset_values)
+        .zip(valid_rows.iter())
+        .enumerate()
+    {
+        if !valid {
+            continue;
+        }
+        let child_index = child_indices[usize::from(*type_id)]
             .ok_or_else(|| vortex_err!("DenseUnion contains unknown type ID {type_id}"))?;
-        let offset = u32::try_from(offset_values[row]).map_err(|_| {
-            vortex_err!(
-                "DenseUnion contains negative offset {} at row {row}",
-                offset_values[row]
-            )
+        let offset = u32::try_from(*offset).map_err(|_| {
+            vortex_err!("DenseUnion contains negative offset {offset} at row {row}")
         })?;
         let child_len = child_lengths[child_index];
         vortex_ensure!(
             (offset as usize) < child_len,
             "DenseUnion offset {offset} is out of bounds for child {child_index} of length {child_len}"
         );
-        let codes = codes_by_child[child_index].get_or_insert_with(|| BufferMut::zeroed(len));
-        codes[row] = offset;
-        Ok(())
-    };
-
-    match valid_rows.indices() {
-        AllOr::All => {
-            for row in 0..len {
-                assign_row(row)?;
-            }
-        }
-        AllOr::None => {}
-        AllOr::Some(rows) => {
-            for &row in rows {
-                assign_row(row)?;
-            }
-        }
+        codes_by_child[child_index].get_or_insert_with(|| BufferMut::zeroed(len))[row] = offset;
     }
 
     let sparse_children = array
         .iter_children()
         .zip(codes_by_child)
         .map(|(child, codes)| {
-            let codes = match codes {
-                Some(codes) => {
-                    PrimitiveArray::new(codes.freeze(), Validity::NonNullable).into_array()
-                }
-                None => ConstantArray::new(0u32, len).into_array(),
-            };
-            let values = if child.is_empty() {
-                ConstantArray::new(Scalar::default_value(child.dtype()), 1).into_array()
-            } else {
-                child.clone()
+            // Codes are only recorded for a variant some valid row selects at an in-bounds
+            // offset, so a variant without them is unreachable and needs no values at all.
+            let Some(codes) = codes else {
+                return Ok(
+                    ConstantArray::new(Scalar::default_value(child.dtype()), len).into_array(),
+                );
             };
-            DictArray::try_new(codes, values).map(IntoArray::into_array)
+            let codes = PrimitiveArray::new(codes.freeze(), Validity::NonNullable).into_array();
+            DictArray::try_new(codes, child.clone()).map(IntoArray::into_array)
         })
         .collect::<VortexResult<Vec<_>>>()?;
 

vortex-spatial/src/dense_union/compact.rs

@@ -0,0 +1,123 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright the Vortex contributors
+
+use vortex_array::ArrayRef;
+use vortex_array::ExecutionCtx;
+use vortex_array::IntoArray;
+use vortex_array::arrays::PrimitiveArray;
+use vortex_array::dtype::UnionVariants;
+use vortex_error::VortexResult;
+use vortex_error::vortex_bail;
+use vortex_error::vortex_err;
+use vortex_mask::Mask;
+
+use super::tag_lookup;
+
+/// A dense union compacted into a directly Arrow-exportable layout.
+pub(crate) struct CompactUnion {
+    variants: UnionVariants,
+    /// The row-aligned Arrow type IDs.
+    pub type_ids: Vec<i8>,
+    /// The row-aligned offsets into [`Self::children`], increasing within each child.
+    pub offsets: Vec<i32>,
+    /// The compacted children in variant order, each holding exactly the rows that select it.
+    children: Vec<ArrayRef>,
+    /// The row validity, which Arrow instead expresses through the selected child.
+    pub validity: Mask,
+}
+
+impl CompactUnion {
+    /// Return the compacted child selected by a data-level type tag.
+    pub(crate) fn child(&self, tag: u8) -> Option<&ArrayRef> {
+        self.variants
+            .tag_to_child_index(tag)
+            .and_then(|child_index| self.children.get(child_index))
+    }
+}
+
+/// Compact a dense union so that its offsets increase within each child.
+///
+/// Vortex's selector-only operations reorder and repeat per-child offsets and retain unselected
+/// child rows, so the compact children are not an Arrow dense-union layout as they stand. This
+/// gathers each child down to exactly the rows that select it, in row order, and rebases the
+/// offsets onto the result.
+///
+/// A row's nullity moves from the union onto the selected child, matching Arrow's dense union,
+/// which has no validity of its own.
+///
+/// # Errors
+///
+/// Returns an error for unknown type IDs on valid rows, for offsets that exceed `i32`, or when
+/// gathering a child fails.
+pub(crate) fn compact_for_arrow(
+    variants: UnionVariants,
+    type_ids: &PrimitiveArray,
+    offsets: &PrimitiveArray,
+    children: Vec<ArrayRef>,
+    ctx: &mut ExecutionCtx,
+) -> VortexResult<CompactUnion> {
+    let validity = type_ids.validity()?.execute_mask(type_ids.len(), ctx)?;
+    let all_valid = validity.all_true();
+    let type_id_values = type_ids.as_slice::<u8>();
+    let offset_values = offsets.as_slice::<i32>();
+
+    let child_indices = tag_lookup(&variants);
+    // A null row is free to carry a type ID no variant declares: canonicalization and `scalar_at`
+    // both skip a null row's selectors without reading them. Arrow has no such slack, so park
+    // those rows on the first variant, where the null index below makes the value null.
+    let fallback_child = variants
+        .type_ids()
+        .first()
+        .and_then(|tag| child_indices[usize::from(*tag)])
+        .ok_or_else(|| vortex_err!("DenseUnion has no variants"))?;
+
+    let mut selections = vec![Vec::<Option<i32>>::new(); children.len()];
+    let mut arrow_type_ids = Vec::with_capacity(type_id_values.len());
+    let mut arrow_offsets = Vec::with_capacity(offset_values.len());
+    for (row, ((type_id, offset), valid)) in type_id_values
+        .iter()
+        .zip(offset_values)
+        .zip(validity.iter())
+        .enumerate()
+    {
+        let child_index = match (child_indices[usize::from(*type_id)], valid) {
+            (Some(child_index), _) => child_index,
+            (None, false) => fallback_child,
+            (None, true) => vortex_bail!("DenseUnion row has unknown type ID {type_id}"),
+        };
+        let arrow_type_id = variants.child_index_to_tag(child_index);
+        arrow_type_ids.push(
+            i8::try_from(arrow_type_id)
+                .map_err(|_| vortex_err!("DenseUnion type ID {arrow_type_id} exceeds i8"))?,
+        );
+        arrow_offsets.push(
+            i32::try_from(selections[child_index].len())
+                .map_err(|_| vortex_err!("DenseUnion child offset exceeds i32 at row {row}"))?,
+        );
+        selections[child_index].push(valid.then_some(*offset));
+    }
+
+    let children = children
+        .into_iter()
+        .zip(selections)
+        .map(|(child, selection)| {
+            // A null index gathers to a null value, which is how the row's nullity reaches the
+            // child. Keeping the child non-nullable when nothing selected it while null lets the
+            // Arrow export match a non-nullable child field.
+            let indices = if all_valid || selection.iter().all(Option::is_some) {
+                PrimitiveArray::from_iter(selection.into_iter().flatten()).into_array()
+            } else {
+                PrimitiveArray::from_option_iter(selection).into_array()
+            };
+            child.take(indices)
+        })
+        .collect::<VortexResult<Vec<_>>>()?;
+
+    Ok(CompactUnion {
+        variants,
+        type_ids: arrow_type_ids,
+        offsets: arrow_offsets,
+        children,
+        validity,
+    })
+}

vortex-spatial/src/dense_union/compute/filter.rs

@@ -3,23 +3,20 @@
 
 use vortex_array::ArrayRef;
 use vortex_array::ArrayView;
-use vortex_array::IntoArray;
 use vortex_array::arrays::filter::FilterReduce;
 use vortex_error::VortexResult;
 use vortex_mask::Mask;
 
+use super::with_selectors;
 use crate::dense_union::DenseUnion;
-use crate::dense_union::DenseUnionArrayExt;
 use crate::dense_union::DenseUnionArraySlotsExt;
 
 impl FilterReduce for DenseUnion {
     fn filter(array: ArrayView<'_, Self>, mask: &Mask) -> VortexResult<Option<ArrayRef>> {
-        DenseUnion::try_new(
+        with_selectors(
+            array,
             array.type_ids().filter(mask.clone())?,
             array.offsets().filter(mask.clone())?,
-            array.variants().clone(),
-            array.iter_children().cloned(),
         )
-        .map(|array| Some(array.into_array()))
     }
 }

vortex-spatial/src/dense_union/compute/mask.rs

@@ -3,23 +3,20 @@
 
 use vortex_array::ArrayRef;
 use vortex_array::ArrayView;
-use vortex_array::IntoArray;
 use vortex_array::builtins::ArrayBuiltins;
 use vortex_array::scalar_fn::fns::mask::MaskReduce;
 use vortex_error::VortexResult;
 
+use super::with_selectors;
 use crate::dense_union::DenseUnion;
-use crate::dense_union::DenseUnionArrayExt;
 use crate::dense_union::DenseUnionArraySlotsExt;
 
 impl MaskReduce for DenseUnion {
     fn mask(array: ArrayView<'_, Self>, mask: &ArrayRef) -> VortexResult<Option<ArrayRef>> {
-        DenseUnion::try_new(
+        with_selectors(
+            array,
             array.type_ids().clone().mask(mask.clone())?,
             array.offsets().clone(),
-            array.variants().clone(),
-            array.iter_children().cloned(),
         )
-        .map(|array| Some(array.into_array()))
     }
 }

vortex-spatial/src/dense_union/compute/mod.rs

@@ -2,8 +2,35 @@
 // SPDX-FileCopyrightText: Copyright the Vortex contributors
 
 //! Compute kernels for dense unions.
+//!
+//! Every kernel here is selector-only: it transforms the type IDs and offsets and retains each
+//! compact child in full. That is O(selected rows) rather than O(child values), at the cost of
+//! leaving unselected values behind and reordering per-child offsets.
 
 mod filter;
 mod mask;
 mod slice;
 mod take;
+
+use vortex_array::ArrayRef;
+use vortex_array::ArrayView;
+use vortex_array::IntoArray;
+use vortex_error::VortexResult;
+
+use super::DenseUnion;
+use super::DenseUnionArrayExt;
+
+/// Rebuild a dense union around transformed row selectors, retaining its compact children.
+fn with_selectors(
+    array: ArrayView<'_, DenseUnion>,
+    type_ids: ArrayRef,
+    offsets: ArrayRef,
+) -> VortexResult<Option<ArrayRef>> {
+    DenseUnion::try_new(
+        type_ids,
+        offsets,
+        array.variants().clone(),
+        array.iter_children().cloned(),
+    )
+    .map(|array| Some(array.into_array()))
+}

vortex-spatial/src/dense_union/compute/slice.rs

@@ -5,22 +5,19 @@ use std::ops::Range;
 
 use vortex_array::ArrayRef;
 use vortex_array::ArrayView;
-use vortex_array::IntoArray;
 use vortex_array::arrays::slice::SliceReduce;
 use vortex_error::VortexResult;
 
+use super::with_selectors;
 use crate::dense_union::DenseUnion;
-use crate::dense_union::DenseUnionArrayExt;
 use crate::dense_union::DenseUnionArraySlotsExt;
 
 impl SliceReduce for DenseUnion {
     fn slice(array: ArrayView<'_, Self>, range: Range<usize>) -> VortexResult<Option<ArrayRef>> {
-        DenseUnion::try_new(
+        with_selectors(
+            array,
             array.type_ids().slice(range.clone())?,
             array.offsets().slice(range)?,
-            array.variants().clone(),
-            array.iter_children().cloned(),
         )
-        .map(|array| Some(array.into_array()))
     }
 }