From ad600f730b6490996c06a7cdb44926f01a95d08a Mon Sep 17 00:00:00 2001 From: Nemo Yu Date: Wed, 5 Aug 2026 17:07:06 -0400 Subject: [PATCH 1/7] 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 22127a050536780f1fb6a2101b40bce94b7fdbdb Mon Sep 17 00:00:00 2001 From: Nemo Yu Date: Wed, 5 Aug 2026 17:07:15 -0400 Subject: [PATCH 2/7] 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 c3cf312504c4141a8f8b4d74865232abdbc99d78 Mon Sep 17 00:00:00 2001 From: Nemo Yu Date: Thu, 6 Aug 2026 16:15:13 -0400 Subject: [PATCH 3/7] 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 f972ca28b14dab6ebf4dc650799124b646ea9836 Mon Sep 17 00:00:00 2001 From: Nemo Yu Date: Fri, 7 Aug 2026 17:07:58 -0400 Subject: [PATCH 4/7] 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 de907e5531fcbe422285d729aba53dd1111be163 Mon Sep 17 00:00:00 2001 From: Nemo Yu Date: Fri, 7 Aug 2026 17:08:12 -0400 Subject: [PATCH 5/7] 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)]; From 7c3d42a343d0846e06b45d94cbbc58ce35ee2b46 Mon Sep 17 00:00:00 2001 From: Nemo Yu Date: Mon, 10 Aug 2026 14:35:46 -0400 Subject: [PATCH 6/7] refactor(vortex-spatial): clarify collect list semantics Signed-off-by: Nemo Yu --- vortex-spatial/benches/collect.rs | 16 +++--- vortex-spatial/src/scalar_fn/collect.rs | 73 ++++++++++++++----------- 2 files changed, 49 insertions(+), 40 deletions(-) diff --git a/vortex-spatial/benches/collect.rs b/vortex-spatial/benches/collect.rs index e07ed00fde3..6e0ec2e3686 100644 --- a/vortex-spatial/benches/collect.rs +++ b/vortex-spatial/benches/collect.rs @@ -1,10 +1,12 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors -//! Microbenchmarks for native `ST_Collect` over homogeneous geometry lists. +//! Microbenchmarks for the per-row scalar `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. +//! The benchmark inputs are already list-valued, modeling the output of a preceding `ARRAY_AGG` or +//! `list` aggregate rather than measuring aggregation itself. 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`. @@ -113,8 +115,8 @@ fn polygon_lists() -> ArrayRef { geometry_lists(polygons, POLYGONS_PER_ROW) } -fn collect(input: &ArrayRef, ctx: &mut ExecutionCtx) -> ArrayRef { - SpatialCollect::try_new_array(input.clone()) +fn collect_list_rows(geometry_lists: &ArrayRef, ctx: &mut ExecutionCtx) -> ArrayRef { + SpatialCollect::try_new_array(geometry_lists.clone()) .unwrap() .into_array() .execute::(ctx) @@ -122,11 +124,11 @@ fn collect(input: &ArrayRef, ctx: &mut ExecutionCtx) -> ArrayRef { .into_array() } -fn bench_collect(bencher: Bencher, input: ArrayRef) { +fn bench_collect(bencher: Bencher, geometry_lists: ArrayRef) { let mut ctx = SESSION.create_execution_ctx(); bencher .counter(ItemsCount::new(ROWS)) - .bench_local(|| collect(&input, &mut ctx)); + .bench_local(|| collect_list_rows(&geometry_lists, &mut ctx)); } #[divan::bench] diff --git a/vortex-spatial/src/scalar_fn/collect.rs b/vortex-spatial/src/scalar_fn/collect.rs index 2f80401904e..4d0b6049c4d 100644 --- a/vortex-spatial/src/scalar_fn/collect.rs +++ b/vortex-spatial/src/scalar_fn/collect.rs @@ -1,7 +1,12 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors -//! `ST_Collect`: collect homogeneous native geometries into their native multi-geometry type. +//! Per-row scalar `ST_Collect` over homogeneous native geometry lists. +//! +//! Each input row is one `List`, `List`, or `List` and produces one +//! corresponding multi-geometry row. This function does not aggregate geometry values across +//! rows; a SQL query that starts with individual geometry rows must first group them with an +//! aggregate such as `ARRAY_AGG` or `list`. use std::sync::Arc; @@ -102,36 +107,36 @@ fn valid_count(mask: &Mask, start: usize, end: usize) -> usize { } } -/// Rewrap a homogeneous geometry list as its corresponding multi-geometry array. +/// Rewrap each homogeneous geometry-list row as its corresponding multi-geometry row. /// /// 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. 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, +fn collect_list_rows( + mut lists: ListViewArray, validity: Validity, output_dtype: &ExtDTypeRef, ctx: &mut ExecutionCtx, ) -> VortexResult { - let mut element_valid = list + let mut element_valid = lists .elements() .validity()? - .execute_mask(list.elements().len(), ctx)?; + .execute_mask(lists.elements().len(), ctx)?; if !element_valid.all_true() { - list = list.rebuild(ListViewRebuildMode::MakeExact, ctx)?; - element_valid = list + lists = lists.rebuild(ListViewRebuildMode::MakeExact, ctx)?; + element_valid = lists .elements() .validity()? - .execute_mask(list.elements().len(), ctx)?; + .execute_mask(lists.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(); + // zero-copy to a `ListArray` exactly when `lists` is, which `MakeExact` above has already + // ensured for every list array that reaches compaction. + let zero_copy_to_list = lists.is_zero_copy_to_list(); + let parts = lists.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") @@ -196,29 +201,29 @@ fn collect_list( 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( +/// Execute per-row list-to-multi-geometry conversion after shared unary shape and null dispatch. +fn execute_collect_list_rows( execution: Execution<1, Validity>, output_dtype: &ExtDTypeRef, ctx: &mut ExecutionCtx, ) -> VortexResult { match execution.operands { [Operand::Constant(scalar)] => { - let one = ConstantArray::new(scalar, 1) + let one_list = ConstantArray::new(scalar, 1) .into_array() .execute::(ctx)?; - let collected = collect_list( - one, + let collected = collect_list_rows( + one_list, Validity::from_mask(Mask::new_true(1), execution.nullability), output_dtype, ctx, )?; Ok(ConstantArray::new(collected.execute_scalar(0, ctx)?, execution.len).into_array()) } - [Operand::Column(array)] => { + [Operand::Column(geometry_lists)] => { let valid = execution.valid.execute_mask(execution.len, ctx)?; - collect_list( - array.execute::(ctx)?, + collect_list_rows( + geometry_lists.execute::(ctx)?, Validity::from_mask(valid, execution.nullability), output_dtype, ctx, @@ -227,19 +232,21 @@ fn execute_collect( } } -/// 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. +/// Scalar `ST_Collect` over list-valued rows of native geometries. +/// +/// Each input list row produces one `MultiPoint`, `MultiLineString`, or `MultiPolygon` row. This is +/// distinct from an aggregate function: it does not combine geometry values from different rows. +/// 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 { + /// Create a lazy scalar array that converts each geometry-list row to one multi-geometry row. + pub fn try_new_array(geometry_lists: ArrayRef) -> VortexResult { ScalarFnArray::try_new( TypedScalarFnInstance::new(SpatialCollect, EmptyOptions).erased(), - vec![array], + vec![geometry_lists], ) } } @@ -266,7 +273,7 @@ impl ScalarFnVTable for SpatialCollect { fn child_name(&self, _: &Self::Options, child_idx: usize) -> ChildName { match child_idx { - 0 => ChildName::from("geometries"), + 0 => ChildName::from("geometry_list"), _ => unreachable!("collect has exactly one child"), } } @@ -281,12 +288,12 @@ impl ScalarFnVTable for SpatialCollect { args: &dyn ExecutionArgs, ctx: &mut ExecutionCtx, ) -> VortexResult { - let input = args.get(0)?; - let output_dtype = collect_dtype(std::slice::from_ref(input.dtype()))?; + let geometry_lists = args.get(0)?; + let output_dtype = collect_dtype(std::slice::from_ref(geometry_lists.dtype()))?; dispatch_unary( - &input, + &geometry_lists, DType::Extension(output_dtype.clone()), - |execution, ctx| execute_collect(execution, &output_dtype, ctx), + |execution, ctx| execute_collect_list_rows(execution, &output_dtype, ctx), ctx, ) } From 8153fd859b07a6f88b21151b1b421a24d0ab9be2 Mon Sep 17 00:00:00 2001 From: Nemo Yu Date: Mon, 10 Aug 2026 15:05:01 -0400 Subject: [PATCH 7/7] refactor(vortex-spatial): simplify collect row compaction Collapse the two parallel `compact_elements` branches in `collect_list_rows` into one, which removes the derived flag, a mask clone, and a duplicated cast. Extract the row-view rebuild into `compact_row_views` and drop its read of the incoming offsets: compaction only runs after `MakeExact`, which leaves the views a gapless in-order cover, so each row starts where the previous one ended. That saves a cast and a full array execution per null-bearing batch. Rename `execute_collect_list_rows` to `execute_collect` to match the sibling `execute_envelope` and `execute_make_line`, and use one name, `element_mask`, for the element validity throughout. Add a test for a non-exact list view, the `MakeExact` path the offsets derivation relies on and the only path no test covered. Signed-off-by: Nemo Yu --- vortex-spatial/benches/collect.rs | 14 +- vortex-spatial/src/scalar_fn/collect.rs | 239 +++++++++++++----------- 2 files changed, 137 insertions(+), 116 deletions(-) diff --git a/vortex-spatial/benches/collect.rs b/vortex-spatial/benches/collect.rs index 6e0ec2e3686..3aee9e25a76 100644 --- a/vortex-spatial/benches/collect.rs +++ b/vortex-spatial/benches/collect.rs @@ -3,10 +3,9 @@ //! Microbenchmarks for the per-row scalar `ST_Collect` over homogeneous geometry lists. //! -//! The benchmark inputs are already list-valued, modeling the output of a preceding `ARRAY_AGG` or -//! `list` aggregate rather than measuring aggregation itself. 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. +//! Inputs are already list-valued, standing in for a preceding `ARRAY_AGG` rather than measuring +//! the aggregate itself. The cases cover each overload plus the null-element path, and execute to +//! canonical form so the whole multi-geometry construction is timed. //! //! Run with `cargo bench -p vortex-spatial --bench collect`. @@ -153,10 +152,9 @@ fn nullable_points(bencher: Bencher) { /// 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. +/// The cases above stop at [`Canonical`], whose list form is a `ListViewArray`, so they cannot see +/// whether the output still reports itself zero-copy to a list. `ST_Envelope` reaches that path via +/// `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() diff --git a/vortex-spatial/src/scalar_fn/collect.rs b/vortex-spatial/src/scalar_fn/collect.rs index 4d0b6049c4d..a2d373c2b43 100644 --- a/vortex-spatial/src/scalar_fn/collect.rs +++ b/vortex-spatial/src/scalar_fn/collect.rs @@ -3,10 +3,10 @@ //! Per-row scalar `ST_Collect` over homogeneous native geometry lists. //! -//! Each input row is one `List`, `List`, or `List` and produces one -//! corresponding multi-geometry row. This function does not aggregate geometry values across -//! rows; a SQL query that starts with individual geometry rows must first group them with an -//! aggregate such as `ARRAY_AGG` or `list`. +//! One `List`, `List`, or `List` row in; one `MultiPoint`, +//! `MultiLineString`, or `MultiPolygon` row out. This is not an aggregate: it never combines +//! geometries from different rows, so a query over individual geometry rows must group them +//! first with `ARRAY_AGG` or `list`. use std::sync::Arc; @@ -38,6 +38,7 @@ use vortex_array::scalar_fn::TypedScalarFnInstance; use vortex_array::validity::Validity; use vortex_buffer::Buffer; use vortex_buffer::BufferMut; +use vortex_error::VortexExpect; use vortex_error::VortexResult; use vortex_error::vortex_bail; use vortex_error::vortex_ensure; @@ -67,8 +68,7 @@ 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]); }; - // Multi-geometries cannot contain null components. Null list elements are ignored during - // execution, so their storage is non-nullable in the result. + // Execution ignores null list elements, so the result's element storage is non-nullable. let multi_storage = |element: &ExtDTypeRef| { DType::List( Arc::new(element.storage_dtype().as_nonnullable()), @@ -98,111 +98,119 @@ fn collect_dtype(dtypes: &[DType]) -> VortexResult { } } -/// 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() { +/// Count valid elements in `start..end` with one range popcount, not per-element lookups. +fn count_valid(element_mask: &Mask, start: usize, end: usize) -> usize { + match element_mask.bit_buffer() { AllOr::All => end - start, AllOr::None => 0, AllOr::Some(bits) => bits.count_range(start, end), } } -/// Rewrap each homogeneous geometry-list row as its corresponding multi-geometry row. +/// Re-address the rows once null elements are filtered out of the payload. /// -/// 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. 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. +/// Takes the `row_sizes` of an exact list view and the mask over its elements; returns the new +/// `(offsets, sizes)`. The old offsets are redundant: `MakeExact` leaves the views a gapless +/// in-order cover, so each row starts where the previous one ended. Neither running sum can +/// exceed the element count. +fn compact_row_views( + row_sizes: &ArrayRef, + element_mask: &Mask, + ctx: &mut ExecutionCtx, +) -> VortexResult<(ArrayRef, ArrayRef)> { + let sizes = row_sizes + .cast(DType::Primitive(PType::U64, Nullability::NonNullable))? + .execute::>(ctx)?; + let mut compact_offsets = BufferMut::::with_capacity(sizes.len()); + let mut compact_sizes = BufferMut::::with_capacity(sizes.len()); + let mut start = 0usize; + let mut offset = 0_u64; + + for &size in sizes.iter() { + let end = usize::try_from(size) + .ok() + .and_then(|row_len| start.checked_add(row_len)) + .filter(|end| *end <= element_mask.len()) + .ok_or_else(|| { + vortex_err!( + "spatial: collect row at element {start} exceeds the {} list elements", + element_mask.len() + ) + })?; + let valid = u64::try_from(count_valid(element_mask, start, end)) + .map_err(|_| vortex_err!("spatial: collect valid element count exceeds u64"))?; + compact_offsets.push(offset); + compact_sizes.push(valid); + offset += valid; + start = end; + } + Ok((compact_offsets.into_array(), compact_sizes.into_array())) +} + +/// Rewrap each geometry-list row as one multi-geometry row. +/// +/// All-valid rows reuse the payload and views untouched. Null elements must be ignored (DuckDB +/// semantics), so that path makes the views exact, filters the payload, and re-addresses the rows. +/// Both paths forward the input's zero-copy-to-list flag. fn collect_list_rows( mut lists: ListViewArray, validity: Validity, output_dtype: &ExtDTypeRef, ctx: &mut ExecutionCtx, ) -> VortexResult { - let mut element_valid = lists + let mut element_mask = lists .elements() .validity()? .execute_mask(lists.elements().len(), ctx)?; - if !element_valid.all_true() { + if !element_mask.all_true() { + // Filtering addresses elements by position, so make the views exact first. `MakeExact` + // re-gathers the elements, so the mask must be taken again. lists = lists.rebuild(ListViewRebuildMode::MakeExact, ctx)?; - element_valid = lists + element_mask = lists .elements() .validity()? .execute_mask(lists.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 `lists` is, which `MakeExact` above has already - // ensured for every list array that reaches compaction. + // Both paths keep the views exact — one forwards them, the other rebuilds them as a running + // sum in the same element order — so the output is zero-copy to a `ListArray` whenever the + // input is. let zero_copy_to_list = lists.is_zero_copy_to_list(); let parts = lists.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()) + let (element_storage, offsets, sizes) = if element_mask.all_true() { + (elements.storage_array().clone(), parts.offsets, parts.sizes) } else { - (parts.offsets, parts.sizes) + let (offsets, sizes) = compact_row_views(&parts.sizes, &element_mask, ctx)?; + ( + elements.storage_array().filter(element_mask)?, + offsets, + sizes, + ) }; - 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 output_element_dtype = output_dtype + .storage_dtype() + .as_list_element_opt() + .vortex_expect("collect output storage is always a list") + .as_ref() + .clone(); + let storage = ListViewArray::try_new( + element_storage.cast(output_element_dtype)?, + offsets, + sizes, + validity, + )?; + // SAFETY: the views were either forwarded unchanged or rebuilt as a gapless, non-overlapping + // running sum over the same elements, so the flag still holds. Forwarding it matters: + // `list_from_list_view` re-gathers the whole payload when it reads `false`. 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()) } -/// Execute per-row list-to-multi-geometry conversion after shared unary shape and null dispatch. -fn execute_collect_list_rows( +/// Apply [`collect_list_rows`] to a constant or column, after shared unary null dispatch. +fn execute_collect( execution: Execution<1, Validity>, output_dtype: &ExtDTypeRef, ctx: &mut ExecutionCtx, @@ -232,12 +240,11 @@ fn execute_collect_list_rows( } } -/// Scalar `ST_Collect` over list-valued rows of native geometries. +/// Scalar `ST_Collect`: one `List`, `List`, or `List` row in, one +/// `MultiPoint`, `MultiLineString`, or `MultiPolygon` row out. /// -/// Each input list row produces one `MultiPoint`, `MultiLineString`, or `MultiPolygon` row. This is -/// distinct from an aggregate function: it does not combine geometry values from different rows. -/// Null geometry elements are ignored. Mixed geometry lists are rejected by the list element dtype -/// rather than represented as a geometry union. +/// Not an aggregate — it never combines rows. Null elements are ignored, and mixed geometry lists +/// are rejected by the element dtype rather than widened to a union. #[derive(Debug, Clone, Default, PartialEq, Eq, Hash)] pub struct SpatialCollect; @@ -293,7 +300,7 @@ impl ScalarFnVTable for SpatialCollect { dispatch_unary( &geometry_lists, DType::Extension(output_dtype.clone()), - |execution, ctx| execute_collect_list_rows(execution, &output_dtype, ctx), + |execution, ctx| execute_collect(execution, &output_dtype, ctx), ctx, ) } @@ -366,16 +373,47 @@ mod tests { list_with_validity(elements, offsets, Validity::NonNullable) } + /// Assert that `ST_Collect` of `input` equals `expected`. + fn assert_collects(input: ArrayRef, expected: ArrayRef) -> VortexResult<()> { + let mut ctx = vortex_array::array_session().create_execution_ctx(); + let result = SpatialCollect::try_new_array(input)?.into_array(); + + assert_arrays_eq!(result, expected, &mut ctx); + Ok(()) + } + #[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(()) + assert_collects(input, expected) + } + + /// Out-of-order rows make the view non-exact. `compact_row_views` derives offsets from `sizes` + /// alone, which only holds after `MakeExact` reorders the elements. + #[test] + fn collects_out_of_order_list_views() -> VortexResult<()> { + let points = nullable_point_column(vec![ + Some((0.0, 4.0)), + None, + Some((2.0, 6.0)), + Some((3.0, 7.0)), + ])?; + let views = ListViewArray::try_new( + points, + PrimitiveArray::from_iter([2u32, 0]).into_array(), + PrimitiveArray::from_iter([2u32, 2]).into_array(), + Validity::NonNullable, + )?; + assert!( + !views.is_zero_copy_to_list(), + "out-of-order row views are not zero-copy to a list" + ); + let expected = multipoint_column(vec![vec![(2.0, 6.0), (3.0, 7.0)], vec![(0.0, 4.0)]])?; + + assert_collects(views.into_array(), expected) } #[test] @@ -443,11 +481,8 @@ mod tests { &[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(()) + assert_collects(input, expected) } #[test] @@ -464,11 +499,8 @@ mod tests { &[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(()) + assert_collects(input, expected) } #[test] @@ -502,22 +534,16 @@ mod tests { 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(()) + assert_collects(input, expected) } #[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(()) + assert_collects(input, expected) } #[test] @@ -532,11 +558,8 @@ mod tests { 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(()) + assert_collects(input, expected) } #[test]