diff --git a/Cargo.lock b/Cargo.lock index 8d3e9b7..01a4d79 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -298,6 +298,18 @@ dependencies = [ "regex-syntax", ] +[[package]] +name = "as-slice" +version = "0.1.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "45403b49e3954a4b8428a0ac21a4b7afadccf92bfd96273f1a58cd4812496ae0" +dependencies = [ + "generic-array 0.12.4", + "generic-array 0.13.3", + "generic-array 0.14.9", + "stable_deref_trait", +] + [[package]] name = "async-stream" version = "0.3.6" @@ -340,6 +352,15 @@ dependencies = [ "num-traits", ] +[[package]] +name = "atomic-polyfill" +version = "1.0.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8cf2bce30dfe09ef0bfaef228b9d414faaf7e563035494d7fe092dba54b300f4" +dependencies = [ + "critical-section", +] + [[package]] name = "atomic-waker" version = "1.1.2" @@ -544,6 +565,12 @@ dependencies = [ "cfg-if", ] +[[package]] +name = "critical-section" +version = "1.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "790eea4361631c5e7d22598ecd5723ff611904e3344ce8720784c93e3d83d40b" + [[package]] name = "crossbeam-deque" version = "0.8.6" @@ -1541,6 +1568,34 @@ dependencies = [ "slab", ] +[[package]] +name = "generic-array" +version = "0.12.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ffdf9f34f1447443d37393cc6c2b8313aebddcd96906caf34e54c68d8e57d7bd" +dependencies = [ + "typenum", +] + +[[package]] +name = "generic-array" +version = "0.13.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f797e67af32588215eaaab8327027ee8e71b9dd0b2b26996aedf20c030fce309" +dependencies = [ + "typenum", +] + +[[package]] +name = "generic-array" +version = "0.14.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4bb6743198531e02858aeaea5398fcc883e71851fcbcb5a2f773e2fb6cb1edf2" +dependencies = [ + "typenum", + "version_check", +] + [[package]] name = "geo" version = "0.31.0" @@ -1555,7 +1610,7 @@ dependencies = [ "log", "num-traits", "robust", - "rstar", + "rstar 0.12.2", "spade", ] @@ -1570,14 +1625,18 @@ dependencies = [ [[package]] name = "geo-types" -version = "0.7.17" +version = "0.7.19" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "75a4dcd69d35b2c87a7c83bce9af69fd65c9d68d3833a0ded568983928f3fc99" +checksum = "94776032c45f950d30a13af6113c2ad5625316c9abfbccee4dd5a6695f8fe0f5" dependencies = [ "approx", "num-traits", "rayon", - "rstar", + "rstar 0.10.0", + "rstar 0.11.0", + "rstar 0.12.2", + "rstar 0.8.4", + "rstar 0.9.3", "serde", ] @@ -1682,6 +1741,7 @@ dependencies = [ "approx", "arrow-arith", "arrow-array", + "arrow-buffer", "arrow-schema", "datafusion", "geo", @@ -1931,6 +1991,24 @@ dependencies = [ "zerocopy", ] +[[package]] +name = "hash32" +version = "0.1.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d4041af86e63ac4298ce40e5cca669066e75b6f1aa3390fe2561ffa5e1d9f4cc" +dependencies = [ + "byteorder", +] + +[[package]] +name = "hash32" +version = "0.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b0c35f58762feb77d74ebe43bdbc3210f09be9fe6742234d573bacc26ed92b67" +dependencies = [ + "byteorder", +] + [[package]] name = "hash32" version = "0.3.1" @@ -1974,13 +2052,38 @@ dependencies = [ "foldhash 0.2.0", ] +[[package]] +name = "heapless" +version = "0.6.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "634bd4d29cbf24424d0a4bfcbf80c6960129dc24424752a7d1d1390607023422" +dependencies = [ + "as-slice", + "generic-array 0.14.9", + "hash32 0.1.1", + "stable_deref_trait", +] + +[[package]] +name = "heapless" +version = "0.7.17" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cdc6457c0eb62c71aac4bc17216026d8410337c4126773b9c5daba343f17964f" +dependencies = [ + "atomic-polyfill", + "hash32 0.2.1", + "rustc_version", + "spin", + "stable_deref_trait", +] + [[package]] name = "heapless" version = "0.8.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "0bfb9eb618601c89945a70e254898da93b13be0388091d42117462b265bb3fad" dependencies = [ - "hash32", + "hash32 0.3.1", "stable_deref_trait", ] @@ -2136,9 +2239,9 @@ checksum = "9190f86706ca38ac8add223b2aed8b1330002b5cdbbce28fb58b10914d38fc27" [[package]] name = "i_overlay" -version = "4.0.6" +version = "4.0.7" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "0fcccbd4e4274e0f80697f5fbc6540fdac533cce02f2081b328e68629cce24f9" +checksum = "413183068e6e0289e18d7d0a1f661b81546e6918d5453a44570b9ab30cbed1b3" dependencies = [ "i_float", "i_key_sort", @@ -2736,6 +2839,12 @@ version = "1.0.15" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "57c0d7b74b563b49d38dae00a0c37d4d6de9b432382b2892f0574ddcae73fd0a" +[[package]] +name = "pdqselect" +version = "0.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4ec91767ecc0a0bbe558ce8c9da33c068066c57ecc8bb8477ef8c1ad3ef77c27" + [[package]] name = "percent-encoding" version = "2.3.2" @@ -3123,14 +3232,64 @@ version = "1.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "4e27ee8bb91ca0adcf0ecb116293afa12d393f9c2b9b9cd54d33e8078fe19839" +[[package]] +name = "rstar" +version = "0.8.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3a45c0e8804d37e4d97e55c6f258bc9ad9c5ee7b07437009dd152d764949a27c" +dependencies = [ + "heapless 0.6.1", + "num-traits", + "pdqselect", + "serde", + "smallvec", +] + +[[package]] +name = "rstar" +version = "0.9.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b40f1bfe5acdab44bc63e6699c28b74f75ec43afb59f3eda01e145aff86a25fa" +dependencies = [ + "heapless 0.7.17", + "num-traits", + "serde", + "smallvec", +] + +[[package]] +name = "rstar" +version = "0.10.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1f39465655a1e3d8ae79c6d9e007f4953bfc5d55297602df9dc38f9ae9f1359a" +dependencies = [ + "heapless 0.7.17", + "num-traits", + "serde", + "smallvec", +] + +[[package]] +name = "rstar" +version = "0.11.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "73111312eb7a2287d229f06c00ff35b51ddee180f017ab6dec1f69d62ac098d6" +dependencies = [ + "heapless 0.7.17", + "num-traits", + "serde", + "smallvec", +] + [[package]] name = "rstar" version = "0.12.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "421400d13ccfd26dfa5858199c30a5d76f9c54e0dba7575273025b43c5175dbb" dependencies = [ - "heapless", + "heapless 0.8.0", "num-traits", + "serde", "smallvec", ] @@ -3454,6 +3613,15 @@ dependencies = [ "smallvec", ] +[[package]] +name = "spin" +version = "0.9.8" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6980e8d7511241f8acf4aebddbb1ff938df5eebe98691418c4468d0b72a96a67" +dependencies = [ + "lock_api", +] + [[package]] name = "sqlparser" version = "0.61.0" @@ -3851,6 +4019,12 @@ version = "2.1.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "9ea3136b675547379c4bd395ca6b938e5ad3c3d20fad76e7fe85f9e0d011419c" +[[package]] +name = "typenum" +version = "1.20.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b6f5e870be6c3b371b77fe0ee0bafb859fa4964b4404c27de1d380043c4dda20" + [[package]] name = "unicode-ident" version = "1.0.19" diff --git a/Cargo.toml b/Cargo.toml index 5b49396..d14f377 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -21,6 +21,7 @@ categories = ["science::geo"] approx = "0.5.1" arrow-arith = "58.1" arrow-array = "58.1" +arrow-buffer = "58.1" arrow-schema = "58.1" async-trait = "0.1" datafusion = { version = "53", default-features = false } diff --git a/README.md b/README.md index 282432b..782f17d 100644 --- a/README.md +++ b/README.md @@ -48,7 +48,7 @@ Functions are explicitly modeled after the [PostGIS API](https://postgis.net/doc | ST_BoundingDiagonal | | Returns the diagonal of a geometry's bounding box. | | ST_CoordDim | ✅ | Return the coordinate dimension of a geometry. | | ST_Dimension | | Returns the topological dimension of a geometry. | -| ST_Dump | | Returns a set of geometry_dump rows for the components of a geometry. | +| ST_Dump | ✅ | Returns a set of geometry_dump rows for the components of a geometry. | | ST_DumpPoints | | Returns a set of geometry_dump rows for the coordinates in a geometry. | | ST_DumpSegments | | Returns a set of geometry_dump rows for the segments in a geometry. | | ST_DumpRings | | Returns a set of geometry_dump rows for the exterior and interior rings of a Polygon. | diff --git a/rust/geodatafusion/Cargo.toml b/rust/geodatafusion/Cargo.toml index 0432792..13b003a 100644 --- a/rust/geodatafusion/Cargo.toml +++ b/rust/geodatafusion/Cargo.toml @@ -21,6 +21,7 @@ geos-3_11 = ["geos/v3_11_0"] [dependencies] arrow-arith = { workspace = true } arrow-array = { workspace = true } +arrow-buffer = { workspace = true } arrow-schema = { workspace = true } datafusion = { workspace = true } geo = { workspace = true } diff --git a/rust/geodatafusion/src/udf/native/accessors/dump.rs b/rust/geodatafusion/src/udf/native/accessors/dump.rs new file mode 100644 index 0000000..9fe6dfe --- /dev/null +++ b/rust/geodatafusion/src/udf/native/accessors/dump.rs @@ -0,0 +1,468 @@ +use std::any::Any; +use std::sync::{Arc, OnceLock}; + +use arrow_array::builder::{ArrayBuilder, Int32Builder, ListBuilder}; +use arrow_array::{ArrayRef, ListArray, StructArray}; +use arrow_buffer::{NullBufferBuilder, OffsetBuffer}; +use arrow_schema::{DataType, Field, FieldRef, Fields}; +use datafusion::error::{DataFusionError, Result}; +use datafusion::logical_expr::scalar_doc_sections::DOC_SECTION_OTHER; +use datafusion::logical_expr::{ + ColumnarValue, Documentation, ReturnFieldArgs, ScalarFunctionArgs, ScalarUDFImpl, Signature, +}; +use geo_traits::{ + GeometryCollectionTrait, GeometryTrait, MultiLineStringTrait, MultiPointTrait, + MultiPolygonTrait, +}; +use geoarrow_array::array::from_arrow_array; +use geoarrow_array::builder::GeometryBuilder; +use geoarrow_array::{GeoArrowArray, GeoArrowArrayAccessor, downcast_geoarrow_array}; +use geoarrow_schema::{CoordType, GeometryType, Metadata}; + +use crate::data_types::any_single_geometry_type_input; +use crate::error::GeoDataFusionResult; +use crate::udf::native::accessors::is_empty::is_geometry_topologically_empty; + +/// Decomposes a geometry into its atomic components (POINT, LINESTRING, POLYGON). +#[derive(Debug, Eq, PartialEq, Hash)] +pub struct Dump { + coord_type: CoordType, +} + +impl Dump { + pub fn new(coord_type: CoordType) -> Self { + Self { coord_type } + } +} + +impl Default for Dump { + fn default() -> Self { + Self::new(Default::default()) + } +} + +static DOCUMENTATION: OnceLock = OnceLock::new(); + +impl ScalarUDFImpl for Dump { + fn as_any(&self) -> &dyn Any { + self + } + + fn name(&self) -> &str { + "st_dump" + } + + fn signature(&self) -> &Signature { + any_single_geometry_type_input() + } + + fn return_type(&self, _arg_types: &[DataType]) -> Result { + Err(DataFusionError::Internal( + "Return field is computed from metadata in return_field_from_args.".to_string(), + )) + } + + fn return_field_from_args(&self, args: ReturnFieldArgs) -> Result { + let metadata = + Arc::new(Metadata::try_from(args.arg_fields[0].as_ref()).unwrap_or_default()); + let output_type = GeometryType::new(metadata).with_coord_type(self.coord_type); + Ok(output_field(&output_type)) + } + + fn invoke_with_args(&self, args: ScalarFunctionArgs) -> Result { + Ok(dump_impl(args, self.coord_type)?) + } + + fn documentation(&self) -> Option<&Documentation> { + Some(DOCUMENTATION.get_or_init(|| { + Documentation::builder( + DOC_SECTION_OTHER, + "Decomposes a geometry into its atomic components (Point, LineString, Polygon). \ + Multi-geometries and GeometryCollections are split (recursively) into their parts. \ + Atomic geometries are returned unchanged. This includes Polygons, which are returned whole. \ + Returns, per input row, a list of `(path, geom)` structs where \ + `path` is the 1-based navigation path to the component (empty for an atomic input). \ + Use `unnest` if you want to expand into a row per component instead of row per input geom. \ + Empty inputs (recursively, per the same definition as ST_Empty) produce zero components.", + "ST_Dump(geom)", + ) + .with_argument("geom", "geometry") + .build() + })) + } +} + +fn path_values_field() -> FieldRef { + Arc::new(Field::new("item", DataType::Int32, false)) +} + +fn path_field() -> Field { + Field::new("path", DataType::List(path_values_field()), false) +} + +fn item_field(geom_type: &GeometryType) -> FieldRef { + Arc::new(Field::new( + "item", + DataType::Struct(item_struct_fields(geom_type)), + false, + )) +} + +fn item_struct_fields(geom_type: &GeometryType) -> Fields { + Fields::from(vec![path_field(), geom_type.to_field("geom", true)]) +} + +fn output_field(geom_type: &GeometryType) -> FieldRef { + Arc::new(Field::new("", DataType::List(item_field(geom_type)), true)) +} + +fn dump_impl( + args: ScalarFunctionArgs, + coord_type: CoordType, +) -> GeoDataFusionResult { + let array = ColumnarValue::values_to_arrays(&args.args)? + .into_iter() + .next() + .expect("should have one input argument or else planning will fail"); + let geo_array = from_arrow_array(&array, &args.arg_fields[0])?; + let geom_type = + GeometryType::new(geo_array.data_type().metadata().clone()).with_coord_type(coord_type); + + let geo_array_ref = geo_array.as_ref(); + let list_array = downcast_geoarrow_array!(geo_array_ref, dump_array, &geom_type)?; + + Ok(ColumnarValue::Array(Arc::new(list_array))) +} + +/// The main decomposition loop. +fn dump_array<'a>( + array: &'a impl GeoArrowArrayAccessor<'a>, + geom_type: &GeometryType, +) -> GeoDataFusionResult { + let mut geom_builder = GeometryBuilder::new(geom_type.clone()); + let mut path_builder = ListBuilder::new(Int32Builder::new()).with_field(path_values_field()); + let mut path_stack = Vec::new(); + let mut row_lengths = Vec::with_capacity(array.len()); + let mut row_validity = NullBufferBuilder::new(array.len()); + + for maybe_geom in array.iter() { + let start = path_builder.len(); + match maybe_geom.transpose()? { + Some(geom) => { + row_validity.append_non_null(); + // Topologically empty geometries produce an empty, non-null component list. + if !is_geometry_topologically_empty(&geom) { + dump_geometry(&geom, &mut geom_builder, &mut path_builder, &mut path_stack)?; + } + } + None => row_validity.append_null(), + } + row_lengths.push(path_builder.len() - start); + } + + let path_array = path_builder.finish(); + let geom_array = geom_builder.finish().into_array_ref(); + let struct_array = StructArray::new( + item_struct_fields(geom_type), + vec![Arc::new(path_array) as ArrayRef, geom_array], + None, + ); + Ok(ListArray::new( + item_field(geom_type), + OffsetBuffer::::from_lengths(row_lengths), + Arc::new(struct_array), + row_validity.finish(), + )) +} + +/// Visitor for every non-topologically-empty geometry +/// which emits atomic geoms directly, and recursively dumps children in containers. +fn dump_geometry( + geom: &impl GeometryTrait, + geom_builder: &mut GeometryBuilder, + path_builder: &mut ListBuilder, + path_stack: &mut Vec, +) -> GeoDataFusionResult<()> { + use geo_traits::GeometryType::*; + + match geom.as_type() { + Point(_) | LineString(_) | Polygon(_) | Rect(_) | Triangle(_) | Line(_) => { + geom_builder.push_geometry(Some(geom))?; + path_builder.append_value(path_stack.iter().map(|&i| Some(i))); + } + MultiPoint(g) => dump_multi_children(g.points(), geom_builder, path_builder, path_stack)?, + MultiLineString(g) => { + dump_multi_children(g.line_strings(), geom_builder, path_builder, path_stack)? + } + MultiPolygon(g) => { + dump_multi_children(g.polygons(), geom_builder, path_builder, path_stack)? + } + GeometryCollection(g) => { + dump_multi_children(g.geometries(), geom_builder, path_builder, path_stack)? + } + } + Ok(()) +} + +/// Process children of a multi-geometry container. +fn dump_multi_children( + children: impl Iterator>, + geom_builder: &mut GeometryBuilder, + path_builder: &mut ListBuilder, + path_stack: &mut Vec, +) -> GeoDataFusionResult<()> { + for (i, child) in children.enumerate() { + path_stack.push(i as i32 + 1); + dump_geometry(&child, geom_builder, path_builder, path_stack)?; + path_stack.pop(); + } + Ok(()) +} + +#[cfg(test)] +mod test { + use arrow_array::Array; + use arrow_array::cast::AsArray; + use arrow_array::types::Int32Type; + use arrow_schema::DataType; + use datafusion::prelude::SessionContext; + use geoarrow_array::cast::to_wkt; + + use super::*; + use crate::udf::native::io::{AsText, GeomFromText}; + + fn ctx() -> SessionContext { + let ctx = SessionContext::new(); + ctx.register_udf(Dump::default().into()); + ctx.register_udf(GeomFromText::default().into()); + ctx.register_udf(AsText.into()); + ctx + } + + /// Run `ST_Dump` on a single WKT input and return the result decoded as WKT. + async fn dump_rows(ctx: &SessionContext, wkt: &str) -> Vec<(Vec, String)> { + let sql = format!("SELECT ST_Dump(ST_GeomFromText('{wkt}'))"); + let df = ctx.sql(&sql).await.unwrap(); + let batch = df.collect().await.unwrap().into_iter().next().unwrap(); + let list = batch.column(0).as_list::(); + assert!(!list.is_null(0), "unexpected NULL dump for {wkt}"); + + let components = list.value(0); + let structs = components.as_struct(); + let path_col = structs.column(0).as_list::(); + + let DataType::Struct(fields) = structs.data_type() else { + panic!("ST_Dump must return a list of structs"); + }; + let geo = from_arrow_array(structs.column(1).as_ref(), fields[1].as_ref()).unwrap(); + let wkt_ref = to_wkt::(geo.as_ref()).unwrap().into_array_ref(); + let wkts = wkt_ref.as_string::(); + + (0..structs.len()) + .map(|i| { + let path = path_col + .value(i) + .as_primitive::() + .values() + .to_vec(); + (path, wkts.value(i).to_string()) + }) + .collect() + } + + #[tokio::test] + async fn test_atomic_round_trips() { + let ctx = ctx(); + assert_eq!( + dump_rows(&ctx, "POINT(1 2)").await, + vec![(vec![], "POINT(1 2)".to_string())], + "Point is atomic and should round-trip to itself.", + ); + + assert_eq!( + dump_rows(&ctx, "LINESTRING(0 0,1 1)").await, + vec![(vec![], "LINESTRING(0 0,1 1)".to_string())], + "LineString is atomic and should round-trip to itself.", + ); + + assert_eq!( + dump_rows(&ctx, "POLYGON((0 0,0 3,3 3,3 0,0 0),(1 1,1 2,2 2,2 1,1 1))").await, + vec![( + vec![], + "POLYGON((0 0,0 3,3 3,3 0,0 0),(1 1,1 2,2 2,2 1,1 1))".to_string() + )], + "Polygon is atomic and should round-trip to itself.", + ); + } + + #[tokio::test] + async fn test_multi_geometries_decompose() { + let ctx = ctx(); + assert_eq!( + dump_rows(&ctx, "MULTIPOINT(0 0,1 1)").await, + vec![ + (vec![1], "POINT(0 0)".to_string()), + (vec![2], "POINT(1 1)".to_string()), + ], + ); + + assert_eq!( + dump_rows(&ctx, "MULTILINESTRING((0 0,1 1),(2 2,3 3))").await, + vec![ + (vec![1], "LINESTRING(0 0,1 1)".to_string()), + (vec![2], "LINESTRING(2 2,3 3)".to_string()), + ], + ); + + assert_eq!( + dump_rows( + &ctx, + "MULTIPOLYGON(((0 0,0 1,1 1,1 0,0 0)),((10 10,10 20,20 20,20 10,10 10)))" + ) + .await, + vec![ + (vec![1], "POLYGON((0 0,0 1,1 1,1 0,0 0))".to_string()), + ( + vec![2], + "POLYGON((10 10,10 20,20 20,20 10,10 10))".to_string() + ), + ], + ); + } + + #[tokio::test] + async fn test_collection_decomposition() { + let ctx = ctx(); + assert_eq!( + dump_rows(&ctx, "GEOMETRYCOLLECTION(POINT(0 0),LINESTRING(1 1,2 2))").await, + vec![ + (vec![1], "POINT(0 0)".to_string()), + (vec![2], "LINESTRING(1 1,2 2)".to_string()), + ], + "GeometryCollection should be decomposed into components" + ); + + assert_eq!( + dump_rows( + &ctx, + "GEOMETRYCOLLECTION(MULTIPOLYGON(((0 0,0 1,1 1,1 0,0 0)),((10 10,10 20,20 20,20 10,10 10))),POINT(5 5))" + ) + .await, + vec![ + (vec![1, 1], "POLYGON((0 0,0 1,1 1,1 0,0 0))".to_string()), + (vec![1, 2], "POLYGON((10 10,10 20,20 20,20 10,10 10))".to_string()), + (vec![2], "POINT(5 5)".to_string()), + ], + "Nesting a MultiPolygon in a GeometryCollection should unroll into atomic Polygons" + ); + } + + #[tokio::test] + async fn test_dimension_preservation() { + let ctx = ctx(); + assert_eq!( + dump_rows(&ctx, "MULTIPOINT Z (0 0 1,1 1 2)").await, + vec![ + (vec![1], "POINT Z(0 0 1)".to_string()), + (vec![2], "POINT Z(1 1 2)".to_string()), + ], + ); + assert_eq!( + dump_rows(&ctx, "MULTIPOINT ZM (0 0 1 5,1 1 2 6)").await, + vec![ + (vec![1], "POINT ZM(0 0 1 5)".to_string()), + (vec![2], "POINT ZM(1 1 2 6)".to_string()), + ], + ); + } + + #[tokio::test] + async fn test_dump_empty_inputs_yield_zero_rows() { + let ctx = ctx(); + assert_eq!(dump_rows(&ctx, "POINT EMPTY").await, vec![]); + assert_eq!(dump_rows(&ctx, "LINESTRING EMPTY").await, vec![]); + assert_eq!(dump_rows(&ctx, "POLYGON EMPTY").await, vec![]); + assert_eq!(dump_rows(&ctx, "MULTIPOINT EMPTY").await, vec![]); + assert_eq!(dump_rows(&ctx, "MULTILINESTRING EMPTY").await, vec![]); + assert_eq!(dump_rows(&ctx, "MULTIPOLYGON EMPTY").await, vec![]); + assert_eq!(dump_rows(&ctx, "GEOMETRYCOLLECTION EMPTY").await, vec![]); + } + + #[tokio::test] + async fn test_dump_recursive_empty_collection_yields_zero_rows() { + let ctx = ctx(); + let wkt = "GEOMETRYCOLLECTION(POINT EMPTY, LINESTRING EMPTY, POLYGON EMPTY, MULTIPOINT EMPTY, MULTILINESTRING EMPTY, MULTIPOLYGON EMPTY)"; + assert_eq!(dump_rows(&ctx, wkt).await, vec![]); + } + + #[tokio::test] + async fn test_dump_mixed_empty_and_nonempty_children_emits_empty_leaves() { + let ctx = ctx(); + assert_eq!( + dump_rows(&ctx, "GEOMETRYCOLLECTION(POINT(0 0), POLYGON EMPTY)").await, + vec![ + (vec![1], "POINT(0 0)".to_string()), + (vec![2], "POLYGON EMPTY".to_string()), + ], + "A single non-empty geometry in a collection should make the collection non-empty, including empty children" + ); + } + + #[tokio::test] + async fn test_dump_null_yields_null() { + let ctx = ctx(); + let df = ctx + .sql("SELECT ST_Dump(ST_GeomFromText(NULL))") + .await + .unwrap(); + let batch = df.collect().await.unwrap().into_iter().next().unwrap(); + let list = batch.column(0).as_list::(); + assert!(list.is_null(0), "ST_Dump(NULL) should be NULL"); + } + + #[tokio::test] + async fn test_unnest_expands_to_multiple_rows() { + let ctx = ctx(); + let df = ctx + .sql("SELECT unnest(ST_Dump(ST_GeomFromText('MULTIPOINT(0 0,1 1)'))) AS comp") + .await + .unwrap(); + let batches = df.collect().await.unwrap(); + let rows: usize = batches.iter().map(|b| b.num_rows()).sum(); + assert_eq!( + rows, 2, + "unnest should expand a 2-component dump into 2 rows" + ); + assert!(matches!( + batches[0].column(0).data_type(), + DataType::Struct(_) + )); + } + + #[tokio::test] + async fn test_st_dump_field_access_sql() { + let ctx = ctx(); + let df = ctx + .sql( + // This SQL illustrates how to get a PostGIS-style shape with multiple rows per non-atomic geometry. + "SELECT ST_AsText(comp.geom) AS wkt, comp.path AS path \ + FROM (SELECT unnest(ST_Dump(ST_GeomFromText('MULTIPOINT(0 0,1 1)'))) AS comp)", + ) + .await + .unwrap(); + let batch = df.collect().await.unwrap().into_iter().next().unwrap(); + + // `comp['geom']` keeps the geometry's extension type, so ST_AsText accepts it directly. + let wkt = batch.column_by_name("wkt").unwrap().as_string::(); + assert_eq!( + (0..wkt.len()).map(|i| wkt.value(i)).collect::>(), + vec!["POINT(0 0)", "POINT(1 1)"], + ); + + // `comp['path']` extracts the navigation-path list (a single 1-based index for each part). + let path = batch.column_by_name("path").unwrap().as_list::(); + assert_eq!(path.len(), 2); + assert_eq!(path.value(0).as_primitive::().values(), &[1]); + assert_eq!(path.value(1).as_primitive::().values(), &[2]); + } +} diff --git a/rust/geodatafusion/src/udf/native/accessors/mod.rs b/rust/geodatafusion/src/udf/native/accessors/mod.rs index 057f12f..aa8c960 100644 --- a/rust/geodatafusion/src/udf/native/accessors/mod.rs +++ b/rust/geodatafusion/src/udf/native/accessors/mod.rs @@ -1,4 +1,5 @@ mod coord_dim; +mod dump; mod geometry_type; mod is_closed; mod is_empty; @@ -8,6 +9,7 @@ mod num_interior_rings; mod point; pub use coord_dim::{CoordDim, NDims}; +pub use dump::Dump; pub use geometry_type::{GeometryType, ST_GeometryType}; pub use is_closed::IsClosed; pub use is_empty::IsEmpty; @@ -23,6 +25,7 @@ pub fn register(session_context: &datafusion::prelude::SessionContext) { session_context.register_udf(ST_GeometryType.into()); session_context.register_udf(IsClosed.into()); session_context.register_udf(IsEmpty.into()); + session_context.register_udf(Dump::default().into()); session_context.register_udf(EndPoint::default().into()); session_context.register_udf(StartPoint::default().into()); session_context.register_udf(NPoints.into());