diff --git a/vortex-duckdb/src/convert/dtype.rs b/vortex-duckdb/src/convert/dtype.rs index 816fa7868f6..863a41cd869 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 abf0b68c3b4..b2c763e0613 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; @@ -107,7 +110,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(()) +}