diff --git a/Cargo.lock b/Cargo.lock index 4e58b02f838..24bc0c29ce9 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -10607,6 +10607,7 @@ name = "vortex-spatial" version = "0.1.0" dependencies = [ "arrow-array 58.4.0", + "arrow-buffer 58.4.0", "arrow-schema 58.4.0", "codspeed-divan-compat", "geo", @@ -10620,6 +10621,7 @@ dependencies = [ "vortex-array", "vortex-arrow", "vortex-buffer", + "vortex-dense-union", "vortex-error", "vortex-layout", "vortex-mask", diff --git a/vortex-spatial/Cargo.toml b/vortex-spatial/Cargo.toml index a758087da03..fcd7ce3a751 100644 --- a/vortex-spatial/Cargo.toml +++ b/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 } @@ -25,6 +26,7 @@ prost = { workspace = true } vortex-array = { workspace = true } vortex-arrow = { workspace = true } vortex-buffer = { workspace = true } +vortex-dense-union = { workspace = true } vortex-error = { workspace = true } vortex-mask = { workspace = true } vortex-session = { workspace = true } diff --git a/vortex-spatial/src/aggregate_fn/aabb.rs b/vortex-spatial/src/aggregate_fn/aabb.rs index 7c3bcaf7659..3cd0927210a 100644 --- a/vortex-spatial/src/aggregate_fn/aabb.rs +++ b/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; @@ -23,12 +24,14 @@ use vortex_error::vortex_err; use vortex_session::VortexSession; use vortex_session::registry::CachedId; +use crate::extension::Geometry; use crate::extension::Rect; use crate::extension::SpatialMetadata; 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_native_geometry; @@ -205,6 +208,18 @@ 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 array + .dtype() + .as_extension_opt() + .is_some_and(|ext| ext.is::()) + { + 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). diff --git a/vortex-spatial/src/extension/geometry.rs b/vortex-spatial/src/extension/geometry.rs new file mode 100644 index 00000000000..759254b55f3 --- /dev/null +++ b/vortex-spatial/src/extension/geometry.rs @@ -0,0 +1,633 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! The [`Geometry`] extension type (`vortex.st.geometry`), represented logically as a union +//! of native geometry extension arrays and mapped to `geoarrow.geometry`. Arrow imports retain +//! the compact layout through the external DenseUnion physical encoding. + +use std::sync::Arc; + +use arrow_array::Array as _; +use arrow_array::ArrayRef as ArrowArrayRef; +use arrow_array::UnionArray as ArrowUnionArray; +use arrow_array::new_empty_array; +use arrow_buffer::ScalarBuffer; +use arrow_schema::DataType; +use arrow_schema::Field; +use arrow_schema::UnionMode; +use arrow_schema::extension::ExtensionType; +use geo_traits::to_geo::ToGeoGeometry; +use geo_types::Geometry as GeoGeometry; +use geoarrow::array::GeoArrowArrayAccessor; +use geoarrow::array::GeometryArray as GeoArrowGeometryArray; +use geoarrow::datatypes::CoordType; +use geoarrow::datatypes::GeometryType as GeoArrowGeometryType; +use prost::Message; +use vortex_array::ArrayRef; +use vortex_array::ExecutionCtx; +use vortex_array::IntoArray; +use vortex_array::arrays::ExtensionArray; +use vortex_array::arrays::PrimitiveArray; +use vortex_array::arrays::UnionArray; +use vortex_array::arrays::extension::ExtensionArrayExt; +use vortex_array::arrays::union::UnionArrayExt; +use vortex_array::arrays::union::UnionArraySlotsExt; +use vortex_array::builtins::ArrayBuiltins; +use vortex_array::dtype::DType; +use vortex_array::dtype::FieldNames; +use vortex_array::dtype::Nullability; +use vortex_array::dtype::UnionVariants; +use vortex_array::dtype::extension::ExtDType; +use vortex_array::dtype::extension::ExtId; +use vortex_array::dtype::extension::ExtVTable; +use vortex_array::scalar::Scalar; +use vortex_array::scalar::ScalarValue; +use vortex_array::validity::Validity; +use vortex_arrow::ArrowExport; +use vortex_arrow::ArrowExportVTable; +use vortex_arrow::ArrowImport; +use vortex_arrow::ArrowImportVTable; +use vortex_arrow::ArrowSession; +use vortex_arrow::ArrowSessionExt; +use vortex_arrow::FromArrowArray; +use vortex_dense_union::DenseUnion; +use vortex_dense_union::DenseUnionArrayExt; +use vortex_dense_union::DenseUnionArraySlotsExt; +use vortex_error::VortexError; +use vortex_error::VortexResult; +use vortex_error::vortex_bail; +use vortex_error::vortex_ensure; +use vortex_error::vortex_err; +use vortex_session::registry::CachedId; +use vortex_session::registry::Id; + +use super::LineString; +use super::MultiLineString; +use super::MultiPoint; +use super::MultiPolygon; +use super::Point; +use super::Polygon; +use super::SpatialMetadata; +use super::coordinate::Dimension; +use super::coordinate::coordinate_dimension; +use super::geoarrow_metadata; +use super::geoarrow_to_wkb; +use super::linestring_dimension; +use super::multilinestring_dimension; +use super::multipoint_dimension; +use super::multipolygon_dimension; +use super::polygon_dimension; +use super::spatial_metadata_from_arrow; + +/// A mixed native geometry column whose logical union variants are Point, LineString, Polygon, +/// MultiPoint, MultiLineString, and MultiPolygon extension arrays. GeoArrow GeometryCollection +/// fields may be present in the Arrow schema, but selected GeometryCollection values are not yet +/// supported. +#[derive(Debug, Clone, Default, PartialEq, Eq, Hash)] +pub struct Geometry; + +#[derive(Clone, Copy, Debug, Eq, PartialEq)] +enum GeoArrowGeometryKind { + Point, + LineString, + Polygon, + MultiPoint, + MultiLineString, + MultiPolygon, + GeometryCollection, +} + +fn geoarrow_type_id_parts(type_id: u8) -> VortexResult<(GeoArrowGeometryKind, Dimension)> { + // GeoArrow assigns the ones digit to geometry kind and the tens digit to coordinate dimension. + let dimension = match type_id / 10 { + 0 => Dimension::Xy, + 1 => Dimension::Xyz, + 2 => Dimension::Xym, + 3 => Dimension::Xyzm, + _ => vortex_bail!("unsupported GeoArrow geometry type ID {type_id}"), + }; + let kind = match type_id % 10 { + 1 => GeoArrowGeometryKind::Point, + 2 => GeoArrowGeometryKind::LineString, + 3 => GeoArrowGeometryKind::Polygon, + 4 => GeoArrowGeometryKind::MultiPoint, + 5 => GeoArrowGeometryKind::MultiLineString, + 6 => GeoArrowGeometryKind::MultiPolygon, + 7 => GeoArrowGeometryKind::GeometryCollection, + _ => vortex_bail!("unsupported GeoArrow geometry type ID {type_id}"), + }; + Ok((kind, dimension)) +} + +fn validate_variant_dtype(type_id: u8, dtype: &DType) -> VortexResult<()> { + let (kind, expected_dimension) = geoarrow_type_id_parts(type_id)?; + let ext = dtype + .as_extension_opt() + .ok_or_else(|| vortex_err!("GeoArrow geometry variant {type_id} must be an extension"))?; + let actual_dimension = match kind { + GeoArrowGeometryKind::Point => { + vortex_ensure!( + ext.is::(), + "type ID {type_id} must contain Point values" + ); + coordinate_dimension(ext.storage_dtype())? + } + GeoArrowGeometryKind::LineString => { + vortex_ensure!( + ext.is::(), + "type ID {type_id} must contain LineString values" + ); + linestring_dimension(ext.storage_dtype())? + } + GeoArrowGeometryKind::Polygon => { + vortex_ensure!( + ext.is::(), + "type ID {type_id} must contain Polygon values" + ); + polygon_dimension(ext.storage_dtype())? + } + GeoArrowGeometryKind::MultiPoint => { + vortex_ensure!( + ext.is::(), + "type ID {type_id} must contain MultiPoint values" + ); + multipoint_dimension(ext.storage_dtype())? + } + GeoArrowGeometryKind::MultiLineString => { + vortex_ensure!( + ext.is::(), + "type ID {type_id} must contain MultiLineString values" + ); + multilinestring_dimension(ext.storage_dtype())? + } + GeoArrowGeometryKind::MultiPolygon => { + vortex_ensure!( + ext.is::(), + "type ID {type_id} must contain MultiPolygon values" + ); + multipolygon_dimension(ext.storage_dtype())? + } + GeoArrowGeometryKind::GeometryCollection => { + vortex_bail!("GeoArrow GeometryCollection values are not supported yet") + } + }; + vortex_ensure!( + actual_dimension == expected_dimension, + "GeoArrow geometry type ID {type_id} requires {expected_dimension:?} storage, got {actual_dimension:?}" + ); + Ok(()) +} + +fn native_child_dtype( + type_id: u8, + metadata: &SpatialMetadata, + storage_dtype: DType, +) -> VortexResult { + let (kind, _) = geoarrow_type_id_parts(type_id)?; + let dtype = match kind { + GeoArrowGeometryKind::Point => { + DType::Extension(ExtDType::::try_new(metadata.clone(), storage_dtype)?.erased()) + } + GeoArrowGeometryKind::LineString => DType::Extension( + ExtDType::::try_new(metadata.clone(), storage_dtype)?.erased(), + ), + GeoArrowGeometryKind::Polygon => DType::Extension( + ExtDType::::try_new(metadata.clone(), storage_dtype)?.erased(), + ), + GeoArrowGeometryKind::MultiPoint => DType::Extension( + ExtDType::::try_new(metadata.clone(), storage_dtype)?.erased(), + ), + GeoArrowGeometryKind::MultiLineString => DType::Extension( + ExtDType::::try_new(metadata.clone(), storage_dtype)?.erased(), + ), + GeoArrowGeometryKind::MultiPolygon => DType::Extension( + ExtDType::::try_new(metadata.clone(), storage_dtype)?.erased(), + ), + GeoArrowGeometryKind::GeometryCollection => { + vortex_bail!("GeoArrow GeometryCollection values are not supported yet") + } + }; + validate_variant_dtype(type_id, &dtype)?; + Ok(dtype) +} + +fn geometry_variants(dtype: &DType) -> VortexResult<(&UnionVariants, Nullability)> { + let DType::Union(variants, nullability) = dtype else { + vortex_bail!("Geometry storage must be a union, got {dtype}"); + }; + Ok((variants, *nullability)) +} + +struct GeometryUnionParts { + variants: UnionVariants, + type_ids: PrimitiveArray, + offsets: PrimitiveArray, + children: Vec, +} + +impl GeometryUnionParts { + fn child(&self, type_id: u8) -> Option<&ArrayRef> { + self.variants + .tag_to_child_index(type_id) + .and_then(|child_index| self.children.get(child_index)) + } +} + +fn geometry_union_parts( + storage: ArrayRef, + ctx: &mut ExecutionCtx, +) -> VortexResult { + if let Some(dense) = storage.as_opt::() { + return Ok(GeometryUnionParts { + variants: dense.variants().clone(), + type_ids: dense.type_ids().clone().execute::(ctx)?, + offsets: dense.offsets().clone().execute::(ctx)?, + children: dense.iter_children().cloned().collect(), + }); + } + + // Other physical encodings execute to the canonical sparse union. Its row-aligned children + // can use row indices as dense-union offsets, so export does not need to compact them. + let sparse = storage.execute::(ctx)?; + let offsets = (0..sparse.len()) + .map(|offset| { + i32::try_from(offset) + .map_err(|_| vortex_err!("Geometry row offset {offset} exceeds i32")) + }) + .collect::>>()?; + let type_ids = sparse.type_ids().clone().execute::(ctx)?; + let outer_validity = type_ids + .validity()? + .execute_mask(type_ids.len(), ctx)? + .into_array(); + let children = sparse + .iter_children() + .cloned() + .map(|child| child.mask(outer_validity.clone())) + .collect::>>()?; + Ok(GeometryUnionParts { + variants: sparse.variants().clone(), + type_ids, + offsets: PrimitiveArray::from_iter(offsets), + children, + }) +} + +impl ExtVTable for Geometry { + type Metadata = SpatialMetadata; + type NativeValue<'a> = Scalar; + + fn id(&self) -> ExtId { + static ID: CachedId = CachedId::new("vortex.st.geometry"); + *ID + } + + fn serialize_metadata(&self, metadata: &Self::Metadata) -> VortexResult> { + Ok(metadata.encode_to_vec()) + } + + fn deserialize_metadata(&self, metadata: &[u8]) -> VortexResult { + Ok(SpatialMetadata::decode(metadata)?) + } + + fn validate_dtype(ext_dtype: &ExtDType) -> VortexResult<()> { + let (variants, _) = geometry_variants(ext_dtype.storage_dtype())?; + for (type_id, dtype) in variants.type_ids().iter().zip(variants.variants()) { + validate_variant_dtype(*type_id, &dtype)?; + } + Ok(()) + } + + fn unpack_native<'a>( + ext_dtype: &'a ExtDType, + storage_value: &'a ScalarValue, + ) -> VortexResult> { + Scalar::try_new( + ext_dtype.storage_dtype().clone(), + Some(storage_value.clone()), + ) + } +} + +static GEOARROW_GEOMETRY: CachedId = CachedId::new(GeoArrowGeometryType::NAME); + +fn geoarrow_geometry_type(metadata: &SpatialMetadata) -> GeoArrowGeometryType { + GeoArrowGeometryType::new(geoarrow_metadata(metadata)).with_coord_type(CoordType::Separated) +} + +/// A materialized mixed geometry extension array. +pub struct GeometryData(ExtensionArray); + +impl TryFrom for GeometryData { + type Error = VortexError; + + fn try_from(ext: ExtensionArray) -> Result { + vortex_ensure!( + ext.ext_dtype().is::(), + "expected a Geometry extension array" + ); + Ok(Self(ext)) + } +} + +impl GeometryData { + /// Serialize mixed geometries to WKB. + pub fn to_wkb(&self, ctx: &mut ExecutionCtx) -> VortexResult { + geoarrow_to_wkb(&geoarrow_geometry_array(&self.0.clone().into_array(), ctx)?) + } +} + +fn geoarrow_geometry_array( + array: &ArrayRef, + ctx: &mut ExecutionCtx, +) -> VortexResult { + let session = ctx.session().clone(); + let field = session.arrow().to_arrow_field("", array.dtype())?; + let geometry_type = field + .try_extension_type::() + .map_err(|e| vortex_err!("failed to construct GeoArrow GeometryType: {e}"))?; + let arrow = session + .arrow() + .execute_arrow(array.clone(), Some(&field), ctx)?; + GeoArrowGeometryArray::try_from((arrow.as_ref(), geometry_type)) + .map_err(|e| vortex_err!("failed to construct GeoArrow GeometryArray: {e}")) +} + +pub(crate) fn decode_mixed_geometries( + array: &ArrayRef, + ctx: &mut ExecutionCtx, +) -> VortexResult>> { + geoarrow_geometry_array(array, ctx)? + .iter() + .map(|geometry| -> VortexResult> { + Ok(geometry + .ok_or_else(|| vortex_err!("spatial: null geometry is not supported"))? + .map_err(|e| vortex_err!("spatial: geometry access failed: {e}"))? + .to_geometry()) + }) + .collect() +} + +impl ArrowExportVTable for Geometry { + fn arrow_ext_id(&self) -> Id { + *GEOARROW_GEOMETRY + } + + fn vortex_id(&self) -> Id { + self.id() + } + + fn to_arrow_field( + &self, + name: &str, + dtype: &DType, + _session: &ArrowSession, + ) -> VortexResult> { + let ext_dtype = dtype.as_extension(); + let metadata = ext_dtype.metadata::(); + let (_, nullability) = geometry_variants(ext_dtype.storage_dtype())?; + Ok(Some( + geoarrow_geometry_type(metadata).to_field(name, nullability.is_nullable()), + )) + } + + fn execute_arrow( + &self, + array: ArrayRef, + target: &Field, + ctx: &mut ExecutionCtx, + ) -> VortexResult { + let Some(ext_dtype) = array.dtype().as_extension_opt() else { + return Ok(ArrowExport::Unsupported(array)); + }; + if !ext_dtype.is::() { + return Ok(ArrowExport::Unsupported(array)); + } + let Ok(target_type) = target.try_extension_type::() else { + return Ok(ArrowExport::Unsupported(array)); + }; + if target_type.coord_type() != CoordType::Separated { + return Ok(ArrowExport::Unsupported(array)); + } + let DataType::Union(target_fields, UnionMode::Dense) = target.data_type() else { + return Ok(ArrowExport::Unsupported(array)); + }; + + let extension = array.execute::(ctx)?; + let parts = geometry_union_parts(extension.storage_array().clone(), ctx)?; + let mut known_type_ids = [false; 256]; + for source_id in parts.variants.type_ids() { + known_type_ids[usize::from(*source_id)] = true; + let source_id = i8::try_from(*source_id) + .map_err(|_| vortex_err!("GeoArrow geometry type ID {source_id} exceeds i8"))?; + vortex_ensure!( + target_fields + .iter() + .any(|(target_id, _)| target_id == source_id), + "target GeoArrow geometry union is missing type ID {source_id}" + ); + } + + let validity = parts + .type_ids + .validity()? + .execute_mask(parts.type_ids.len(), ctx)?; + let fallback_type_id = parts + .variants + .type_ids() + .first() + .copied() + .ok_or_else(|| vortex_err!("Geometry union has no variants"))?; + let arrow_type_ids = parts + .type_ids + .as_slice::() + .iter() + .zip(validity.iter()) + .map(|(type_id, valid)| { + let type_id = match (known_type_ids[usize::from(*type_id)], valid) { + (true, _) => *type_id, + (false, false) => fallback_type_id, + (false, true) => { + vortex_bail!("Geometry row has unknown type ID {type_id}") + } + }; + i8::try_from(type_id) + .map_err(|_| vortex_err!("GeoArrow geometry type ID {type_id} exceeds i8")) + }) + .collect::>>()?; + let arrow_offsets = parts.offsets.as_slice::().to_vec(); + + let session = ctx.session().clone(); + let mut arrow_children = Vec::with_capacity(target_fields.len()); + for (type_id, child_field) in target_fields.iter() { + let source_id = u8::try_from(type_id) + .map_err(|_| vortex_err!("GeoArrow geometry type ID {type_id} is negative"))?; + let Some(child) = parts.child(source_id) else { + arrow_children.push(new_empty_array(child_field.data_type())); + continue; + }; + let child = child.clone().execute::(ctx)?; + arrow_children.push(session.arrow().execute_arrow( + child.storage_array().clone(), + Some(child_field.as_ref()), + ctx, + )?); + } + + let union = ArrowUnionArray::try_new( + target_fields.clone(), + ScalarBuffer::from(arrow_type_ids), + Some(ScalarBuffer::from(arrow_offsets)), + arrow_children, + ) + .map_err(|e| vortex_err!("failed to construct Arrow dense union: {e}"))?; + + if !target.is_nullable() { + vortex_ensure!( + validity.all_true(), + "cannot export nullable Geometry values to a non-nullable Arrow field" + ); + } + for (row, _) in validity.iter().enumerate().filter(|(_, valid)| !valid) { + let type_id = union.type_id(row); + let offset = union.value_offset(row); + vortex_ensure!( + union.child(type_id).is_null(offset), + "GeoArrow Geometry null at row {row} is not represented by a null selected child" + ); + } + + Ok(ArrowExport::Exported(Arc::new(union))) + } +} + +impl ArrowImportVTable for Geometry { + fn arrow_ext_id(&self) -> Id { + *GEOARROW_GEOMETRY + } + + fn from_arrow_field( + &self, + field: &Field, + session: &ArrowSession, + ) -> VortexResult> { + let Ok(geometry_type) = field.try_extension_type::() else { + return Ok(None); + }; + vortex_ensure!( + geometry_type.coord_type() == CoordType::Separated, + "geoarrow.geometry with interleaved coordinates is not supported; re-encode with separated coordinates" + ); + let DataType::Union(fields, UnionMode::Dense) = field.data_type() else { + vortex_bail!("geoarrow.geometry requires dense union storage"); + }; + let metadata = spatial_metadata_from_arrow(geometry_type.metadata()); + let mut names = Vec::new(); + let mut dtypes = Vec::new(); + let mut type_ids = Vec::new(); + for (type_id, child_field) in fields.iter() { + let type_id = u8::try_from(type_id) + .map_err(|_| vortex_err!("GeoArrow geometry type ID {type_id} is negative"))?; + let (kind, _) = geoarrow_type_id_parts(type_id)?; + if kind == GeoArrowGeometryKind::GeometryCollection { + continue; + } + let storage_dtype = session + .from_arrow_datatype(child_field.data_type(), child_field.is_nullable().into())?; + names.push(child_field.name().as_str()); + dtypes.push(native_child_dtype(type_id, &metadata, storage_dtype)?); + type_ids.push(type_id); + } + let variants = UnionVariants::try_new(FieldNames::from(names), dtypes, type_ids)?; + let storage_dtype = DType::Union(variants, field.is_nullable().into()); + Ok(Some(DType::Extension( + ExtDType::::try_new(metadata, storage_dtype)?.erased(), + ))) + } + + fn from_arrow_array( + &self, + array: ArrowArrayRef, + field: &Field, + dtype: &DType, + ) -> VortexResult { + let Some(ext_dtype) = dtype.as_extension_opt() else { + return Ok(ArrowImport::Unsupported(array)); + }; + if !ext_dtype.is::() { + return Ok(ArrowImport::Unsupported(array)); + } + let Some(union) = array.as_any().downcast_ref::() else { + return Ok(ArrowImport::Unsupported(array)); + }; + let DataType::Union(fields, UnionMode::Dense) = union.data_type() else { + return Ok(ArrowImport::Unsupported(array)); + }; + let offsets = union + .offsets() + .ok_or_else(|| vortex_err!("geoarrow.geometry requires dense union offsets"))?; + let type_ids = union + .type_ids() + .iter() + .map(|type_id| { + let type_id = u8::try_from(*type_id) + .map_err(|_| vortex_err!("GeoArrow geometry type ID {type_id} is negative"))?; + let (kind, _) = geoarrow_type_id_parts(type_id)?; + vortex_ensure!( + kind != GeoArrowGeometryKind::GeometryCollection, + "GeoArrow GeometryCollection values are not supported yet" + ); + Ok(type_id) + }) + .collect::>>()?; + + let (variants, nullability) = geometry_variants(ext_dtype.storage_dtype())?; + let mut children = Vec::with_capacity(variants.len()); + for (type_id, child_dtype) in variants.type_ids().iter().zip(variants.variants()) { + let arrow_id = i8::try_from(*type_id) + .map_err(|_| vortex_err!("GeoArrow geometry type ID {type_id} exceeds i8"))?; + let child_field = fields + .iter() + .find_map(|(candidate, field)| (candidate == arrow_id).then_some(field)) + .ok_or_else(|| vortex_err!("missing GeoArrow geometry child {type_id}"))?; + let storage = + ArrayRef::from_arrow(union.child(arrow_id).as_ref(), child_field.is_nullable())?; + let child_ext = child_dtype.as_extension(); + children.push(ExtensionArray::try_new(child_ext.clone(), storage)?.into_array()); + } + + let row_validity = union + .type_ids() + .iter() + .zip(offsets.iter()) + .enumerate() + .map(|(row, (type_id, offset))| { + let offset = usize::try_from(*offset) + .map_err(|_| vortex_err!("negative GeoArrow union offset at row {row}"))?; + let child = union.child(*type_id); + vortex_ensure!( + offset < child.len(), + "GeoArrow union offset {offset} is out of bounds at row {row}" + ); + Ok(child.is_valid(offset)) + }) + .collect::>>()?; + if !field.is_nullable() { + vortex_ensure!( + row_validity.iter().all(|valid| *valid), + "non-nullable geoarrow.geometry field contains null values" + ); + } + let validity = match nullability { + Nullability::NonNullable => Validity::NonNullable, + Nullability::Nullable => row_validity.into_iter().collect(), + }; + let type_ids = PrimitiveArray::new(type_ids, validity).into_array(); + let offsets = PrimitiveArray::from_iter(offsets.iter().copied()).into_array(); + let dense = + DenseUnion::try_new(type_ids, offsets, variants.clone(), children)?.into_array(); + Ok(ArrowImport::Imported( + ExtensionArray::try_new(ext_dtype.clone(), dense)?.into_array(), + )) + } +} diff --git a/vortex-spatial/src/extension/mod.rs b/vortex-spatial/src/extension/mod.rs index d1e2c37ebf4..bfac5db583f 100644 --- a/vortex-spatial/src/extension/mod.rs +++ b/vortex-spatial/src/extension/mod.rs @@ -2,6 +2,7 @@ // SPDX-FileCopyrightText: Copyright the Vortex contributors pub(crate) mod coordinate; +mod geometry; mod linestring; mod multilinestring; mod multipoint; @@ -16,7 +17,7 @@ use std::sync::Arc; use ::wkb::reader::GeometryType; use arrow_array::BinaryArray; -use geo_types::Geometry; +use geo_types::Geometry as GeoGeometry; use geoarrow::array::GenericWkbArray; use geoarrow::array::GeoArrowArray; use geoarrow::datatypes::CoordType; @@ -32,6 +33,7 @@ use geoarrow::datatypes::PointType; use geoarrow::datatypes::PolygonType; use geoarrow::datatypes::WkbType; use geoarrow_cast::cast::cast; +pub use geometry::*; pub use linestring::*; pub use multilinestring::*; pub use multipoint::*; @@ -73,6 +75,7 @@ pub(crate) fn is_native_geometry(dtype: &DType) -> bool { || ext.is::() || ext.is::() || ext.is::() + || ext.is::() || ext.is::() }) } @@ -143,17 +146,20 @@ pub(crate) fn flatten_row_offsets( Ok((row_offsets, level.execute::(ctx)?)) } -/// Decode a native geometry column to `geo_types`. A non-geometry operand is an error. -pub(crate) fn geometries( +/// Decode a native geometry column to the row-oriented `geo_types` representation. +pub(crate) fn decode_geometries( array: &ArrayRef, ctx: &mut ExecutionCtx, -) -> VortexResult>> { +) -> VortexResult>> { let Some(ext) = array.dtype().as_extension_opt() else { vortex_bail!( "spatial: operand is not a geometry extension type, was {}", array.dtype() ); }; + if ext.is::() { + return decode_mixed_geometries(array, ctx); + } let storage = array .clone() .execute::(ctx)? @@ -180,12 +186,12 @@ pub(crate) fn geometries( /// Decode a constant operand scalar to one geometry, a constant of any /// supported geometry type is decoded exactly like a column. -pub(crate) fn single_geometry( +pub(crate) fn decode_geometry_scalar( scalar: &Scalar, ctx: &mut ExecutionCtx, -) -> VortexResult> { +) -> VortexResult> { let array = ConstantArray::new(scalar.clone(), 1).into_array(); - geometries(&array, ctx)? + decode_geometries(&array, ctx)? .pop() .ok_or_else(|| vortex_err!("spatial: constant operand decoded to no geometry")) } diff --git a/vortex-spatial/src/lib.rs b/vortex-spatial/src/lib.rs index 8d89c2f322f..a73bc4c96fb 100644 --- a/vortex-spatial/src/lib.rs +++ b/vortex-spatial/src/lib.rs @@ -11,6 +11,7 @@ use vortex_arrow::ArrowSessionExt; use vortex_session::VortexSession; use crate::aggregate_fn::GeometryAabb; +use crate::extension::Geometry; use crate::extension::LineString; use crate::extension::MultiLineString; use crate::extension::MultiPoint; @@ -40,6 +41,8 @@ mod tests; /// Set up a session with support for spatial extension types, encodings and layouts. pub fn initialize(session: &VortexSession) { + vortex_dense_union::initialize(session); + // Register the spatial extension types. session.dtypes().register(WellKnownBinary); session.arrow().register_exporter(Arc::new(WellKnownBinary)); @@ -65,6 +68,9 @@ pub fn initialize(session: &VortexSession) { session.dtypes().register(Rect); session.arrow().register_exporter(Arc::new(Rect)); session.arrow().register_importer(Arc::new(Rect)); + session.dtypes().register(Geometry); + session.arrow().register_exporter(Arc::new(Geometry)); + session.arrow().register_importer(Arc::new(Geometry)); // Register the geometry scalar functions. session.scalar_fns().register(SpatialArea); diff --git a/vortex-spatial/src/prune/mod.rs b/vortex-spatial/src/prune/mod.rs index 32c2058b384..d3455a817e7 100644 --- a/vortex-spatial/src/prune/mod.rs +++ b/vortex-spatial/src/prune/mod.rs @@ -41,8 +41,8 @@ use vortex_array::stats::rewrite::StatsRewriteCtx; use vortex_error::VortexResult; use crate::aggregate_fn::GeometryAabb; +use crate::extension::decode_geometry_scalar; use crate::extension::is_native_geometry; -use crate::extension::single_geometry; /// Splits a symmetric two-operand spatial predicate into the scope-rooted geometry column and the /// constant operand's scalar. @@ -91,7 +91,7 @@ fn query_aabb( // Decoding the constant into a concrete geometry runs through the compute stack, which needs // an execution context. let mut exec = ctx.session().create_execution_ctx(); - Ok(single_geometry(constant, &mut exec)?.bounding_rect()) + Ok(decode_geometry_scalar(constant, &mut exec)?.bounding_rect()) } /// The chunk's AABB statistic, as the storage struct with `xmin`/`ymin`/`xmax`/`ymax` fields. diff --git a/vortex-spatial/src/scalar_fn/envelope.rs b/vortex-spatial/src/scalar_fn/envelope.rs index f1e8ca33b1a..63c0c84f1fc 100644 --- a/vortex-spatial/src/scalar_fn/envelope.rs +++ b/vortex-spatial/src/scalar_fn/envelope.rs @@ -6,6 +6,7 @@ //! A row-oriented consumer (e.g. bulk-loading an in-memory R-tree in a spatial-join operator) //! reads the resulting box column back row by row. +use geo::BoundingRect; use vortex_array::ArrayRef; use vortex_array::ExecutionCtx; use vortex_array::IntoArray; @@ -37,6 +38,7 @@ use vortex_mask::Mask; use vortex_session::VortexSession; use vortex_session::registry::CachedId; +use crate::extension::Geometry; use crate::extension::Rect; use crate::extension::SpatialMetadata; use crate::extension::box_field_names; @@ -45,6 +47,7 @@ use crate::extension::build_rect_array; 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_row_offsets; use crate::extension::is_native_geometry; use crate::scalar_fn::execute::Execution; @@ -138,6 +141,55 @@ fn row_boxes( )) } +fn mixed_geometry_boxes( + array: ArrayRef, + valid: Mask, + output_dtype: &ExtDType, + ctx: &mut ExecutionCtx, +) -> VortexResult { + let len = array.len(); + let decoded = decode_mixed_geometries(&array.filter(valid.clone())?, ctx)?; + let mut decoded = decoded.into_iter(); + let mut xmins = BufferMut::zeroed(len); + let mut ymins = BufferMut::zeroed(len); + let mut xmaxs = BufferMut::zeroed(len); + let mut ymaxs = BufferMut::zeroed(len); + let mut output_validity = vec![false; len]; + + for (row, valid) in valid.iter().enumerate() { + if !valid { + continue; + } + let geometry = decoded + .next() + .ok_or_else(|| vortex_err!("Geometry decode returned too few rows"))?; + let Some(rect) = geometry.bounding_rect() else { + continue; + }; + xmins[row] = rect.min().x; + ymins[row] = rect.min().y; + xmaxs[row] = rect.max().x; + ymaxs[row] = rect.max().y; + output_validity[row] = true; + } + vortex_ensure!( + decoded.next().is_none(), + "Geometry decode returned too many rows" + ); + + build_rect_array( + output_dtype, + vec![ + xmins.freeze().into_array(), + ymins.freeze().into_array(), + xmaxs.freeze().into_array(), + ymaxs.freeze().into_array(), + ], + len, + Validity::from_iter(output_validity), + ) +} + /// Compute boxes directly over a non-constant native geometry column. fn envelope_array( array: ArrayRef, @@ -146,11 +198,15 @@ fn envelope_array( ctx: &mut ExecutionCtx, ) -> VortexResult { let len = array.len(); - let is_rect = array + let ext_dtype = array .dtype() .as_extension_opt() - .ok_or_else(|| vortex_err!("spatial: envelope operand is not a geometry extension type"))? - .is::(); + .ok_or_else(|| vortex_err!("spatial: envelope operand is not a geometry extension type"))?; + if ext_dtype.is::() { + let valid = validity.execute_mask(len, ctx)?; + return mixed_geometry_boxes(array, valid, output_dtype, ctx); + } + let is_rect = ext_dtype.is::(); let storage = array .execute::(ctx)? .storage_array() @@ -237,9 +293,9 @@ impl ScalarFnVTable for SpatialEnvelope { Ok(DType::Extension(output_box_dtype()?.erased())) } - /// Compute each row's box directly over the native coordinate storage — no decode to - /// `geo_types`, no Arrow round-trip. A null row, or a valid row that owns no coordinate (an - /// empty geometry), yields a null box. + /// Compute each row's box directly over homogeneous native coordinate storage. Geometry + /// unions use the row-oriented fallback. A null row, or a valid row that owns no coordinate + /// (an empty geometry), yields a null box. fn execute( &self, _: &Self::Options, diff --git a/vortex-spatial/src/scalar_fn/execute/binary.rs b/vortex-spatial/src/scalar_fn/execute/binary.rs index f2c03bd1beb..fda6b806487 100644 --- a/vortex-spatial/src/scalar_fn/execute/binary.rs +++ b/vortex-spatial/src/scalar_fn/execute/binary.rs @@ -4,7 +4,7 @@ //! Binary pairwise dispatch, plus an adapter for row-oriented `geo_types` kernels. use geo::BoundingRect; -use geo_types::Geometry; +use geo_types::Geometry as GeoGeometry; use geo_types::Rect; use vortex_array::ArrayRef; use vortex_array::ExecutionCtx; @@ -22,7 +22,7 @@ use super::Operand; use super::geo_types::GeoTypesOutput; use super::geo_types::eval_column; use super::geo_types::eval_column_pair; -use crate::extension::single_geometry; +use crate::extension::decode_geometry_scalar; /// Dispatch a binary strict geometry kernel over constants and columns. /// @@ -112,7 +112,7 @@ pub(crate) fn execute_binary_geo_types( ) -> VortexResult where T: GeoTypesOutput, - F: Fn(&Geometry, &Geometry) -> T + Copy, + F: Fn(&GeoGeometry, &GeoGeometry) -> T + Copy, { let nullability = Nullability::from(left.dtype().is_nullable() || right.dtype().is_nullable()); dispatch_binary( @@ -121,8 +121,8 @@ where T::dtype(nullability), |execution, ctx| match execution.operands { [Operand::Constant(left), Operand::Constant(right)] => { - let left = single_geometry(&left, ctx)?; - let right = single_geometry(&right, ctx)?; + let left = decode_geometry_scalar(&left, ctx)?; + let right = decode_geometry_scalar(&right, ctx)?; Ok(ConstantArray::new( compute(&left, &right).into_scalar(execution.nullability), execution.len, @@ -130,7 +130,7 @@ where .into_array()) } [Operand::Constant(left), Operand::Column(right)] => { - let left = single_geometry(&left, ctx)?; + let left = decode_geometry_scalar(&left, ctx)?; let prescreen = bbox_precheck.zip(left.bounding_rect()); eval_column( &right, @@ -145,7 +145,7 @@ where ) } [Operand::Column(left), Operand::Constant(right)] => { - let right = single_geometry(&right, ctx)?; + let right = decode_geometry_scalar(&right, ctx)?; let prescreen = bbox_precheck.zip(right.bounding_rect()); eval_column( &left, diff --git a/vortex-spatial/src/scalar_fn/execute/geo_types.rs b/vortex-spatial/src/scalar_fn/execute/geo_types.rs index 038aca46502..7133775cc83 100644 --- a/vortex-spatial/src/scalar_fn/execute/geo_types.rs +++ b/vortex-spatial/src/scalar_fn/execute/geo_types.rs @@ -6,7 +6,7 @@ //! `geo_types` is the row representation consumed by the kernel. These helpers always construct //! and return Vortex arrays; they do not expose `geo_types` values as scalar-function outputs. -use geo_types::Geometry; +use geo_types::Geometry as GeoGeometry; use vortex_array::ArrayRef; use vortex_array::ExecutionCtx; use vortex_array::IntoArray; @@ -22,7 +22,7 @@ use vortex_error::VortexResult; use vortex_mask::AllOr; use vortex_mask::Mask; -use crate::extension::geometries; +use crate::extension::decode_geometries; /// A primitive result produced after kernel inputs are decoded to `geo_types`. pub(crate) trait GeoTypesOutput: Copy { @@ -111,10 +111,10 @@ pub(super) fn eval_column( ) -> VortexResult where T: GeoTypesOutput, - F: Fn(&Geometry) -> T, + F: Fn(&GeoGeometry) -> T, { let len = column.len(); - let decoded = geometries(&column.filter(valid.clone())?, ctx)?; + let decoded = decode_geometries(&column.filter(valid.clone())?, ctx)?; let values = decoded.iter().map(compute).collect(); Ok(T::build_array(len, valid, values, nullability)) } @@ -130,11 +130,11 @@ pub(super) fn eval_column_pair( ) -> VortexResult where T: GeoTypesOutput, - F: Fn(&Geometry, &Geometry) -> T, + F: Fn(&GeoGeometry, &GeoGeometry) -> T, { let len = left.len(); - let left = geometries(&left.filter(valid.clone())?, ctx)?; - let right = geometries(&right.filter(valid.clone())?, ctx)?; + let left = decode_geometries(&left.filter(valid.clone())?, ctx)?; + let right = decode_geometries(&right.filter(valid.clone())?, ctx)?; let values = left .iter() .zip(&right) diff --git a/vortex-spatial/src/scalar_fn/execute/unary.rs b/vortex-spatial/src/scalar_fn/execute/unary.rs index a8bb74f850c..63aaae2a677 100644 --- a/vortex-spatial/src/scalar_fn/execute/unary.rs +++ b/vortex-spatial/src/scalar_fn/execute/unary.rs @@ -3,7 +3,7 @@ //! Unary operand dispatch, plus an adapter for row-oriented `geo_types` kernels. -use geo_types::Geometry; +use geo_types::Geometry as GeoGeometry; use vortex_array::ArrayRef; use vortex_array::ExecutionCtx; use vortex_array::IntoArray; @@ -18,7 +18,7 @@ use super::Execution; use super::Operand; use super::geo_types::GeoTypesOutput; use super::geo_types::eval_column; -use crate::extension::single_geometry; +use crate::extension::decode_geometry_scalar; /// Dispatch a unary strict geometry kernel over a constant or column. /// @@ -77,7 +77,7 @@ pub(crate) fn execute_unary_geo_types( ) -> VortexResult where T: GeoTypesOutput, - F: Fn(&Geometry) -> T, + F: Fn(&GeoGeometry) -> T, { let nullability = array.dtype().nullability(); dispatch_unary( @@ -85,7 +85,7 @@ where T::dtype(nullability), |execution, ctx| match execution.operands { [Operand::Constant(scalar)] => { - let geometry = single_geometry(&scalar, ctx)?; + let geometry = decode_geometry_scalar(&scalar, ctx)?; Ok(ConstantArray::new( compute(&geometry).into_scalar(execution.nullability), execution.len, diff --git a/vortex-spatial/src/scalar_fn/make_line.rs b/vortex-spatial/src/scalar_fn/make_line.rs index c32d4c76691..75f1bc17b77 100644 --- a/vortex-spatial/src/scalar_fn/make_line.rs +++ b/vortex-spatial/src/scalar_fn/make_line.rs @@ -270,8 +270,8 @@ mod tests { use crate::extension::coordinate::Dimension; use crate::extension::coordinate::coordinate_dimension; use crate::extension::coordinate::ordinates; + use crate::extension::decode_geometries; use crate::extension::flatten_coordinates; - use crate::extension::geometries; use crate::test_harness::point_column; fn dimensional_point( @@ -319,7 +319,7 @@ mod tests { let lines = SpatialMakeLine::try_new_array(starts, ends)?.into_array(); assert!(lines.dtype().as_extension().is::()); assert_eq!( - geometries(&lines, &mut ctx)?, + decode_geometries(&lines, &mut ctx)?, vec![ Geometry::LineString(GeoLineString::new(vec![ Coord { x: 0.0, y: 0.0 }, @@ -351,7 +351,7 @@ mod tests { }; assert_eq!(lines.len(), 3); assert_eq!( - geometries(&lines.into_array(), &mut ctx)?, + decode_geometries(&lines.into_array(), &mut ctx)?, vec![ Geometry::LineString(GeoLineString::new(vec![ Coord { x: 0.0, y: 0.0 }, @@ -391,7 +391,7 @@ mod tests { })) }) .collect::>(); - assert_eq!(geometries(&lines, &mut ctx)?, expected); + assert_eq!(decode_geometries(&lines, &mut ctx)?, expected); Ok(()) } diff --git a/vortex-spatial/src/tests/geometry.rs b/vortex-spatial/src/tests/geometry.rs new file mode 100644 index 00000000000..fcd5bfe958b --- /dev/null +++ b/vortex-spatial/src/tests/geometry.rs @@ -0,0 +1,335 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Arrow interop for the mixed `vortex.st.geometry` extension (`geoarrow.geometry`). + +use std::sync::Arc; + +use arrow_array::ArrayRef as ArrowArrayRef; +use arrow_array::UnionArray as ArrowUnionArray; +use arrow_schema::Field; +use geo_traits::to_geo::ToGeoGeometry; +use geo_types::Geometry as GeoGeometry; +use geo_types::GeometryCollection; +use geo_types::LineString; +use geo_types::MultiLineString; +use geo_types::MultiPoint; +use geo_types::MultiPolygon; +use geo_types::Point; +use geo_types::Polygon; +use geoarrow::array::GeoArrowArrayAccessor; +use geoarrow::array::GeometryArray as GeoArrowGeometryArray; +use geoarrow::array::GeometryBuilder; +use geoarrow::array::IntoArrow; +use geoarrow::datatypes::CoordType; +use geoarrow::datatypes::Crs; +use geoarrow::datatypes::GeometryType as GeoArrowGeometryType; +use geoarrow::datatypes::Metadata; +use vortex_array::IntoArray; +use vortex_array::VortexSessionExecute; +use vortex_array::aggregate_fn::Accumulator; +use vortex_array::aggregate_fn::DynAccumulator; +use vortex_array::aggregate_fn::EmptyOptions; +use vortex_array::arrays::ConstantArray; +use vortex_array::arrays::ExtensionArray; +use vortex_array::arrays::UnionArray as SparseUnionArray; +use vortex_array::arrays::extension::ExtensionArrayExt; +use vortex_array::assert_arrays_eq; +use vortex_array::dtype::DType; +use vortex_array::scalar::Scalar; +use vortex_arrow::ArrowSessionExt; +use vortex_dense_union::DenseUnion; +use vortex_error::VortexResult; +use vortex_error::vortex_err; + +use super::SESSION; +use crate::aggregate_fn::GeometryAabb; +use crate::extension::Geometry; +use crate::extension::decode_geometries; +use crate::scalar_fn::envelope::SpatialEnvelope; +use crate::test_harness::nullable_rect_column; + +fn polygon(points: &[(f64, f64)]) -> Polygon { + Polygon::new(LineString::from(points.to_vec()), vec![]) +} + +fn supported_geometries() -> Vec>> { + let line = LineString::from(vec![(3.0, -4.0), (8.0, 7.0)]); + let polygon = polygon(&[(0.0, 0.0), (4.0, 0.0), (4.0, 3.0), (0.0, 0.0)]); + vec![ + Some(GeoGeometry::Point(Point::new(1.0, 2.0))), + Some(GeoGeometry::LineString(line.clone())), + Some(GeoGeometry::Polygon(polygon.clone())), + Some(GeoGeometry::MultiPoint(MultiPoint::new(vec![ + Point::new(-2.0, 5.0), + Point::new(4.0, 9.0), + ]))), + Some(GeoGeometry::MultiLineString(MultiLineString::new(vec![ + line, + ]))), + Some(GeoGeometry::MultiPolygon(MultiPolygon::new(vec![polygon]))), + None, + ] +} + +fn geoarrow_geometry_type() -> GeoArrowGeometryType { + let crs = Crs::from_unknown_crs_type("EPSG:4326".to_string()); + GeoArrowGeometryType::new(Arc::new(Metadata::new(crs, None))) + .with_coord_type(CoordType::Separated) +} + +fn arrow_fixture(geometries: &[Option>]) -> VortexResult<(ArrowArrayRef, Field)> { + let geometry_type = geoarrow_geometry_type(); + let array = GeometryBuilder::from_nullable_geometries(geometries, geometry_type.clone()) + .map_err(|e| vortex_err!("failed to build GeoArrow geometry array: {e}"))? + .finish(); + let field = geometry_type.to_field("geom", true); + Ok((Arc::new(array.into_arrow()), field)) +} + +fn decode_arrow( + array: &ArrowArrayRef, + field: &Field, +) -> VortexResult>>> { + let geometries = GeoArrowGeometryArray::try_from((array.as_ref(), field)) + .map_err(|e| vortex_err!("failed to decode exported GeoArrow geometry array: {e}"))?; + geometries + .iter() + .map(|geometry| match geometry { + None => Ok(None), + Some(Ok(geometry)) => Ok(Some(geometry.to_geometry())), + Some(Err(e)) => Err(vortex_err!("failed to access GeoArrow geometry: {e}")), + }) + .collect() +} + +fn aabb(result: &Scalar) -> VortexResult<(f64, f64, f64, f64)> { + let storage = result.as_extension().to_storage_scalar(); + let fields = storage.as_struct(); + let read = |name: &str| -> VortexResult { + f64::try_from( + &fields + .field(name) + .ok_or_else(|| vortex_err!("AABB result is missing {name}"))?, + ) + }; + Ok((read("xmin")?, read("ymin")?, read("xmax")?, read("ymax")?)) +} + +#[test] +fn imports_as_geometry_over_logical_union() -> VortexResult<()> { + let field = geoarrow_geometry_type().to_field("geom", true); + let dtype = SESSION.arrow().from_arrow_field(&field)?; + + let DType::Extension(ext) = &dtype else { + return Err(vortex_err!( + "expected Geometry extension dtype, got {dtype}" + )); + }; + assert!(ext.is::()); + assert_eq!(ext.metadata::().crs.as_deref(), Some("EPSG:4326")); + + let DType::Union(variants, nullability) = ext.storage_dtype() else { + return Err(vortex_err!( + "expected logical Union storage, got {}", + ext.storage_dtype() + )); + }; + assert!(nullability.is_nullable()); + assert_eq!(variants.len(), 24); + assert!(variants.type_ids().iter().all(|type_id| type_id % 10 != 7)); + Ok(()) +} + +#[test] +fn infers_canonical_geoarrow_field() -> VortexResult<()> { + let expected_type = geoarrow_geometry_type(); + let dtype = SESSION + .arrow() + .from_arrow_field(&expected_type.to_field("input", true))?; + let field = SESSION.arrow().to_arrow_field("geom", &dtype)?; + + assert_eq!(field.name(), "geom"); + assert!(field.is_nullable()); + assert_eq!(field.data_type(), &expected_type.data_type()); + let actual_type = field + .try_extension_type::() + .map_err(|e| vortex_err!("failed to read inferred GeoArrow GeometryType: {e}"))?; + assert_eq!(actual_type, expected_type); + Ok(()) +} + +#[test] +fn roundtrips_supported_kinds_and_nulls() -> VortexResult<()> { + let expected = supported_geometries(); + let (arrow, field) = arrow_fixture(&expected)?; + let imported = SESSION.arrow().from_arrow_array(arrow, &field)?; + + let mut ctx = SESSION.create_execution_ctx(); + let extension = imported.clone().execute::(&mut ctx)?; + assert!(extension.storage_array().is::()); + assert_eq!( + imported + .validity()? + .execute_mask(imported.len(), &mut ctx)? + .iter() + .collect::>(), + vec![true, true, true, true, true, true, false] + ); + + let exported = SESSION + .arrow() + .execute_arrow(imported, Some(&field), &mut ctx)?; + assert_eq!(decode_arrow(&exported, &field)?, expected); + Ok(()) +} + +#[test] +fn preserves_non_first_type_id_for_nulls() -> VortexResult<()> { + let expected = vec![ + Some(GeoGeometry::LineString(LineString::from(vec![ + (0.0, 1.0), + (2.0, 3.0), + ]))), + None, + ]; + let (arrow, field) = arrow_fixture(&expected)?; + let input = arrow + .as_any() + .downcast_ref::() + .ok_or_else(|| vortex_err!("GeoArrow fixture must use a UnionArray"))?; + assert_eq!(input.type_ids().as_ref(), &[2, 2]); + + let imported = SESSION.arrow().from_arrow_array(arrow, &field)?; + let mut ctx = SESSION.create_execution_ctx(); + let exported = SESSION + .arrow() + .execute_arrow(imported, Some(&field), &mut ctx)?; + let output = exported + .as_any() + .downcast_ref::() + .ok_or_else(|| vortex_err!("exported GeoArrow Geometry must use a UnionArray"))?; + assert_eq!(output.type_ids().as_ref(), &[2, 2]); + assert_eq!(decode_arrow(&exported, &field)?, expected); + Ok(()) +} + +#[test] +fn slice_preserves_dense_union_storage() -> VortexResult<()> { + let expected = supported_geometries(); + let (arrow, field) = arrow_fixture(&expected)?; + let imported = SESSION.arrow().from_arrow_array(arrow, &field)?; + let sliced = imported.slice(1..5)?; + + let mut ctx = SESSION.create_execution_ctx(); + let extension = sliced.clone().execute::(&mut ctx)?; + assert!(extension.storage_array().is::()); + + let exported = SESSION + .arrow() + .execute_arrow(sliced, Some(&field), &mut ctx)?; + assert_eq!(decode_arrow(&exported, &field)?, expected[1..5]); + Ok(()) +} + +#[test] +fn exports_canonical_sparse_union() -> VortexResult<()> { + let expected = supported_geometries(); + let (arrow, field) = arrow_fixture(&expected)?; + let imported = SESSION.arrow().from_arrow_array(arrow, &field)?; + let mut ctx = SESSION.create_execution_ctx(); + + let extension = imported.execute::(&mut ctx)?; + let sparse = extension + .storage_array() + .clone() + .execute::(&mut ctx)?; + let sparse_geometry = + ExtensionArray::try_new(extension.ext_dtype().clone(), sparse.into_array())?.into_array(); + let exported = SESSION + .arrow() + .execute_arrow(sparse_geometry, Some(&field), &mut ctx)?; + assert_eq!(decode_arrow(&exported, &field)?, expected); + Ok(()) +} + +#[test] +fn exports_constant_nulls() -> VortexResult<()> { + let expected = supported_geometries(); + let (arrow, field) = arrow_fixture(&expected)?; + let imported = SESSION.arrow().from_arrow_array(arrow, &field)?; + let mut ctx = SESSION.create_execution_ctx(); + let nulls = ConstantArray::new(Scalar::null(imported.dtype().clone()), 2).into_array(); + let exported = SESSION + .arrow() + .execute_arrow(nulls, Some(&field), &mut ctx)?; + assert_eq!(decode_arrow(&exported, &field)?, vec![None, None]); + Ok(()) +} + +#[test] +fn decodes_supported_variants() -> VortexResult<()> { + let expected = supported_geometries(); + let (arrow, field) = arrow_fixture(&expected)?; + let imported = SESSION.arrow().from_arrow_array(arrow, &field)?; + let mut ctx = SESSION.create_execution_ctx(); + + let non_null = imported.slice(0..6)?; + let expected_non_null = expected[..6] + .iter() + .filter_map(Clone::clone) + .collect::>(); + assert_eq!(decode_geometries(&non_null, &mut ctx)?, expected_non_null); + Ok(()) +} + +#[test] +fn computes_aabb_across_variants() -> VortexResult<()> { + let expected = supported_geometries(); + let (arrow, field) = arrow_fixture(&expected)?; + let imported = SESSION.arrow().from_arrow_array(arrow, &field)?; + let mut ctx = SESSION.create_execution_ctx(); + let mut accumulator = + Accumulator::try_new(GeometryAabb, EmptyOptions, imported.dtype().clone())?; + accumulator.accumulate(&imported, &mut ctx)?; + assert_eq!(aabb(&accumulator.finish()?)?, (-2.0, -4.0, 8.0, 9.0)); + Ok(()) +} + +#[test] +fn computes_envelopes_across_variants() -> VortexResult<()> { + let expected = supported_geometries(); + let (arrow, field) = arrow_fixture(&expected)?; + let imported = SESSION.arrow().from_arrow_array(arrow, &field)?; + let mut ctx = SESSION.create_execution_ctx(); + let envelopes = SpatialEnvelope::try_new_array(imported)? + .into_array() + .execute::(&mut ctx)?; + let expected_envelopes = nullable_rect_column(vec![ + Some((1.0, 2.0, 1.0, 2.0)), + Some((3.0, -4.0, 8.0, 7.0)), + Some((0.0, 0.0, 4.0, 3.0)), + Some((-2.0, 5.0, 4.0, 9.0)), + Some((3.0, -4.0, 8.0, 7.0)), + Some((0.0, 0.0, 4.0, 3.0)), + None, + ])?; + assert_arrays_eq!(envelopes, expected_envelopes, &mut ctx); + Ok(()) +} + +#[test] +fn rejects_selected_geometry_collection() -> VortexResult<()> { + let geometries = vec![Some(GeoGeometry::GeometryCollection( + GeometryCollection::new_from(vec![ + GeoGeometry::Point(Point::new(1.0, 2.0)), + GeoGeometry::Point(Point::new(3.0, 4.0)), + ]), + ))]; + let (arrow, field) = arrow_fixture(&geometries)?; + let Err(error) = SESSION.arrow().from_arrow_array(arrow, &field) else { + return Err(vortex_err!("selected GeometryCollection must be rejected")); + }; + assert!(error.to_string().contains("GeometryCollection")); + Ok(()) +} diff --git a/vortex-spatial/src/tests/mod.rs b/vortex-spatial/src/tests/mod.rs index c83a30d476d..4734d433508 100644 --- a/vortex-spatial/src/tests/mod.rs +++ b/vortex-spatial/src/tests/mod.rs @@ -4,6 +4,7 @@ //! Arrow interop tests for the spatial extension types, exercising the session wiring set up //! by [`crate::initialize`]. +mod geometry; mod linestring; mod multilinestring; mod multipoint; diff --git a/vortex-spatial/src/tests/rect.rs b/vortex-spatial/src/tests/rect.rs index c1f752cfb21..0797158b680 100644 --- a/vortex-spatial/src/tests/rect.rs +++ b/vortex-spatial/src/tests/rect.rs @@ -110,7 +110,8 @@ fn roundtrips_through_arrow() -> VortexResult<()> { Ok(()) } -/// The existing spatial scalar functions run on a `Rect` operand via the shared `geometries()` decode, +/// The existing spatial scalar functions run on a `Rect` operand via the shared +/// `decode_geometries()` path, /// producing the same results as the equivalent polygon: a box `(0,0)-(10,10)` against interior /// point `(5,5)` and exterior point `(20,20)`. #[test]