From c09aa3a01c58810665dac3a7e4fa40fc23c90c60 Mon Sep 17 00:00:00 2001 From: Nemo Yu Date: Tue, 7 Jul 2026 16:53:16 -0400 Subject: [PATCH] feat(geo): add native LineString, MultiPoint, MultiLineString types Complete the set of OGC simple-feature geometries alongside the existing Point/Polygon/MultiPolygon: - vortex.geo.linestring - List> - vortex.geo.multipoint - List> - vortex.geo.multilinestring - List>> Each mirrors the existing Polygon/MultiPolygon extension types: GeoArrow import/export with separated (struct) coordinates, WKB literal decoding, geo_types conversion for the geo scalar functions, and registration in initialize(). LineString/MultiPoint share a storage shape (as do MultiLineString/Polygon); they are disambiguated by their GeoArrow extension name and Vortex extension id, not by shape. Wire them through the DuckDB integration so they surface as GEOMETRY like the existing native types: the dtype -> GEOMETRY mapping, per-type WKB exporters, and native-geometry-column recognition for predicate pushdown. Tests cover dtype validation across all four dimensions, Arrow field import/export (including that the shape-identical types resolve by name), and WKB-literal decode round-trips. Signed-off-by: Nemo Yu --- vortex-duckdb/src/convert/dtype.rs | 8 +- vortex-duckdb/src/convert/expr.rs | 12 +- vortex-duckdb/src/exporter/extension.rs | 18 + vortex-duckdb/src/exporter/geo.rs | 35 ++ vortex-geo/src/extension/linestring.rs | 341 +++++++++++++++++++ vortex-geo/src/extension/mod.rs | 110 ++++++ vortex-geo/src/extension/multilinestring.rs | 359 ++++++++++++++++++++ vortex-geo/src/extension/multipoint.rs | 336 ++++++++++++++++++ vortex-geo/src/lib.rs | 12 + vortex-geo/src/tests/linestring.rs | 88 +++++ vortex-geo/src/tests/mod.rs | 3 + vortex-geo/src/tests/multilinestring.rs | 92 +++++ vortex-geo/src/tests/multipoint.rs | 87 +++++ 13 files changed, 1499 insertions(+), 2 deletions(-) create mode 100644 vortex-geo/src/extension/linestring.rs create mode 100644 vortex-geo/src/extension/multilinestring.rs create mode 100644 vortex-geo/src/extension/multipoint.rs create mode 100644 vortex-geo/src/tests/linestring.rs create mode 100644 vortex-geo/src/tests/multilinestring.rs create mode 100644 vortex-geo/src/tests/multipoint.rs diff --git a/vortex-duckdb/src/convert/dtype.rs b/vortex-duckdb/src/convert/dtype.rs index d05fdb9f22e..6837bd778c2 100644 --- a/vortex-duckdb/src/convert/dtype.rs +++ b/vortex-duckdb/src/convert/dtype.rs @@ -58,6 +58,9 @@ use vortex::extension::datetime::Time; use vortex::extension::datetime::TimeUnit; use vortex::extension::datetime::Timestamp; use vortex_geo::extension::GeoMetadata; +use vortex_geo::extension::LineString; +use vortex_geo::extension::MultiLineString; +use vortex_geo::extension::MultiPoint; use vortex_geo::extension::MultiPolygon; use vortex_geo::extension::Point; use vortex_geo::extension::Polygon; @@ -248,10 +251,13 @@ impl TryFrom<&DType> for LogicalType { return temporal_to_duckdb(temporal); } - // Native Point/Polygon and WKB all surface to DuckDB as GEOMETRY so `ST_*` bind. + // Native geometry types and WKB all surface to DuckDB as GEOMETRY so `ST_*` bind. if let Some(geo) = ext_dtype .metadata_opt::() + .or_else(|| ext_dtype.metadata_opt::()) + .or_else(|| ext_dtype.metadata_opt::()) .or_else(|| ext_dtype.metadata_opt::()) + .or_else(|| ext_dtype.metadata_opt::()) .or_else(|| ext_dtype.metadata_opt::()) .or_else(|| ext_dtype.metadata_opt::()) { diff --git a/vortex-duckdb/src/convert/expr.rs b/vortex-duckdb/src/convert/expr.rs index 32778870845..7e05a633c19 100644 --- a/vortex-duckdb/src/convert/expr.rs +++ b/vortex-duckdb/src/convert/expr.rs @@ -38,6 +38,9 @@ use vortex::scalar_fn::fns::like::Like; use vortex::scalar_fn::fns::like::LikeOptions; use vortex::scalar_fn::fns::literal::Literal; use vortex::scalar_fn::fns::operators::Operator; +use vortex_geo::extension::LineString; +use vortex_geo::extension::MultiLineString; +use vortex_geo::extension::MultiPoint; use vortex_geo::extension::MultiPolygon; use vortex_geo::extension::Point; use vortex_geo::extension::Polygon; @@ -105,7 +108,14 @@ fn is_native_geo_column(fields: Option<&[DuckdbField]>, name: &str) -> bool { .flatten() .filter(|field| field.name == name) .any(|field| match field.dtype.as_extension_opt() { - Some(ext) => ext.is::() || ext.is::() || ext.is::(), + Some(ext) => { + ext.is::() + || ext.is::() + || ext.is::() + || ext.is::() + || ext.is::() + || ext.is::() + } None => false, }) } diff --git a/vortex-duckdb/src/exporter/extension.rs b/vortex-duckdb/src/exporter/extension.rs index 30eb3715cc2..c07dc25a886 100644 --- a/vortex-duckdb/src/exporter/extension.rs +++ b/vortex-duckdb/src/exporter/extension.rs @@ -8,6 +8,12 @@ use vortex::array::arrays::extension::ExtensionArrayExt; use vortex::array::extension::datetime::AnyTemporal; use vortex::error::VortexResult; use vortex::error::vortex_bail; +use vortex_geo::extension::LineString; +use vortex_geo::extension::LineStringData; +use vortex_geo::extension::MultiLineString; +use vortex_geo::extension::MultiLineStringData; +use vortex_geo::extension::MultiPoint; +use vortex_geo::extension::MultiPointData; use vortex_geo::extension::MultiPolygon; use vortex_geo::extension::MultiPolygonData; use vortex_geo::extension::Point; @@ -37,10 +43,22 @@ pub(crate) fn new_exporter( return geo::new_point_exporter(PointData::try_from(ext)?, ctx); } + if ext.ext_dtype().is::() { + return geo::new_linestring_exporter(LineStringData::try_from(ext)?, ctx); + } + + if ext.ext_dtype().is::() { + return geo::new_multipoint_exporter(MultiPointData::try_from(ext)?, ctx); + } + if ext.ext_dtype().is::() { return geo::new_polygon_exporter(PolygonData::try_from(ext)?, ctx); } + if ext.ext_dtype().is::() { + return geo::new_multilinestring_exporter(MultiLineStringData::try_from(ext)?, ctx); + } + if ext.ext_dtype().is::() { return geo::new_multipolygon_exporter(MultiPolygonData::try_from(ext)?, ctx); } diff --git a/vortex-duckdb/src/exporter/geo.rs b/vortex-duckdb/src/exporter/geo.rs index fc6634c121c..ccc1893ea18 100644 --- a/vortex-duckdb/src/exporter/geo.rs +++ b/vortex-duckdb/src/exporter/geo.rs @@ -4,6 +4,9 @@ use vortex::array::ExecutionCtx; use vortex::array::arrays::VarBinViewArray; use vortex::error::VortexResult; +use vortex_geo::extension::LineStringData; +use vortex_geo::extension::MultiLineStringData; +use vortex_geo::extension::MultiPointData; use vortex_geo::extension::MultiPolygonData; use vortex_geo::extension::PointData; use vortex_geo::extension::PolygonData; @@ -51,3 +54,35 @@ pub(crate) fn new_multipolygon_exporter( let values = multipolygon.to_wkb(ctx)?.execute::(ctx)?; new_exporter(values, ctx) } + +/// Create an exporter for a native `LineString` column, serialized to WKB via +/// [`LineStringData::to_wkb`] (see [`new_point_exporter`]). +pub(crate) fn new_linestring_exporter( + linestring: LineStringData, + ctx: &mut ExecutionCtx, +) -> VortexResult> { + let values = linestring.to_wkb(ctx)?.execute::(ctx)?; + new_exporter(values, ctx) +} + +/// Create an exporter for a native `MultiPoint` column, serialized to WKB via +/// [`MultiPointData::to_wkb`] (see [`new_point_exporter`]). +pub(crate) fn new_multipoint_exporter( + multipoint: MultiPointData, + ctx: &mut ExecutionCtx, +) -> VortexResult> { + let values = multipoint.to_wkb(ctx)?.execute::(ctx)?; + new_exporter(values, ctx) +} + +/// Create an exporter for a native `MultiLineString` column, serialized to WKB via +/// [`MultiLineStringData::to_wkb`] (see [`new_point_exporter`]). +pub(crate) fn new_multilinestring_exporter( + multilinestring: MultiLineStringData, + ctx: &mut ExecutionCtx, +) -> VortexResult> { + let values = multilinestring + .to_wkb(ctx)? + .execute::(ctx)?; + new_exporter(values, ctx) +} diff --git a/vortex-geo/src/extension/linestring.rs b/vortex-geo/src/extension/linestring.rs new file mode 100644 index 00000000000..f7d8fcff1de --- /dev/null +++ b/vortex-geo/src/extension/linestring.rs @@ -0,0 +1,341 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! The [`LineString`] geometry extension type (`vortex.geo.linestring`): an ordered path of the +//! [`Point`](super::Point) coordinate struct, stored as `List>` and tagged +//! with [`GeoMetadata`] (CRS). + +use std::sync::Arc; + +use arrow_array::ArrayRef as ArrowArrayRef; +use arrow_schema::DataType; +use arrow_schema::Field; +use arrow_schema::extension::ExtensionType; +use geo_traits::to_geo::ToGeoGeometry; +use geo_types::Geometry; +use geoarrow::array::GeoArrowArrayAccessor; +use geoarrow::array::IntoArrow; +use geoarrow::array::LineStringArray; +use geoarrow::datatypes::CoordType; +use geoarrow::datatypes::LineStringType; +use prost::Message; +use vortex_array::ArrayRef; +use vortex_array::ExecutionCtx; +use vortex_array::IntoArray; +use vortex_array::arrays::ExtensionArray; +use vortex_array::arrays::extension::ExtensionArrayExt; +use vortex_array::arrow::ArrowExport; +use vortex_array::arrow::ArrowExportVTable; +use vortex_array::arrow::ArrowImport; +use vortex_array::arrow::ArrowImportVTable; +use vortex_array::arrow::ArrowSession; +use vortex_array::arrow::ArrowSessionExt; +use vortex_array::arrow::FromArrowArray; +use vortex_array::dtype::DType; +use vortex_array::dtype::Nullability; +use vortex_array::dtype::arrow::FromArrowType; +use vortex_array::dtype::extension::ExtDType; +use vortex_array::dtype::extension::ExtId; +use vortex_array::dtype::extension::ExtVTable; +use vortex_array::scalar::ScalarValue; +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::GeoMetadata; +use super::coordinate::Dimension; +use super::coordinate::coordinate_dimension; +use super::coordinate::coordinate_storage_dtype; +use super::geo_metadata_from_arrow; +use super::geoarrow_metadata; +use super::geoarrow_to_wkb; + +/// A line string: `geoarrow.linestring`, stored as `List>` (an ordered path +/// of vertices). +#[derive(Debug, Clone, Default, PartialEq, Eq, Hash)] +pub struct LineString; + +impl ExtVTable for LineString { + type Metadata = GeoMetadata; + // No cheap owned value like Point's `Coordinate`; expose the raw storage scalar. + type NativeValue<'a> = &'a ScalarValue; + + fn id(&self) -> ExtId { + static ID: CachedId = CachedId::new("vortex.geo.linestring"); + *ID + } + + fn serialize_metadata(&self, metadata: &Self::Metadata) -> VortexResult> { + Ok(metadata.encode_to_vec()) + } + + fn deserialize_metadata(&self, metadata: &[u8]) -> VortexResult { + Ok(GeoMetadata::decode(metadata)?) + } + + fn validate_dtype(ext_dtype: &ExtDType) -> VortexResult<()> { + linestring_dimension(ext_dtype.storage_dtype()).map(|_| ()) + } + + fn unpack_native<'a>( + _ext_dtype: &'a ExtDType, + storage_value: &'a ScalarValue, + ) -> VortexResult<&'a ScalarValue> { + Ok(storage_value) + } +} + +/// Canonical line-string storage: a list of the coordinate `Struct`. +pub(crate) fn linestring_storage_dtype(dim: Dimension, nullability: Nullability) -> DType { + let coords = coordinate_storage_dtype(dim, Nullability::NonNullable); + DType::List(Arc::new(coords), nullability) +} + +/// Validate `dtype` is `List` and return its [`Dimension`]. +pub(crate) fn linestring_dimension(dtype: &DType) -> VortexResult { + let DType::List(coords, _) = dtype else { + vortex_bail!("linestring storage must be a List of coordinates, was {dtype}"); + }; + coordinate_dimension(coords) +} + +static ARROW_LINESTRING: CachedId = CachedId::new(LineStringType::NAME); + +/// The `geoarrow.linestring` extension type for `dimension`, with separated (struct) coordinates +/// matching `LineString` storage. +fn linestring_type(geo_metadata: &GeoMetadata, dimension: Dimension) -> LineStringType { + LineStringType::new(dimension.into(), geoarrow_metadata(geo_metadata)) +} + +/// Decode `LineString` storage (`List`) to `geo_types` line strings, for the geo scalar +/// functions. CRS does not affect planar geometry ops, so default metadata is used. +pub(crate) fn linestring_geometries( + storage: &ArrayRef, + ctx: &mut ExecutionCtx, +) -> VortexResult>> { + linestring_array(storage, ctx)? + .iter() + .map(|geometry| -> VortexResult> { + Ok(geometry + .ok_or_else(|| vortex_err!("geo: null geometry is not supported"))? + .map_err(|e| vortex_err!("geo: geometry access failed: {e}"))? + .to_geometry()) + }) + .collect() +} + +/// Build a geoarrow `LineStringArray` from a `LineString`'s `List` storage. +fn linestring_array(storage: &ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult { + let linestring_type = linestring_type( + &GeoMetadata::default(), + linestring_dimension(storage.dtype())?, + ); + let session = ctx.session().clone(); + let arrow = session.arrow().execute_arrow(storage.clone(), None, ctx)?; + LineStringArray::try_from((arrow.as_ref(), linestring_type)) + .map_err(|e| vortex_err!("failed to construct LineStringArray: {e}")) +} + +/// A validated `LineString` array (`try_from` checks the extension type). +pub struct LineStringData(ExtensionArray); + +impl TryFrom for LineStringData { + type Error = VortexError; + + fn try_from(ext: ExtensionArray) -> Result { + vortex_ensure!( + ext.ext_dtype().is::(), + "expected a LineString extension array" + ); + Ok(LineStringData(ext)) + } +} + +impl LineStringData { + /// Serialize line strings to WKB (a view array) — the form DuckDB `GEOMETRY` takes. + pub fn to_wkb(&self, ctx: &mut ExecutionCtx) -> VortexResult { + geoarrow_to_wkb(&linestring_array(self.0.storage_array(), ctx)?) + } +} + +impl ArrowExportVTable for LineString { + fn arrow_ext_id(&self) -> Id { + *ARROW_LINESTRING + } + + fn vortex_id(&self) -> Id { + self.id() + } + + fn to_arrow_field( + &self, + name: &str, + dtype: &DType, + session: &ArrowSession, + ) -> VortexResult> { + let ext_type = dtype.as_extension(); + let geo_metadata = ext_type.metadata::(); + let dimension = linestring_dimension(ext_type.storage_dtype())?; + + let mut field = session.to_arrow_field(name, ext_type.storage_dtype())?; + field.try_with_extension_type(linestring_type(geo_metadata, dimension))?; + + Ok(Some(field)) + } + + fn execute_arrow( + &self, + array: ArrayRef, + target: &Field, + ctx: &mut ExecutionCtx, + ) -> VortexResult { + let is_linestring = array + .dtype() + .as_extension_opt() + .map(|ext| ext.is::()) + .unwrap_or(false); + if !is_linestring { + return Ok(ArrowExport::Unsupported(array)); + } + + let Ok(linestring_meta) = target.try_extension_type::() else { + return Ok(ArrowExport::Unsupported(array)); + }; + if linestring_meta.coord_type() != CoordType::Separated { + return Ok(ArrowExport::Unsupported(array)); + } + + let executed = array.execute::(ctx)?; + let storage = executed.storage_array().clone(); + + let storage_field = Field::new( + String::new(), + target.data_type().clone(), + target.is_nullable(), + ); + let session = ctx.session().clone(); + let arrow_storage = session + .arrow() + .execute_arrow(storage, Some(&storage_field), ctx)?; + + // Round-trip through GeoArrow's line-string array; `into_arrow` is concrete, so wrap in `Arc`. + let linestrings = LineStringArray::try_from((arrow_storage.as_ref(), linestring_meta)) + .map_err(|e| vortex_err!("failed to construct LineStringArray: {e}"))?; + + Ok(ArrowExport::Exported(Arc::new(linestrings.into_arrow()))) + } +} + +impl ArrowImportVTable for LineString { + fn arrow_ext_id(&self) -> Id { + *ARROW_LINESTRING + } + + /// Import a `geoarrow.linestring` field as the [`LineString`] dtype. Keyed off the standard + /// GeoArrow name, so any producer resolves here. Accepts the full `LineStringType` extension, or + /// — for a metadata-less geometry literal — the name alone, inferring the dimension from the + /// coordinate field names. + fn from_arrow_field(&self, field: &Field) -> VortexResult> { + let (dimension, metadata) = + if let Ok(linestring_meta) = field.try_extension_type::() { + vortex_ensure!( + linestring_meta.coord_type() == CoordType::Separated, + "geoarrow.linestring with interleaved coordinates is not supported; \ + re-encode with separated (struct) coordinates" + ); + ( + linestring_meta.dimension().into(), + geo_metadata_from_arrow(linestring_meta.metadata()), + ) + } else { + // Literal: peel the `List` layer to the coordinate struct and read its dimension from + // the field names (the canonical check rejects nullable coordinates). + if field.extension_type_name() != Some(LineStringType::NAME) { + return Ok(None); + } + let DType::List(coords, _) = DType::from_arrow(field) else { + return Ok(None); + }; + let DType::Struct(fields, _) = coords.as_ref() else { + return Ok(None); + }; + let Ok(dimension) = Dimension::from_field_names(fields.names()) else { + return Ok(None); + }; + (dimension, GeoMetadata::default()) + }; + + let storage_dtype = linestring_storage_dtype(dimension, field.is_nullable().into()); + Ok(Some(DType::Extension( + ExtDType::try_with_vtable(LineString, 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::() + || field.try_extension_type::().is_err() + || !matches!(array.data_type(), DataType::List(_)) + { + return Ok(ArrowImport::Unsupported(array)); + } + + let storage = ArrayRef::from_arrow(array.as_ref(), field.is_nullable())?; + Ok(ArrowImport::Imported( + ExtensionArray::try_new(ext_dtype.clone(), storage)?.into_array(), + )) + } +} + +#[cfg(test)] +mod tests { + use rstest::rstest; + use vortex_array::dtype::DType; + use vortex_array::dtype::Nullability; + use vortex_array::dtype::PType; + use vortex_array::dtype::extension::ExtDType; + use vortex_error::VortexResult; + + use super::LineString; + use super::linestring_storage_dtype; + use crate::extension::GeoMetadata; + use crate::extension::coordinate::Dimension; + + fn geo_meta() -> GeoMetadata { + GeoMetadata { + crs: Some("EPSG:4326".to_string()), + } + } + + /// `LineString` accepts the canonical `List` storage of every dimension. + #[rstest] + #[case::xy(Dimension::Xy)] + #[case::xyz(Dimension::Xyz)] + #[case::xym(Dimension::Xym)] + #[case::xyzm(Dimension::Xyzm)] + fn linestring_validates_every_dimension(#[case] dim: Dimension) -> VortexResult<()> { + let storage = linestring_storage_dtype(dim, Nullability::NonNullable); + ExtDType::::try_new(geo_meta(), storage)?; + Ok(()) + } + + /// Non-linestring storage is rejected at dtype construction: a bare coordinate struct (point) is + /// not a list of coordinates. + #[test] + fn linestring_rejects_invalid_storage() -> VortexResult<()> { + let primitive = DType::Primitive(PType::F64, Nullability::NonNullable); + assert!(ExtDType::::try_new(geo_meta(), primitive).is_err()); + Ok(()) + } +} diff --git a/vortex-geo/src/extension/mod.rs b/vortex-geo/src/extension/mod.rs index 4e4896aa5e5..9da09020dbf 100644 --- a/vortex-geo/src/extension/mod.rs +++ b/vortex-geo/src/extension/mod.rs @@ -2,6 +2,9 @@ // SPDX-FileCopyrightText: Copyright the Vortex contributors pub(crate) mod coordinate; +mod linestring; +mod multilinestring; +mod multipoint; mod multipolygon; mod point; mod polygon; @@ -19,12 +22,18 @@ use geoarrow::datatypes::CoordType; use geoarrow::datatypes::Crs; use geoarrow::datatypes::Dimension; use geoarrow::datatypes::GeoArrowType; +use geoarrow::datatypes::LineStringType; use geoarrow::datatypes::Metadata; +use geoarrow::datatypes::MultiLineStringType; +use geoarrow::datatypes::MultiPointType; use geoarrow::datatypes::MultiPolygonType; use geoarrow::datatypes::PointType; use geoarrow::datatypes::PolygonType; use geoarrow::datatypes::WkbType; use geoarrow_cast::cast::cast; +pub use linestring::*; +pub use multilinestring::*; +pub use multipoint::*; pub use multipolygon::*; pub use point::*; pub use polygon::*; @@ -61,8 +70,14 @@ pub(crate) fn geometries( .clone(); if ext.is::() { point_geometries(&storage, ctx) + } else if ext.is::() { + linestring_geometries(&storage, ctx) + } else if ext.is::() { + multipoint_geometries(&storage, ctx) } else if ext.is::() { polygon_geometries(&storage, ctx) + } else if ext.is::() { + multilinestring_geometries(&storage, ctx) } else if ext.is::() { multipolygon_geometries(&storage, ctx) } else { @@ -107,12 +122,31 @@ pub fn native_geometry_scalar_from_wkb(bytes: &[u8]) -> VortexResult { + let target = GeoArrowType::LineString( + LineStringType::new(Dimension::XY, metadata).with_coord_type(CoordType::Separated), + ); + geo_ext_scalar(LineString, to_storage(&target)?)? + } GeometryType::Polygon => { let target = GeoArrowType::Polygon( PolygonType::new(Dimension::XY, metadata).with_coord_type(CoordType::Separated), ); geo_ext_scalar(Polygon, to_storage(&target)?)? } + GeometryType::MultiPoint => { + let target = GeoArrowType::MultiPoint( + MultiPointType::new(Dimension::XY, metadata).with_coord_type(CoordType::Separated), + ); + geo_ext_scalar(MultiPoint, to_storage(&target)?)? + } + GeometryType::MultiLineString => { + let target = GeoArrowType::MultiLineString( + MultiLineStringType::new(Dimension::XY, metadata) + .with_coord_type(CoordType::Separated), + ); + geo_ext_scalar(MultiLineString, to_storage(&target)?)? + } GeometryType::MultiPolygon => { let target = GeoArrowType::MultiPolygon( MultiPolygonType::new(Dimension::XY, metadata) @@ -201,6 +235,9 @@ mod tests { use vortex_error::VortexResult; use vortex_error::vortex_err; + use super::LineString; + use super::MultiLineString; + use super::MultiPoint; use super::Point; use super::Polygon; use super::native_geometry_scalar_from_wkb; @@ -256,4 +293,77 @@ mod tests { assert!(ext.is::()); Ok(()) } + + /// A little-endian WKB `LINESTRING` literal decodes to the native `LineString` extension scalar. + #[test] + fn decodes_wkb_linestring_to_native() -> VortexResult<()> { + let points = [(0.0, 0.0), (1.0, 1.0)]; + let mut wkb = vec![1u8]; // little-endian byte order + wkb.extend_from_slice(&2u32.to_le_bytes()); // geometry type: linestring + let len = u32::try_from(points.len()).map_err(|e| vortex_err!("{e}"))?; + wkb.extend_from_slice(&len.to_le_bytes()); + for (x, y) in points { + wkb.extend_from_slice(&f64::to_le_bytes(x)); + wkb.extend_from_slice(&f64::to_le_bytes(y)); + } + + let scalar = native_geometry_scalar_from_wkb(&wkb)?.expect("a linestring scalar"); + let DType::Extension(ext) = scalar.dtype() else { + panic!("expected an extension dtype, got {}", scalar.dtype()); + }; + assert!(ext.is::()); + Ok(()) + } + + /// A little-endian WKB `MULTIPOINT` literal decodes to the native `MultiPoint` extension scalar. + #[test] + fn decodes_wkb_multipoint_to_native() -> VortexResult<()> { + let points = [(0.0, 0.0), (1.0, 1.0)]; + let mut wkb = vec![1u8]; // little-endian byte order + wkb.extend_from_slice(&4u32.to_le_bytes()); // geometry type: multipoint + let len = u32::try_from(points.len()).map_err(|e| vortex_err!("{e}"))?; + wkb.extend_from_slice(&len.to_le_bytes()); + for (x, y) in points { + // each member is a full WKB point + wkb.push(1u8); + wkb.extend_from_slice(&1u32.to_le_bytes()); + wkb.extend_from_slice(&f64::to_le_bytes(x)); + wkb.extend_from_slice(&f64::to_le_bytes(y)); + } + + let scalar = native_geometry_scalar_from_wkb(&wkb)?.expect("a multipoint scalar"); + let DType::Extension(ext) = scalar.dtype() else { + panic!("expected an extension dtype, got {}", scalar.dtype()); + }; + assert!(ext.is::()); + Ok(()) + } + + /// A little-endian WKB `MULTILINESTRING` literal decodes to the native `MultiLineString` scalar. + #[test] + fn decodes_wkb_multilinestring_to_native() -> VortexResult<()> { + let lines = [[(0.0, 0.0), (1.0, 1.0)], [(2.0, 2.0), (3.0, 3.0)]]; + let mut wkb = vec![1u8]; // little-endian byte order + wkb.extend_from_slice(&5u32.to_le_bytes()); // geometry type: multilinestring + let num_lines = u32::try_from(lines.len()).map_err(|e| vortex_err!("{e}"))?; + wkb.extend_from_slice(&num_lines.to_le_bytes()); + for line in lines { + // each member is a full WKB linestring + wkb.push(1u8); + wkb.extend_from_slice(&2u32.to_le_bytes()); + let len = u32::try_from(line.len()).map_err(|e| vortex_err!("{e}"))?; + wkb.extend_from_slice(&len.to_le_bytes()); + for (x, y) in line { + wkb.extend_from_slice(&f64::to_le_bytes(x)); + wkb.extend_from_slice(&f64::to_le_bytes(y)); + } + } + + let scalar = native_geometry_scalar_from_wkb(&wkb)?.expect("a multilinestring scalar"); + let DType::Extension(ext) = scalar.dtype() else { + panic!("expected an extension dtype, got {}", scalar.dtype()); + }; + assert!(ext.is::()); + Ok(()) + } } diff --git a/vortex-geo/src/extension/multilinestring.rs b/vortex-geo/src/extension/multilinestring.rs new file mode 100644 index 00000000000..edfb0d3554c --- /dev/null +++ b/vortex-geo/src/extension/multilinestring.rs @@ -0,0 +1,359 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! The [`MultiLineString`] extension type (`vortex.geo.multilinestring`), stored as +//! `List>>` (line strings → coordinates) and tagged with +//! [`GeoMetadata`]. The storage layout matches [`Polygon`](super::Polygon); the two are +//! distinguished by their GeoArrow extension name, not their shape. + +use std::sync::Arc; + +use arrow_array::ArrayRef as ArrowArrayRef; +use arrow_schema::DataType; +use arrow_schema::Field; +use arrow_schema::extension::ExtensionType; +use geo_traits::to_geo::ToGeoGeometry; +use geo_types::Geometry; +use geoarrow::array::GeoArrowArrayAccessor; +use geoarrow::array::IntoArrow; +use geoarrow::array::MultiLineStringArray; +use geoarrow::datatypes::CoordType; +use geoarrow::datatypes::MultiLineStringType; +use prost::Message; +use vortex_array::ArrayRef; +use vortex_array::ExecutionCtx; +use vortex_array::IntoArray; +use vortex_array::arrays::ExtensionArray; +use vortex_array::arrays::extension::ExtensionArrayExt; +use vortex_array::arrow::ArrowExport; +use vortex_array::arrow::ArrowExportVTable; +use vortex_array::arrow::ArrowImport; +use vortex_array::arrow::ArrowImportVTable; +use vortex_array::arrow::ArrowSession; +use vortex_array::arrow::ArrowSessionExt; +use vortex_array::arrow::FromArrowArray; +use vortex_array::dtype::DType; +use vortex_array::dtype::Nullability; +use vortex_array::dtype::arrow::FromArrowType; +use vortex_array::dtype::extension::ExtDType; +use vortex_array::dtype::extension::ExtId; +use vortex_array::dtype::extension::ExtVTable; +use vortex_array::scalar::ScalarValue; +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::GeoMetadata; +use super::coordinate::Dimension; +use super::coordinate::coordinate_dimension; +use super::coordinate::coordinate_storage_dtype; +use super::geo_metadata_from_arrow; +use super::geoarrow_metadata; +use super::geoarrow_to_wkb; + +/// A multilinestring: `geoarrow.multilinestring`, stored as `List>>` +/// (line strings of vertices). +#[derive(Debug, Clone, Default, PartialEq, Eq, Hash)] +pub struct MultiLineString; + +impl ExtVTable for MultiLineString { + type Metadata = GeoMetadata; + // No cheap owned value like Point's `Coordinate`; expose the raw storage scalar. + type NativeValue<'a> = &'a ScalarValue; + + fn id(&self) -> ExtId { + static ID: CachedId = CachedId::new("vortex.geo.multilinestring"); + *ID + } + + fn serialize_metadata(&self, metadata: &Self::Metadata) -> VortexResult> { + Ok(metadata.encode_to_vec()) + } + + fn deserialize_metadata(&self, metadata: &[u8]) -> VortexResult { + Ok(GeoMetadata::decode(metadata)?) + } + + fn validate_dtype(ext_dtype: &ExtDType) -> VortexResult<()> { + multilinestring_dimension(ext_dtype.storage_dtype()).map(|_| ()) + } + + fn unpack_native<'a>( + _ext_dtype: &'a ExtDType, + storage_value: &'a ScalarValue, + ) -> VortexResult<&'a ScalarValue> { + Ok(storage_value) + } +} + +/// Canonical multilinestring storage: an outer list of line strings, each a list of the coordinate +/// `Struct`. +pub(crate) fn multilinestring_storage_dtype(dim: Dimension, nullability: Nullability) -> DType { + let coords = coordinate_storage_dtype(dim, Nullability::NonNullable); + let line = DType::List(Arc::new(coords), Nullability::NonNullable); + DType::List(Arc::new(line), nullability) +} + +/// Validate `dtype` is `List>` and return its [`Dimension`]. +pub(crate) fn multilinestring_dimension(dtype: &DType) -> VortexResult { + let DType::List(line, _) = dtype else { + vortex_bail!("multilinestring storage must be a List of line strings, was {dtype}"); + }; + let DType::List(coords, _) = line.as_ref() else { + vortex_bail!("multilinestring line storage must be a List of coordinates, was {line}"); + }; + coordinate_dimension(coords) +} + +static ARROW_MULTILINESTRING: CachedId = CachedId::new(MultiLineStringType::NAME); + +/// The `geoarrow.multilinestring` type for `dimension`, with separated (struct) coordinates. +fn multilinestring_type(geo_metadata: &GeoMetadata, dimension: Dimension) -> MultiLineStringType { + MultiLineStringType::new(dimension.into(), geoarrow_metadata(geo_metadata)) +} + +/// Decode storage to `geo_types` for the geo scalar functions (CRS is irrelevant to planar ops). +pub(crate) fn multilinestring_geometries( + storage: &ArrayRef, + ctx: &mut ExecutionCtx, +) -> VortexResult>> { + multilinestring_array(storage, ctx)? + .iter() + .map(|geometry| -> VortexResult> { + Ok(geometry + .ok_or_else(|| vortex_err!("geo: null geometry is not supported"))? + .map_err(|e| vortex_err!("geo: geometry access failed: {e}"))? + .to_geometry()) + }) + .collect() +} + +/// Build a geoarrow `MultiLineStringArray` from the `MultiLineString` storage. +fn multilinestring_array( + storage: &ArrayRef, + ctx: &mut ExecutionCtx, +) -> VortexResult { + let multilinestring_type = multilinestring_type( + &GeoMetadata::default(), + multilinestring_dimension(storage.dtype())?, + ); + let session = ctx.session().clone(); + let arrow = session.arrow().execute_arrow(storage.clone(), None, ctx)?; + MultiLineStringArray::try_from((arrow.as_ref(), multilinestring_type)) + .map_err(|e| vortex_err!("failed to construct MultiLineStringArray: {e}")) +} + +/// A validated `MultiLineString` array (`try_from` checks the extension type). +pub struct MultiLineStringData(ExtensionArray); + +impl TryFrom for MultiLineStringData { + type Error = VortexError; + + fn try_from(ext: ExtensionArray) -> Result { + vortex_ensure!( + ext.ext_dtype().is::(), + "expected a MultiLineString extension array" + ); + Ok(MultiLineStringData(ext)) + } +} + +impl MultiLineStringData { + /// Serialize multilinestrings to WKB (a view array) — the form DuckDB `GEOMETRY` takes. + pub fn to_wkb(&self, ctx: &mut ExecutionCtx) -> VortexResult { + geoarrow_to_wkb(&multilinestring_array(self.0.storage_array(), ctx)?) + } +} + +impl ArrowExportVTable for MultiLineString { + fn arrow_ext_id(&self) -> Id { + *ARROW_MULTILINESTRING + } + + fn vortex_id(&self) -> Id { + self.id() + } + + fn to_arrow_field( + &self, + name: &str, + dtype: &DType, + session: &ArrowSession, + ) -> VortexResult> { + let ext_type = dtype.as_extension(); + let geo_metadata = ext_type.metadata::(); + let dimension = multilinestring_dimension(ext_type.storage_dtype())?; + + let mut field = session.to_arrow_field(name, ext_type.storage_dtype())?; + field.try_with_extension_type(multilinestring_type(geo_metadata, dimension))?; + + Ok(Some(field)) + } + + fn execute_arrow( + &self, + array: ArrayRef, + target: &Field, + ctx: &mut ExecutionCtx, + ) -> VortexResult { + let is_multilinestring = array + .dtype() + .as_extension_opt() + .map(|ext| ext.is::()) + .unwrap_or(false); + if !is_multilinestring { + return Ok(ArrowExport::Unsupported(array)); + } + + let Ok(multilinestring_meta) = target.try_extension_type::() else { + return Ok(ArrowExport::Unsupported(array)); + }; + if multilinestring_meta.coord_type() != CoordType::Separated { + return Ok(ArrowExport::Unsupported(array)); + } + + let executed = array.execute::(ctx)?; + let storage = executed.storage_array().clone(); + + let storage_field = Field::new( + String::new(), + target.data_type().clone(), + target.is_nullable(), + ); + let session = ctx.session().clone(); + let arrow_storage = session + .arrow() + .execute_arrow(storage, Some(&storage_field), ctx)?; + + let multilinestrings = + MultiLineStringArray::try_from((arrow_storage.as_ref(), multilinestring_meta)) + .map_err(|e| vortex_err!("failed to construct MultiLineStringArray: {e}"))?; + + Ok(ArrowExport::Exported(Arc::new( + multilinestrings.into_arrow(), + ))) + } +} + +impl ArrowImportVTable for MultiLineString { + fn arrow_ext_id(&self) -> Id { + *ARROW_MULTILINESTRING + } + + /// Import a `geoarrow.multilinestring` field (matched by GeoArrow name). Accepts the full + /// `MultiLineStringType`, or a metadata-less literal (name only), inferring the dimension. + fn from_arrow_field(&self, field: &Field) -> VortexResult> { + let (dimension, metadata) = + if let Ok(multilinestring_meta) = field.try_extension_type::() { + vortex_ensure!( + multilinestring_meta.coord_type() == CoordType::Separated, + "geoarrow.multilinestring with interleaved coordinates is not supported; \ + re-encode with separated (struct) coordinates" + ); + ( + multilinestring_meta.dimension().into(), + geo_metadata_from_arrow(multilinestring_meta.metadata()), + ) + } else { + // Literal: peel the two `List` layers to the coordinate struct and read its dimension + // from the field names (the canonical check rejects nullable coordinates). + if field.extension_type_name() != Some(MultiLineStringType::NAME) { + return Ok(None); + } + let DType::List(line, _) = DType::from_arrow(field) else { + return Ok(None); + }; + let DType::List(coords, _) = line.as_ref() else { + return Ok(None); + }; + let DType::Struct(fields, _) = coords.as_ref() else { + return Ok(None); + }; + let Ok(dimension) = Dimension::from_field_names(fields.names()) else { + return Ok(None); + }; + (dimension, GeoMetadata::default()) + }; + + let storage_dtype = multilinestring_storage_dtype(dimension, field.is_nullable().into()); + Ok(Some(DType::Extension( + ExtDType::try_with_vtable(MultiLineString, 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::() + || field.try_extension_type::().is_err() + || !matches!(array.data_type(), DataType::List(_)) + { + return Ok(ArrowImport::Unsupported(array)); + } + + let storage = ArrayRef::from_arrow(array.as_ref(), field.is_nullable())?; + Ok(ArrowImport::Imported( + ExtensionArray::try_new(ext_dtype.clone(), storage)?.into_array(), + )) + } +} + +#[cfg(test)] +mod tests { + use std::sync::Arc; + + use rstest::rstest; + use vortex_array::dtype::DType; + use vortex_array::dtype::Nullability; + use vortex_array::dtype::PType; + use vortex_array::dtype::extension::ExtDType; + use vortex_error::VortexResult; + + use super::MultiLineString; + use super::multilinestring_storage_dtype; + use crate::extension::GeoMetadata; + use crate::extension::coordinate::Dimension; + use crate::extension::coordinate::coordinate_storage_dtype; + + fn geo_meta() -> GeoMetadata { + GeoMetadata { + crs: Some("EPSG:4326".to_string()), + } + } + + /// `MultiLineString` accepts the canonical `List>` storage of every + /// dimension. + #[rstest] + #[case::xy(Dimension::Xy)] + #[case::xyz(Dimension::Xyz)] + #[case::xym(Dimension::Xym)] + #[case::xyzm(Dimension::Xyzm)] + fn multilinestring_validates_every_dimension(#[case] dim: Dimension) -> VortexResult<()> { + let storage = multilinestring_storage_dtype(dim, Nullability::NonNullable); + ExtDType::::try_new(geo_meta(), storage)?; + Ok(()) + } + + /// Non-multilinestring storage is rejected: a bare coordinate struct (point) fails. + #[test] + fn multilinestring_rejects_invalid_storage() -> VortexResult<()> { + let primitive = DType::Primitive(PType::F64, Nullability::NonNullable); + assert!(ExtDType::::try_new(geo_meta(), primitive).is_err()); + + // A bare list of coordinates is a single line string, not a multilinestring. + let coords = coordinate_storage_dtype(Dimension::Xy, Nullability::NonNullable); + let line = DType::List(Arc::new(coords), Nullability::NonNullable); + assert!(ExtDType::::try_new(geo_meta(), line).is_err()); + Ok(()) + } +} diff --git a/vortex-geo/src/extension/multipoint.rs b/vortex-geo/src/extension/multipoint.rs new file mode 100644 index 00000000000..3e8778744e5 --- /dev/null +++ b/vortex-geo/src/extension/multipoint.rs @@ -0,0 +1,336 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! The [`MultiPoint`] geometry extension type (`vortex.geo.multipoint`): an unordered set of the +//! [`Point`](super::Point) coordinate struct, stored as `List>` and tagged +//! with [`GeoMetadata`] (CRS). The storage layout matches [`LineString`](super::LineString); the +//! two are distinguished by their GeoArrow extension name, not their shape. + +use std::sync::Arc; + +use arrow_array::ArrayRef as ArrowArrayRef; +use arrow_schema::DataType; +use arrow_schema::Field; +use arrow_schema::extension::ExtensionType; +use geo_traits::to_geo::ToGeoGeometry; +use geo_types::Geometry; +use geoarrow::array::GeoArrowArrayAccessor; +use geoarrow::array::IntoArrow; +use geoarrow::array::MultiPointArray; +use geoarrow::datatypes::CoordType; +use geoarrow::datatypes::MultiPointType; +use prost::Message; +use vortex_array::ArrayRef; +use vortex_array::ExecutionCtx; +use vortex_array::IntoArray; +use vortex_array::arrays::ExtensionArray; +use vortex_array::arrays::extension::ExtensionArrayExt; +use vortex_array::arrow::ArrowExport; +use vortex_array::arrow::ArrowExportVTable; +use vortex_array::arrow::ArrowImport; +use vortex_array::arrow::ArrowImportVTable; +use vortex_array::arrow::ArrowSession; +use vortex_array::arrow::ArrowSessionExt; +use vortex_array::arrow::FromArrowArray; +use vortex_array::dtype::DType; +use vortex_array::dtype::Nullability; +use vortex_array::dtype::arrow::FromArrowType; +use vortex_array::dtype::extension::ExtDType; +use vortex_array::dtype::extension::ExtId; +use vortex_array::dtype::extension::ExtVTable; +use vortex_array::scalar::ScalarValue; +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::GeoMetadata; +use super::coordinate::Dimension; +use super::coordinate::coordinate_dimension; +use super::coordinate::coordinate_storage_dtype; +use super::geo_metadata_from_arrow; +use super::geoarrow_metadata; +use super::geoarrow_to_wkb; + +/// A multipoint: `geoarrow.multipoint`, stored as `List>` (a set of points). +#[derive(Debug, Clone, Default, PartialEq, Eq, Hash)] +pub struct MultiPoint; + +impl ExtVTable for MultiPoint { + type Metadata = GeoMetadata; + // No cheap owned value like Point's `Coordinate`; expose the raw storage scalar. + type NativeValue<'a> = &'a ScalarValue; + + fn id(&self) -> ExtId { + static ID: CachedId = CachedId::new("vortex.geo.multipoint"); + *ID + } + + fn serialize_metadata(&self, metadata: &Self::Metadata) -> VortexResult> { + Ok(metadata.encode_to_vec()) + } + + fn deserialize_metadata(&self, metadata: &[u8]) -> VortexResult { + Ok(GeoMetadata::decode(metadata)?) + } + + fn validate_dtype(ext_dtype: &ExtDType) -> VortexResult<()> { + multipoint_dimension(ext_dtype.storage_dtype()).map(|_| ()) + } + + fn unpack_native<'a>( + _ext_dtype: &'a ExtDType, + storage_value: &'a ScalarValue, + ) -> VortexResult<&'a ScalarValue> { + Ok(storage_value) + } +} + +/// Canonical multipoint storage: a list of the coordinate `Struct`. +pub(crate) fn multipoint_storage_dtype(dim: Dimension, nullability: Nullability) -> DType { + let coords = coordinate_storage_dtype(dim, Nullability::NonNullable); + DType::List(Arc::new(coords), nullability) +} + +/// Validate `dtype` is `List` and return its [`Dimension`]. +pub(crate) fn multipoint_dimension(dtype: &DType) -> VortexResult { + let DType::List(coords, _) = dtype else { + vortex_bail!("multipoint storage must be a List of coordinates, was {dtype}"); + }; + coordinate_dimension(coords) +} + +static ARROW_MULTIPOINT: CachedId = CachedId::new(MultiPointType::NAME); + +/// The `geoarrow.multipoint` extension type for `dimension`, with separated (struct) coordinates. +fn multipoint_type(geo_metadata: &GeoMetadata, dimension: Dimension) -> MultiPointType { + MultiPointType::new(dimension.into(), geoarrow_metadata(geo_metadata)) +} + +/// Decode `MultiPoint` storage (`List`) to `geo_types`, for the geo scalar functions. +pub(crate) fn multipoint_geometries( + storage: &ArrayRef, + ctx: &mut ExecutionCtx, +) -> VortexResult>> { + multipoint_array(storage, ctx)? + .iter() + .map(|geometry| -> VortexResult> { + Ok(geometry + .ok_or_else(|| vortex_err!("geo: null geometry is not supported"))? + .map_err(|e| vortex_err!("geo: geometry access failed: {e}"))? + .to_geometry()) + }) + .collect() +} + +/// Build a geoarrow `MultiPointArray` from a `MultiPoint`'s `List` storage. +fn multipoint_array(storage: &ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult { + let multipoint_type = multipoint_type( + &GeoMetadata::default(), + multipoint_dimension(storage.dtype())?, + ); + let session = ctx.session().clone(); + let arrow = session.arrow().execute_arrow(storage.clone(), None, ctx)?; + MultiPointArray::try_from((arrow.as_ref(), multipoint_type)) + .map_err(|e| vortex_err!("failed to construct MultiPointArray: {e}")) +} + +/// A validated `MultiPoint` array (`try_from` checks the extension type). +pub struct MultiPointData(ExtensionArray); + +impl TryFrom for MultiPointData { + type Error = VortexError; + + fn try_from(ext: ExtensionArray) -> Result { + vortex_ensure!( + ext.ext_dtype().is::(), + "expected a MultiPoint extension array" + ); + Ok(MultiPointData(ext)) + } +} + +impl MultiPointData { + /// Serialize multipoints to WKB (a view array) — the form DuckDB `GEOMETRY` takes. + pub fn to_wkb(&self, ctx: &mut ExecutionCtx) -> VortexResult { + geoarrow_to_wkb(&multipoint_array(self.0.storage_array(), ctx)?) + } +} + +impl ArrowExportVTable for MultiPoint { + fn arrow_ext_id(&self) -> Id { + *ARROW_MULTIPOINT + } + + fn vortex_id(&self) -> Id { + self.id() + } + + fn to_arrow_field( + &self, + name: &str, + dtype: &DType, + session: &ArrowSession, + ) -> VortexResult> { + let ext_type = dtype.as_extension(); + let geo_metadata = ext_type.metadata::(); + let dimension = multipoint_dimension(ext_type.storage_dtype())?; + + let mut field = session.to_arrow_field(name, ext_type.storage_dtype())?; + field.try_with_extension_type(multipoint_type(geo_metadata, dimension))?; + + Ok(Some(field)) + } + + fn execute_arrow( + &self, + array: ArrayRef, + target: &Field, + ctx: &mut ExecutionCtx, + ) -> VortexResult { + let is_multipoint = array + .dtype() + .as_extension_opt() + .map(|ext| ext.is::()) + .unwrap_or(false); + if !is_multipoint { + return Ok(ArrowExport::Unsupported(array)); + } + + let Ok(multipoint_meta) = target.try_extension_type::() else { + return Ok(ArrowExport::Unsupported(array)); + }; + if multipoint_meta.coord_type() != CoordType::Separated { + return Ok(ArrowExport::Unsupported(array)); + } + + let executed = array.execute::(ctx)?; + let storage = executed.storage_array().clone(); + + let storage_field = Field::new( + String::new(), + target.data_type().clone(), + target.is_nullable(), + ); + let session = ctx.session().clone(); + let arrow_storage = session + .arrow() + .execute_arrow(storage, Some(&storage_field), ctx)?; + + // Round-trip through GeoArrow's multipoint array; `into_arrow` is concrete, so wrap in `Arc`. + let multipoints = MultiPointArray::try_from((arrow_storage.as_ref(), multipoint_meta)) + .map_err(|e| vortex_err!("failed to construct MultiPointArray: {e}"))?; + + Ok(ArrowExport::Exported(Arc::new(multipoints.into_arrow()))) + } +} + +impl ArrowImportVTable for MultiPoint { + fn arrow_ext_id(&self) -> Id { + *ARROW_MULTIPOINT + } + + /// Import a `geoarrow.multipoint` field as the [`MultiPoint`] dtype (matched by GeoArrow name). + /// Accepts the full `MultiPointType`, or a metadata-less literal (name only), inferring the + /// dimension from the coordinate field names. + fn from_arrow_field(&self, field: &Field) -> VortexResult> { + let (dimension, metadata) = + if let Ok(multipoint_meta) = field.try_extension_type::() { + vortex_ensure!( + multipoint_meta.coord_type() == CoordType::Separated, + "geoarrow.multipoint with interleaved coordinates is not supported; \ + re-encode with separated (struct) coordinates" + ); + ( + multipoint_meta.dimension().into(), + geo_metadata_from_arrow(multipoint_meta.metadata()), + ) + } else { + if field.extension_type_name() != Some(MultiPointType::NAME) { + return Ok(None); + } + let DType::List(coords, _) = DType::from_arrow(field) else { + return Ok(None); + }; + let DType::Struct(fields, _) = coords.as_ref() else { + return Ok(None); + }; + let Ok(dimension) = Dimension::from_field_names(fields.names()) else { + return Ok(None); + }; + (dimension, GeoMetadata::default()) + }; + + let storage_dtype = multipoint_storage_dtype(dimension, field.is_nullable().into()); + Ok(Some(DType::Extension( + ExtDType::try_with_vtable(MultiPoint, 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::() + || field.try_extension_type::().is_err() + || !matches!(array.data_type(), DataType::List(_)) + { + return Ok(ArrowImport::Unsupported(array)); + } + + let storage = ArrayRef::from_arrow(array.as_ref(), field.is_nullable())?; + Ok(ArrowImport::Imported( + ExtensionArray::try_new(ext_dtype.clone(), storage)?.into_array(), + )) + } +} + +#[cfg(test)] +mod tests { + use rstest::rstest; + use vortex_array::dtype::DType; + use vortex_array::dtype::Nullability; + use vortex_array::dtype::PType; + use vortex_array::dtype::extension::ExtDType; + use vortex_error::VortexResult; + + use super::MultiPoint; + use super::multipoint_storage_dtype; + use crate::extension::GeoMetadata; + use crate::extension::coordinate::Dimension; + + fn geo_meta() -> GeoMetadata { + GeoMetadata { + crs: Some("EPSG:4326".to_string()), + } + } + + /// `MultiPoint` accepts the canonical `List` storage of every dimension. + #[rstest] + #[case::xy(Dimension::Xy)] + #[case::xyz(Dimension::Xyz)] + #[case::xym(Dimension::Xym)] + #[case::xyzm(Dimension::Xyzm)] + fn multipoint_validates_every_dimension(#[case] dim: Dimension) -> VortexResult<()> { + let storage = multipoint_storage_dtype(dim, Nullability::NonNullable); + ExtDType::::try_new(geo_meta(), storage)?; + Ok(()) + } + + /// Non-multipoint storage is rejected at dtype construction: a bare coordinate struct (point) is + /// not a list of coordinates. + #[test] + fn multipoint_rejects_invalid_storage() -> VortexResult<()> { + let primitive = DType::Primitive(PType::F64, Nullability::NonNullable); + assert!(ExtDType::::try_new(geo_meta(), primitive).is_err()); + Ok(()) + } +} diff --git a/vortex-geo/src/lib.rs b/vortex-geo/src/lib.rs index 2cc8004efc5..c2637ecac55 100644 --- a/vortex-geo/src/lib.rs +++ b/vortex-geo/src/lib.rs @@ -8,6 +8,9 @@ use vortex_array::dtype::session::DTypeSessionExt; use vortex_array::scalar_fn::session::ScalarFnSessionExt; use vortex_session::VortexSession; +use crate::extension::LineString; +use crate::extension::MultiLineString; +use crate::extension::MultiPoint; use crate::extension::MultiPolygon; use crate::extension::Point; use crate::extension::Polygon; @@ -30,9 +33,18 @@ pub fn initialize(session: &VortexSession) { session.dtypes().register(Point); session.arrow().register_exporter(Arc::new(Point)); session.arrow().register_importer(Arc::new(Point)); + session.dtypes().register(LineString); + session.arrow().register_exporter(Arc::new(LineString)); + session.arrow().register_importer(Arc::new(LineString)); + session.dtypes().register(MultiPoint); + session.arrow().register_exporter(Arc::new(MultiPoint)); + session.arrow().register_importer(Arc::new(MultiPoint)); session.dtypes().register(Polygon); session.arrow().register_exporter(Arc::new(Polygon)); session.arrow().register_importer(Arc::new(Polygon)); + session.dtypes().register(MultiLineString); + session.arrow().register_exporter(Arc::new(MultiLineString)); + session.arrow().register_importer(Arc::new(MultiLineString)); session.dtypes().register(MultiPolygon); session.arrow().register_exporter(Arc::new(MultiPolygon)); session.arrow().register_importer(Arc::new(MultiPolygon)); diff --git a/vortex-geo/src/tests/linestring.rs b/vortex-geo/src/tests/linestring.rs new file mode 100644 index 00000000000..a1a9b3a4845 --- /dev/null +++ b/vortex-geo/src/tests/linestring.rs @@ -0,0 +1,88 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Arrow interop for the `vortex.geo.linestring` extension type (`geoarrow.linestring`). + +use std::sync::Arc; + +use arrow_schema::DataType; +use arrow_schema::Field; +use arrow_schema::extension::ExtensionType as _; +use geoarrow::datatypes::CoordType; +use geoarrow::datatypes::Crs; +use geoarrow::datatypes::Dimension as GeoArrowDimension; +use geoarrow::datatypes::LineStringType; +use geoarrow::datatypes::Metadata; +use vortex_array::arrow::ArrowSessionExt; +use vortex_array::dtype::DType; +use vortex_array::dtype::Nullability; +use vortex_error::VortexResult; + +use super::SESSION; +use crate::extension::LineString; + +/// A `geoarrow.linestring` Arrow field with separated (struct) XY coordinates. +fn linestring_field(name: &str, nullable: bool, crs: Option<&str>) -> Field { + let crs = crs + .map(|crs| Crs::from_unknown_crs_type(crs.to_string())) + .unwrap_or_default(); + let metadata = Arc::new(Metadata::new(crs, None)); + LineStringType::new(GeoArrowDimension::XY, metadata).to_field(name, nullable) +} + +/// An imported `geoarrow.linestring` field maps to the LineString extension dtype, recovering the +/// CRS, the `List>` storage, and nullability. +#[test] +fn import_field_recovers_extension() -> VortexResult<()> { + let field = linestring_field("geom", true, Some("EPSG:4326")); + let dtype = SESSION.arrow().from_arrow_field(&field)?; + + let DType::Extension(ext) = &dtype else { + panic!("expected Extension dtype, got {dtype}"); + }; + assert!(ext.is::()); + assert_eq!( + ext.metadata::().crs.as_deref(), + Some("EPSG:4326") + ); + + // Storage peels one List layer (linestring → coordinates) to the coordinate struct. + let DType::List(coords, nullability) = ext.storage_dtype() else { + panic!("expected List storage, got {}", ext.storage_dtype()); + }; + assert_eq!(*nullability, Nullability::Nullable); + let DType::Struct(fields, _) = coords.as_ref() else { + panic!("expected coordinate Struct"); + }; + let names: Vec<&str> = fields.names().iter().map(|n| n.as_ref()).collect(); + assert_eq!(names, vec!["x", "y"]); + Ok(()) +} + +/// A field with interleaved (`FixedSizeList`) coordinates fails to import. +#[test] +fn import_interleaved_field_fails() { + let linestring_type = LineStringType::new(GeoArrowDimension::XY, Default::default()) + .with_coord_type(CoordType::Interleaved); + let field = linestring_type.to_field("geom", false); + assert!(SESSION.arrow().from_arrow_field(&field).is_err()); +} + +/// A field imported to the LineString dtype and exported back carries the `geoarrow.linestring` +/// extension over its `List` storage. +#[test] +fn export_field_carries_extension() -> VortexResult<()> { + let imported = + SESSION + .arrow() + .from_arrow_field(&linestring_field("geom", false, Some("EPSG:4326")))?; + let field = SESSION.arrow().to_arrow_field("geom", &imported)?; + + assert_eq!(field.extension_type_name(), Some(LineStringType::NAME)); + assert!( + matches!(field.data_type(), DataType::List(_)), + "expected List storage, got {}", + field.data_type() + ); + Ok(()) +} diff --git a/vortex-geo/src/tests/mod.rs b/vortex-geo/src/tests/mod.rs index 87b25ed1293..0d4de3ae6a3 100644 --- a/vortex-geo/src/tests/mod.rs +++ b/vortex-geo/src/tests/mod.rs @@ -4,6 +4,9 @@ //! Arrow interop tests for the geospatial extension types, exercising the session wiring set up //! by [`crate::initialize`]. +mod linestring; +mod multilinestring; +mod multipoint; mod multipolygon; mod point; mod wkb; diff --git a/vortex-geo/src/tests/multilinestring.rs b/vortex-geo/src/tests/multilinestring.rs new file mode 100644 index 00000000000..cf5b5204681 --- /dev/null +++ b/vortex-geo/src/tests/multilinestring.rs @@ -0,0 +1,92 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Arrow interop for the `vortex.geo.multilinestring` extension type (`geoarrow.multilinestring`). + +use std::sync::Arc; + +use arrow_schema::DataType; +use arrow_schema::Field; +use arrow_schema::extension::ExtensionType as _; +use geoarrow::datatypes::CoordType; +use geoarrow::datatypes::Crs; +use geoarrow::datatypes::Dimension as GeoArrowDimension; +use geoarrow::datatypes::Metadata; +use geoarrow::datatypes::MultiLineStringType; +use vortex_array::arrow::ArrowSessionExt; +use vortex_array::dtype::DType; +use vortex_array::dtype::Nullability; +use vortex_error::VortexResult; + +use super::SESSION; +use crate::extension::MultiLineString; + +/// A `geoarrow.multilinestring` Arrow field with separated (struct) XY coordinates. +fn multilinestring_field(name: &str, nullable: bool, crs: Option<&str>) -> Field { + let crs = crs + .map(|crs| Crs::from_unknown_crs_type(crs.to_string())) + .unwrap_or_default(); + let metadata = Arc::new(Metadata::new(crs, None)); + MultiLineStringType::new(GeoArrowDimension::XY, metadata).to_field(name, nullable) +} + +/// An imported `geoarrow.multilinestring` field maps to the MultiLineString extension dtype (not +/// Polygon, despite the identical `List>>` storage), recovering CRS and shape. +#[test] +fn import_field_recovers_extension() -> VortexResult<()> { + let field = multilinestring_field("geom", true, Some("EPSG:4326")); + let dtype = SESSION.arrow().from_arrow_field(&field)?; + + let DType::Extension(ext) = &dtype else { + panic!("expected Extension dtype, got {dtype}"); + }; + assert!(ext.is::()); + assert_eq!( + ext.metadata::().crs.as_deref(), + Some("EPSG:4326") + ); + + // Storage peels two List layers (multilinestring → line strings) to the coordinate struct. + let DType::List(lines, nullability) = ext.storage_dtype() else { + panic!("expected List storage, got {}", ext.storage_dtype()); + }; + assert_eq!(*nullability, Nullability::Nullable); + let DType::List(coords, _) = lines.as_ref() else { + panic!("expected List of line strings"); + }; + let DType::Struct(fields, _) = coords.as_ref() else { + panic!("expected coordinate Struct"); + }; + let names: Vec<&str> = fields.names().iter().map(|n| n.as_ref()).collect(); + assert_eq!(names, vec!["x", "y"]); + Ok(()) +} + +/// A field with interleaved (`FixedSizeList`) coordinates fails to import. +#[test] +fn import_interleaved_field_fails() { + let multilinestring_type = MultiLineStringType::new(GeoArrowDimension::XY, Default::default()) + .with_coord_type(CoordType::Interleaved); + let field = multilinestring_type.to_field("geom", false); + assert!(SESSION.arrow().from_arrow_field(&field).is_err()); +} + +/// A field imported to the MultiLineString dtype and exported back carries the +/// `geoarrow.multilinestring` extension over its `List` storage. +#[test] +fn export_field_carries_extension() -> VortexResult<()> { + let imported = SESSION.arrow().from_arrow_field(&multilinestring_field( + "geom", + false, + Some("EPSG:4326"), + ))?; + let field = SESSION.arrow().to_arrow_field("geom", &imported)?; + + assert_eq!(field.extension_type_name(), Some(MultiLineStringType::NAME)); + assert!( + matches!(field.data_type(), DataType::List(_)), + "expected List storage, got {}", + field.data_type() + ); + Ok(()) +} diff --git a/vortex-geo/src/tests/multipoint.rs b/vortex-geo/src/tests/multipoint.rs new file mode 100644 index 00000000000..c1de946f0ba --- /dev/null +++ b/vortex-geo/src/tests/multipoint.rs @@ -0,0 +1,87 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Arrow interop for the `vortex.geo.multipoint` extension type (`geoarrow.multipoint`). + +use std::sync::Arc; + +use arrow_schema::DataType; +use arrow_schema::Field; +use arrow_schema::extension::ExtensionType as _; +use geoarrow::datatypes::CoordType; +use geoarrow::datatypes::Crs; +use geoarrow::datatypes::Dimension as GeoArrowDimension; +use geoarrow::datatypes::Metadata; +use geoarrow::datatypes::MultiPointType; +use vortex_array::arrow::ArrowSessionExt; +use vortex_array::dtype::DType; +use vortex_array::dtype::Nullability; +use vortex_error::VortexResult; + +use super::SESSION; +use crate::extension::MultiPoint; + +/// A `geoarrow.multipoint` Arrow field with separated (struct) XY coordinates. +fn multipoint_field(name: &str, nullable: bool, crs: Option<&str>) -> Field { + let crs = crs + .map(|crs| Crs::from_unknown_crs_type(crs.to_string())) + .unwrap_or_default(); + let metadata = Arc::new(Metadata::new(crs, None)); + MultiPointType::new(GeoArrowDimension::XY, metadata).to_field(name, nullable) +} + +/// An imported `geoarrow.multipoint` field maps to the MultiPoint extension dtype (not LineString, +/// despite the identical `List>` storage), recovering the CRS and nullability. +#[test] +fn import_field_recovers_extension() -> VortexResult<()> { + let field = multipoint_field("geom", true, Some("EPSG:4326")); + let dtype = SESSION.arrow().from_arrow_field(&field)?; + + let DType::Extension(ext) = &dtype else { + panic!("expected Extension dtype, got {dtype}"); + }; + assert!(ext.is::()); + assert_eq!( + ext.metadata::().crs.as_deref(), + Some("EPSG:4326") + ); + + let DType::List(coords, nullability) = ext.storage_dtype() else { + panic!("expected List storage, got {}", ext.storage_dtype()); + }; + assert_eq!(*nullability, Nullability::Nullable); + let DType::Struct(fields, _) = coords.as_ref() else { + panic!("expected coordinate Struct"); + }; + let names: Vec<&str> = fields.names().iter().map(|n| n.as_ref()).collect(); + assert_eq!(names, vec!["x", "y"]); + Ok(()) +} + +/// A field with interleaved (`FixedSizeList`) coordinates fails to import. +#[test] +fn import_interleaved_field_fails() { + let multipoint_type = MultiPointType::new(GeoArrowDimension::XY, Default::default()) + .with_coord_type(CoordType::Interleaved); + let field = multipoint_type.to_field("geom", false); + assert!(SESSION.arrow().from_arrow_field(&field).is_err()); +} + +/// A field imported to the MultiPoint dtype and exported back carries the `geoarrow.multipoint` +/// extension over its `List` storage. +#[test] +fn export_field_carries_extension() -> VortexResult<()> { + let imported = + SESSION + .arrow() + .from_arrow_field(&multipoint_field("geom", false, Some("EPSG:4326")))?; + let field = SESSION.arrow().to_arrow_field("geom", &imported)?; + + assert_eq!(field.extension_type_name(), Some(MultiPointType::NAME)); + assert!( + matches!(field.data_type(), DataType::List(_)), + "expected List storage, got {}", + field.data_type() + ); + Ok(()) +}