diff --git a/vortex-spatial/Cargo.toml b/vortex-spatial/Cargo.toml index 8b790da3c05..a6c9cbfbfc3 100644 --- a/vortex-spatial/Cargo.toml +++ b/vortex-spatial/Cargo.toml @@ -58,5 +58,9 @@ harness = false name = "distance" harness = false +[[bench]] +name = "make_line" +harness = false + [lints] workspace = true diff --git a/vortex-spatial/benches/make_line.rs b/vortex-spatial/benches/make_line.rs new file mode 100644 index 00000000000..0d87df3dc59 --- /dev/null +++ b/vortex-spatial/benches/make_line.rs @@ -0,0 +1,122 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Microbenchmarks for native `ST_MakeLine`. +//! +//! The cases cover the normal paired-column operation, a broadcast point constant, and strict +//! null propagation. They execute the result to its canonical representation so the benchmark +//! includes construction of the two-vertex line storage. +//! +//! `ROWS` keeps each case near the roughly 1 ms iteration budget recommended for CodSpeed. +//! +//! Run with `cargo bench -p vortex-spatial --bench make_line`. + +#![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::ConstantArray; +use vortex_session::VortexSession; +use vortex_spatial::scalar_fn::make_line::SpatialMakeLine; +use vortex_spatial::test_harness::nullable_point_column; +use vortex_spatial::test_harness::point_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(); +} + +/// Deterministic pseudo-random value in `[0, 1)`. +fn unit(i: usize) -> f64 { + ((i.wrapping_mul(2_654_435_761) >> 8) % 10_000) as f64 / 10_000.0 +} + +fn points(offset: usize) -> ArrayRef { + let xs = (0..ROWS) + .map(|i| 300.0 * unit(i + offset) - 150.0) + .collect(); + let ys = (0..ROWS) + .map(|i| 300.0 * unit(i + offset + 1) - 150.0) + .collect(); + point_column(xs, ys).unwrap() +} + +fn nullable_points(offset: usize, null_every: usize) -> ArrayRef { + nullable_point_column( + (0..ROWS) + .map(|i| { + (!i.is_multiple_of(null_every)).then(|| { + ( + 300.0 * unit(i + offset) - 150.0, + 300.0 * unit(i + offset + 1) - 150.0, + ) + }) + }) + .collect(), + ) + .unwrap() +} + +fn point_constant(ctx: &mut ExecutionCtx) -> ArrayRef { + let scalar = point_column(vec![0.0], vec![0.0]) + .unwrap() + .execute_scalar(0, ctx) + .unwrap(); + ConstantArray::new(scalar, ROWS).into_array() +} + +fn make_lines(starts: &ArrayRef, ends: &ArrayRef, ctx: &mut ExecutionCtx) -> ArrayRef { + SpatialMakeLine::try_new_array(starts.clone(), ends.clone()) + .unwrap() + .into_array() + .execute::(ctx) + .unwrap() + .into_array() +} + +#[divan::bench] +fn column_x_column(bencher: Bencher) { + let starts = points(0); + let ends = points(97); + let mut ctx = SESSION.create_execution_ctx(); + bencher + .counter(ItemsCount::new(ROWS)) + .bench_local(|| make_lines(&starts, &ends, &mut ctx)); +} + +#[divan::bench] +fn column_x_constant(bencher: Bencher) { + let starts = points(0); + let mut ctx = SESSION.create_execution_ctx(); + let end = point_constant(&mut ctx); + bencher + .counter(ItemsCount::new(ROWS)) + .bench_local(|| make_lines(&starts, &end, &mut ctx)); +} + +#[divan::bench] +fn nullable_columns(bencher: Bencher) { + let starts = nullable_points(0, 8); + let ends = nullable_points(97, 11); + let mut ctx = SESSION.create_execution_ctx(); + bencher + .counter(ItemsCount::new(ROWS)) + .bench_local(|| make_lines(&starts, &ends, &mut ctx)); +} diff --git a/vortex-spatial/src/extension/coordinate.rs b/vortex-spatial/src/extension/coordinate.rs index dc3537cfb30..43599a705b7 100644 --- a/vortex-spatial/src/extension/coordinate.rs +++ b/vortex-spatial/src/extension/coordinate.rs @@ -72,6 +72,21 @@ impl Dimension { Dimension::Xyzm => &["x", "y", "z", "m"], } } + + /// Promote two coordinate dimensions to the smallest dimension that represents both. + /// + /// Missing `z`/`m` ordinates are materialized as zero when values are converted to this + /// dimension, matching DuckDB Spatial's `ST_MakeLine` promotion. + pub(crate) fn promote(self, other: Self) -> Self { + match (self, other) { + (Self::Xyzm, _) | (_, Self::Xyzm) | (Self::Xyz, Self::Xym) | (Self::Xym, Self::Xyz) => { + Self::Xyzm + } + (Self::Xyz, _) | (_, Self::Xyz) => Self::Xyz, + (Self::Xym, _) | (_, Self::Xym) => Self::Xym, + (Self::Xy, Self::Xy) => Self::Xy, + } + } } impl From for Dimension { diff --git a/vortex-spatial/src/extension/linestring.rs b/vortex-spatial/src/extension/linestring.rs index 8de97ed386e..4a820b4dc06 100644 --- a/vortex-spatial/src/extension/linestring.rs +++ b/vortex-spatial/src/extension/linestring.rs @@ -22,14 +22,22 @@ use prost::Message; 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::InterleaveArray; +use vortex_array::arrays::ListArray; +use vortex_array::arrays::PrimitiveArray; +use vortex_array::arrays::StructArray; use vortex_array::arrays::extension::ExtensionArrayExt; +use vortex_array::arrays::struct_::StructArrayExt; use vortex_array::dtype::DType; +use vortex_array::dtype::FieldNames; use vortex_array::dtype::Nullability; use vortex_array::dtype::extension::ExtDType; use vortex_array::dtype::extension::ExtId; use vortex_array::dtype::extension::ExtVTable; use vortex_array::scalar::ScalarValue; +use vortex_array::validity::Validity; use vortex_arrow::ArrowExport; use vortex_arrow::ArrowExportVTable; use vortex_arrow::ArrowImport; @@ -38,10 +46,12 @@ use vortex_arrow::ArrowSession; use vortex_arrow::ArrowSessionExt; use vortex_arrow::FromArrowArray; use vortex_arrow::FromArrowType; +use vortex_buffer::Buffer; use vortex_error::VortexError; use vortex_error::VortexResult; use vortex_error::vortex_bail; use vortex_error::vortex_ensure; +use vortex_error::vortex_ensure_eq; use vortex_error::vortex_err; use vortex_session::registry::CachedId; use vortex_session::registry::Id; @@ -103,6 +113,76 @@ pub(crate) fn linestring_dimension(dtype: &DType) -> VortexResult { coordinate_dimension(coords) } +/// Return one coordinate ordinate, filling an ordinate absent from the point dimension with zero. +fn point_ordinate( + points: &StructArray, + dimension: Dimension, + name: &str, +) -> VortexResult { + if dimension.field_names().contains(&name) { + points.unmasked_field_by_name(name).cloned() + } else { + Ok(ConstantArray::new(0.0f64, points.len()).into_array()) + } +} + +/// Build one native [`LineString`] per corresponding pair of point coordinate rows. +pub(crate) fn linestring_array_from_point_pairs( + ext_dtype: &ExtDType, + starts: &StructArray, + ends: &StructArray, + validity: Validity, +) -> VortexResult { + let len = starts.len(); + vortex_ensure_eq!( + len, + ends.len(), + "spatial: line string point columns must have equal lengths" + ); + let vertex_count = len + .checked_mul(2) + .ok_or_else(|| vortex_err!("spatial: two-vertex line string length overflow"))?; + let last_offset = i32::try_from(vertex_count) + .map_err(|_| vortex_err!("spatial: two-vertex line string offset overflow"))?; + let row_count = u32::try_from(len) + .map_err(|_| vortex_err!("spatial: two-vertex line string row count overflow"))?; + let dimension = linestring_dimension(ext_dtype.storage_dtype())?; + let start_dimension = coordinate_dimension(starts.dtype())?; + let end_dimension = coordinate_dimension(ends.dtype())?; + + let array_indices = PrimitiveArray::from_iter((0..len).flat_map(|_| [0u8, 1])).into_array(); + let row_indices = Buffer::from_iter((0..row_count).flat_map(|row| [row; 2])).into_array(); + + let ordinates = dimension + .field_names() + .iter() + .map(|name| { + Ok(InterleaveArray::try_new( + vec![ + point_ordinate(starts, start_dimension, name)?, + point_ordinate(ends, end_dimension, name)?, + ], + array_indices.clone(), + row_indices.clone(), + )? + .into_array()) + }) + .collect::>>()?; + + let vertices = StructArray::try_new( + FieldNames::from(dimension.field_names()), + ordinates, + vertex_count, + Validity::NonNullable, + )? + .into_array(); + + let offsets = Buffer::from_iter((0..=last_offset).step_by(2)).into_array(); + let storage = ListArray::try_new(vertices, offsets, validity)?.into_array(); + + Ok(ExtensionArray::try_new(ext_dtype.clone().erased(), storage)?.into_array()) +} + static ARROW_LINESTRING: CachedId = CachedId::new(LineStringType::NAME); /// The `geoarrow.linestring` extension type for `dimension`, with separated (struct) coordinates diff --git a/vortex-spatial/src/lib.rs b/vortex-spatial/src/lib.rs index 70067f57b16..318e6736e8f 100644 --- a/vortex-spatial/src/lib.rs +++ b/vortex-spatial/src/lib.rs @@ -25,6 +25,7 @@ use crate::scalar_fn::contains::SpatialContains; use crate::scalar_fn::distance::SpatialDistance; use crate::scalar_fn::envelope::SpatialEnvelope; use crate::scalar_fn::intersects::SpatialIntersects; +use crate::scalar_fn::make_line::SpatialMakeLine; pub mod aggregate_fn; pub mod extension; @@ -68,6 +69,7 @@ pub fn initialize(session: &VortexSession) { session.scalar_fns().register(SpatialContains); session.scalar_fns().register(SpatialDistance); session.scalar_fns().register(SpatialIntersects); + session.scalar_fns().register(SpatialMakeLine); // The axis-aligned bounding-box (AABB) aggregate; self-declares as a per-chunk zone stat for // geometry columns. diff --git a/vortex-spatial/src/scalar_fn/execute.rs b/vortex-spatial/src/scalar_fn/execute.rs index ca5b4018249..c83e03474b1 100644 --- a/vortex-spatial/src/scalar_fn/execute.rs +++ b/vortex-spatial/src/scalar_fn/execute.rs @@ -3,18 +3,19 @@ //! Shared execution for native geometry scalar functions. //! -//! [`dispatch_unary`] and the binary dispatcher handle constant/column operands and strict null +//! [`dispatch_unary`] and [`dispatch_binary`] handle constant/column operands and strict null //! propagation without prescribing how a kernel represents geometries or builds its output. -//! Native columnar kernels such as `ST_Envelope` use the unary dispatcher directly. +//! Native columnar kernels such as `ST_MakeLine` use these dispatchers directly. //! -//! [`execute_binary_geo_types`] adapts row-oriented algorithms from the `geo` ecosystem. It decodes -//! valid inputs into `geo_types::Geometry`; the final output is still a Vortex [`ArrayRef`], such -//! as an `f64` or boolean array. +//! [`execute_binary_geo_types`] is a convenience adapter for row-oriented algorithms from the +//! `geo` ecosystem. It decodes valid inputs into `geo_types::Geometry`; the final output is still +//! a Vortex [`ArrayRef`], such as an `f64` or boolean array. mod binary; mod geo_types; mod unary; +pub(crate) use binary::dispatch_binary; pub(crate) use binary::execute_binary_geo_types; pub(crate) use unary::dispatch_unary; use vortex_array::ArrayRef; diff --git a/vortex-spatial/src/scalar_fn/make_line.rs b/vortex-spatial/src/scalar_fn/make_line.rs new file mode 100644 index 00000000000..c32d4c76691 --- /dev/null +++ b/vortex-spatial/src/scalar_fn/make_line.rs @@ -0,0 +1,542 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! `ST_MakeLine`: construct a native line string between two native points. + +use vortex_array::ArrayRef; +use vortex_array::ExecutionCtx; +use vortex_array::IntoArray; +use vortex_array::arrays::ConstantArray; +use vortex_array::arrays::ScalarFnArray; +use vortex_array::arrays::StructArray; +use vortex_array::dtype::DType; +use vortex_array::dtype::Nullability; +use vortex_array::dtype::extension::ExtDType; +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_error::VortexResult; +use vortex_error::vortex_ensure; +use vortex_error::vortex_err; +use vortex_mask::Mask; +use vortex_session::VortexSession; +use vortex_session::registry::CachedId; + +use crate::extension::LineString; +use crate::extension::Point; +use crate::extension::SpatialMetadata; +use crate::extension::coordinate::coordinate_dimension; +use crate::extension::flatten_coordinates; +use crate::extension::linestring_array_from_point_pairs; +use crate::extension::linestring_storage_dtype; +use crate::scalar_fn::execute::Execution; +use crate::scalar_fn::execute::Operand; +use crate::scalar_fn::execute::dispatch_binary; + +/// Validate the two point operands accepted by `ST_MakeLine`. +fn validate_make_line_operands(dtypes: &[DType]) -> VortexResult<()> { + vortex_ensure!( + dtypes.len() == 2, + "spatial: make_line requires exactly two point operands, got {}", + dtypes.len() + ); + for dtype in dtypes { + vortex_ensure!( + dtype + .as_extension_opt() + .is_some_and(|extension| extension.is::()), + "spatial: make_line operand {dtype} is not a native point" + ); + } + Ok(()) +} + +/// Resolve DuckDB's `ST_MakeLine` CRS propagation for two geometry operands. +fn make_line_metadata( + left: &SpatialMetadata, + right: &SpatialMetadata, +) -> VortexResult { + match (&left.crs, &right.crs) { + (Some(left_crs), Some(right_crs)) => { + vortex_ensure!( + left_crs == right_crs, + "spatial: make_line operands have different coordinate reference systems: \ + {left_crs} and {right_crs}" + ); + Ok(left.clone()) + } + (Some(_), None) => Ok(left.clone()), + (None, Some(_)) => Ok(right.clone()), + (None, None) => Ok(SpatialMetadata::default()), + } +} + +/// The native `LineString` dtype emitted by `ST_MakeLine`. +fn make_line_dtype(dtypes: &[DType]) -> VortexResult> { + validate_make_line_operands(dtypes)?; + let left = dtypes[0].as_extension(); + let right = dtypes[1].as_extension(); + let dimension = coordinate_dimension(left.storage_dtype())? + .promote(coordinate_dimension(right.storage_dtype())?); + let metadata = make_line_metadata(left.metadata::(), right.metadata::())?; + let nullability = Nullability::from(dtypes.iter().any(DType::is_nullable)); + ExtDType::try_new(metadata, linestring_storage_dtype(dimension, nullability)) +} + +/// Expose a point operand as coordinate fields without expanding a constant into value buffers. +fn point_coordinates( + operand: Operand, + len: usize, + ctx: &mut ExecutionCtx, +) -> VortexResult { + match operand { + Operand::Column(points) => flatten_coordinates(&points, ctx), + Operand::Constant(point) => { + let storage = point.as_extension().to_storage_scalar(); + let fields = storage.as_struct(); + let names = fields.names().clone(); + let arrays = names + .iter() + .map(|name| { + fields + .field(name) + .map(|value| ConstantArray::new(value, len).into_array()) + .ok_or_else(|| vortex_err!("spatial: point coordinate missing {name}")) + }) + .collect::>>()?; + StructArray::try_new(names, arrays, len, Validity::NonNullable) + } + } +} + +/// Build a native line-string column from dispatched point operands. +fn build_make_lines( + operands: [Operand; 2], + len: usize, + valid: Mask, + output_dtype: &ExtDType, + ctx: &mut ExecutionCtx, +) -> VortexResult { + let [start, end] = operands; + let starts = point_coordinates(start, len, ctx)?; + let ends = point_coordinates(end, len, ctx)?; + linestring_array_from_point_pairs( + output_dtype, + &starts, + &ends, + Validity::from_mask(valid, output_dtype.storage_dtype().nullability()), + ) +} + +/// Execute `ST_MakeLine` after shared constant/column and null dispatch. +fn execute_make_line( + execution: Execution<2>, + output_dtype: &ExtDType, + ctx: &mut ExecutionCtx, +) -> VortexResult { + match execution.operands { + [Operand::Constant(start), Operand::Constant(end)] => { + let one = build_make_lines( + [Operand::Constant(start), Operand::Constant(end)], + 1, + Mask::new_true(1), + output_dtype, + ctx, + )?; + Ok(ConstantArray::new(one.execute_scalar(0, ctx)?, execution.len).into_array()) + } + operands => build_make_lines(operands, execution.len, execution.valid, output_dtype, ctx), + } +} + +/// Construct `LineString`s from paired native point operands. The output's vertices preserve all +/// coordinate ordinates (`x`, `y`, and any `z`/`m`) and appear in operand order. When the points +/// have different dimensions, it promotes to their union and fills absent `z`/`m` ordinates with +/// zero. +#[derive(Debug, Clone, Default, PartialEq, Eq, Hash)] +pub struct SpatialMakeLine; + +impl SpatialMakeLine { + /// A lazy `ScalarFnArray` constructing one two-vertex line string per pair of point operands. + pub fn try_new_array(a: ArrayRef, b: ArrayRef) -> VortexResult { + ScalarFnArray::try_new( + TypedScalarFnInstance::new(SpatialMakeLine, EmptyOptions).erased(), + vec![a, b], + ) + } +} + +impl ScalarFnVTable for SpatialMakeLine { + type Options = EmptyOptions; + + fn id(&self) -> ScalarFnId { + static ID: CachedId = CachedId::new("vortex.st.make_line"); + *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(2) + } + + fn child_name(&self, _: &Self::Options, child_idx: usize) -> ChildName { + match child_idx { + 0 => ChildName::from("start"), + 1 => ChildName::from("end"), + _ => unreachable!("make_line has exactly two children"), + } + } + + fn return_dtype(&self, _: &Self::Options, dtypes: &[DType]) -> VortexResult { + Ok(DType::Extension(make_line_dtype(dtypes)?.erased())) + } + + fn execute( + &self, + _: &Self::Options, + args: &dyn ExecutionArgs, + ctx: &mut ExecutionCtx, + ) -> VortexResult { + let a = args.get(0)?; + let b = args.get(1)?; + let output_dtype = make_line_dtype(&[a.dtype().clone(), b.dtype().clone()])?; + dispatch_binary( + &a, + &b, + DType::Extension(output_dtype.clone().erased()), + |execution, ctx| execute_make_line(execution, &output_dtype, 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 geo_types::Coord; + use geo_types::Geometry; + use geo_types::LineString as GeoLineString; + use rstest::rstest; + use vortex_array::ArrayRef; + use vortex_array::Columnar; + use vortex_array::ExecutionCtx; + use vortex_array::IntoArray; + use vortex_array::VortexSessionExecute; + use vortex_array::arrays::ConstantArray; + use vortex_array::arrays::ExtensionArray; + use vortex_array::arrays::PrimitiveArray; + use vortex_array::arrays::StructArray; + use vortex_array::dtype::DType; + use vortex_array::dtype::extension::ExtDType; + use vortex_array::scalar::Scalar; + use vortex_array::scalar_fn::EmptyOptions; + use vortex_array::scalar_fn::ScalarFnVTable; + use vortex_error::VortexResult; + use vortex_error::vortex_err; + + use super::SpatialMakeLine; + use crate::extension::LineString; + use crate::extension::Point; + use crate::extension::SpatialMetadata; + use crate::extension::coordinate::Dimension; + use crate::extension::coordinate::coordinate_dimension; + use crate::extension::coordinate::ordinates; + use crate::extension::flatten_coordinates; + use crate::extension::geometries; + use crate::test_harness::point_column; + + fn dimensional_point( + dimension: Dimension, + coordinate: [f64; 4], + crs: Option<&str>, + ) -> VortexResult { + let mut fields = vec![ + ("x", PrimitiveArray::from_iter([coordinate[0]]).into_array()), + ("y", PrimitiveArray::from_iter([coordinate[1]]).into_array()), + ]; + if matches!(dimension, Dimension::Xyz | Dimension::Xyzm) { + fields.push(("z", PrimitiveArray::from_iter([coordinate[2]]).into_array())); + } + if matches!(dimension, Dimension::Xym | Dimension::Xyzm) { + fields.push(("m", PrimitiveArray::from_iter([coordinate[3]]).into_array())); + } + let storage = StructArray::from_fields(&fields)?.into_array(); + let dtype = ExtDType::::try_new( + SpatialMetadata { + crs: crs.map(str::to_owned), + }, + storage.dtype().clone(), + )?; + Ok(ExtensionArray::try_new(dtype.erased(), storage)?.into_array()) + } + + fn point_constant( + x: f64, + y: f64, + len: usize, + ctx: &mut ExecutionCtx, + ) -> VortexResult { + let scalar = point_column(vec![x], vec![y])?.execute_scalar(0, ctx)?; + Ok(ConstantArray::new(scalar, len).into_array()) + } + + #[test] + fn connects_paired_points_in_operand_order() -> VortexResult<()> { + let session = vortex_array::array_session(); + let mut ctx = session.create_execution_ctx(); + let starts = point_column(vec![0.0, 3.0], vec![0.0, 4.0])?; + let ends = point_column(vec![3.0, 0.0], vec![4.0, 0.0])?; + + let lines = SpatialMakeLine::try_new_array(starts, ends)?.into_array(); + assert!(lines.dtype().as_extension().is::()); + assert_eq!( + geometries(&lines, &mut ctx)?, + vec![ + Geometry::LineString(GeoLineString::new(vec![ + Coord { x: 0.0, y: 0.0 }, + Coord { x: 3.0, y: 4.0 }, + ])), + Geometry::LineString(GeoLineString::new(vec![ + Coord { x: 3.0, y: 4.0 }, + Coord { x: 0.0, y: 0.0 }, + ])), + ] + ); + Ok(()) + } + + #[test] + fn two_constants_are_built_once_and_remain_constant() -> VortexResult<()> { + let session = vortex_array::array_session(); + let mut ctx = session.create_execution_ctx(); + let starts = point_constant(0.0, 0.0, 3, &mut ctx)?; + let ends = point_constant(3.0, 4.0, 3, &mut ctx)?; + + let result = SpatialMakeLine::try_new_array(starts, ends)? + .into_array() + .execute::(&mut ctx)?; + let Columnar::Constant(lines) = result else { + return Err(vortex_err!( + "make_line of two constants should remain constant" + )); + }; + assert_eq!(lines.len(), 3); + assert_eq!( + geometries(&lines.into_array(), &mut ctx)?, + vec![ + Geometry::LineString(GeoLineString::new(vec![ + Coord { x: 0.0, y: 0.0 }, + Coord { x: 3.0, y: 4.0 }, + ])); + 3 + ] + ); + Ok(()) + } + + #[rstest] + #[case::constant_start(true)] + #[case::constant_end(false)] + fn constant_and_column_are_paired_by_row(#[case] constant_start: bool) -> VortexResult<()> { + let session = vortex_array::array_session(); + let mut ctx = session.create_execution_ctx(); + let constant = point_constant(0.0, 0.0, 2, &mut ctx)?; + let column = point_column(vec![3.0, 6.0], vec![4.0, 8.0])?; + let (starts, ends) = if constant_start { + (constant, column) + } else { + (column, constant) + }; + + let lines = SpatialMakeLine::try_new_array(starts, ends)?.into_array(); + let endpoints = [(3.0, 4.0), (6.0, 8.0)]; + let expected = endpoints + .into_iter() + .map(|(x, y)| { + let constant = Coord { x: 0.0, y: 0.0 }; + let column = Coord { x, y }; + Geometry::LineString(GeoLineString::new(if constant_start { + vec![constant, column] + } else { + vec![column, constant] + })) + }) + .collect::>(); + assert_eq!(geometries(&lines, &mut ctx)?, expected); + Ok(()) + } + + #[test] + fn null_constant_is_all_null() -> VortexResult<()> { + let session = vortex_array::array_session(); + let mut ctx = session.create_execution_ctx(); + let point_dtype = point_column(vec![0.0], vec![0.0])?.dtype().as_nullable(); + let starts = ConstantArray::new(Scalar::null(point_dtype), 2).into_array(); + let ends = point_column(vec![3.0, 6.0], vec![4.0, 8.0])?; + + let result = SpatialMakeLine::try_new_array(starts, ends)? + .into_array() + .execute::(&mut ctx)?; + let Columnar::Constant(lines) = result else { + return Err(vortex_err!( + "make_line with a null constant should remain constant" + )); + }; + assert_eq!(lines.len(), 2); + assert!(lines.scalar().is_null()); + assert!(lines.dtype().as_extension().is::()); + Ok(()) + } + + #[rstest] + #[case::xy_xyz( + Dimension::Xy, + [1.0, 2.0, 0.0, 0.0], + Dimension::Xyz, + [3.0, 4.0, 5.0, 0.0], + Dimension::Xyz, + Some([0.0, 5.0]), + None + )] + #[case::xyz_xy( + Dimension::Xyz, + [1.0, 2.0, 5.0, 0.0], + Dimension::Xy, + [3.0, 4.0, 0.0, 0.0], + Dimension::Xyz, + Some([5.0, 0.0]), + None + )] + #[case::xym_xyz( + Dimension::Xym, + [1.0, 2.0, 0.0, 6.0], + Dimension::Xyz, + [3.0, 4.0, 5.0, 0.0], + Dimension::Xyzm, + Some([0.0, 5.0]), + Some([6.0, 0.0]) + )] + #[case::xyz_xym( + Dimension::Xyz, + [1.0, 2.0, 5.0, 0.0], + Dimension::Xym, + [3.0, 4.0, 0.0, 6.0], + Dimension::Xyzm, + Some([5.0, 0.0]), + Some([0.0, 6.0]) + )] + fn promotes_mixed_point_dimensions( + #[case] start_dimension: Dimension, + #[case] start: [f64; 4], + #[case] end_dimension: Dimension, + #[case] end: [f64; 4], + #[case] expected_dimension: Dimension, + #[case] expected_z: Option<[f64; 2]>, + #[case] expected_m: Option<[f64; 2]>, + ) -> VortexResult<()> { + let session = vortex_array::array_session(); + let mut ctx = session.create_execution_ctx(); + let expected_x = [start[0], end[0]]; + let expected_y = [start[1], end[1]]; + let start = dimensional_point(start_dimension, start, None)?; + let end = dimensional_point(end_dimension, end, None)?; + + let lines = SpatialMakeLine::try_new_array(start, end)?.into_array(); + let vertices = flatten_coordinates(&lines, &mut ctx)?; + assert_eq!(coordinate_dimension(vertices.dtype())?, expected_dimension); + for (name, expected) in [("x", expected_x), ("y", expected_y)] { + assert_eq!(ordinates(&vertices, name, &mut ctx)?.as_slice(), expected); + } + for (name, expected) in [("z", expected_z), ("m", expected_m)] { + if let Some(expected) = expected { + assert_eq!(ordinates(&vertices, name, &mut ctx)?.as_slice(), expected); + } + } + Ok(()) + } + + #[rstest] + #[case::matching(Some("EPSG:4326"), Some("EPSG:4326"), Some("EPSG:4326"))] + #[case::left_only(Some("EPSG:4326"), None, Some("EPSG:4326"))] + #[case::right_only(None, Some("EPSG:3857"), Some("EPSG:3857"))] + #[case::unreferenced(None, None, None)] + fn propagates_compatible_crs( + #[case] left_crs: Option<&str>, + #[case] right_crs: Option<&str>, + #[case] expected: Option<&str>, + ) -> VortexResult<()> { + let left = dimensional_point(Dimension::Xy, [0.0; 4], left_crs)?; + let right = dimensional_point(Dimension::Xy, [1.0; 4], right_crs)?; + + let dtype = SpatialMakeLine.return_dtype( + &EmptyOptions, + &[left.dtype().clone(), right.dtype().clone()], + )?; + assert_eq!( + dtype.as_extension().metadata::().crs.as_deref(), + expected + ); + Ok(()) + } + + #[test] + fn rejects_mismatched_crs() -> VortexResult<()> { + let left = dimensional_point(Dimension::Xy, [0.0; 4], Some("EPSG:4326"))?; + let right = dimensional_point(Dimension::Xy, [1.0; 4], Some("EPSG:3857"))?; + assert!( + SpatialMakeLine + .return_dtype( + &EmptyOptions, + &[left.dtype().clone(), right.dtype().clone()] + ) + .is_err() + ); + Ok(()) + } + + #[test] + fn rejects_non_point_dtype() -> VortexResult<()> { + let points = point_column(vec![0.0], vec![0.0])?; + assert!( + SpatialMakeLine + .return_dtype(&EmptyOptions, std::slice::from_ref(points.dtype())) + .is_err() + ); + let non_point = DType::Bool(vortex_array::dtype::Nullability::NonNullable); + assert!( + SpatialMakeLine + .return_dtype(&EmptyOptions, &[points.dtype().clone(), non_point]) + .is_err() + ); + Ok(()) + } +} diff --git a/vortex-spatial/src/scalar_fn/mod.rs b/vortex-spatial/src/scalar_fn/mod.rs index e6770be4fff..fda570a2c1e 100644 --- a/vortex-spatial/src/scalar_fn/mod.rs +++ b/vortex-spatial/src/scalar_fn/mod.rs @@ -8,3 +8,4 @@ pub mod distance; pub mod envelope; mod execute; pub mod intersects; +pub mod make_line;