diff --git a/Cargo.lock b/Cargo.lock index cd2662d2372..22d3da8c8e2 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -10603,6 +10603,7 @@ dependencies = [ "mimalloc", "prost 0.14.4", "rstest", + "serde_json", "vortex-array", "vortex-arrow", "vortex-buffer", diff --git a/vortex-bench/src/conversions.rs b/vortex-bench/src/conversions.rs index 0086d321e89..59e890a6f83 100644 --- a/vortex-bench/src/conversions.rs +++ b/vortex-bench/src/conversions.rs @@ -356,10 +356,7 @@ fn geoparquet_columns(metadata: &ParquetMetaData) -> HashSet { /// The erased `vortex.st.wkb` extension dtype over a binary `storage` dtype. fn wkb_ext_dtype(storage: &DType) -> VortexResult { - Ok( - ExtDType::::try_new(SpatialMetadata { crs: None }, storage.clone())? - .erased(), - ) + Ok(ExtDType::::try_new(SpatialMetadata::default(), storage.clone())?.erased()) } /// Re-type the named binary columns of a struct `dtype` as `vortex.st.wkb`, so the column diff --git a/vortex-duckdb/src/convert/dtype.rs b/vortex-duckdb/src/convert/dtype.rs index b00dee81a02..5dbb29a4724 100644 --- a/vortex-duckdb/src/convert/dtype.rs +++ b/vortex-duckdb/src/convert/dtype.rs @@ -173,7 +173,10 @@ impl FromLogicalType for DType { let crs = logical_type.geometry_crs().map(|crs| crs.to_string()); DType::Extension( ExtDType::::try_new( - SpatialMetadata { crs }, + SpatialMetadata { + crs, + ..Default::default() + }, DType::Binary(nullability), )? .erased(), @@ -616,6 +619,7 @@ mod tests { ExtDType::::try_new( SpatialMetadata { crs: Some("EPSG:4326".to_string()), + ..Default::default() }, DType::Binary(Nullability::NonNullable), )? diff --git a/vortex-duckdb/src/convert/scalar.rs b/vortex-duckdb/src/convert/scalar.rs index 804db2b05b5..a44bd5d923d 100644 --- a/vortex-duckdb/src/convert/scalar.rs +++ b/vortex-duckdb/src/convert/scalar.rs @@ -482,6 +482,7 @@ mod tests { Scalar::extension::( SpatialMetadata { crs: crs.map(str::to_string), + ..Default::default() }, Scalar::binary(bytes.to_vec(), Nullability::Nullable), ) @@ -533,6 +534,7 @@ mod tests { let dtype = ExtDType::::try_new( SpatialMetadata { crs: Some("EPSG:4326".to_string()), + ..Default::default() }, DType::Binary(Nullability::Nullable), ) diff --git a/vortex-duckdb/src/convert/vector.rs b/vortex-duckdb/src/convert/vector.rs index 4a5bd42cda5..eba70c6c9c7 100644 --- a/vortex-duckdb/src/convert/vector.rs +++ b/vortex-duckdb/src/convert/vector.rs @@ -262,7 +262,10 @@ pub fn flat_vector_to_vortex(vector: &VectorRef, len: usize) -> VortexResult::try_new( - SpatialMetadata { crs }, + SpatialMetadata { + crs, + ..Default::default() + }, DType::Binary(Nullability::Nullable), )? .erased(); diff --git a/vortex-duckdb/src/e2e_test/spatial_pushdown_test.rs b/vortex-duckdb/src/e2e_test/spatial_pushdown_test.rs index 39301da81b4..0bed547c516 100644 --- a/vortex-duckdb/src/e2e_test/spatial_pushdown_test.rs +++ b/vortex-duckdb/src/e2e_test/spatial_pushdown_test.rs @@ -59,9 +59,8 @@ fn native_point_file() -> NamedTempFile { let storage = StructArray::from_fields(&[("x", xs), ("y", ys)]) .unwrap() .into_array(); - let dtype = - ExtDType::::try_new(SpatialMetadata { crs: None }, storage.dtype().clone()) - .unwrap(); + let dtype = ExtDType::::try_new(SpatialMetadata::default(), storage.dtype().clone()) + .unwrap(); let points = ExtensionArray::new(dtype.erased(), storage).into_array(); let file = NamedTempFile::with_suffix(".vortex").unwrap(); diff --git a/vortex-duckdb/src/e2e_test/vortex_scan_test.rs b/vortex-duckdb/src/e2e_test/vortex_scan_test.rs index 0876be1ca4c..ef82549df5a 100644 --- a/vortex-duckdb/src/e2e_test/vortex_scan_test.rs +++ b/vortex-duckdb/src/e2e_test/vortex_scan_test.rs @@ -994,6 +994,7 @@ fn test_geometry() { ExtDType::::try_new( SpatialMetadata { crs: Some("EPSG:32600".to_string()), + ..Default::default() }, geometry.dtype().clone(), ) diff --git a/vortex-spatial/Cargo.toml b/vortex-spatial/Cargo.toml index a758087da03..9f43bfa78bc 100644 --- a/vortex-spatial/Cargo.toml +++ b/vortex-spatial/Cargo.toml @@ -22,6 +22,7 @@ geo-types = { workspace = true } geoarrow = { workspace = true } geoarrow-cast = { workspace = true } prost = { workspace = true } +serde_json = { workspace = true } vortex-array = { workspace = true } vortex-arrow = { workspace = true } vortex-buffer = { workspace = true } diff --git a/vortex-spatial/src/extension/geometry.rs b/vortex-spatial/src/extension/geometry.rs new file mode 100644 index 00000000000..a80c0561ca1 --- /dev/null +++ b/vortex-spatial/src/extension/geometry.rs @@ -0,0 +1,145 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Native geometry helpers and conversion to [`geo_types::Geometry`] for `geo` algorithms. + +use geo_types::Geometry; +use vortex_array::ArrayRef; +use vortex_array::ExecutionCtx; +use vortex_array::IntoArray; +use vortex_array::arrays::ConstantArray; +use vortex_array::arrays::ExtensionArray; +use vortex_array::arrays::ListViewArray; +use vortex_array::arrays::StructArray; +use vortex_array::arrays::extension::ExtensionArrayExt; +use vortex_array::arrays::list::ListArraySlotsExt; +use vortex_array::arrays::listview::ListViewArraySlotsExt; +use vortex_array::arrays::listview::list_from_list_view; +use vortex_array::builtins::ArrayBuiltins; +use vortex_array::dtype::DType; +use vortex_array::dtype::Nullability; +use vortex_array::dtype::PType; +use vortex_array::scalar::Scalar; +use vortex_buffer::Buffer; +use vortex_error::VortexResult; +use vortex_error::vortex_bail; +use vortex_error::vortex_err; + +use super::LineString; +use super::MultiLineString; +use super::MultiPoint; +use super::MultiPolygon; +use super::Point; +use super::Polygon; +use super::Rect; +use super::linestring_geometries; +use super::multilinestring_geometries; +use super::multipoint_geometries; +use super::multipolygon_geometries; +use super::point_geometries; +use super::polygon_geometries; +use super::rect_geometries; + +/// Whether `dtype` is a native geometry extension type. +pub(crate) fn is_native_geometry(dtype: &DType) -> bool { + dtype.as_extension_opt().is_some_and(|ext| { + ext.is::() + || ext.is::() + || ext.is::() + || ext.is::() + || ext.is::() + || ext.is::() + || ext.is::() + }) +} + +/// Flatten a native geometry column to its coordinates. +pub(crate) fn flatten_coordinates( + array: &ArrayRef, + ctx: &mut ExecutionCtx, +) -> VortexResult { + if !is_native_geometry(array.dtype()) { + vortex_bail!( + "spatial: operand is not a native geometry extension type, was {}", + array.dtype() + ); + } + let mut node = array + .clone() + .execute::(ctx)? + .storage_array() + .clone(); + while node.dtype().is_list() { + node = node.execute::(ctx)?.elements().clone(); + } + node.execute::(ctx) +} + +/// Flatten native geometry storage and return each row's coordinate offsets. +pub(crate) fn flatten_row_offsets( + storage: ArrayRef, + ctx: &mut ExecutionCtx, +) -> VortexResult<(Vec, StructArray)> { + let mut row_offsets: Vec = (0..=storage.len()).collect(); + let mut level = storage; + while level.dtype().is_list() { + let list = list_from_list_view(level.execute::(ctx)?, ctx)?; + let offsets = list + .offsets() + .clone() + .cast(DType::Primitive(PType::U64, Nullability::NonNullable))? + .execute::>(ctx)?; + for row_offset in &mut row_offsets { + *row_offset = usize::try_from(offsets[*row_offset]) + .map_err(|_| vortex_err!("spatial: list offset exceeds usize"))?; + } + level = list.elements().clone(); + } + Ok((row_offsets, level.execute::(ctx)?)) +} + +/// Decode a native geometry column to `geo_types`. +pub(crate) fn geometries( + array: &ArrayRef, + ctx: &mut ExecutionCtx, +) -> VortexResult>> { + let Some(ext) = array.dtype().as_extension_opt() else { + vortex_bail!( + "spatial: operand is not a geometry extension type, was {}", + array.dtype() + ); + }; + let storage = array + .clone() + .execute::(ctx)? + .storage_array() + .clone(); + if ext.is::() { + point_geometries(&storage, ext.metadata::(), ctx) + } else if ext.is::() { + linestring_geometries(&storage, ext.metadata::(), ctx) + } else if ext.is::() { + multipoint_geometries(&storage, ext.metadata::(), ctx) + } else if ext.is::() { + polygon_geometries(&storage, ext.metadata::(), ctx) + } else if ext.is::() { + multilinestring_geometries(&storage, ext.metadata::(), ctx) + } else if ext.is::() { + multipolygon_geometries(&storage, ext.metadata::(), ctx) + } else if ext.is::() { + rect_geometries(&storage, ext.metadata::(), ctx) + } else { + vortex_bail!("spatial: unsupported geometry extension {}", array.dtype()) + } +} + +/// Decode a constant operand to one geometry. +pub(crate) fn single_geometry( + scalar: &Scalar, + ctx: &mut ExecutionCtx, +) -> VortexResult> { + let array = ConstantArray::new(scalar.clone(), 1).into_array(); + geometries(&array, ctx)? + .pop() + .ok_or_else(|| vortex_err!("spatial: constant operand decoded to no geometry")) +} diff --git a/vortex-spatial/src/extension/linestring.rs b/vortex-spatial/src/extension/linestring.rs index 2c6a46ab06a..7a56f3805ac 100644 --- a/vortex-spatial/src/extension/linestring.rs +++ b/vortex-spatial/src/extension/linestring.rs @@ -3,7 +3,7 @@ //! The [`LineString`] geometry extension type (`vortex.st.linestring`): an ordered path of the //! [`Point`](super::Point) coordinate struct, stored as `List>` and tagged -//! with [`SpatialMetadata`] (CRS). +//! with [`SpatialMetadata`]. use std::sync::Arc; @@ -18,7 +18,6 @@ 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; @@ -59,9 +58,9 @@ use super::SpatialMetadata; use super::coordinate::Dimension; use super::coordinate::coordinate_dimension; use super::coordinate::coordinate_storage_dtype; -use super::geoarrow_metadata; use super::geoarrow_to_wkb; -use super::spatial_metadata_from_arrow; +use super::metadata::from_geoarrow; +use super::metadata::to_geoarrow; /// A line string: `geoarrow.linestring`, stored as `List>` (an ordered path /// of vertices). @@ -79,11 +78,11 @@ impl ExtVTable for LineString { } fn serialize_metadata(&self, metadata: &Self::Metadata) -> VortexResult> { - Ok(metadata.encode_to_vec()) + Ok(metadata.serialize()) } fn deserialize_metadata(&self, metadata: &[u8]) -> VortexResult { - Ok(SpatialMetadata::decode(metadata)?) + SpatialMetadata::deserialize(metadata) } fn validate_dtype(ext_dtype: &ExtDType) -> VortexResult<()> { @@ -186,17 +185,23 @@ 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(spatial_metadata: &SpatialMetadata, dimension: Dimension) -> LineStringType { - LineStringType::new(dimension.into(), geoarrow_metadata(spatial_metadata)) +fn linestring_type( + metadata: &SpatialMetadata, + dimension: Dimension, +) -> VortexResult { + Ok(LineStringType::new( + dimension.into(), + to_geoarrow(metadata)?, + )) } -/// Decode `LineString` storage (`List`) to `geo_types` line strings, for the spatial scalar -/// functions. CRS does not affect planar geometry ops, so default metadata is used. +/// Decode line string storage to `geo_types`. pub(crate) fn linestring_geometries( storage: &ArrayRef, + metadata: &SpatialMetadata, ctx: &mut ExecutionCtx, ) -> VortexResult>> { - linestring_array(storage, ctx)? + linestring_array(storage, metadata, ctx)? .iter() .map(|geometry| -> VortexResult> { Ok(geometry @@ -208,11 +213,12 @@ pub(crate) fn linestring_geometries( } /// Build a geoarrow `LineStringArray` from a `LineString`'s `List` storage. -fn linestring_array(storage: &ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult { - let linestring_type = linestring_type( - &SpatialMetadata::default(), - linestring_dimension(storage.dtype())?, - ); +fn linestring_array( + storage: &ArrayRef, + metadata: &SpatialMetadata, + ctx: &mut ExecutionCtx, +) -> VortexResult { + let linestring_type = linestring_type(metadata, 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)) @@ -237,7 +243,11 @@ impl TryFrom for LineStringData { 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)?) + geoarrow_to_wkb(&linestring_array( + self.0.storage_array(), + self.0.ext_dtype().metadata::(), + ctx, + )?) } } @@ -257,11 +267,11 @@ impl ArrowExportVTable for LineString { session: &ArrowSession, ) -> VortexResult> { let ext_type = dtype.as_extension(); - let spatial_metadata = ext_type.metadata::(); + let 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(spatial_metadata, dimension))?; + field.try_with_extension_type(linestring_type(metadata, dimension)?)?; Ok(Some(field)) } @@ -332,7 +342,7 @@ impl ArrowImportVTable for LineString { ); ( linestring_meta.dimension().into(), - spatial_metadata_from_arrow(linestring_meta.metadata()), + from_geoarrow(linestring_meta.metadata()), ) } else { // Literal: peel the `List` layer to the coordinate struct and read its dimension from @@ -400,6 +410,7 @@ mod tests { fn spatial_meta() -> SpatialMetadata { SpatialMetadata { crs: Some("EPSG:4326".to_string()), + ..Default::default() } } diff --git a/vortex-spatial/src/extension/literal.rs b/vortex-spatial/src/extension/literal.rs new file mode 100644 index 00000000000..117fbaffad3 --- /dev/null +++ b/vortex-spatial/src/extension/literal.rs @@ -0,0 +1,232 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! WKB geometry literal conversion. + +use std::sync::Arc; + +use ::wkb::reader::GeometryType; +use arrow_array::BinaryArray; +use geoarrow::array::GenericWkbArray; +use geoarrow::datatypes::CoordType; +use geoarrow::datatypes::Dimension; +use geoarrow::datatypes::GeoArrowType; +use geoarrow::datatypes::LineStringType; +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; +use vortex_array::ArrayRef; +use vortex_array::IntoArray; +use vortex_array::arrays::ExtensionArray; +use vortex_array::dtype::extension::ExtDType; +use vortex_array::dtype::extension::ExtVTable; +use vortex_array::scalar::Scalar; +use vortex_arrow::FromArrowArray; +use vortex_error::VortexResult; +use vortex_error::vortex_err; + +use super::LineString; +use super::MultiLineString; +use super::MultiPoint; +use super::MultiPolygon; +use super::Point; +use super::Polygon; +use super::SpatialMetadata; +use super::Wkb; +use super::metadata::to_geoarrow; + +/// Decode a WKB literal to a native geometry scalar. +pub fn native_geometry_scalar_from_wkb(bytes: &[u8]) -> VortexResult> { + let metadata = to_geoarrow(&SpatialMetadata::default())?; + let binary = BinaryArray::from(vec![Some(bytes)]); + let wkb = GenericWkbArray::::try_from(( + &binary as &dyn arrow_array::Array, + WkbType::new(Arc::clone(&metadata)), + )) + .map_err(|e| vortex_err!("failed to read WKB literal: {e}"))?; + + let to_storage = |target: &GeoArrowType| -> VortexResult { + let native = + cast(&wkb, target).map_err(|e| vortex_err!("failed to cast WKB literal: {e}"))?; + ArrayRef::from_arrow(native.to_array_ref().as_ref(), false) + }; + + let scalar = match Wkb::try_from_bytes(bytes)?.geometry_type() { + GeometryType::Point => { + let target = GeoArrowType::Point( + PointType::new(Dimension::XY, metadata).with_coord_type(CoordType::Separated), + ); + spatial_ext_scalar(Point, to_storage(&target)?)? + } + GeometryType::LineString => { + let target = GeoArrowType::LineString( + LineStringType::new(Dimension::XY, metadata).with_coord_type(CoordType::Separated), + ); + spatial_ext_scalar(LineString, to_storage(&target)?)? + } + GeometryType::Polygon => { + let target = GeoArrowType::Polygon( + PolygonType::new(Dimension::XY, metadata).with_coord_type(CoordType::Separated), + ); + spatial_ext_scalar(Polygon, to_storage(&target)?)? + } + GeometryType::MultiPoint => { + let target = GeoArrowType::MultiPoint( + MultiPointType::new(Dimension::XY, metadata).with_coord_type(CoordType::Separated), + ); + spatial_ext_scalar(MultiPoint, to_storage(&target)?)? + } + GeometryType::MultiLineString => { + let target = GeoArrowType::MultiLineString( + MultiLineStringType::new(Dimension::XY, metadata) + .with_coord_type(CoordType::Separated), + ); + spatial_ext_scalar(MultiLineString, to_storage(&target)?)? + } + GeometryType::MultiPolygon => { + let target = GeoArrowType::MultiPolygon( + MultiPolygonType::new(Dimension::XY, metadata) + .with_coord_type(CoordType::Separated), + ); + spatial_ext_scalar(MultiPolygon, to_storage(&target)?)? + } + _ => return Ok(None), + }; + Ok(Some(scalar)) +} + +// `scalar_at` is deprecated, but literal conversion has no execution context. +#[allow(deprecated)] +fn spatial_ext_scalar>( + vtable: V, + storage: ArrayRef, +) -> VortexResult { + let ext = + ExtDType::try_with_vtable(vtable, SpatialMetadata::default(), storage.dtype().clone())? + .erased(); + ExtensionArray::try_new(ext, storage)? + .into_array() + .scalar_at(0) +} + +#[cfg(test)] +mod tests { + use vortex_array::dtype::DType; + 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; + + #[test] + fn decodes_wkb_point_to_native() -> VortexResult<()> { + let mut wkb = vec![1u8]; + wkb.extend_from_slice(&1u32.to_le_bytes()); + wkb.extend_from_slice(&1.0f64.to_le_bytes()); + wkb.extend_from_slice(&2.0f64.to_le_bytes()); + + let scalar = native_geometry_scalar_from_wkb(&wkb)?.expect("a point scalar"); + let DType::Extension(ext) = scalar.dtype() else { + panic!("expected an extension dtype, got {}", scalar.dtype()); + }; + assert!(ext.is::()); + Ok(()) + } + + #[test] + fn decodes_wkb_polygon_to_native() -> VortexResult<()> { + let ring = [(0.0, 0.0), (1.0, 0.0), (0.0, 1.0), (0.0, 0.0)]; + let mut wkb = vec![1u8]; + wkb.extend_from_slice(&3u32.to_le_bytes()); + wkb.extend_from_slice(&1u32.to_le_bytes()); + let ring_len = u32::try_from(ring.len()).map_err(|e| vortex_err!("{e}"))?; + wkb.extend_from_slice(&ring_len.to_le_bytes()); + for (x, y) in ring { + 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 polygon scalar"); + let DType::Extension(ext) = scalar.dtype() else { + panic!("expected an extension dtype, got {}", scalar.dtype()); + }; + assert!(ext.is::()); + Ok(()) + } + + #[test] + fn decodes_wkb_linestring_to_native() -> VortexResult<()> { + let points = [(0.0, 0.0), (1.0, 1.0)]; + let mut wkb = vec![1u8]; + wkb.extend_from_slice(&2u32.to_le_bytes()); + 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(()) + } + + #[test] + fn decodes_wkb_multipoint_to_native() -> VortexResult<()> { + let points = [(0.0, 0.0), (1.0, 1.0)]; + let mut wkb = vec![1u8]; + wkb.extend_from_slice(&4u32.to_le_bytes()); + 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.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(()) + } + + #[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]; + wkb.extend_from_slice(&5u32.to_le_bytes()); + 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 { + 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-spatial/src/extension/metadata.rs b/vortex-spatial/src/extension/metadata.rs new file mode 100644 index 00000000000..8d6afecef05 --- /dev/null +++ b/vortex-spatial/src/extension/metadata.rs @@ -0,0 +1,296 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Spatial metadata and GeoArrow conversion. + +use std::fmt::Display; +use std::sync::Arc; + +use geoarrow::datatypes::Crs; +use geoarrow::datatypes::CrsType as ArrowCrsType; +use geoarrow::datatypes::Edges as ArrowEdges; +use geoarrow::datatypes::Metadata; +use prost::Message; +use vortex_error::VortexError; +use vortex_error::VortexResult; +use vortex_error::vortex_bail; +use vortex_error::vortex_ensure; +use vortex_error::vortex_err; + +/// CRS serialization format. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, prost::Enumeration)] +#[repr(i32)] +pub enum CrsType { + /// A PROJJSON object. + Projjson = 0, + /// A WKT2:2019 string. + Wkt2_2019 = 1, + /// An `AUTHORITY:CODE` identifier. + AuthorityCode = 2, + /// An opaque, producer-specific identifier. + Srid = 3, +} + +impl From for ArrowCrsType { + fn from(value: CrsType) -> Self { + match value { + CrsType::Projjson => Self::Projjson, + CrsType::Wkt2_2019 => Self::Wkt2_2019, + CrsType::AuthorityCode => Self::AuthorityCode, + CrsType::Srid => Self::Srid, + } + } +} + +impl From for CrsType { + fn from(value: ArrowCrsType) -> Self { + match value { + ArrowCrsType::Projjson => Self::Projjson, + ArrowCrsType::Wkt2_2019 => Self::Wkt2_2019, + ArrowCrsType::AuthorityCode => Self::AuthorityCode, + ArrowCrsType::Srid => Self::Srid, + } + } +} + +/// Edge interpretation. An omitted value means planar edges. +#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, prost::Enumeration)] +#[repr(i32)] +pub enum Edges { + /// Andoyer geodesics on the CRS ellipsoid. + Andoyer = 0, + /// Karney geodesics on the CRS ellipsoid. + Karney = 1, + /// Great-circle paths. + Spherical = 2, + /// Thomas geodesics on the CRS ellipsoid. + Thomas = 3, + /// Vincenty geodesics on the CRS ellipsoid. + Vincenty = 4, +} + +impl From for ArrowEdges { + fn from(value: Edges) -> Self { + match value { + Edges::Andoyer => Self::Andoyer, + Edges::Karney => Self::Karney, + Edges::Spherical => Self::Spherical, + Edges::Thomas => Self::Thomas, + Edges::Vincenty => Self::Vincenty, + } + } +} + +impl From for Edges { + fn from(value: ArrowEdges) -> Self { + match value { + ArrowEdges::Andoyer => Self::Andoyer, + ArrowEdges::Karney => Self::Karney, + ArrowEdges::Spherical => Self::Spherical, + ArrowEdges::Thomas => Self::Thomas, + ArrowEdges::Vincenty => Self::Vincenty, + } + } +} + +/// Metadata shared by all spatial extension types. +#[derive(Debug, Default, Clone, PartialEq, Eq, Hash)] +pub struct SpatialMetadata { + /// Coordinate reference system, if known. + pub crs: Option, + /// Serialization format of [`Self::crs`]. + pub crs_type: Option, + /// Edge interpretation. `None` means planar edges. + pub edges: Option, +} + +/// Protobuf representation of [`SpatialMetadata`]. +#[derive(Clone, PartialEq, Message)] +struct SpatialMetadataProto { + #[prost(optional, string, tag = "1")] + crs: Option, + #[prost(enumeration = "CrsType", optional, tag = "2")] + crs_type: Option, + #[prost(enumeration = "Edges", optional, tag = "3")] + edges: Option, +} + +impl From<&SpatialMetadata> for SpatialMetadataProto { + fn from(metadata: &SpatialMetadata) -> Self { + Self { + crs: metadata.crs.clone(), + crs_type: metadata.crs_type.map(Into::into), + edges: metadata.edges.map(Into::into), + } + } +} + +impl TryFrom for SpatialMetadata { + type Error = VortexError; + + fn try_from(proto: SpatialMetadataProto) -> VortexResult { + let crs_type = proto + .crs_type + .map(CrsType::try_from) + .transpose() + .map_err(|error| vortex_err!("spatial: invalid CRS type: {error}"))?; + let edges = proto + .edges + .map(Edges::try_from) + .transpose() + .map_err(|error| vortex_err!("spatial: invalid edges value: {error}"))?; + Ok(Self { + crs: proto.crs, + crs_type, + edges, + }) + } +} + +impl SpatialMetadata { + pub(super) fn serialize(&self) -> Vec { + SpatialMetadataProto::from(self).encode_to_vec() + } + + pub(super) fn deserialize(bytes: &[u8]) -> VortexResult { + SpatialMetadataProto::decode(bytes)?.try_into() + } +} + +impl Display for SpatialMetadata { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + match self.crs.as_ref() { + Some(crs) => write!(f, "Geometry(crs={crs})"), + None => write!(f, "Geometry(unreferenced)"), + } + } +} + +/// Convert spatial metadata to GeoArrow. +pub(crate) fn to_geoarrow(metadata: &SpatialMetadata) -> VortexResult> { + let crs = match (metadata.crs.as_deref(), metadata.crs_type) { + (None, None) => Crs::default(), + (None, Some(crs_type)) => { + vortex_bail!("spatial: CRS type {crs_type:?} requires a CRS value") + } + (Some(crs), None) => Crs::from_unknown_crs_type(crs.to_owned()), + (Some(crs), Some(CrsType::Projjson)) => { + let value: serde_json::Value = serde_json::from_str(crs) + .map_err(|error| vortex_err!("spatial: invalid PROJJSON CRS: {error}"))?; + vortex_ensure!( + value.is_object(), + "spatial: PROJJSON CRS must be a JSON object" + ); + Crs::from_projjson(value) + } + (Some(crs), Some(CrsType::Wkt2_2019)) => Crs::from_wkt2_2019(crs.to_owned()), + (Some(crs), Some(CrsType::AuthorityCode)) => { + vortex_ensure!( + crs.contains(':'), + "spatial: authority-code CRS must have the form AUTHORITY:CODE" + ); + Crs::from_authority_code(crs.to_owned()) + } + (Some(crs), Some(CrsType::Srid)) => Crs::from_srid(crs.to_owned()), + }; + Ok(Arc::new(Metadata::new(crs, metadata.edges.map(Into::into)))) +} + +/// Convert GeoArrow metadata to [`SpatialMetadata`]. +pub(crate) fn from_geoarrow(metadata: &Metadata) -> SpatialMetadata { + let arrow_crs = metadata.crs(); + let value = arrow_crs.crs_value(); + let crs = value.map(|value| { + value + .as_str() + .map_or_else(|| value.to_string(), str::to_owned) + }); + // GeoArrow defines an object-valued CRS as PROJJSON, even when `crs_type` is omitted. + let crs_type = arrow_crs.crs_type().map(Into::into).or_else(|| { + value + .filter(|value| value.is_object()) + .map(|_| CrsType::Projjson) + }); + + SpatialMetadata { + crs, + crs_type, + edges: metadata.edges().map(Into::into), + } +} + +#[cfg(test)] +mod tests { + use rstest::rstest; + use vortex_error::VortexResult; + + use super::CrsType; + use super::Edges; + use super::SpatialMetadata; + use super::from_geoarrow; + use super::to_geoarrow; + + #[test] + fn metadata_roundtrips_serialization() -> VortexResult<()> { + let metadata = SpatialMetadata { + crs: Some("EPSG:4326".to_string()), + crs_type: Some(CrsType::AuthorityCode), + edges: Some(Edges::Spherical), + }; + + assert_eq!(metadata.to_string(), "Geometry(crs=EPSG:4326)"); + let decoded = SpatialMetadata::deserialize(&metadata.serialize())?; + assert_eq!(decoded, metadata); + Ok(()) + } + + #[test] + fn decodes_legacy_crs_only_bytes() -> VortexResult<()> { + let legacy = b"\x0a\x09EPSG:4326"; + assert_eq!( + SpatialMetadata::deserialize(legacy)?, + SpatialMetadata { + crs: Some("EPSG:4326".to_string()), + crs_type: None, + edges: None, + } + ); + Ok(()) + } + + #[rstest] + #[case::projjson(CrsType::Projjson, r#"{"type":"GeographicCRS"}"#)] + #[case::wkt(CrsType::Wkt2_2019, "GEOGCRS[\"WGS 84\"]")] + #[case::authority_code(CrsType::AuthorityCode, "EPSG:4326")] + #[case::srid(CrsType::Srid, "database-crs-42")] + fn geoarrow_crs_type_roundtrips( + #[case] crs_type: CrsType, + #[case] crs: &str, + ) -> VortexResult<()> { + let metadata = SpatialMetadata { + crs: Some(crs.to_string()), + crs_type: Some(crs_type), + edges: None, + }; + let geoarrow = to_geoarrow(&metadata)?; + assert_eq!(from_geoarrow(&geoarrow), metadata); + Ok(()) + } + + #[rstest] + #[case::andoyer(Edges::Andoyer)] + #[case::karney(Edges::Karney)] + #[case::spherical(Edges::Spherical)] + #[case::thomas(Edges::Thomas)] + #[case::vincenty(Edges::Vincenty)] + fn geoarrow_edges_roundtrip(#[case] edges: Edges) -> VortexResult<()> { + let metadata = SpatialMetadata { + crs: None, + crs_type: None, + edges: Some(edges), + }; + let geoarrow = to_geoarrow(&metadata)?; + assert_eq!(from_geoarrow(&geoarrow), metadata); + Ok(()) + } +} diff --git a/vortex-spatial/src/extension/mod.rs b/vortex-spatial/src/extension/mod.rs index d1e2c37ebf4..fb3b2344192 100644 --- a/vortex-spatial/src/extension/mod.rs +++ b/vortex-spatial/src/extension/mod.rs @@ -2,7 +2,10 @@ // SPDX-FileCopyrightText: Copyright the Vortex contributors pub(crate) mod coordinate; +mod geometry; mod linestring; +mod literal; +mod metadata; mod multilinestring; mod multipoint; mod multipolygon; @@ -11,454 +14,20 @@ mod polygon; mod rect; mod wkb; -use std::fmt::Display; -use std::sync::Arc; - -use ::wkb::reader::GeometryType; -use arrow_array::BinaryArray; -use geo_types::Geometry; -use geoarrow::array::GenericWkbArray; -use geoarrow::array::GeoArrowArray; -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(crate) use geometry::flatten_coordinates; +pub(crate) use geometry::flatten_row_offsets; +pub(crate) use geometry::geometries; +pub(crate) use geometry::is_native_geometry; +pub(crate) use geometry::single_geometry; pub use linestring::*; +pub use literal::native_geometry_scalar_from_wkb; +pub use metadata::CrsType; +pub use metadata::Edges; +pub use metadata::SpatialMetadata; pub use multilinestring::*; pub use multipoint::*; pub use multipolygon::*; pub use point::*; pub use polygon::*; pub use rect::*; -use vortex_array::ArrayRef; -use vortex_array::ExecutionCtx; -use vortex_array::IntoArray; -use vortex_array::arrays::ConstantArray; -use vortex_array::arrays::ExtensionArray; -use vortex_array::arrays::ListViewArray; -use vortex_array::arrays::StructArray; -use vortex_array::arrays::extension::ExtensionArrayExt; -use vortex_array::arrays::list::ListArraySlotsExt; -use vortex_array::arrays::listview::ListViewArraySlotsExt; -use vortex_array::arrays::listview::list_from_list_view; -use vortex_array::builtins::ArrayBuiltins; -use vortex_array::dtype::DType; -use vortex_array::dtype::Nullability; -use vortex_array::dtype::PType; -use vortex_array::dtype::extension::ExtDType; -use vortex_array::dtype::extension::ExtVTable; -use vortex_array::scalar::Scalar; -use vortex_arrow::FromArrowArray; -use vortex_buffer::Buffer; -use vortex_error::VortexResult; -use vortex_error::vortex_bail; -use vortex_error::vortex_err; pub use wkb::*; - -/// Whether `dtype` is one of the native geometry extension types the spatial kernels operate on. -pub(crate) fn is_native_geometry(dtype: &DType) -> bool { - dtype.as_extension_opt().is_some_and(|ext| { - ext.is::() - || ext.is::() - || ext.is::() - || ext.is::() - || ext.is::() - || ext.is::() - || ext.is::() - }) -} - -/// Flatten a native geometry column into a single coordinate `Struct` containing -/// every vertex of every geometry. -pub(crate) fn flatten_coordinates( - array: &ArrayRef, - ctx: &mut ExecutionCtx, -) -> VortexResult { - if !is_native_geometry(array.dtype()) { - vortex_bail!( - "spatial: operand is not a native geometry extension type, was {}", - array.dtype() - ); - } - let mut node = array - .clone() - .execute::(ctx)? - .storage_array() - .clone(); - while node.dtype().is_list() { - node = node.execute::(ctx)?.elements().clone(); - } - node.execute::(ctx) -} - -/// Flatten a native geometry `storage` array to its leaf coordinates, keeping track of which -/// row owns which coordinates. -/// -/// Returns the flat coordinate `Struct` plus `row_offsets` (one entry per row plus a final -/// cap): row `r`'s coordinates are `coordinates[row_offsets[r]..row_offsets[r + 1]]`, an empty -/// range if the row has none. -/// -/// Row boundaries are pushed down one `List` level at a time. A boundary at list `e` moves to -/// `offsets[e]`, where that list's children start. For example, a 2-row `MultiPolygon` column — -/// row 0 = one polygon of two rings (3 + 4 vertices), row 1 = one polygon of one ring (5 -/// vertices): -/// -/// ```text -/// level offsets row_offsets after the level -/// rows → polygons [0,1,2] [0,1,2] row 1 starts at polygon 1 -/// polygons → rings [0,2,3] [0,2,3] row 1 starts at ring 2 -/// rings → vertices [0,3,7,12] [0,7,12] row 1 starts at vertex 7 -/// ``` -pub(crate) fn flatten_row_offsets( - storage: ArrayRef, - ctx: &mut ExecutionCtx, -) -> VortexResult<(Vec, StructArray)> { - let len = storage.len(); - - // At the outermost level, row `r` starts at element `r`; the extra entry caps the last row. - let mut row_offsets: Vec = (0..=len).collect(); - let mut level = storage; - while level.dtype().is_list() { - let list = list_from_list_view(level.execute::(ctx)?, ctx)?; - let offsets = list - .offsets() - .clone() - .cast(DType::Primitive(PType::U64, Nullability::NonNullable))? - .execute::>(ctx)?; - for row_offset in &mut row_offsets { - *row_offset = usize::try_from(offsets[*row_offset]) - .map_err(|_| vortex_err!("spatial: list offset exceeds usize"))?; - } - level = list.elements().clone(); - } - Ok((row_offsets, level.execute::(ctx)?)) -} - -/// Decode a native geometry column to `geo_types`. A non-geometry operand is an error. -pub(crate) fn geometries( - array: &ArrayRef, - ctx: &mut ExecutionCtx, -) -> VortexResult>> { - let Some(ext) = array.dtype().as_extension_opt() else { - vortex_bail!( - "spatial: operand is not a geometry extension type, was {}", - array.dtype() - ); - }; - let storage = array - .clone() - .execute::(ctx)? - .storage_array() - .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 if ext.is::() { - rect_geometries(&storage, ctx) - } else { - vortex_bail!("spatial: unsupported geometry extension {}", array.dtype()) - } -} - -/// Decode a constant operand scalar to one geometry, a constant of any -/// supported geometry type is decoded exactly like a column. -pub(crate) fn single_geometry( - scalar: &Scalar, - ctx: &mut ExecutionCtx, -) -> VortexResult> { - let array = ConstantArray::new(scalar.clone(), 1).into_array(); - geometries(&array, ctx)? - .pop() - .ok_or_else(|| vortex_err!("spatial: constant operand decoded to no geometry")) -} - -/// Decode a WKB geometry literal (DuckDB's wire form for `GEOMETRY` constants) to its native -/// `Point`/`Polygon`/`MultiPolygon` scalar. `None` for unsupported types. Plan-time, one value only. -pub fn native_geometry_scalar_from_wkb(bytes: &[u8]) -> VortexResult> { - let metadata = geoarrow_metadata(&SpatialMetadata::default()); - let binary = BinaryArray::from(vec![Some(bytes)]); - let wkb = GenericWkbArray::::try_from(( - &binary as &dyn arrow_array::Array, - WkbType::new(Arc::clone(&metadata)), - )) - .map_err(|e| vortex_err!("failed to read WKB literal: {e}"))?; - - // Cast the WKB value to `target`, import its native storage as a Vortex array. - let to_storage = |target: &GeoArrowType| -> VortexResult { - let native = - cast(&wkb, target).map_err(|e| vortex_err!("failed to cast WKB literal: {e}"))?; - ArrayRef::from_arrow(native.to_array_ref().as_ref(), false) - }; - - let scalar = match Wkb::try_from_bytes(bytes)?.geometry_type() { - GeometryType::Point => { - let target = GeoArrowType::Point( - PointType::new(Dimension::XY, metadata).with_coord_type(CoordType::Separated), - ); - spatial_ext_scalar(Point, to_storage(&target)?)? - } - GeometryType::LineString => { - let target = GeoArrowType::LineString( - LineStringType::new(Dimension::XY, metadata).with_coord_type(CoordType::Separated), - ); - spatial_ext_scalar(LineString, to_storage(&target)?)? - } - GeometryType::Polygon => { - let target = GeoArrowType::Polygon( - PolygonType::new(Dimension::XY, metadata).with_coord_type(CoordType::Separated), - ); - spatial_ext_scalar(Polygon, to_storage(&target)?)? - } - GeometryType::MultiPoint => { - let target = GeoArrowType::MultiPoint( - MultiPointType::new(Dimension::XY, metadata).with_coord_type(CoordType::Separated), - ); - spatial_ext_scalar(MultiPoint, to_storage(&target)?)? - } - GeometryType::MultiLineString => { - let target = GeoArrowType::MultiLineString( - MultiLineStringType::new(Dimension::XY, metadata) - .with_coord_type(CoordType::Separated), - ); - spatial_ext_scalar(MultiLineString, to_storage(&target)?)? - } - GeometryType::MultiPolygon => { - let target = GeoArrowType::MultiPolygon( - MultiPolygonType::new(Dimension::XY, metadata) - .with_coord_type(CoordType::Separated), - ); - spatial_ext_scalar(MultiPolygon, to_storage(&target)?)? - } - _ => return Ok(None), - }; - Ok(Some(scalar)) -} - -/// Wrap cast-from-WKB `storage` in its `vtable` extension type and pull out the single scalar. -// `scalar_at` is deprecated for `execute_scalar`, but there is no execution context at plan time. -#[allow(deprecated)] -fn spatial_ext_scalar>( - vtable: V, - storage: ArrayRef, -) -> VortexResult { - let ext = - ExtDType::try_with_vtable(vtable, SpatialMetadata::default(), storage.dtype().clone())? - .erased(); - ExtensionArray::try_new(ext, storage)? - .into_array() - .scalar_at(0) -} - -/// Extension metadata that is common to all the spatial extension types. -/// -/// Currently, this is just the coordinate reference system (CRS). -/// We may wish to add a second field for edges interpretation in the future similar to -/// the GeoArrow standard. -#[derive(Clone, PartialEq, Eq, Hash, prost::Message)] -pub struct SpatialMetadata { - #[prost(optional, string, tag = "1")] - pub crs: Option, -} - -impl Display for SpatialMetadata { - fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - match self.crs.as_ref() { - Some(crs) => write!(f, "Geometry(crs={crs})"), - None => write!(f, "Geometry(unreferenced)"), - } - } -} - -/// The GeoArrow [`Metadata`] equivalent of `spatial_metadata`. -pub(crate) fn geoarrow_metadata(spatial_metadata: &SpatialMetadata) -> Arc { - Arc::new(Metadata::new( - spatial_metadata - .crs - .as_ref() - .map(|crs| Crs::from_unknown_crs_type(crs.to_string())) - .unwrap_or_default(), - None, - )) -} - -/// Serialize a native geometry array to WKB (a `WkbView` array) via geoarrow's cast. -/// Shared by the `to_wkb` methods on the geometry extension types. -pub(crate) fn geoarrow_to_wkb(geoarrow_array: &dyn GeoArrowArray) -> VortexResult { - let wkb_type = - GeoArrowType::WkbView(WkbType::new(geoarrow_metadata(&SpatialMetadata::default()))); - let wkb = cast(geoarrow_array, &wkb_type) - .map_err(|e| vortex_err!("failed to cast geometry to WKB: {e}"))?; - ArrayRef::from_arrow(wkb.to_array_ref().as_ref(), false) -} - -/// Recover [`SpatialMetadata`] from GeoArrow metadata. -pub(crate) fn spatial_metadata_from_arrow(metadata: &Metadata) -> SpatialMetadata { - let crs = metadata.crs().crs_value().map(|value| { - // `Crs::from_unknown_crs_type` stores the user's string verbatim as a JSON string - // value, so prefer the raw string when available to round-trip cleanly. For other - // CRS encodings (PROJJSON object, etc.), fall back to the JSON-encoded form. - value - .as_str() - .map(str::to_string) - .unwrap_or_else(|| value.to_string()) - }); - SpatialMetadata { crs } -} - -#[cfg(test)] -mod tests { - use prost::Message; - use vortex_array::dtype::DType; - 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; - use crate::extension::SpatialMetadata; - - #[test] - fn test_metadata() { - let meta = SpatialMetadata { - crs: Some("EPSG:4326".to_string()), - }; - - assert_eq!(meta.to_string(), "Geometry(crs=EPSG:4326)"); - // round trip - let bytes = meta.encode_to_vec(); - let decoded = SpatialMetadata::decode(bytes.as_slice()).unwrap(); - assert_eq!(decoded, meta); - } - - /// A little-endian WKB `POINT` literal decodes to the native `Point` extension scalar. - #[test] - fn decodes_wkb_point_to_native() -> VortexResult<()> { - let mut wkb = vec![1u8]; // little-endian byte order - wkb.extend_from_slice(&1u32.to_le_bytes()); // geometry type: point - wkb.extend_from_slice(&1.0f64.to_le_bytes()); // x - wkb.extend_from_slice(&2.0f64.to_le_bytes()); // y - - let scalar = native_geometry_scalar_from_wkb(&wkb)?.expect("a point scalar"); - let DType::Extension(ext) = scalar.dtype() else { - panic!("expected an extension dtype, got {}", scalar.dtype()); - }; - assert!(ext.is::()); - Ok(()) - } - - /// A little-endian WKB `POLYGON` literal decodes to the native `Polygon` extension scalar. - #[test] - fn decodes_wkb_polygon_to_native() -> VortexResult<()> { - let ring = [(0.0, 0.0), (1.0, 0.0), (0.0, 1.0), (0.0, 0.0)]; - let mut wkb = vec![1u8]; // little-endian byte order - wkb.extend_from_slice(&3u32.to_le_bytes()); // geometry type: polygon - wkb.extend_from_slice(&1u32.to_le_bytes()); // one ring - let ring_len = u32::try_from(ring.len()).map_err(|e| vortex_err!("{e}"))?; - wkb.extend_from_slice(&ring_len.to_le_bytes()); - for (x, y) in ring { - 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 polygon scalar"); - let DType::Extension(ext) = scalar.dtype() else { - panic!("expected an extension dtype, got {}", scalar.dtype()); - }; - 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-spatial/src/extension/multilinestring.rs b/vortex-spatial/src/extension/multilinestring.rs index a7e7104927a..81f6f4e31be 100644 --- a/vortex-spatial/src/extension/multilinestring.rs +++ b/vortex-spatial/src/extension/multilinestring.rs @@ -19,7 +19,6 @@ 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; @@ -50,9 +49,9 @@ use super::SpatialMetadata; use super::coordinate::Dimension; use super::coordinate::coordinate_dimension; use super::coordinate::coordinate_storage_dtype; -use super::geoarrow_metadata; use super::geoarrow_to_wkb; -use super::spatial_metadata_from_arrow; +use super::metadata::from_geoarrow; +use super::metadata::to_geoarrow; /// A multilinestring: `geoarrow.multilinestring`, stored as `List>>` /// (line strings of vertices). @@ -70,11 +69,11 @@ impl ExtVTable for MultiLineString { } fn serialize_metadata(&self, metadata: &Self::Metadata) -> VortexResult> { - Ok(metadata.encode_to_vec()) + Ok(metadata.serialize()) } fn deserialize_metadata(&self, metadata: &[u8]) -> VortexResult { - Ok(SpatialMetadata::decode(metadata)?) + SpatialMetadata::deserialize(metadata) } fn validate_dtype(ext_dtype: &ExtDType) -> VortexResult<()> { @@ -112,18 +111,22 @@ static ARROW_MULTILINESTRING: CachedId = CachedId::new(MultiLineStringType::NAME /// The `geoarrow.multilinestring` type for `dimension`, with separated (struct) coordinates. fn multilinestring_type( - spatial_metadata: &SpatialMetadata, + metadata: &SpatialMetadata, dimension: Dimension, -) -> MultiLineStringType { - MultiLineStringType::new(dimension.into(), geoarrow_metadata(spatial_metadata)) +) -> VortexResult { + Ok(MultiLineStringType::new( + dimension.into(), + to_geoarrow(metadata)?, + )) } -/// Decode storage to `geo_types` for the spatial scalar functions (CRS is irrelevant to planar ops). +/// Decode multi-line string storage to `geo_types`. pub(crate) fn multilinestring_geometries( storage: &ArrayRef, + metadata: &SpatialMetadata, ctx: &mut ExecutionCtx, ) -> VortexResult>> { - multilinestring_array(storage, ctx)? + multilinestring_array(storage, metadata, ctx)? .iter() .map(|geometry| -> VortexResult> { Ok(geometry @@ -137,12 +140,11 @@ pub(crate) fn multilinestring_geometries( /// Build a geoarrow `MultiLineStringArray` from the `MultiLineString` storage. fn multilinestring_array( storage: &ArrayRef, + metadata: &SpatialMetadata, ctx: &mut ExecutionCtx, ) -> VortexResult { - let multilinestring_type = multilinestring_type( - &SpatialMetadata::default(), - multilinestring_dimension(storage.dtype())?, - ); + let multilinestring_type = + multilinestring_type(metadata, 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)) @@ -167,7 +169,11 @@ impl TryFrom for MultiLineStringData { 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)?) + geoarrow_to_wkb(&multilinestring_array( + self.0.storage_array(), + self.0.ext_dtype().metadata::(), + ctx, + )?) } } @@ -187,11 +193,11 @@ impl ArrowExportVTable for MultiLineString { session: &ArrowSession, ) -> VortexResult> { let ext_type = dtype.as_extension(); - let spatial_metadata = ext_type.metadata::(); + let 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(spatial_metadata, dimension))?; + field.try_with_extension_type(multilinestring_type(metadata, dimension)?)?; Ok(Some(field)) } @@ -262,7 +268,7 @@ impl ArrowImportVTable for MultiLineString { ); ( multilinestring_meta.dimension().into(), - spatial_metadata_from_arrow(multilinestring_meta.metadata()), + from_geoarrow(multilinestring_meta.metadata()), ) } else { // Literal: peel the two `List` layers to the coordinate struct and read its dimension @@ -336,6 +342,7 @@ mod tests { fn spatial_meta() -> SpatialMetadata { SpatialMetadata { crs: Some("EPSG:4326".to_string()), + ..Default::default() } } diff --git a/vortex-spatial/src/extension/multipoint.rs b/vortex-spatial/src/extension/multipoint.rs index 46597ccdd5c..f0d9199e81b 100644 --- a/vortex-spatial/src/extension/multipoint.rs +++ b/vortex-spatial/src/extension/multipoint.rs @@ -3,8 +3,8 @@ //! The [`MultiPoint`] geometry extension type (`vortex.st.multipoint`): an unordered set of the //! [`Point`](super::Point) coordinate struct, stored as `List>` and tagged -//! with [`SpatialMetadata`] (CRS). The storage layout matches [`LineString`](super::LineString); the -//! two are distinguished by their GeoArrow extension name, not their shape. +//! with [`SpatialMetadata`]. The storage layout matches [`LineString`](super::LineString); the two +//! are distinguished by their GeoArrow extension name, not their shape. use std::sync::Arc; @@ -19,7 +19,6 @@ 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; @@ -50,9 +49,9 @@ use super::SpatialMetadata; use super::coordinate::Dimension; use super::coordinate::coordinate_dimension; use super::coordinate::coordinate_storage_dtype; -use super::geoarrow_metadata; use super::geoarrow_to_wkb; -use super::spatial_metadata_from_arrow; +use super::metadata::from_geoarrow; +use super::metadata::to_geoarrow; /// A multipoint: `geoarrow.multipoint`, stored as `List>` (a set of points). #[derive(Debug, Clone, Default, PartialEq, Eq, Hash)] @@ -69,11 +68,11 @@ impl ExtVTable for MultiPoint { } fn serialize_metadata(&self, metadata: &Self::Metadata) -> VortexResult> { - Ok(metadata.encode_to_vec()) + Ok(metadata.serialize()) } fn deserialize_metadata(&self, metadata: &[u8]) -> VortexResult { - Ok(SpatialMetadata::decode(metadata)?) + SpatialMetadata::deserialize(metadata) } fn validate_dtype(ext_dtype: &ExtDType) -> VortexResult<()> { @@ -105,16 +104,23 @@ pub(crate) fn multipoint_dimension(dtype: &DType) -> VortexResult { static ARROW_MULTIPOINT: CachedId = CachedId::new(MultiPointType::NAME); /// The `geoarrow.multipoint` extension type for `dimension`, with separated (struct) coordinates. -fn multipoint_type(spatial_metadata: &SpatialMetadata, dimension: Dimension) -> MultiPointType { - MultiPointType::new(dimension.into(), geoarrow_metadata(spatial_metadata)) +fn multipoint_type( + metadata: &SpatialMetadata, + dimension: Dimension, +) -> VortexResult { + Ok(MultiPointType::new( + dimension.into(), + to_geoarrow(metadata)?, + )) } -/// Decode `MultiPoint` storage (`List`) to `geo_types`, for the spatial scalar functions. +/// Decode multipoint storage to `geo_types`. pub(crate) fn multipoint_geometries( storage: &ArrayRef, + metadata: &SpatialMetadata, ctx: &mut ExecutionCtx, ) -> VortexResult>> { - multipoint_array(storage, ctx)? + multipoint_array(storage, metadata, ctx)? .iter() .map(|geometry| -> VortexResult> { Ok(geometry @@ -126,11 +132,12 @@ pub(crate) fn multipoint_geometries( } /// Build a geoarrow `MultiPointArray` from a `MultiPoint`'s `List` storage. -fn multipoint_array(storage: &ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult { - let multipoint_type = multipoint_type( - &SpatialMetadata::default(), - multipoint_dimension(storage.dtype())?, - ); +fn multipoint_array( + storage: &ArrayRef, + metadata: &SpatialMetadata, + ctx: &mut ExecutionCtx, +) -> VortexResult { + let multipoint_type = multipoint_type(metadata, 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)) @@ -155,7 +162,11 @@ impl TryFrom for MultiPointData { 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)?) + geoarrow_to_wkb(&multipoint_array( + self.0.storage_array(), + self.0.ext_dtype().metadata::(), + ctx, + )?) } } @@ -175,11 +186,11 @@ impl ArrowExportVTable for MultiPoint { session: &ArrowSession, ) -> VortexResult> { let ext_type = dtype.as_extension(); - let spatial_metadata = ext_type.metadata::(); + let 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(spatial_metadata, dimension))?; + field.try_with_extension_type(multipoint_type(metadata, dimension)?)?; Ok(Some(field)) } @@ -249,7 +260,7 @@ impl ArrowImportVTable for MultiPoint { ); ( multipoint_meta.dimension().into(), - spatial_metadata_from_arrow(multipoint_meta.metadata()), + from_geoarrow(multipoint_meta.metadata()), ) } else { if field.extension_type_name() != Some(MultiPointType::NAME) { @@ -315,6 +326,7 @@ mod tests { fn spatial_meta() -> SpatialMetadata { SpatialMetadata { crs: Some("EPSG:4326".to_string()), + ..Default::default() } } diff --git a/vortex-spatial/src/extension/multipolygon.rs b/vortex-spatial/src/extension/multipolygon.rs index 80078a4b07e..e2bc4538589 100644 --- a/vortex-spatial/src/extension/multipolygon.rs +++ b/vortex-spatial/src/extension/multipolygon.rs @@ -18,7 +18,6 @@ use geoarrow::array::IntoArrow; use geoarrow::array::MultiPolygonArray; use geoarrow::datatypes::CoordType; use geoarrow::datatypes::MultiPolygonType; -use prost::Message; use vortex_array::ArrayRef; use vortex_array::ExecutionCtx; use vortex_array::IntoArray; @@ -49,9 +48,9 @@ use super::SpatialMetadata; use super::coordinate::Dimension; use super::coordinate::coordinate_dimension; use super::coordinate::coordinate_storage_dtype; -use super::geoarrow_metadata; use super::geoarrow_to_wkb; -use super::spatial_metadata_from_arrow; +use super::metadata::from_geoarrow; +use super::metadata::to_geoarrow; /// A multipolygon (`geoarrow.multipolygon`); a single `Polygon` is a one-element multipolygon. #[derive(Debug, Clone, Default, PartialEq, Eq, Hash)] @@ -68,11 +67,11 @@ impl ExtVTable for MultiPolygon { } fn serialize_metadata(&self, metadata: &Self::Metadata) -> VortexResult> { - Ok(metadata.encode_to_vec()) + Ok(metadata.serialize()) } fn deserialize_metadata(&self, metadata: &[u8]) -> VortexResult { - Ok(SpatialMetadata::decode(metadata)?) + SpatialMetadata::deserialize(metadata) } fn validate_dtype(ext_dtype: &ExtDType) -> VortexResult<()> { @@ -112,16 +111,23 @@ pub(crate) fn multipolygon_dimension(dtype: &DType) -> VortexResult { static ARROW_MULTIPOLYGON: CachedId = CachedId::new(MultiPolygonType::NAME); /// The `geoarrow.multipolygon` type for `dimension`, with separated (struct) coordinates. -fn multipolygon_type(spatial_metadata: &SpatialMetadata, dimension: Dimension) -> MultiPolygonType { - MultiPolygonType::new(dimension.into(), geoarrow_metadata(spatial_metadata)) +fn multipolygon_type( + metadata: &SpatialMetadata, + dimension: Dimension, +) -> VortexResult { + Ok(MultiPolygonType::new( + dimension.into(), + to_geoarrow(metadata)?, + )) } -/// Decode storage to `geo_types` for the spatial scalar functions (CRS is irrelevant to planar ops). +/// Decode multipolygon storage to `geo_types`. pub(crate) fn multipolygon_geometries( storage: &ArrayRef, + metadata: &SpatialMetadata, ctx: &mut ExecutionCtx, ) -> VortexResult>> { - multipolygon_array(storage, ctx)? + multipolygon_array(storage, metadata, ctx)? .iter() .map(|geometry| -> VortexResult> { Ok(geometry @@ -135,12 +141,10 @@ pub(crate) fn multipolygon_geometries( /// Build a geoarrow `MultiPolygonArray` from the `MultiPolygon` storage. fn multipolygon_array( storage: &ArrayRef, + metadata: &SpatialMetadata, ctx: &mut ExecutionCtx, ) -> VortexResult { - let multipolygon_type = multipolygon_type( - &SpatialMetadata::default(), - multipolygon_dimension(storage.dtype())?, - ); + let multipolygon_type = multipolygon_type(metadata, multipolygon_dimension(storage.dtype())?)?; let session = ctx.session().clone(); let arrow = session.arrow().execute_arrow(storage.clone(), None, ctx)?; MultiPolygonArray::try_from((arrow.as_ref(), multipolygon_type)) @@ -165,7 +169,11 @@ impl TryFrom for MultiPolygonData { impl MultiPolygonData { /// Serialize multipolygons to WKB (a view array) — the form DuckDB `GEOMETRY` takes. pub fn to_wkb(&self, ctx: &mut ExecutionCtx) -> VortexResult { - geoarrow_to_wkb(&multipolygon_array(self.0.storage_array(), ctx)?) + geoarrow_to_wkb(&multipolygon_array( + self.0.storage_array(), + self.0.ext_dtype().metadata::(), + ctx, + )?) } } @@ -185,11 +193,11 @@ impl ArrowExportVTable for MultiPolygon { session: &ArrowSession, ) -> VortexResult> { let ext_type = dtype.as_extension(); - let spatial_metadata = ext_type.metadata::(); + let metadata = ext_type.metadata::(); let dimension = multipolygon_dimension(ext_type.storage_dtype())?; let mut field = session.to_arrow_field(name, ext_type.storage_dtype())?; - field.try_with_extension_type(multipolygon_type(spatial_metadata, dimension))?; + field.try_with_extension_type(multipolygon_type(metadata, dimension)?)?; Ok(Some(field)) } @@ -258,7 +266,7 @@ impl ArrowImportVTable for MultiPolygon { ); ( multipolygon_meta.dimension().into(), - spatial_metadata_from_arrow(multipolygon_meta.metadata()), + from_geoarrow(multipolygon_meta.metadata()), ) } else { // Literal: peel the three `List` layers to the coordinate struct and read its @@ -335,6 +343,7 @@ mod tests { fn spatial_meta() -> SpatialMetadata { SpatialMetadata { crs: Some("EPSG:4326".to_string()), + ..Default::default() } } diff --git a/vortex-spatial/src/extension/point.rs b/vortex-spatial/src/extension/point.rs index 5dfbbb99949..d1658fdfe8c 100644 --- a/vortex-spatial/src/extension/point.rs +++ b/vortex-spatial/src/extension/point.rs @@ -3,8 +3,8 @@ //! The [`Point`] geometry extension type (`vortex.st.point`): a location stored columnarly as //! `Struct` of non-nullable `f64` — the four GeoArrow dimensions XY, XYZ, XYM, -//! XYZM — tagged with [`SpatialMetadata`] (CRS). `z` is an optional elevation and `m` an optional -//! measure: an arbitrary per-point value such as distance along a route or a timestamp. +//! XYZM — tagged with [`SpatialMetadata`]. `z` is an optional elevation and `m` an optional measure: +//! an arbitrary per-point value such as distance along a route or a timestamp. use arrow_array::ArrayRef as ArrowArrayRef; use arrow_schema::DataType; @@ -17,7 +17,6 @@ use geoarrow::array::IntoArrow; use geoarrow::array::PointArray; use geoarrow::datatypes::CoordType; use geoarrow::datatypes::PointType; -use prost::Message; use vortex_array::ArrayRef; use vortex_array::ExecutionCtx; use vortex_array::IntoArray; @@ -49,9 +48,9 @@ use super::coordinate::Dimension; use super::coordinate::coordinate_dimension; use super::coordinate::coordinate_from_struct; use super::coordinate::coordinate_storage_dtype; -use super::geoarrow_metadata; use super::geoarrow_to_wkb; -use super::spatial_metadata_from_arrow; +use super::metadata::from_geoarrow; +use super::metadata::to_geoarrow; /// A single location: `geoarrow.point`, stored as `Struct` of non-nullable `f64`. #[derive(Debug, Clone, Default, PartialEq, Eq, Hash)] @@ -67,11 +66,11 @@ impl ExtVTable for Point { } fn serialize_metadata(&self, metadata: &Self::Metadata) -> VortexResult> { - Ok(metadata.encode_to_vec()) + Ok(metadata.serialize()) } fn deserialize_metadata(&self, metadata: &[u8]) -> VortexResult { - Ok(SpatialMetadata::decode(metadata)?) + SpatialMetadata::deserialize(metadata) } fn validate_dtype(ext_dtype: &ExtDType) -> VortexResult<()> { @@ -94,8 +93,8 @@ static ARROW_POINT: CachedId = CachedId::new(PointType::NAME); /// The `geoarrow.point` extension type for `dimension`, with separated (struct) coordinates /// matching `Point` storage. -fn point_type(spatial_metadata: &SpatialMetadata, dimension: Dimension) -> PointType { - PointType::new(dimension.into(), geoarrow_metadata(spatial_metadata)) +fn point_type(metadata: &SpatialMetadata, dimension: Dimension) -> VortexResult { + Ok(PointType::new(dimension.into(), to_geoarrow(metadata)?)) } pub struct PointData(ExtensionArray); @@ -115,29 +114,35 @@ impl TryFrom for PointData { impl PointData { /// Serialize points to WKB (a view array) — the form DuckDB `GEOMETRY` takes. pub fn to_wkb(&self, ctx: &mut ExecutionCtx) -> VortexResult { - geoarrow_to_wkb(&point_array(self.0.storage_array(), ctx)?) + geoarrow_to_wkb(&point_array( + self.0.storage_array(), + self.0.ext_dtype().metadata::(), + ctx, + )?) } } /// Build a geoarrow `PointArray` from a `Point`'s `Struct` storage, shared by WKB export /// and `geo_types` decoding. -fn point_array(storage: &ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult { - let point_type = point_type( - &SpatialMetadata::default(), - coordinate_dimension(storage.dtype())?, - ); +fn point_array( + storage: &ArrayRef, + metadata: &SpatialMetadata, + ctx: &mut ExecutionCtx, +) -> VortexResult { + let point_type = point_type(metadata, coordinate_dimension(storage.dtype())?)?; let session = ctx.session().clone(); let arrow = session.arrow().execute_arrow(storage.clone(), None, ctx)?; PointArray::try_from((arrow.as_ref(), point_type)) .map_err(|e| vortex_err!("failed to construct PointArray: {e}")) } -/// Decode `Point` storage to `geo_types` points, for the spatial scalar functions. +/// Decode point storage to `geo_types`. pub(crate) fn point_geometries( storage: &ArrayRef, + metadata: &SpatialMetadata, ctx: &mut ExecutionCtx, ) -> VortexResult>> { - point_array(storage, ctx)? + point_array(storage, metadata, ctx)? .iter() .map(|geometry| -> VortexResult> { Ok(geometry @@ -164,11 +169,11 @@ impl ArrowExportVTable for Point { session: &ArrowSession, ) -> VortexResult> { let ext_type = dtype.as_extension(); - let spatial_metadata = ext_type.metadata::(); + let metadata = ext_type.metadata::(); let dimension = coordinate_dimension(ext_type.storage_dtype())?; let mut field = session.to_arrow_field(name, ext_type.storage_dtype())?; - field.try_with_extension_type(point_type(spatial_metadata, dimension))?; + field.try_with_extension_type(point_type(metadata, dimension)?)?; Ok(Some(field)) } @@ -240,7 +245,7 @@ impl ArrowImportVTable for Point { ); ( point_meta.dimension().into(), - spatial_metadata_from_arrow(point_meta.metadata()), + from_geoarrow(point_meta.metadata()), ) } else { // Infer the dimension from the field names, not the canonical storage check: a literal's @@ -312,6 +317,7 @@ mod tests { fn spatial_meta() -> SpatialMetadata { SpatialMetadata { crs: Some("EPSG:4326".to_string()), + ..Default::default() } } diff --git a/vortex-spatial/src/extension/polygon.rs b/vortex-spatial/src/extension/polygon.rs index e727a7ab3cb..ccc7dac435b 100644 --- a/vortex-spatial/src/extension/polygon.rs +++ b/vortex-spatial/src/extension/polygon.rs @@ -2,8 +2,8 @@ // SPDX-FileCopyrightText: Copyright the Vortex contributors //! The [`Polygon`] geometry extension type (`vortex.st.polygon`): rings of the -//! [`Point`](super::Point) coordinate struct, stored as `List>>` and tagged with -//! [`SpatialMetadata`] (CRS). The first ring is the exterior boundary; the rest are holes. +//! [`Point`](super::Point) coordinate struct, stored as `List>>` and +//! tagged with [`SpatialMetadata`]. The first ring is the exterior boundary; the rest are holes. use std::sync::Arc; @@ -18,7 +18,6 @@ use geoarrow::array::IntoArrow; use geoarrow::array::PolygonArray; use geoarrow::datatypes::CoordType; use geoarrow::datatypes::PolygonType; -use prost::Message; use vortex_array::ArrayRef; use vortex_array::ExecutionCtx; use vortex_array::IntoArray; @@ -49,9 +48,9 @@ use super::SpatialMetadata; use super::coordinate::Dimension; use super::coordinate::coordinate_dimension; use super::coordinate::coordinate_storage_dtype; -use super::geoarrow_metadata; use super::geoarrow_to_wkb; -use super::spatial_metadata_from_arrow; +use super::metadata::from_geoarrow; +use super::metadata::to_geoarrow; /// A polygon: `geoarrow.polygon`, stored as `List>>` (rings of vertices). #[derive(Debug, Clone, Default, PartialEq, Eq, Hash)] @@ -68,11 +67,11 @@ impl ExtVTable for Polygon { } fn serialize_metadata(&self, metadata: &Self::Metadata) -> VortexResult> { - Ok(metadata.encode_to_vec()) + Ok(metadata.serialize()) } fn deserialize_metadata(&self, metadata: &[u8]) -> VortexResult { - Ok(SpatialMetadata::decode(metadata)?) + SpatialMetadata::deserialize(metadata) } fn validate_dtype(ext_dtype: &ExtDType) -> VortexResult<()> { @@ -109,17 +108,17 @@ static ARROW_POLYGON: CachedId = CachedId::new(PolygonType::NAME); /// The `geoarrow.polygon` extension type for `dimension`, with separated (struct) coordinates /// matching `Polygon` storage. -fn polygon_type(spatial_metadata: &SpatialMetadata, dimension: Dimension) -> PolygonType { - PolygonType::new(dimension.into(), geoarrow_metadata(spatial_metadata)) +fn polygon_type(metadata: &SpatialMetadata, dimension: Dimension) -> VortexResult { + Ok(PolygonType::new(dimension.into(), to_geoarrow(metadata)?)) } -/// Decode `Polygon` storage (`List>`) to `geo_types` polygons, for the spatial scalar -/// functions. CRS does not affect planar geometry ops, so default metadata is used. +/// Decode polygon storage to `geo_types`. pub(crate) fn polygon_geometries( storage: &ArrayRef, + metadata: &SpatialMetadata, ctx: &mut ExecutionCtx, ) -> VortexResult>> { - polygon_array(storage, ctx)? + polygon_array(storage, metadata, ctx)? .iter() .map(|geometry| -> VortexResult> { Ok(geometry @@ -131,11 +130,12 @@ pub(crate) fn polygon_geometries( } /// Build a geoarrow `PolygonArray` from a `Polygon`'s `List>` storage. -fn polygon_array(storage: &ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult { - let polygon_type = polygon_type( - &SpatialMetadata::default(), - polygon_dimension(storage.dtype())?, - ); +fn polygon_array( + storage: &ArrayRef, + metadata: &SpatialMetadata, + ctx: &mut ExecutionCtx, +) -> VortexResult { + let polygon_type = polygon_type(metadata, polygon_dimension(storage.dtype())?)?; let session = ctx.session().clone(); let arrow = session.arrow().execute_arrow(storage.clone(), None, ctx)?; PolygonArray::try_from((arrow.as_ref(), polygon_type)) @@ -160,7 +160,11 @@ impl TryFrom for PolygonData { impl PolygonData { /// Serialize polygons to WKB (a view array) — the form DuckDB `GEOMETRY` takes. pub fn to_wkb(&self, ctx: &mut ExecutionCtx) -> VortexResult { - geoarrow_to_wkb(&polygon_array(self.0.storage_array(), ctx)?) + geoarrow_to_wkb(&polygon_array( + self.0.storage_array(), + self.0.ext_dtype().metadata::(), + ctx, + )?) } } @@ -180,11 +184,11 @@ impl ArrowExportVTable for Polygon { session: &ArrowSession, ) -> VortexResult> { let ext_type = dtype.as_extension(); - let spatial_metadata = ext_type.metadata::(); + let metadata = ext_type.metadata::(); let dimension = polygon_dimension(ext_type.storage_dtype())?; let mut field = session.to_arrow_field(name, ext_type.storage_dtype())?; - field.try_with_extension_type(polygon_type(spatial_metadata, dimension))?; + field.try_with_extension_type(polygon_type(metadata, dimension)?)?; Ok(Some(field)) } @@ -255,7 +259,7 @@ impl ArrowImportVTable for Polygon { ); ( polygon_meta.dimension().into(), - spatial_metadata_from_arrow(polygon_meta.metadata()), + from_geoarrow(polygon_meta.metadata()), ) } else { // Infer the dimension from the field names, not the canonical storage check: a literal's @@ -330,6 +334,7 @@ mod tests { fn spatial_meta() -> SpatialMetadata { SpatialMetadata { crs: Some("EPSG:4326".to_string()), + ..Default::default() } } diff --git a/vortex-spatial/src/extension/rect.rs b/vortex-spatial/src/extension/rect.rs index ff25b4ba6a3..2bb26c46523 100644 --- a/vortex-spatial/src/extension/rect.rs +++ b/vortex-spatial/src/extension/rect.rs @@ -4,7 +4,7 @@ //! The [`Rect`] bounding-box extension type (`vortex.st.box`): an axis-aligned envelope stored //! columnarly as `Struct` of non-nullable //! `f64` — the lower corner's ordinates followed by the upper corner's — tagged with -//! [`SpatialMetadata`] (CRS). Its GeoArrow wire type is `geoarrow.box`. +//! [`SpatialMetadata`]. Its GeoArrow wire type is `geoarrow.box`. //! //! Decoding to `geo_types` yields a 2D [`Geometry::Rect`]; any `z`/`m` //! bounds are dropped, as for the other geometry types. @@ -21,7 +21,6 @@ use geoarrow::array::GeoArrowArrayAccessor; use geoarrow::array::IntoArrow; use geoarrow::array::RectArray; use geoarrow::datatypes::BoxType; -use prost::Message; use vortex_array::ArrayRef; use vortex_array::ExecutionCtx; use vortex_array::IntoArray; @@ -54,8 +53,8 @@ use vortex_session::registry::Id; use super::SpatialMetadata; use super::coordinate::Dimension; -use super::geoarrow_metadata; -use super::spatial_metadata_from_arrow; +use super::metadata::from_geoarrow; +use super::metadata::to_geoarrow; /// An axis-aligned bounding box (`geoarrow.box`), stored as `Struct`. // Named `Rect`, not `Box`: matches `geo::Rect` / geoarrow-rs `RectArray`, and `Box` is a std name. @@ -73,11 +72,11 @@ impl ExtVTable for Rect { } fn serialize_metadata(&self, metadata: &Self::Metadata) -> VortexResult> { - Ok(metadata.encode_to_vec()) + Ok(metadata.serialize()) } fn deserialize_metadata(&self, metadata: &[u8]) -> VortexResult { - Ok(SpatialMetadata::decode(metadata)?) + SpatialMetadata::deserialize(metadata) } fn validate_dtype(ext_dtype: &ExtDType) -> VortexResult<()> { @@ -177,25 +176,30 @@ pub(crate) fn build_rect_array( static ARROW_BOX: CachedId = CachedId::new(BoxType::NAME); /// The `geoarrow.box` extension type for `dimension`. -fn box_type(spatial_metadata: &SpatialMetadata, dimension: Dimension) -> BoxType { - BoxType::new(dimension.into(), geoarrow_metadata(spatial_metadata)) +fn box_type(metadata: &SpatialMetadata, dimension: Dimension) -> VortexResult { + Ok(BoxType::new(dimension.into(), to_geoarrow(metadata)?)) } /// Build a geoarrow `RectArray` from a `Rect`'s box `Struct` storage. -fn rect_array(storage: &ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult { - let box_type = box_type(&SpatialMetadata::default(), box_dimension(storage.dtype())?); +fn rect_array( + storage: &ArrayRef, + metadata: &SpatialMetadata, + ctx: &mut ExecutionCtx, +) -> VortexResult { + let box_type = box_type(metadata, box_dimension(storage.dtype())?)?; let session = ctx.session().clone(); let arrow = session.arrow().execute_arrow(storage.clone(), None, ctx)?; RectArray::try_from((arrow.as_ref(), box_type)) .map_err(|e| vortex_err!("failed to construct RectArray: {e}")) } -/// Decode `Rect` storage to `geo_types` (2D [`Geometry::Rect`]), for the spatial scalar functions. +/// Decode box storage to 2D [`Geometry::Rect`] values. pub(crate) fn rect_geometries( storage: &ArrayRef, + metadata: &SpatialMetadata, ctx: &mut ExecutionCtx, ) -> VortexResult>> { - rect_array(storage, ctx)? + rect_array(storage, metadata, ctx)? .iter() .map(|geometry| -> VortexResult> { Ok(geometry @@ -222,11 +226,11 @@ impl ArrowExportVTable for Rect { session: &ArrowSession, ) -> VortexResult> { let ext_type = dtype.as_extension(); - let spatial_metadata = ext_type.metadata::(); + let metadata = ext_type.metadata::(); let dimension = box_dimension(ext_type.storage_dtype())?; let mut field = session.to_arrow_field(name, ext_type.storage_dtype())?; - field.try_with_extension_type(box_type(spatial_metadata, dimension))?; + field.try_with_extension_type(box_type(metadata, dimension)?)?; Ok(Some(field)) } @@ -287,7 +291,7 @@ impl ArrowImportVTable for Rect { let (dimension, metadata) = if let Ok(box_meta) = field.try_extension_type::() { ( box_meta.dimension().into(), - spatial_metadata_from_arrow(box_meta.metadata()), + from_geoarrow(box_meta.metadata()), ) } else { if field.extension_type_name() != Some(BoxType::NAME) { @@ -358,6 +362,7 @@ mod tests { fn spatial_meta() -> SpatialMetadata { SpatialMetadata { crs: Some("EPSG:4326".to_string()), + ..Default::default() } } @@ -405,12 +410,13 @@ mod tests { let mut ctx = session.create_execution_ctx(); let rects = crate::test_harness::rect_column(vec![(0.0, 0.0, 2.0, 3.0)])?; + let metadata = rects.dtype().as_extension().metadata::().clone(); let storage = rects .execute::(&mut ctx)? .storage_array() .clone(); - let geometries = rect_geometries(&storage, &mut ctx)?; + let geometries = rect_geometries(&storage, &metadata, &mut ctx)?; assert!(matches!(geometries.as_slice(), [Geometry::Rect(_)])); Ok(()) } diff --git a/vortex-spatial/src/extension/wkb.rs b/vortex-spatial/src/extension/wkb.rs index 977e2855d51..ab09401f3e2 100644 --- a/vortex-spatial/src/extension/wkb.rs +++ b/vortex-spatial/src/extension/wkb.rs @@ -10,10 +10,12 @@ use arrow_schema::DataType; use arrow_schema::Field; use arrow_schema::extension::ExtensionType; use geoarrow::array::GenericWkbArray; +use geoarrow::array::GeoArrowArray; use geoarrow::array::IntoArrow; use geoarrow::array::WkbViewArray; +use geoarrow::datatypes::GeoArrowType; use geoarrow::datatypes::WkbType; -use prost::Message; +use geoarrow_cast::cast::cast; use vortex_array::ArrayRef; use vortex_array::ExecutionCtx; use vortex_array::IntoArray; @@ -40,8 +42,8 @@ use vortex_session::registry::Id; use wkb::reader::GeometryType; use crate::extension::SpatialMetadata; -use crate::extension::geoarrow_metadata; -use crate::extension::spatial_metadata_from_arrow; +use crate::extension::metadata::from_geoarrow; +use crate::extension::metadata::to_geoarrow; /// A typed handle to an [`ExtensionArray`] that contains WKB-encoded data. /// @@ -141,11 +143,11 @@ impl ExtVTable for WellKnownBinary { } fn serialize_metadata(&self, metadata: &Self::Metadata) -> VortexResult> { - Ok(metadata.encode_to_vec()) + Ok(metadata.serialize()) } fn deserialize_metadata(&self, metadata: &[u8]) -> VortexResult { - Ok(SpatialMetadata::decode(metadata)?) + SpatialMetadata::deserialize(metadata) } fn validate_dtype(ext_dtype: &ExtDType) -> VortexResult<()> { @@ -184,10 +186,10 @@ impl ArrowExportVTable for WellKnownBinary { session: &ArrowSession, ) -> VortexResult> { let ext_type = dtype.as_extension(); - let spatial_metadata = ext_type.metadata::(); + let metadata = ext_type.metadata::(); let mut field = session.to_arrow_field(name, ext_type.storage_dtype())?; - field.try_with_extension_type(wkb_type(spatial_metadata))?; + field.try_with_extension_type(wkb_type(metadata)?)?; Ok(Some(field)) } @@ -266,8 +268,12 @@ impl ArrowImportVTable for WellKnownBinary { let storage_dtype = DType::Binary(field.is_nullable().into()); Ok(Some(DType::Extension( - ExtDType::try_with_vtable(WellKnownBinary, spatial_metadata(&wkb_meta), storage_dtype)? - .erased(), + ExtDType::try_with_vtable( + WellKnownBinary, + from_geoarrow(wkb_meta.metadata()), + storage_dtype, + )? + .erased(), ))) } @@ -297,10 +303,15 @@ impl ArrowImportVTable for WellKnownBinary { } } -fn wkb_type(spatial_metadata: &SpatialMetadata) -> WkbType { - WkbType::new(geoarrow_metadata(spatial_metadata)) +fn wkb_type(metadata: &SpatialMetadata) -> VortexResult { + Ok(WkbType::new(to_geoarrow(metadata)?)) } -fn spatial_metadata(wkb_type: &WkbType) -> SpatialMetadata { - spatial_metadata_from_arrow(wkb_type.metadata()) +/// Serialize a native geometry array to WKB. +pub(crate) fn geoarrow_to_wkb(geoarrow_array: &dyn GeoArrowArray) -> VortexResult { + let data_type = geoarrow_array.data_type(); + let wkb_type = GeoArrowType::WkbView(WkbType::new(Arc::clone(data_type.metadata()))); + let wkb = cast(geoarrow_array, &wkb_type) + .map_err(|e| vortex_err!("failed to cast geometry to WKB: {e}"))?; + ArrayRef::from_arrow(wkb.to_array_ref().as_ref(), false) } diff --git a/vortex-spatial/src/scalar_fn/make_line.rs b/vortex-spatial/src/scalar_fn/make_line.rs index c32d4c76691..3c21b20d066 100644 --- a/vortex-spatial/src/scalar_fn/make_line.rs +++ b/vortex-spatial/src/scalar_fn/make_line.rs @@ -293,6 +293,7 @@ mod tests { let dtype = ExtDType::::try_new( SpatialMetadata { crs: crs.map(str::to_owned), + ..Default::default() }, storage.dtype().clone(), )?; diff --git a/vortex-spatial/src/test_harness.rs b/vortex-spatial/src/test_harness.rs index 7b471bdf2c4..16b2df1e85a 100644 --- a/vortex-spatial/src/test_harness.rs +++ b/vortex-spatial/src/test_harness.rs @@ -49,6 +49,7 @@ pub fn spatial_session() -> VortexSession { fn wgs84() -> SpatialMetadata { SpatialMetadata { crs: Some("EPSG:4326".to_string()), + ..Default::default() } } diff --git a/vortex-spatial/src/tests/wkb.rs b/vortex-spatial/src/tests/wkb.rs index 40477a464c9..c74af6bf83b 100644 --- a/vortex-spatial/src/tests/wkb.rs +++ b/vortex-spatial/src/tests/wkb.rs @@ -20,6 +20,8 @@ use geo_types::Geometry; use geo_types::LineString; use geo_types::Polygon; use geoarrow::datatypes::Crs; +use geoarrow::datatypes::CrsType as GeoArrowCrsType; +use geoarrow::datatypes::Edges as GeoArrowEdges; use geoarrow::datatypes::Metadata; use geoarrow::datatypes::WkbType; use rstest::rstest; @@ -37,6 +39,8 @@ use vortex_error::vortex_err; use wkb::writer::WriteOptions; use super::SESSION; +use crate::extension::CrsType; +use crate::extension::Edges; use crate::extension::SpatialMetadata; use crate::extension::WellKnownBinary; @@ -71,6 +75,7 @@ fn wkb_extension_array() -> VortexResult<(Vec, vortex_array::ArrayRef)> { let dtype = ExtDType::::try_new( SpatialMetadata { crs: Some("EPSG:4326".to_string()), + ..Default::default() }, DType::Binary(Nullability::NonNullable), )?; @@ -107,6 +112,7 @@ fn scalar_unpacks_to_geometry() -> VortexResult<()> { let dtype = ExtDType::::try_new( SpatialMetadata { crs: Some("EPSG:4326".to_string()), + ..Default::default() }, DType::Binary(Nullability::NonNullable), )?; @@ -168,6 +174,92 @@ fn import_field_recovers_extension() -> VortexResult<()> { Ok(()) } +/// GeoArrow CRS representation and edge semantics survive Arrow → Vortex → Arrow conversion. +#[test] +fn roundtrips_geoarrow_metadata() -> VortexResult<()> { + let projjson = serde_json::json!({ + "type": "GeographicCRS", + "name": "WGS 84" + }); + let metadata = Arc::new(Metadata::new( + Crs::from_projjson(projjson.clone()), + Some(GeoArrowEdges::Spherical), + )); + let field = Field::new("geom", DataType::Binary, false) + .with_extension_type(WkbType::new(Arc::clone(&metadata))); + + let dtype = SESSION.arrow().from_arrow_field(&field)?; + let spatial_metadata = dtype.as_extension().metadata::(); + let projjson_string = projjson.to_string(); + assert_eq!( + spatial_metadata.crs.as_deref(), + Some(projjson_string.as_str()) + ); + assert_eq!(spatial_metadata.crs_type, Some(CrsType::Projjson)); + assert_eq!(spatial_metadata.edges, Some(Edges::Spherical)); + + let exported = SESSION.arrow().to_arrow_field("geom", &dtype)?; + let wkb_type = exported.try_extension_type::()?; + assert_eq!(wkb_type.metadata().as_ref(), metadata.as_ref()); + assert_eq!( + wkb_type.metadata().crs().crs_type(), + Some(GeoArrowCrsType::Projjson) + ); + assert_eq!(wkb_type.metadata().crs().crs_value(), Some(&projjson)); + assert_eq!(wkb_type.metadata().edges(), Some(GeoArrowEdges::Spherical)); + Ok(()) +} + +/// A GeoArrow CRS object without the optional `crs_type` remains a PROJJSON object on export. +#[test] +fn infers_projjson_from_crs_object() -> VortexResult<()> { + let projjson = serde_json::json!({ + "type": "GeographicCRS", + "name": "WGS 84" + }); + let metadata = Arc::new( + serde_json::from_value::(serde_json::json!({ + "crs": projjson + })) + .map_err(|error| vortex_err!("failed to build GeoArrow metadata: {error}"))?, + ); + assert_eq!(metadata.crs().crs_type(), None); + + let field = + Field::new("geom", DataType::Binary, false).with_extension_type(WkbType::new(metadata)); + let dtype = SESSION.arrow().from_arrow_field(&field)?; + let spatial_metadata = dtype.as_extension().metadata::(); + assert_eq!(spatial_metadata.crs_type, Some(CrsType::Projjson)); + + let exported = SESSION.arrow().to_arrow_field("geom", &dtype)?; + let wkb_type = exported.try_extension_type::()?; + assert_eq!(wkb_type.metadata().crs().crs_value(), Some(&projjson)); + assert_eq!( + wkb_type.metadata().crs().crs_type(), + Some(GeoArrowCrsType::Projjson) + ); + Ok(()) +} + +/// A CRS declared as PROJJSON must contain a JSON object before Arrow export. +#[test] +fn export_rejects_invalid_projjson_metadata() -> VortexResult<()> { + let dtype = DType::Extension( + ExtDType::::try_new( + SpatialMetadata { + crs: Some("not JSON".to_string()), + crs_type: Some(CrsType::Projjson), + edges: None, + }, + DType::Binary(Nullability::NonNullable), + )? + .erased(), + ); + + assert!(SESSION.arrow().to_arrow_field("geom", &dtype).is_err()); + Ok(()) +} + /// A `geoarrow.wkb` field without a CRS imports as an unreferenced geometry. #[test] fn import_field_without_crs() -> VortexResult<()> {