From b64da03219059fad14499b7564e6f5c94840fae7 Mon Sep 17 00:00:00 2001 From: Nemo Yu Date: Wed, 5 Aug 2026 17:07:06 -0400 Subject: [PATCH 1/5] feat(vortex-geo): add collect scalar function Signed-off-by: Nemo Yu --- vortex-spatial/src/lib.rs | 2 + vortex-spatial/src/scalar_fn/collect.rs | 505 ++++++++++++++++++++++++ vortex-spatial/src/scalar_fn/mod.rs | 1 + 3 files changed, 508 insertions(+) create mode 100644 vortex-spatial/src/scalar_fn/collect.rs diff --git a/vortex-spatial/src/lib.rs b/vortex-spatial/src/lib.rs index 6bf96831c48..9904f5649f4 100644 --- a/vortex-spatial/src/lib.rs +++ b/vortex-spatial/src/lib.rs @@ -22,6 +22,7 @@ use crate::extension::WellKnownBinary; use crate::prune::SpatialDistancePrune; use crate::prune::SpatialIntersectsPrune; use crate::scalar_fn::area::SpatialArea; +use crate::scalar_fn::collect::SpatialCollect; use crate::scalar_fn::contains::SpatialContains; use crate::scalar_fn::distance::SpatialDistance; use crate::scalar_fn::envelope::SpatialEnvelope; @@ -67,6 +68,7 @@ pub fn initialize(session: &VortexSession) { // Register the geometry scalar functions. session.scalar_fns().register(SpatialArea); + session.scalar_fns().register(SpatialCollect); session.scalar_fns().register(SpatialEnvelope); session.scalar_fns().register(SpatialContains); session.scalar_fns().register(SpatialDistance); diff --git a/vortex-spatial/src/scalar_fn/collect.rs b/vortex-spatial/src/scalar_fn/collect.rs new file mode 100644 index 00000000000..b138a661665 --- /dev/null +++ b/vortex-spatial/src/scalar_fn/collect.rs @@ -0,0 +1,505 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! `ST_Collect`: collect homogeneous native geometries into their native multi-geometry type. + +use std::sync::Arc; + +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::ScalarFnArray; +use vortex_array::arrays::extension::ExtensionArrayExt; +use vortex_array::arrays::listview::ListViewArraySlotsExt; +use vortex_array::arrays::listview::ListViewRebuildMode; +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::ExtDTypeRef; +use vortex_array::expr::Expression; +use vortex_array::expr::union_child_validities; +use vortex_array::scalar_fn::Arity; +use vortex_array::scalar_fn::ChildName; +use vortex_array::scalar_fn::EmptyOptions; +use vortex_array::scalar_fn::ExecutionArgs; +use vortex_array::scalar_fn::ScalarFnId; +use vortex_array::scalar_fn::ScalarFnVTable; +use vortex_array::scalar_fn::TypedScalarFnInstance; +use vortex_array::validity::Validity; +use vortex_buffer::Buffer; +use vortex_buffer::BufferMut; +use vortex_error::VortexResult; +use vortex_error::vortex_bail; +use vortex_error::vortex_ensure; +use vortex_error::vortex_err; +use vortex_mask::AllOr; +use vortex_mask::Mask; +use vortex_session::VortexSession; +use vortex_session::registry::CachedId; + +use crate::extension::LineString; +use crate::extension::MultiLineString; +use crate::extension::MultiPoint; +use crate::extension::MultiPolygon; +use crate::extension::Point; +use crate::extension::Polygon; +use crate::scalar_fn::execute::Execution; +use crate::scalar_fn::execute::Operand; +use crate::scalar_fn::execute::dispatch_unary; + +/// Resolve the strict homogeneous `ST_Collect` overload for one list operand. +fn collect_dtype(dtypes: &[DType]) -> VortexResult { + vortex_ensure!( + dtypes.len() == 1, + "spatial: collect requires exactly one list operand, got {}", + dtypes.len() + ); + let DType::List(element_dtype, nullability) = &dtypes[0] else { + vortex_bail!("spatial: collect operand {} is not a list", dtypes[0]); + }; + let Some(element) = element_dtype.as_extension_opt() else { + vortex_bail!( + "spatial: collect list element {} is not a native Point, LineString, or Polygon", + element_dtype + ); + }; + // Multi-geometries cannot contain null components. Null list elements are ignored during + // execution, so their storage is non-nullable in the result. + let storage = DType::List( + Arc::new(element.storage_dtype().as_nonnullable()), + *nullability, + ); + let output = if element.is::() { + ExtDType::::try_new(element.metadata::().clone(), storage)?.erased() + } else if element.is::() { + ExtDType::::try_new(element.metadata::().clone(), storage)? + .erased() + } else if element.is::() { + ExtDType::::try_new(element.metadata::().clone(), storage)?.erased() + } else { + vortex_bail!( + "spatial: collect list element {} is not a native Point, LineString, or Polygon", + element_dtype + ); + }; + Ok(DType::Extension(output)) +} + +/// Count valid elements in an exact list row without per-element mask lookups. +fn valid_count(mask: &Mask, start: usize, end: usize) -> usize { + match mask.bit_buffer() { + AllOr::All => end - start, + AllOr::None => 0, + AllOr::Some(bits) => bits.count_range(start, end), + } +} + +/// Rewrap a homogeneous geometry list as its corresponding multi-geometry array. +/// +/// The all-valid path reuses the geometry payload and list views. If geometry elements are null, +/// DuckDB semantics require ignoring them; that path first makes the views exact, then compacts the +/// payload and rebuilds the row views. +fn collect_list( + mut list: ListViewArray, + validity: Validity, + output_dtype: &ExtDTypeRef, + ctx: &mut ExecutionCtx, +) -> VortexResult { + let mut element_valid = list + .elements() + .validity()? + .execute_mask(list.elements().len(), ctx)?; + if !element_valid.all_true() { + list = list.rebuild(ListViewRebuildMode::MakeExact, ctx)?; + element_valid = list + .elements() + .validity()? + .execute_mask(list.elements().len(), ctx)?; + } + + let parts = list.into_data_parts(); + let elements = parts.elements.execute::(ctx)?; + let DType::List(target_element_storage, _) = output_dtype.storage_dtype() else { + unreachable!("collect output storage is always a list") + }; + let target_element_storage = target_element_storage.as_ref().clone(); + + let compact_elements = !element_valid.all_true(); + let element_storage = if compact_elements { + elements + .storage_array() + .filter(element_valid.clone())? + .cast(target_element_storage)? + } else { + elements.storage_array().cast(target_element_storage)? + }; + + let (offsets, sizes) = if compact_elements { + let old_offsets = parts + .offsets + .cast(DType::Primitive(PType::U64, Nullability::NonNullable))? + .execute::>(ctx)?; + let old_sizes = parts + .sizes + .cast(DType::Primitive(PType::U64, Nullability::NonNullable))? + .execute::>(ctx)?; + let mut offsets = BufferMut::::with_capacity(old_offsets.len()); + let mut sizes = BufferMut::::with_capacity(old_sizes.len()); + let mut next_offset = 0_u64; + + for (&old_offset, &old_size) in old_offsets.iter().zip(old_sizes.iter()) { + let start = usize::try_from(old_offset) + .map_err(|_| vortex_err!("spatial: collect element offset exceeds usize"))?; + let size = usize::try_from(old_size) + .map_err(|_| vortex_err!("spatial: collect element count exceeds usize"))?; + let end = start + .checked_add(size) + .ok_or_else(|| vortex_err!("spatial: collect element range overflows usize"))?; + vortex_ensure!( + end <= element_valid.len(), + "spatial: collect element range {start}..{end} exceeds element length {}", + element_valid.len() + ); + let size = u64::try_from(valid_count(&element_valid, start, end)) + .map_err(|_| vortex_err!("spatial: collect valid element count exceeds u64"))?; + offsets.push(next_offset); + sizes.push(size); + next_offset = next_offset + .checked_add(size) + .ok_or_else(|| vortex_err!("spatial: collect output offset exceeds u64"))?; + } + (offsets.into_array(), sizes.into_array()) + } else { + (parts.offsets, parts.sizes) + }; + + let storage = ListViewArray::try_new(element_storage, offsets, sizes, validity)?.into_array(); + Ok(ExtensionArray::try_new(output_dtype.clone(), storage)?.into_array()) +} + +/// Execute the structural collect kernel after shared unary shape and null dispatch. +fn execute_collect( + execution: Execution<1>, + output_dtype: &ExtDTypeRef, + ctx: &mut ExecutionCtx, +) -> VortexResult { + match execution.operands { + [Operand::Constant(scalar)] => { + let one = ConstantArray::new(scalar, 1) + .into_array() + .execute::(ctx)?; + let collected = collect_list( + one, + Validity::from_mask(Mask::new_true(1), output_dtype.nullability()), + output_dtype, + ctx, + )?; + Ok(ConstantArray::new(collected.execute_scalar(0, ctx)?, execution.len).into_array()) + } + [Operand::Column(array)] => collect_list( + array.execute::(ctx)?, + Validity::from_mask(execution.valid, output_dtype.nullability()), + output_dtype, + ctx, + ), + } +} + +/// Collect a homogeneous list of native `Point`, `LineString`, or `Polygon` values into the +/// corresponding `MultiPoint`, `MultiLineString`, or `MultiPolygon` value. Null geometry elements +/// are ignored. Mixed geometry lists are rejected by the list element dtype rather than represented +/// as a geometry union. +#[derive(Debug, Clone, Default, PartialEq, Eq, Hash)] +pub struct SpatialCollect; + +impl SpatialCollect { + /// A lazy `ScalarFnArray` collecting each list row into one native multi-geometry value. + pub fn try_new_array(array: ArrayRef) -> VortexResult { + ScalarFnArray::try_new( + TypedScalarFnInstance::new(SpatialCollect, EmptyOptions).erased(), + vec![array], + ) + } +} + +impl ScalarFnVTable for SpatialCollect { + type Options = EmptyOptions; + + fn id(&self) -> ScalarFnId { + static ID: CachedId = CachedId::new("vortex.st.collect"); + *ID + } + + fn serialize(&self, _: &Self::Options) -> VortexResult>> { + Ok(Some(vec![])) + } + + fn deserialize(&self, _: &[u8], _: &VortexSession) -> VortexResult { + Ok(EmptyOptions) + } + + fn arity(&self, _: &Self::Options) -> Arity { + Arity::Exact(1) + } + + fn child_name(&self, _: &Self::Options, child_idx: usize) -> ChildName { + match child_idx { + 0 => ChildName::from("geometries"), + _ => unreachable!("collect has exactly one child"), + } + } + + fn return_dtype(&self, _: &Self::Options, dtypes: &[DType]) -> VortexResult { + collect_dtype(dtypes) + } + + fn execute( + &self, + _: &Self::Options, + args: &dyn ExecutionArgs, + ctx: &mut ExecutionCtx, + ) -> VortexResult { + let input = args.get(0)?; + let output_dtype = collect_dtype(std::slice::from_ref(input.dtype()))?; + let output = output_dtype.as_extension().clone(); + dispatch_unary( + &input, + output_dtype, + |execution, ctx| execute_collect(execution, &output, ctx), + ctx, + ) + } + + fn validity( + &self, + _: &Self::Options, + expression: &Expression, + ) -> VortexResult> { + union_child_validities(expression) + } + + fn is_strict(&self, _: &Self::Options) -> bool { + true + } + + fn is_fallible(&self, _: &Self::Options) -> bool { + false + } +} + +#[cfg(test)] +mod tests { + use vortex_array::ArrayRef; + use vortex_array::Columnar; + use vortex_array::IntoArray; + use vortex_array::VortexSessionExecute; + use vortex_array::arrays::ConstantArray; + use vortex_array::arrays::ExtensionArray; + use vortex_array::arrays::ListArray; + use vortex_array::arrays::ListViewArray; + use vortex_array::arrays::MaskedArray; + use vortex_array::arrays::PrimitiveArray; + use vortex_array::arrays::extension::ExtensionArrayExt; + use vortex_array::arrays::listview::ListViewArraySlotsExt; + use vortex_array::assert_arrays_eq; + use vortex_array::dtype::DType; + use vortex_array::dtype::Nullability; + use vortex_array::dtype::PType; + use vortex_array::scalar_fn::EmptyOptions; + use vortex_array::scalar_fn::ScalarFnVTable; + use vortex_array::validity::Validity; + use vortex_error::VortexResult; + use vortex_error::vortex_err; + + use super::SpatialCollect; + use crate::test_harness::linestring_column; + use crate::test_harness::multilinestring_column; + use crate::test_harness::multipoint_column; + use crate::test_harness::multipolygon_column; + use crate::test_harness::nullable_point_column; + use crate::test_harness::point_column; + use crate::test_harness::polygon_column; + + fn list_with_validity( + elements: ArrayRef, + offsets: &[u32], + validity: Validity, + ) -> VortexResult { + Ok(ListArray::try_new( + elements, + PrimitiveArray::from_iter(offsets.iter().copied()).into_array(), + validity, + )? + .into_array()) + } + + fn list(elements: ArrayRef, offsets: &[u32]) -> VortexResult { + list_with_validity(elements, offsets, Validity::NonNullable) + } + + #[test] + fn collects_points_into_multipoints() -> VortexResult<()> { + let points = point_column(vec![0.0, 1.0, 2.0], vec![3.0, 4.0, 5.0])?; + let input = list(points, &[0, 2, 3])?; + let expected = multipoint_column(vec![vec![(0.0, 3.0), (1.0, 4.0)], vec![(2.0, 5.0)]])?; + let result = SpatialCollect::try_new_array(input)?.into_array(); + let mut ctx = vortex_array::array_session().create_execution_ctx(); + + assert_arrays_eq!(result, expected, &mut ctx); + Ok(()) + } + + #[test] + fn all_valid_collect_reuses_geometry_storage() -> VortexResult<()> { + let points = point_column(vec![0.0, 1.0, 2.0], vec![3.0, 4.0, 5.0])?; + let mut ctx = vortex_array::array_session().create_execution_ctx(); + let point_storage = points + .clone() + .execute::(&mut ctx)? + .storage_array() + .clone(); + let input = list(points, &[0, 2, 3])?; + + let result = SpatialCollect::try_new_array(input)? + .into_array() + .execute::(&mut ctx)?; + let result_storage = result + .storage_array() + .clone() + .execute::(&mut ctx)?; + + assert!(ArrayRef::ptr_eq(&point_storage, result_storage.elements())); + Ok(()) + } + + #[test] + fn collects_linestrings_into_multilinestrings() -> VortexResult<()> { + let line_a = vec![(0.0, 0.0), (1.0, 1.0)]; + let line_b = vec![(2.0, 2.0), (3.0, 3.0)]; + let line_c = vec![(4.0, 4.0), (5.0, 5.0)]; + let input = list( + linestring_column(vec![line_a.clone(), line_b.clone(), line_c.clone()])?, + &[0, 2, 3], + )?; + let expected = multilinestring_column(vec![vec![line_a, line_b], vec![line_c]])?; + let result = SpatialCollect::try_new_array(input)?.into_array(); + let mut ctx = vortex_array::array_session().create_execution_ctx(); + + assert_arrays_eq!(result, expected, &mut ctx); + Ok(()) + } + + #[test] + fn collects_polygons_into_multipolygons() -> VortexResult<()> { + let polygon_a = vec![vec![(0.0, 0.0), (2.0, 0.0), (0.0, 2.0), (0.0, 0.0)]]; + let polygon_b = vec![vec![(3.0, 0.0), (5.0, 0.0), (3.0, 2.0), (3.0, 0.0)]]; + let polygon_c = vec![vec![(6.0, 0.0), (8.0, 0.0), (6.0, 2.0), (6.0, 0.0)]]; + let input = list( + polygon_column(vec![ + polygon_a.clone(), + polygon_b.clone(), + polygon_c.clone(), + ])?, + &[0, 2, 3], + )?; + let expected = multipolygon_column(vec![vec![polygon_a, polygon_b], vec![polygon_c]])?; + let result = SpatialCollect::try_new_array(input)?.into_array(); + let mut ctx = vortex_array::array_session().create_execution_ctx(); + + assert_arrays_eq!(result, expected, &mut ctx); + Ok(()) + } + + #[test] + fn constant_list_remains_constant() -> VortexResult<()> { + let input = list( + nullable_point_column(vec![Some((0.0, 2.0)), None, Some((1.0, 3.0))])?, + &[0, 3], + )?; + let mut ctx = vortex_array::array_session().create_execution_ctx(); + let scalar = input.execute_scalar(0, &mut ctx)?; + let input = ConstantArray::new(scalar, 3).into_array(); + + let result = SpatialCollect::try_new_array(input)?.into_array(); + let Columnar::Constant(constant) = result.clone().execute::(&mut ctx)? else { + return Err(vortex_err!( + "collect of a constant list should remain constant" + )); + }; + assert_eq!(constant.len(), 3); + let expected = multipoint_column(vec![ + vec![(0.0, 2.0), (1.0, 3.0)], + vec![(0.0, 2.0), (1.0, 3.0)], + vec![(0.0, 2.0), (1.0, 3.0)], + ])?; + assert_arrays_eq!(result, expected, &mut ctx); + Ok(()) + } + + #[test] + fn ignores_null_geometry_elements() -> VortexResult<()> { + let points = nullable_point_column(vec![Some((0.0, 2.0)), None, Some((1.0, 3.0)), None])?; + let input = list(points, &[0, 2, 4])?; + let expected = multipoint_column(vec![vec![(0.0, 2.0)], vec![(1.0, 3.0)]])?; + let result = SpatialCollect::try_new_array(input)?.into_array(); + let mut ctx = vortex_array::array_session().create_execution_ctx(); + + assert_arrays_eq!(result, expected, &mut ctx); + Ok(()) + } + + #[test] + fn all_null_geometry_elements_produce_empty_multi_geometry() -> VortexResult<()> { + let input = list(nullable_point_column(vec![None, None])?, &[0, 2])?; + let expected = multipoint_column(vec![vec![]])?; + let result = SpatialCollect::try_new_array(input)?.into_array(); + let mut ctx = vortex_array::array_session().create_execution_ctx(); + + assert_arrays_eq!(result, expected, &mut ctx); + Ok(()) + } + + #[test] + fn propagates_null_list_rows() -> VortexResult<()> { + let input = list_with_validity( + point_column(vec![0.0, 1.0], vec![2.0, 3.0])?, + &[0, 1, 2], + Validity::from_iter([true, false]), + )?; + let expected = MaskedArray::try_new( + multipoint_column(vec![vec![(0.0, 2.0)], vec![(1.0, 3.0)]])?, + Validity::from_iter([true, false]), + )? + .into_array(); + let result = SpatialCollect::try_new_array(input)?.into_array(); + let mut ctx = vortex_array::array_session().create_execution_ctx(); + + assert_arrays_eq!(result, expected, &mut ctx); + Ok(()) + } + + #[test] + fn rejects_unsupported_inputs() -> VortexResult<()> { + let point = point_column(vec![0.0], vec![0.0])?; + assert!(SpatialCollect::try_new_array(point).is_err()); + + let multipoints = multipoint_column(vec![vec![(0.0, 0.0)]])?; + assert!(SpatialCollect::try_new_array(list(multipoints, &[0, 1])?).is_err()); + + let primitive = DType::Primitive(PType::F64, Nullability::NonNullable); + assert!( + SpatialCollect + .return_dtype( + &EmptyOptions, + &[DType::List(primitive.into(), Nullability::NonNullable)] + ) + .is_err() + ); + Ok(()) + } +} diff --git a/vortex-spatial/src/scalar_fn/mod.rs b/vortex-spatial/src/scalar_fn/mod.rs index 1dcff7d0b95..fce872e81a6 100644 --- a/vortex-spatial/src/scalar_fn/mod.rs +++ b/vortex-spatial/src/scalar_fn/mod.rs @@ -4,6 +4,7 @@ //! Geometry scalar functions over the native geometry extension types. pub mod area; +pub mod collect; pub mod contains; pub mod distance; pub mod envelope; From 3c5e595ac90b655538debde22dc8d8fd20351050 Mon Sep 17 00:00:00 2001 From: Nemo Yu Date: Wed, 5 Aug 2026 17:07:15 -0400 Subject: [PATCH 2/5] bench(vortex-geo): add collect benchmark Signed-off-by: Nemo Yu --- vortex-spatial/Cargo.toml | 5 + vortex-spatial/benches/collect.rs | 149 ++++++++++++++++++++++++++++++ 2 files changed, 154 insertions(+) create mode 100644 vortex-spatial/benches/collect.rs diff --git a/vortex-spatial/Cargo.toml b/vortex-spatial/Cargo.toml index 8d2b8cfe509..dc4cfe3a51e 100644 --- a/vortex-spatial/Cargo.toml +++ b/vortex-spatial/Cargo.toml @@ -65,5 +65,10 @@ harness = false [[bench]] name = "area" harness = false + +[[bench]] +name = "collect" +harness = false + [lints] workspace = true diff --git a/vortex-spatial/benches/collect.rs b/vortex-spatial/benches/collect.rs new file mode 100644 index 00000000000..f1193b4fc27 --- /dev/null +++ b/vortex-spatial/benches/collect.rs @@ -0,0 +1,149 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Microbenchmarks for native `ST_Collect` over homogeneous geometry lists. +//! +//! The cases cover each strict overload and the inner-null compaction path. They execute the +//! result to its canonical representation so the full multi-geometry construction is measured. +//! +//! Run with `cargo bench -p vortex-spatial --bench collect`. + +#![expect(clippy::unwrap_used)] + +use std::sync::LazyLock; + +use divan::Bencher; +use divan::counter::ItemsCount; +use mimalloc::MiMalloc; +use vortex_array::ArrayRef; +use vortex_array::Canonical; +use vortex_array::ExecutionCtx; +use vortex_array::IntoArray; +use vortex_array::VortexSessionExecute; +use vortex_array::arrays::ListArray; +use vortex_array::arrays::PrimitiveArray; +use vortex_array::validity::Validity; +use vortex_session::VortexSession; +use vortex_spatial::scalar_fn::collect::SpatialCollect; +use vortex_spatial::test_harness::linestring_column; +use vortex_spatial::test_harness::nullable_point_column; +use vortex_spatial::test_harness::point_column; +use vortex_spatial::test_harness::polygon_column; +use vortex_spatial::test_harness::spatial_session; + +// Scalar function execution allocates its output inside the timed region, so use the vendored +// allocator instead of measuring glibc differences between CodSpeed runner images. +#[global_allocator] +static GLOBAL: MiMalloc = MiMalloc; + +static SESSION: LazyLock = LazyLock::new(spatial_session); + +const ROWS: usize = 512; + +fn main() { + divan::main(); +} + +fn geometry_lists(elements: ArrayRef, elements_per_row: usize) -> ArrayRef { + let offsets = PrimitiveArray::from_iter( + (0..=ROWS).map(|row| u64::try_from(row * elements_per_row).unwrap()), + ) + .into_array(); + ListArray::try_new(elements, offsets, Validity::NonNullable) + .unwrap() + .into_array() +} + +fn point_lists(nullable: bool) -> ArrayRef { + const POINTS_PER_ROW: usize = 8; + let len = ROWS * POINTS_PER_ROW; + let points = if nullable { + nullable_point_column( + (0..len) + .map(|i| (!i.is_multiple_of(8)).then_some((i as f64, (i + 1) as f64))) + .collect(), + ) + .unwrap() + } else { + point_column( + (0..len).map(|i| i as f64).collect(), + (0..len).map(|i| (i + 1) as f64).collect(), + ) + .unwrap() + }; + geometry_lists(points, POINTS_PER_ROW) +} + +fn linestring_lists() -> ArrayRef { + const LINES_PER_ROW: usize = 4; + let lines = linestring_column( + (0..ROWS * LINES_PER_ROW) + .map(|line| { + (0..8) + .map(|vertex| { + let value = (line * 8 + vertex) as f64; + (value, value + 1.0) + }) + .collect() + }) + .collect(), + ) + .unwrap(); + geometry_lists(lines, LINES_PER_ROW) +} + +fn polygon_lists() -> ArrayRef { + const POLYGONS_PER_ROW: usize = 2; + let polygons = polygon_column( + (0..ROWS * POLYGONS_PER_ROW) + .map(|polygon| { + let x = polygon as f64; + vec![vec![ + (x, 0.0), + (x + 1.0, 0.0), + (x + 1.0, 1.0), + (x, 1.0), + (x, 0.0), + ]] + }) + .collect(), + ) + .unwrap(); + geometry_lists(polygons, POLYGONS_PER_ROW) +} + +fn collect(input: &ArrayRef, ctx: &mut ExecutionCtx) -> ArrayRef { + SpatialCollect::try_new_array(input.clone()) + .unwrap() + .into_array() + .execute::(ctx) + .unwrap() + .into_array() +} + +fn bench_collect(bencher: Bencher, input: ArrayRef) { + let mut ctx = SESSION.create_execution_ctx(); + bencher + .counter(ItemsCount::new(ROWS)) + .bench_local(|| collect(&input, &mut ctx)); +} + +#[divan::bench] +fn points(bencher: Bencher) { + bench_collect(bencher, point_lists(false)); +} + +#[divan::bench] +fn linestrings(bencher: Bencher) { + bench_collect(bencher, linestring_lists()); +} + +#[divan::bench] +fn polygons(bencher: Bencher) { + bench_collect(bencher, polygon_lists()); +} + +#[divan::bench] +fn nullable_points(bencher: Bencher) { + bench_collect(bencher, point_lists(true)); +} From 91e62f56a826e506d167e6ab44d8df163ad0f585 Mon Sep 17 00:00:00 2001 From: Nemo Yu Date: Thu, 6 Aug 2026 16:15:13 -0400 Subject: [PATCH 3/5] fix(vortex-geo): materialize collect validity Signed-off-by: Nemo Yu --- vortex-spatial/src/scalar_fn/collect.rs | 17 ++++++++++------- 1 file changed, 10 insertions(+), 7 deletions(-) diff --git a/vortex-spatial/src/scalar_fn/collect.rs b/vortex-spatial/src/scalar_fn/collect.rs index b138a661665..dd08382863c 100644 --- a/vortex-spatial/src/scalar_fn/collect.rs +++ b/vortex-spatial/src/scalar_fn/collect.rs @@ -184,7 +184,7 @@ fn collect_list( /// Execute the structural collect kernel after shared unary shape and null dispatch. fn execute_collect( - execution: Execution<1>, + execution: Execution<1, Validity>, output_dtype: &ExtDTypeRef, ctx: &mut ExecutionCtx, ) -> VortexResult { @@ -201,12 +201,15 @@ fn execute_collect( )?; Ok(ConstantArray::new(collected.execute_scalar(0, ctx)?, execution.len).into_array()) } - [Operand::Column(array)] => collect_list( - array.execute::(ctx)?, - Validity::from_mask(execution.valid, output_dtype.nullability()), - output_dtype, - ctx, - ), + [Operand::Column(array)] => { + let valid = execution.valid.execute_mask(execution.len, ctx)?; + collect_list( + array.execute::(ctx)?, + Validity::from_mask(valid, output_dtype.nullability()), + output_dtype, + ctx, + ) + } } } From 420a50e44f97d5fa83a813ed5d482824c41cc52e Mon Sep 17 00:00:00 2001 From: Nemo Yu Date: Fri, 7 Aug 2026 17:07:58 -0400 Subject: [PATCH 4/5] refactor(vortex-spatial): tighten collect dtype plumbing Return `ExtDTypeRef` from `collect_dtype` so `execute` stops unwrapping the extension back out of a `DType` and no longer carries two names for one value, matching how `convex_hull_dtype` resolves its output. Fold the two element-type rejections into one match so the "not a native Point, LineString, or Polygon" message has a single source, and take the output nullability from the `Execution` the dispatcher already populated instead of re-deriving it from the output dtype. Signed-off-by: Nemo Yu --- vortex-spatial/src/scalar_fn/collect.rs | 62 +++++++++++++------------ 1 file changed, 32 insertions(+), 30 deletions(-) diff --git a/vortex-spatial/src/scalar_fn/collect.rs b/vortex-spatial/src/scalar_fn/collect.rs index dd08382863c..172ee622603 100644 --- a/vortex-spatial/src/scalar_fn/collect.rs +++ b/vortex-spatial/src/scalar_fn/collect.rs @@ -53,7 +53,7 @@ use crate::scalar_fn::execute::Operand; use crate::scalar_fn::execute::dispatch_unary; /// Resolve the strict homogeneous `ST_Collect` overload for one list operand. -fn collect_dtype(dtypes: &[DType]) -> VortexResult { +fn collect_dtype(dtypes: &[DType]) -> VortexResult { vortex_ensure!( dtypes.len() == 1, "spatial: collect requires exactly one list operand, got {}", @@ -62,32 +62,35 @@ fn collect_dtype(dtypes: &[DType]) -> VortexResult { let DType::List(element_dtype, nullability) = &dtypes[0] else { vortex_bail!("spatial: collect operand {} is not a list", dtypes[0]); }; - let Some(element) = element_dtype.as_extension_opt() else { - vortex_bail!( - "spatial: collect list element {} is not a native Point, LineString, or Polygon", - element_dtype - ); - }; // Multi-geometries cannot contain null components. Null list elements are ignored during // execution, so their storage is non-nullable in the result. - let storage = DType::List( - Arc::new(element.storage_dtype().as_nonnullable()), - *nullability, - ); - let output = if element.is::() { - ExtDType::::try_new(element.metadata::().clone(), storage)?.erased() - } else if element.is::() { - ExtDType::::try_new(element.metadata::().clone(), storage)? - .erased() - } else if element.is::() { - ExtDType::::try_new(element.metadata::().clone(), storage)?.erased() - } else { - vortex_bail!( - "spatial: collect list element {} is not a native Point, LineString, or Polygon", - element_dtype - ); + let multi_storage = |element: &ExtDTypeRef| { + DType::List( + Arc::new(element.storage_dtype().as_nonnullable()), + *nullability, + ) }; - Ok(DType::Extension(output)) + match element_dtype.as_extension_opt() { + Some(element) if element.is::() => Ok(ExtDType::::try_new( + element.metadata::().clone(), + multi_storage(element), + )? + .erased()), + Some(element) if element.is::() => Ok(ExtDType::::try_new( + element.metadata::().clone(), + multi_storage(element), + )? + .erased()), + Some(element) if element.is::() => Ok(ExtDType::::try_new( + element.metadata::().clone(), + multi_storage(element), + )? + .erased()), + _ => vortex_bail!( + "spatial: collect list element {element_dtype} is not a native Point, LineString, \ + or Polygon" + ), + } } /// Count valid elements in an exact list row without per-element mask lookups. @@ -195,7 +198,7 @@ fn execute_collect( .execute::(ctx)?; let collected = collect_list( one, - Validity::from_mask(Mask::new_true(1), output_dtype.nullability()), + Validity::from_mask(Mask::new_true(1), execution.nullability), output_dtype, ctx, )?; @@ -205,7 +208,7 @@ fn execute_collect( let valid = execution.valid.execute_mask(execution.len, ctx)?; collect_list( array.execute::(ctx)?, - Validity::from_mask(valid, output_dtype.nullability()), + Validity::from_mask(valid, execution.nullability), output_dtype, ctx, ) @@ -258,7 +261,7 @@ impl ScalarFnVTable for SpatialCollect { } fn return_dtype(&self, _: &Self::Options, dtypes: &[DType]) -> VortexResult { - collect_dtype(dtypes) + Ok(DType::Extension(collect_dtype(dtypes)?)) } fn execute( @@ -269,11 +272,10 @@ impl ScalarFnVTable for SpatialCollect { ) -> VortexResult { let input = args.get(0)?; let output_dtype = collect_dtype(std::slice::from_ref(input.dtype()))?; - let output = output_dtype.as_extension().clone(); dispatch_unary( &input, - output_dtype, - |execution, ctx| execute_collect(execution, &output, ctx), + DType::Extension(output_dtype.clone()), + |execution, ctx| execute_collect(execution, &output_dtype, ctx), ctx, ) } From 1f9f067ea1c0a1c7aa088337edafc243e3adfead Mon Sep 17 00:00:00 2001 From: Nemo Yu Date: Fri, 7 Aug 2026 17:08:12 -0400 Subject: [PATCH 5/5] perf(vortex-spatial): keep collect output zero-copy to list `ListViewArray::try_new` always reports `is_zero_copy_to_list` as false, so the list view collect handed back forgot that its views are still exact. The next `list_from_list_view` then re-gathered the entire geometry payload that the all-valid path had just reused, moving the copy one operator later instead of avoiding it. Forward the input's flag instead. The reuse path passes `offsets` and `sizes` through untouched, and the compaction path rebuilds them as a running sum over the same element order, so the zero-copy invariant holds on both; `validate_zctl` checks it under debug assertions. `ST_Envelope(ST_Collect(points))` over 512 rows of 8 points improves from 10.54us to 8.42us fastest and 10.72us to 8.54us median. The existing cases cannot observe this because `Canonical`'s list form is itself a `ListViewArray`, so add one that composes collect with a consumer converting to a `ListArray`. Signed-off-by: Nemo Yu --- vortex-spatial/benches/collect.rs | 28 +++++++++++++++ vortex-spatial/src/scalar_fn/collect.rs | 48 +++++++++++++++++++++++-- 2 files changed, 74 insertions(+), 2 deletions(-) diff --git a/vortex-spatial/benches/collect.rs b/vortex-spatial/benches/collect.rs index f1193b4fc27..e07ed00fde3 100644 --- a/vortex-spatial/benches/collect.rs +++ b/vortex-spatial/benches/collect.rs @@ -25,6 +25,7 @@ use vortex_array::arrays::PrimitiveArray; use vortex_array::validity::Validity; use vortex_session::VortexSession; use vortex_spatial::scalar_fn::collect::SpatialCollect; +use vortex_spatial::scalar_fn::envelope::SpatialEnvelope; use vortex_spatial::test_harness::linestring_column; use vortex_spatial::test_harness::nullable_point_column; use vortex_spatial::test_harness::point_column; @@ -147,3 +148,30 @@ fn polygons(bencher: Bencher) { fn nullable_points(bencher: Bencher) { bench_collect(bencher, point_lists(true)); } + +/// Collect feeding a consumer that converts the result to a `ListArray`. +/// +/// The cases above stop at [`Canonical`], whose list form is a `ListViewArray`, so they cannot +/// observe whether collect's output still reports itself as zero-copy to a list. `ST_Envelope` +/// reaches that path through `flatten_row_offsets`, and re-gathers the whole payload when the +/// flag is missing. +fn envelope_of_collect(input: &ArrayRef, ctx: &mut ExecutionCtx) -> ArrayRef { + let collected = SpatialCollect::try_new_array(input.clone()) + .unwrap() + .into_array(); + SpatialEnvelope::try_new_array(collected) + .unwrap() + .into_array() + .execute::(ctx) + .unwrap() + .into_array() +} + +#[divan::bench] +fn envelope_of_collected_points(bencher: Bencher) { + let input = point_lists(false); + let mut ctx = SESSION.create_execution_ctx(); + bencher + .counter(ItemsCount::new(ROWS)) + .bench_local(|| envelope_of_collect(&input, &mut ctx)); +} diff --git a/vortex-spatial/src/scalar_fn/collect.rs b/vortex-spatial/src/scalar_fn/collect.rs index 172ee622603..2f80401904e 100644 --- a/vortex-spatial/src/scalar_fn/collect.rs +++ b/vortex-spatial/src/scalar_fn/collect.rs @@ -106,7 +106,8 @@ fn valid_count(mask: &Mask, start: usize, end: usize) -> usize { /// /// The all-valid path reuses the geometry payload and list views. If geometry elements are null, /// DuckDB semantics require ignoring them; that path first makes the views exact, then compacts the -/// payload and rebuilds the row views. +/// payload and rebuilds the row views. Either way the output carries the input's zero-copy-to-list +/// flag, so a downstream `ListArray` conversion does not re-gather the reused payload. fn collect_list( mut list: ListViewArray, validity: Validity, @@ -125,6 +126,11 @@ fn collect_list( .execute_mask(list.elements().len(), ctx)?; } + // Both output paths keep the views exact: reuse forwards `offsets` and `sizes` untouched, and + // compaction rebuilds them as a running sum over the same element order. So the result is + // zero-copy to a `ListArray` exactly when `list` is, which `MakeExact` above has already + // ensured for every list that reaches compaction. + let zero_copy_to_list = list.is_zero_copy_to_list(); let parts = list.into_data_parts(); let elements = parts.elements.execute::(ctx)?; let DType::List(target_element_storage, _) = output_dtype.storage_dtype() else { @@ -181,7 +187,12 @@ fn collect_list( (parts.offsets, parts.sizes) }; - let storage = ListViewArray::try_new(element_storage, offsets, sizes, validity)?.into_array(); + let storage = ListViewArray::try_new(element_storage, offsets, sizes, validity)?; + // SAFETY: `zero_copy_to_list` describes views this function either forwarded unchanged or + // replaced with a gapless, non-overlapping running sum over the same elements. Forwarding it + // matters: `list_from_list_view` re-gathers the whole payload for a list view that reports + // `false`, undoing the storage reuse above one operator later. + let storage = unsafe { storage.with_zero_copy_to_list(zero_copy_to_list) }.into_array(); Ok(ExtensionArray::try_new(output_dtype.clone(), storage)?.into_array()) } @@ -299,6 +310,7 @@ impl ScalarFnVTable for SpatialCollect { #[cfg(test)] mod tests { + use rstest::rstest; use vortex_array::ArrayRef; use vortex_array::Columnar; use vortex_array::IntoArray; @@ -382,6 +394,38 @@ mod tests { Ok(()) } + /// A list view that forgets it is zero-copy to a list makes the next + /// `list_from_list_view` re-gather the payload that collect just reused. + #[rstest] + #[case::reused_elements(false)] + #[case::compacted_elements(true)] + fn output_stays_zero_copy_to_list(#[case] null_elements: bool) -> VortexResult<()> { + let points = if null_elements { + nullable_point_column(vec![Some((0.0, 3.0)), None, Some((2.0, 5.0))])? + } else { + point_column(vec![0.0, 1.0, 2.0], vec![3.0, 4.0, 5.0])? + }; + let mut ctx = vortex_array::array_session().create_execution_ctx(); + let input = list(points, &[0, 2, 3])?; + assert!( + input + .clone() + .execute::(&mut ctx)? + .is_zero_copy_to_list(), + "a list column reaches collect as an exact list view" + ); + + let storage = SpatialCollect::try_new_array(input)? + .into_array() + .execute::(&mut ctx)? + .storage_array() + .clone() + .execute::(&mut ctx)?; + + assert!(storage.is_zero_copy_to_list()); + Ok(()) + } + #[test] fn collects_linestrings_into_multilinestrings() -> VortexResult<()> { let line_a = vec![(0.0, 0.0), (1.0, 1.0)];