From 09f16ffee6f2984a6e95f780ecd5e0b73ec55930 Mon Sep 17 00:00:00 2001 From: Nemo Yu Date: Wed, 5 Aug 2026 14:50:22 -0400 Subject: [PATCH] refactor(vortex-geo): generalize scalar execution Signed-off-by: Nemo Yu --- vortex-geo/src/extension/mod.rs | 16 - vortex-geo/src/extension/rect.rs | 20 + vortex-geo/src/scalar_fn/contains.rs | 25 +- vortex-geo/src/scalar_fn/distance.rs | 25 +- vortex-geo/src/scalar_fn/envelope.rs | 173 +++++-- vortex-geo/src/scalar_fn/execute.rs | 470 ++---------------- vortex-geo/src/scalar_fn/execute/binary.rs | 334 +++++++++++++ vortex-geo/src/scalar_fn/execute/geo_types.rs | 144 ++++++ vortex-geo/src/scalar_fn/execute/unary.rs | 62 +++ vortex-geo/src/scalar_fn/intersects.rs | 25 +- 10 files changed, 772 insertions(+), 522 deletions(-) create mode 100644 vortex-geo/src/scalar_fn/execute/binary.rs create mode 100644 vortex-geo/src/scalar_fn/execute/geo_types.rs create mode 100644 vortex-geo/src/scalar_fn/execute/unary.rs diff --git a/vortex-geo/src/extension/mod.rs b/vortex-geo/src/extension/mod.rs index b05069f8ca5..374ddfc466f 100644 --- a/vortex-geo/src/extension/mod.rs +++ b/vortex-geo/src/extension/mod.rs @@ -61,7 +61,6 @@ use vortex_arrow::FromArrowArray; use vortex_buffer::Buffer; use vortex_error::VortexResult; use vortex_error::vortex_bail; -use vortex_error::vortex_ensure; use vortex_error::vortex_err; pub use wkb::*; @@ -78,21 +77,6 @@ pub(crate) fn is_native_geometry(dtype: &DType) -> bool { }) } -/// Validate the operands of a geo scalar function: each must be a native geometry type so the -/// kernel can decode it. The two operands need not share a geometry type — e.g. a `Point` against -/// a `Polygon` is valid, since distance/containment/intersection across types is meaningful. -/// Nullable operands are allowed; the kernels propagate nulls (a null geometry input yields a null -/// result) rather than decoding null rows. -pub(crate) fn validate_geometry_operands(dtypes: &[DType]) -> VortexResult<()> { - for dtype in dtypes { - vortex_ensure!( - is_native_geometry(dtype), - "geo: operand {dtype} is not a native geometry type" - ); - } - Ok(()) -} - /// Flatten a native geometry column into a single coordinate `Struct` containing /// every vertex of every geometry. pub(crate) fn flatten_coordinates( diff --git a/vortex-geo/src/extension/rect.rs b/vortex-geo/src/extension/rect.rs index 4c373a1807f..19f7f9db61e 100644 --- a/vortex-geo/src/extension/rect.rs +++ b/vortex-geo/src/extension/rect.rs @@ -26,6 +26,7 @@ use vortex_array::ArrayRef; use vortex_array::ExecutionCtx; use vortex_array::IntoArray; use vortex_array::arrays::ExtensionArray; +use vortex_array::arrays::StructArray; use vortex_array::arrays::extension::ExtensionArrayExt; use vortex_array::dtype::DType; use vortex_array::dtype::FieldNames; @@ -36,6 +37,7 @@ 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; @@ -155,6 +157,24 @@ pub(crate) fn box_dimension(dtype: &DType) -> VortexResult { .ok_or_else(|| vortex_err!("not a valid geoarrow.box dimension: {:?}", fields.names())) } +/// Build a native [`Rect`] array from canonical min/max ordinate columns. +pub(crate) fn build_rect_array( + ext_dtype: &ExtDType, + corners: Vec, + len: usize, + validity: Validity, +) -> VortexResult { + let dimension = box_dimension(ext_dtype.storage_dtype())?; + let storage = StructArray::try_new( + FieldNames::from(box_field_names(dimension)), + corners, + len, + validity, + )? + .into_array(); + Ok(ExtensionArray::try_new(ext_dtype.clone().erased(), storage)?.into_array()) +} + static ARROW_BOX: CachedId = CachedId::new(BoxType::NAME); /// The `geoarrow.box` extension type for `dimension`. diff --git a/vortex-geo/src/scalar_fn/contains.rs b/vortex-geo/src/scalar_fn/contains.rs index 855a6af1867..a10bd7c4524 100644 --- a/vortex-geo/src/scalar_fn/contains.rs +++ b/vortex-geo/src/scalar_fn/contains.rs @@ -19,11 +19,28 @@ use vortex_array::scalar_fn::ScalarFnId; use vortex_array::scalar_fn::ScalarFnVTable; use vortex_array::scalar_fn::TypedScalarFnInstance; use vortex_error::VortexResult; +use vortex_error::vortex_ensure; use vortex_session::VortexSession; use vortex_session::registry::CachedId; -use crate::extension::validate_geometry_operands; -use crate::scalar_fn::execute::execute_null_propagating; +use crate::extension::is_native_geometry; +use crate::scalar_fn::execute::execute_binary_geo_types; + +/// Validate the two native geometry operands accepted by `ST_Contains`. +fn validate_contains_operands(dtypes: &[DType]) -> VortexResult<()> { + vortex_ensure!( + dtypes.len() == 2, + "geo: contains requires exactly two geometry operands, got {}", + dtypes.len() + ); + for dtype in dtypes { + vortex_ensure!( + is_native_geometry(dtype), + "geo: contains operand {dtype} is not a native geometry type" + ); + } + Ok(()) +} /// OGC `ST_Contains` between two native geometry operands, each a column or a constant /// literal: true where operand `b` lies completely inside operand `a` (boundary contact alone @@ -71,7 +88,7 @@ impl ScalarFnVTable for GeoContains { } fn return_dtype(&self, _: &Self::Options, dtypes: &[DType]) -> VortexResult { - validate_geometry_operands(dtypes)?; + validate_contains_operands(dtypes)?; let nullability = Nullability::from(dtypes.iter().any(DType::is_nullable)); Ok(DType::Bool(nullability)) } @@ -87,7 +104,7 @@ impl ScalarFnVTable for GeoContains { // Containment is not symmetric: `a` is always the container and `b` the contained. A // container's rect must cover the contained's rect (`Rect::contains` is the closed // test), so a contained rect poking outside proves the row false. - execute_null_propagating( + execute_binary_geo_types( &a, &b, |a, b| a.contains(b), diff --git a/vortex-geo/src/scalar_fn/distance.rs b/vortex-geo/src/scalar_fn/distance.rs index dd209a712b4..cd9475f69ba 100644 --- a/vortex-geo/src/scalar_fn/distance.rs +++ b/vortex-geo/src/scalar_fn/distance.rs @@ -21,11 +21,28 @@ use vortex_array::scalar_fn::ScalarFnId; use vortex_array::scalar_fn::ScalarFnVTable; use vortex_array::scalar_fn::TypedScalarFnInstance; use vortex_error::VortexResult; +use vortex_error::vortex_ensure; use vortex_session::VortexSession; use vortex_session::registry::CachedId; -use crate::extension::validate_geometry_operands; -use crate::scalar_fn::execute::execute_null_propagating; +use crate::extension::is_native_geometry; +use crate::scalar_fn::execute::execute_binary_geo_types; + +/// Validate the two native geometry operands accepted by `ST_Distance`. +fn validate_distance_operands(dtypes: &[DType]) -> VortexResult<()> { + vortex_ensure!( + dtypes.len() == 2, + "geo: distance requires exactly two geometry operands, got {}", + dtypes.len() + ); + for dtype in dtypes { + vortex_ensure!( + is_native_geometry(dtype), + "geo: distance operand {dtype} is not a native geometry type" + ); + } + Ok(()) +} /// Planar (Euclidean) `ST_Distance` (no geodesic correction) between two native geometry /// operands, each a column or a constant literal. @@ -72,7 +89,7 @@ impl ScalarFnVTable for GeoDistance { } fn return_dtype(&self, _: &Self::Options, dtypes: &[DType]) -> VortexResult { - validate_geometry_operands(dtypes)?; + validate_distance_operands(dtypes)?; let nullability = Nullability::from(dtypes.iter().any(DType::is_nullable)); Ok(DType::Primitive(PType::F64, nullability)) } @@ -86,7 +103,7 @@ impl ScalarFnVTable for GeoDistance { let a = args.get(0)?; let b = args.get(1)?; // Distance is a value, not a verdict: no bounding-rect test can decide it. - execute_null_propagating(&a, &b, |x, y| Euclidean.distance(x, y), None, ctx) + execute_binary_geo_types(&a, &b, |x, y| Euclidean.distance(x, y), None, ctx) } fn validity( diff --git a/vortex-geo/src/scalar_fn/envelope.rs b/vortex-geo/src/scalar_fn/envelope.rs index 1859bde666c..93c88007959 100644 --- a/vortex-geo/src/scalar_fn/envelope.rs +++ b/vortex-geo/src/scalar_fn/envelope.rs @@ -9,13 +9,13 @@ 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::ScalarFnArray; 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::expr::Expression; @@ -30,6 +30,7 @@ use vortex_array::validity::Validity; use vortex_buffer::BitBuffer; use vortex_buffer::BufferMut; use vortex_error::VortexResult; +use vortex_error::vortex_ensure; use vortex_error::vortex_err; use vortex_mask::Mask; use vortex_session::VortexSession; @@ -39,11 +40,30 @@ use crate::extension::GeoMetadata; use crate::extension::Rect; use crate::extension::box_field_names; use crate::extension::box_storage_dtype; +use crate::extension::build_rect_array; use crate::extension::coordinate::Dimension; use crate::extension::coordinate::box_corners; use crate::extension::coordinate::ordinates; use crate::extension::flatten_row_offsets; -use crate::extension::validate_geometry_operands; +use crate::extension::is_native_geometry; +use crate::scalar_fn::execute::Execution; +use crate::scalar_fn::execute::Operand; +use crate::scalar_fn::execute::dispatch_unary; + +/// Validate the native geometry operand accepted by `envelope`. +fn validate_envelope_operands(dtypes: &[DType]) -> VortexResult<()> { + vortex_ensure!( + dtypes.len() == 1, + "geo: envelope requires exactly one geometry operand, got {}", + dtypes.len() + ); + vortex_ensure!( + is_native_geometry(&dtypes[0]), + "geo: envelope operand {} is not a native geometry type", + dtypes[0] + ); + Ok(()) +} /// `envelope`: the axis-aligned bounding box of each geometry in a native geometry operand (a column /// or a constant literal), as a native 2-D `geoarrow.box` ([`Rect`]) column. @@ -75,9 +95,12 @@ fn output_box_dtype() -> VortexResult> { /// Compute each row's 2-D bounding box: the smallest rectangle covering all of the row's /// coordinates. -fn row_boxes(storage: ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult<(Vec, Validity)> { +fn row_boxes( + storage: ArrayRef, + valid: &Mask, + ctx: &mut ExecutionCtx, +) -> VortexResult<(Vec, Validity)> { let len = storage.len(); - let valid = storage.validity()?.execute_mask(len, ctx)?; let (row_offsets, coords) = flatten_row_offsets(storage, ctx)?; // A row has a box iff it is valid and owns at least one coordinate (an empty geometry has @@ -110,10 +133,73 @@ fn row_boxes(storage: ArrayRef, ctx: &mut ExecutionCtx) -> VortexResult<(Vec, + ctx: &mut ExecutionCtx, +) -> VortexResult { + let len = array.len(); + let is_rect = array + .dtype() + .as_extension_opt() + .ok_or_else(|| vortex_err!("geo: envelope operand is not a geometry extension type"))? + .is::(); + let storage = array + .execute::(ctx)? + .storage_array() + .clone(); + + let (corners, output_validity) = if is_rect { + // A box is its own envelope: project the 2-D corner fields straight out of storage + // (dropping any z/m bounds). A stored box cannot be empty. + let coords = storage.execute::(ctx)?; + let corners = box_field_names(Dimension::Xy) + .iter() + .map(|name| coords.unmasked_field_by_name(name).cloned()) + .collect::>>()?; + ( + corners, + Validity::from_mask(valid.clone(), Nullability::Nullable), + ) + } else if !storage.dtype().is_list() { + // Point storage is the coordinate `Struct` itself: every row owns exactly one + // coordinate, so its box is degenerate and the corner columns are zero-copy projections. + let coords = storage.execute::(ctx)?; + let x = coords.unmasked_field_by_name("x")?.clone(); + let y = coords.unmasked_field_by_name("y")?.clone(); + ( + vec![x.clone(), y.clone(), x, y], + Validity::from_mask(valid.clone(), Nullability::Nullable), + ) + } else { + row_boxes(storage, valid, ctx)? + }; + + build_rect_array(output_dtype, corners, len, output_validity) +} + +/// Execute `envelope` after shared constant/column and null dispatch. +fn execute_envelope( + execution: Execution<1>, + output_dtype: &ExtDType, + ctx: &mut ExecutionCtx, +) -> VortexResult { + match execution.operands { + [Operand::Constant(scalar)] => { + let one = ConstantArray::new(scalar, 1).into_array(); + let output = envelope_array(one, &Mask::new_true(1), output_dtype, ctx)?; + Ok(ConstantArray::new(output.execute_scalar(0, ctx)?, execution.len).into_array()) + } + [Operand::Column(array)] => envelope_array(array, &execution.valid, output_dtype, ctx), + } +} + impl ScalarFnVTable for GeoEnvelope { type Options = EmptyOptions; @@ -142,7 +228,7 @@ impl ScalarFnVTable for GeoEnvelope { } fn return_dtype(&self, _: &Self::Options, dtypes: &[DType]) -> VortexResult { - validate_geometry_operands(dtypes)?; + validate_envelope_operands(dtypes)?; // Always nullable: an empty geometry has no box, so nulls can appear even over a // non-nullable operand. Ok(DType::Extension(output_box_dtype()?.erased())) @@ -158,53 +244,13 @@ impl ScalarFnVTable for GeoEnvelope { ctx: &mut ExecutionCtx, ) -> VortexResult { let array = args.get(0)?; - let len = array.len(); - let ext = array - .dtype() - .as_extension_opt() - .ok_or_else(|| vortex_err!("geo: envelope operand is not a geometry extension type"))?; - let storage = array - .clone() - .execute::(ctx)? - .storage_array() - .clone(); - - let (corners, output_validity) = if ext.is::() { - // A box is its own envelope: project the 2-D corner fields straight out of storage - // (dropping any z/m bounds). A stored box cannot be empty, so the output validity - // is exactly the operand's, kept lazy. - let coords = storage.execute::(ctx)?; - let corners = box_field_names(Dimension::Xy) - .iter() - .map(|name| coords.unmasked_field_by_name(name).cloned()) - .collect::>>()?; - (corners, array.validity()?.into_nullable()) - } else if !storage.dtype().is_list() { - // Point storage is the coordinate `Struct` itself: every row owns exactly one - // coordinate, so its box is degenerate and the corner columns are the ordinate - // arrays, zero-copy. No row can be empty, so the output validity is exactly the - // operand's, kept lazy. - let coords = storage.execute::(ctx)?; - let x = coords.unmasked_field_by_name("x")?.clone(); - let y = coords.unmasked_field_by_name("y")?.clone(); - ( - vec![x.clone(), y.clone(), x, y], - array.validity()?.into_nullable(), - ) - } else { - row_boxes(storage, ctx)? - }; - - // Nullness lives at the box (struct) level: the corner fields stay non-nullable `f64`, - // holding unspecified values under null rows. - let storage = StructArray::try_new( - FieldNames::from(box_field_names(Dimension::Xy)), - corners, - len, - output_validity, - )? - .into_array(); - Ok(ExtensionArray::try_new(output_box_dtype()?.erased(), storage)?.into_array()) + let output_dtype = output_box_dtype()?; + dispatch_unary( + &array, + DType::Extension(output_dtype.clone().erased()), + |execution, ctx| execute_envelope(execution, &output_dtype, ctx), + ctx, + ) } fn validity(&self, _: &Self::Options, _: &Expression) -> VortexResult> { @@ -225,6 +271,7 @@ impl ScalarFnVTable for GeoEnvelope { #[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; @@ -240,6 +287,7 @@ mod tests { use vortex_array::scalar_fn::EmptyOptions; use vortex_array::scalar_fn::ScalarFnVTable; use vortex_error::VortexResult; + use vortex_error::vortex_err; use super::GeoEnvelope; use crate::extension::GeoMetadata; @@ -495,6 +543,25 @@ mod tests { Ok(()) } + /// A non-null constant is boxed once and retained as a constant output. + #[test] + fn constant_point_remains_constant() -> VortexResult<()> { + let session = crate::test_harness::geo_session(); + let mut ctx = session.create_execution_ctx(); + + let scalar = point_column(vec![1.0], vec![2.0])?.execute_scalar(0, &mut ctx)?; + let points = ConstantArray::new(scalar, 3).into_array(); + let result = boxes(points)?.execute::(&mut ctx)?; + let Columnar::Constant(boxes) = result else { + return Err(vortex_err!("envelope of a constant should remain constant")); + }; + assert_eq!(boxes.len(), 3); + + let expected = nullable_rect_column(vec![Some((1.0, 2.0, 1.0, 2.0)); 3])?; + assert_arrays_eq!(boxes.into_array(), expected, &mut ctx); + Ok(()) + } + /// Output is always nullable, even over a non-nullable operand, since an empty geometry has no /// box. #[test] diff --git a/vortex-geo/src/scalar_fn/execute.rs b/vortex-geo/src/scalar_fn/execute.rs index b86f1ef913c..a509e48ebfd 100644 --- a/vortex-geo/src/scalar_fn/execute.rs +++ b/vortex-geo/src/scalar_fn/execute.rs @@ -1,455 +1,43 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright the Vortex contributors -//! Shared execution for the binary geo scalar functions. +//! Shared execution for native geometry scalar functions. //! -//! [`execute_null_propagating`] runs a binary geo kernel (`ST_Distance`, `ST_Intersects`, -//! `ST_Contains`) over its two operands, decoding to `geo_types` and computing per row. Nulls -//! propagate as in SQL — the result is null wherever either operand is null — which the kernels -//! also expose via `vortex_array::expr::union_child_validities` as their `validity()`, so the -//! planner can derive the output null mask without executing them. +//! [`dispatch_unary`] and the binary dispatcher 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. //! -//! When exactly one operand is a constant, a predicate kernel may pass a [`BboxReject`] -//! pre-check: the constant's bounding rect is fixed once per batch, and a row whose own rect -//! already proves the result skips the exact per-row test. +//! [`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. -use geo::BoundingRect; -use geo_types::Geometry; -use geo_types::Rect; +mod binary; +mod geo_types; +mod unary; + +pub(crate) use binary::execute_binary_geo_types; +pub(crate) use unary::dispatch_unary; use vortex_array::ArrayRef; -use vortex_array::ExecutionCtx; -use vortex_array::IntoArray; -use vortex_array::arrays::BoolArray; -use vortex_array::arrays::Constant; -use vortex_array::arrays::ConstantArray; -use vortex_array::arrays::PrimitiveArray; use vortex_array::dtype::Nullability; use vortex_array::scalar::Scalar; -use vortex_array::validity::Validity; -use vortex_buffer::BitBuffer; -use vortex_error::VortexResult; -use vortex_error::vortex_ensure_eq; -use vortex_mask::AllOr; use vortex_mask::Mask; -use crate::extension::geometries; -use crate::extension::single_geometry; - -/// The result type a binary geo kernel produces. Today that is `f64` (for `ST_Distance`) and -/// `bool` (for the `ST_Intersects` / `ST_Contains` predicates), and the trait is implemented for -/// both. A kernel that returns some other type just adds its own `impl GeoOutput`. -pub(crate) trait GeoOutput: Copy { - /// Convert this computed value into a Vortex [`Scalar`] (one typed, nullable value). Used - /// only when both operands are constant: the kernel computes a single result, and this wraps - /// it so a constant array can repeat that one value across every row. - fn into_scalar(self, nullability: Nullability) -> Scalar; - - /// Assemble the `len`-row output: `values` (one per valid row, in row order) land at the set - /// positions of `valid`, and every other row is null. With an empty `valid` this is the - /// all-null output. - fn build_array( - len: usize, - valid: &Mask, - values: Vec, - nullability: Nullability, - ) -> ArrayRef; -} - -impl GeoOutput for f64 { - fn into_scalar(self, nullability: Nullability) -> Scalar { - Scalar::primitive(self, nullability) - } - - fn build_array( - len: usize, - valid: &Mask, - values: Vec, - nullability: Nullability, - ) -> ArrayRef { - let validity = Validity::from_mask(valid.clone(), nullability); - match valid.indices() { - // No nulls: `values` already lines up one-to-one with the rows. - AllOr::All => PrimitiveArray::new(values, validity).into_array(), - // No valid rows: the whole output is null. - AllOr::None => PrimitiveArray::new(vec![0.0f64; len], validity).into_array(), - // Some nulls: scatter each computed value back to the row it came from. - AllOr::Some(rows) => { - let mut data = vec![0.0f64; len]; - for (&row, value) in rows.iter().zip(values) { - data[row] = value; - } - PrimitiveArray::new(data, validity).into_array() - } - } - } -} - -impl GeoOutput for bool { - fn into_scalar(self, nullability: Nullability) -> Scalar { - Scalar::bool(self, nullability) - } - - fn build_array( - len: usize, - valid: &Mask, - values: Vec, - nullability: Nullability, - ) -> ArrayRef { - let validity = Validity::from_mask(valid.clone(), nullability); - match valid.indices() { - // No nulls: `values` already lines up one-to-one with the rows. - AllOr::All => BoolArray::new(BitBuffer::from_iter(values), validity).into_array(), - // No valid rows: the whole output is null. - AllOr::None => BoolArray::new(BitBuffer::new_unset(len), validity).into_array(), - // Some nulls: scatter each computed value back to the row it came from. - AllOr::Some(rows) => { - let mut data = vec![false; len]; - for (&row, value) in rows.iter().zip(values) { - data[row] = value; - } - BoolArray::new(BitBuffer::from_iter(data), validity).into_array() - } - } - } -} - -/// A bounding-rect pre-check for [`execute_null_propagating`]'s one-constant arms. -/// -/// Called per row with the operands' bounding rects in operand order (`a`'s, then `b`'s), it -/// returns `Some(result)` when the rects alone prove the kernel's result — the exact test is -/// skipped — and `None` when they cannot. The proof must be sound, never a guess: disjoint rects -/// prove `ST_Intersects` false, and a container rect not covering the contained rect proves -/// `ST_Contains` false. A kernel a rect cannot decide (`ST_Distance` produces a value, not a -/// verdict) passes `None` for the whole parameter. -pub(crate) type BboxReject = fn(&Rect, &Rect) -> Option; - -/// Run a binary geo kernel over operands `a` and `b`, each a column or a constant literal. -/// -/// The output is null wherever either operand is null, and its type is nullable if either operand -/// is: equivalently, the output validity is the intersection of the operands' validities. -/// -/// The core idea: a geo kernel decodes each operand into a `geo_types` geometry, and a null row -/// has no geometry to decode, so it can't compute over every row and mask the nulls afterwards -/// (the way numeric kernels do). Instead it skips the nulls up front: keep the rows valid in both -/// operands, decode and compute only those, then scatter the results back to their rows and leave -/// every other row null. -/// -/// With exactly one constant operand, `bbox_reject` short-circuits rows from bounding rects -/// alone: the constant's rect is fixed once per batch, each valid row's rect is offered to -/// `bbox_reject` before the exact test, and rows it decides never reach `compute`. An operand -/// without a rect (an empty geometry) always falls through to the exact test. -pub(crate) fn execute_null_propagating( - a: &ArrayRef, - b: &ArrayRef, - compute: F, - bbox_reject: Option>, - ctx: &mut ExecutionCtx, -) -> VortexResult -where - T: GeoOutput, - F: Fn(&Geometry, &Geometry) -> T + Copy, -{ - let len = a.len(); - let nullability = Nullability::from(a.dtype().is_nullable() || b.dtype().is_nullable()); - - // A null constant operand makes every row null (an empty mask builds the all-null output). - for operand in [a, b] { - if operand - .as_opt::() - .is_some_and(|c| c.scalar().is_null()) - { - return Ok(T::build_array( - len, - &Mask::new_false(len), - vec![], - nullability, - )); - } - } - - match (a.as_opt::(), b.as_opt::()) { - // Both constant: compute once and broadcast across every row. - (Some(qa), Some(qb)) => { - let ga = single_geometry(qa.scalar(), ctx)?; - let gb = single_geometry(qb.scalar(), ctx)?; - Ok(ConstantArray::new(compute(&ga, &gb).into_scalar(nullability), len).into_array()) - } - // One constant, one column: fix the constant geometry and evaluate down the column. Its - // bounding rect is also fixed once, so `bbox_reject` can prove rows from their rects - // alone and skip the exact test; `zip` disables the pre-check when the kernel has none - // or the constant has no rect (an empty geometry). The rects go to `bbox_reject` in - // operand order, like the geometries to `compute`. - (Some(qa), None) => { - let ga = single_geometry(qa.scalar(), ctx)?; - let prescreen = bbox_reject.zip(ga.bounding_rect()); - eval_column( - b, - |g| { - prescreen - .and_then(|(reject, fixed)| reject(&fixed, &g.bounding_rect()?)) - .unwrap_or_else(|| compute(&ga, g)) - }, - nullability, - ctx, - ) - } - (None, Some(qb)) => { - let gb = single_geometry(qb.scalar(), ctx)?; - let prescreen = bbox_reject.zip(gb.bounding_rect()); - eval_column( - a, - |g| { - prescreen - .and_then(|(reject, fixed)| reject(&g.bounding_rect()?, &fixed)) - .unwrap_or_else(|| compute(g, &gb)) - }, - nullability, - ctx, - ) - } - // Two columns: evaluate row by row. - (None, None) => { - vortex_ensure_eq!( - a.len(), - b.len(), - "geo binary: operand length mismatch {} vs {}", - a.len(), - b.len() - ); - eval_column_pair(a, b, compute, nullability, ctx) - } - } +/// A non-null operand presented to a geometry kernel. +pub(crate) enum Operand { + /// One scalar value repeated for every row. + Constant(Scalar), + /// A column with one value per row. + Column(ArrayRef), } -/// Evaluate `f` over each valid row of one geometry `column`, propagating the column's nulls. -fn eval_column( - column: &ArrayRef, - f: F, - nullability: Nullability, - ctx: &mut ExecutionCtx, -) -> VortexResult -where - T: GeoOutput, - F: Fn(&Geometry) -> T, -{ - let len = column.len(); - let valid = column.validity()?.execute_mask(len, ctx)?; - // Drop the null rows before decoding, since a null row has no geometry to decode. `filter` - // collapses an all-true mask, so an all-valid column passes through unchanged. - let decoded = geometries(&column.filter(valid.clone())?, ctx)?; - let values = decoded.iter().map(f).collect(); - Ok(T::build_array(len, &valid, values, nullability)) -} - -/// Evaluate `compute` over each row where both geometry columns are valid, propagating the nulls -/// of either column. -fn eval_column_pair( - a: &ArrayRef, - b: &ArrayRef, - compute: F, - nullability: Nullability, - ctx: &mut ExecutionCtx, -) -> VortexResult -where - T: GeoOutput, - F: Fn(&Geometry, &Geometry) -> T, -{ - let len = a.len(); - let a_present = a.validity()?.execute_mask(len, ctx)?; - let b_present = b.validity()?.execute_mask(len, ctx)?; - // A row survives only where both columns are present. - let valid = &a_present & &b_present; - // Keep only the rows valid in both columns, so decoding never sees a null geometry. `filter` - // collapses an all-true mask, so all-valid columns pass through unchanged. - let ag = geometries(&a.filter(valid.clone())?, ctx)?; - let bg = geometries(&b.filter(valid.clone())?, ctx)?; - let values = ag.iter().zip(&bg).map(|(x, y)| compute(x, y)).collect(); - Ok(T::build_array(len, &valid, values, nullability)) -} - -#[cfg(test)] -mod tests { - use std::cell::Cell; - - use geo::Contains; - use geo::Intersects; - use geo_types::Geometry; - use vortex_array::ArrayRef; - use vortex_array::ExecutionCtx; - use vortex_array::IntoArray; - use vortex_array::VortexSessionExecute; - use vortex_array::arrays::BoolArray; - use vortex_array::arrays::ConstantArray; - use vortex_array::assert_arrays_eq; - use vortex_array::validity::Validity; - use vortex_buffer::BitBuffer; - use vortex_error::VortexResult; - - use super::BboxReject; - use super::execute_null_propagating; - use crate::test_harness::linestring_column; - use crate::test_harness::nullable_point_column; - use crate::test_harness::point_column; - use crate::test_harness::polygon_column; - - /// The `ST_Intersects` rejection: disjoint bounding rects prove no intersection. - const DISJOINT_REJECTS: BboxReject = |ra, rb| (!ra.intersects(rb)).then_some(false); - - /// A constant column of length `len`, every row the right triangle `(0,0)-(10,0)-(0,10)`. - /// Its bounding rect is the `[0, 10]` square, so points in the upper-right half of that - /// square are inside the rect but outside the triangle. - fn triangle_constant(len: usize, ctx: &mut ExecutionCtx) -> VortexResult { - let ring = vec![(0.0, 0.0), (10.0, 0.0), (0.0, 10.0), (0.0, 0.0)]; - let single = polygon_column(vec![vec![ring]])?.execute_scalar(0, ctx)?; - Ok(ConstantArray::new(single, len).into_array()) - } - - /// An intersects test that counts how many rows reach the exact per-row computation. - fn counting_intersects( - counter: &Cell, - ) -> impl Fn(&Geometry, &Geometry) -> bool + Copy { - move |x, y| { - counter.set(counter.get() + 1); - x.intersects(y) - } - } - - /// Probes against a constant triangle, one per tier: far outside the bounding rect - /// (short-circuits to false), inside the rect but outside the triangle (exact test says - /// false), and inside the triangle (exact test says true). Only the two in-rect probes - /// reach the exact test. - #[test] - fn bbox_reject_skips_exact_test() -> VortexResult<()> { - let session = vortex_array::array_session(); - let mut ctx = session.create_execution_ctx(); - - let triangle = triangle_constant(3, &mut ctx)?; - let probes = point_column(vec![50.0, 8.0, 2.0], vec![50.0, 8.0, 2.0])?; - - let exact_runs = Cell::new(0); - let result = execute_null_propagating( - &triangle, - &probes, - counting_intersects(&exact_runs), - Some(DISJOINT_REJECTS), - &mut ctx, - )?; - - assert_arrays_eq!(result, BoolArray::from_iter([false, false, true]), &mut ctx); - assert_eq!(exact_runs.get(), 2); - Ok(()) - } - - /// Null rows are untouched by the pre-check: they stay null and never reach the rect test - /// or the exact test; valid rows keep their verdicts. - #[test] - fn bbox_reject_leaves_nulls_alone() -> VortexResult<()> { - let session = vortex_array::array_session(); - let mut ctx = session.create_execution_ctx(); - - let triangle = triangle_constant(3, &mut ctx)?; - let probes = nullable_point_column(vec![Some((50.0, 50.0)), None, Some((2.0, 2.0))])?; - - let exact_runs = Cell::new(0); - let result = execute_null_propagating( - &triangle, - &probes, - counting_intersects(&exact_runs), - Some(DISJOINT_REJECTS), - &mut ctx, - )?; - - let expected = BoolArray::new( - BitBuffer::from_iter([false, false, true]), - Validity::from_iter([true, false, true]), - ) - .into_array(); - assert_arrays_eq!(result, expected, &mut ctx); - // The far probe short-circuits and the null row is filtered before decoding, so only - // the in-triangle probe runs the exact test. - assert_eq!(exact_runs.get(), 1); - Ok(()) - } - - /// The rects reach the rejection in operand order even when the constant is the second - /// operand: a point row's rect never contains the triangle's rect, so every row - /// short-circuits; with the order flipped, the triangle's rect contains the in-rect - /// point's and the exact test would run. - #[test] - fn bbox_reject_sees_rects_in_operand_order() -> VortexResult<()> { - let session = vortex_array::array_session(); - let mut ctx = session.create_execution_ctx(); - - let probes = point_column(vec![2.0, 50.0], vec![2.0, 50.0])?; - let triangle = triangle_constant(2, &mut ctx)?; - - let exact_runs = Cell::new(0); - let counted = |x: &Geometry, y: &Geometry| { - exact_runs.set(exact_runs.get() + 1); - x.contains(y) - }; - let result = execute_null_propagating( - &probes, - &triangle, - counted, - Some(|ra, rb| (!ra.contains(rb)).then_some(false)), - &mut ctx, - )?; - - assert_arrays_eq!(result, BoolArray::from_iter([false, false]), &mut ctx); - assert_eq!(exact_runs.get(), 0); - Ok(()) - } - - /// An empty constant geometry has no bounding rect, so the pre-check is disabled: every - /// valid row falls through to the exact test and none is falsely rejected. - #[test] - fn empty_constant_falls_through_to_exact() -> VortexResult<()> { - let session = vortex_array::array_session(); - let mut ctx = session.create_execution_ctx(); - - let single = linestring_column(vec![vec![]])?.execute_scalar(0, &mut ctx)?; - let empty = ConstantArray::new(single, 2).into_array(); - let probes = point_column(vec![2.0, 50.0], vec![2.0, 50.0])?; - - let exact_runs = Cell::new(0); - let result = execute_null_propagating( - &empty, - &probes, - counting_intersects(&exact_runs), - Some(DISJOINT_REJECTS), - &mut ctx, - )?; - - assert_arrays_eq!(result, BoolArray::from_iter([false, false]), &mut ctx); - assert_eq!(exact_runs.get(), 2); - Ok(()) - } - - /// Property: the pre-check never changes results — a mixed batch (far, in-rect-but-outside, - /// inside, null, boundary, rect corner) computed with the rejection equals the same batch - /// computed with the exact test alone. - #[test] - fn bbox_reject_matches_exact_results() -> VortexResult<()> { - let session = vortex_array::array_session(); - let mut ctx = session.create_execution_ctx(); - - let triangle = triangle_constant(6, &mut ctx)?; - let probes = nullable_point_column(vec![ - Some((50.0, 50.0)), - Some((8.0, 8.0)), - Some((2.0, 2.0)), - None, - Some((0.0, 0.0)), - Some((10.0, 0.0)), - ])?; - let exact = |x: &Geometry, y: &Geometry| x.intersects(y); - - let with_reject = - execute_null_propagating(&triangle, &probes, exact, Some(DISJOINT_REJECTS), &mut ctx)?; - let exact_only = execute_null_propagating(&triangle, &probes, exact, None, &mut ctx)?; - - assert_arrays_eq!(with_reject, exact_only, &mut ctx); - Ok(()) - } +/// Shared batch state presented to a null-propagating geometry kernel with `N` operands. +pub(crate) struct Execution { + /// Constant/column shape of each operand. + pub(crate) operands: [Operand; N], + /// Rows where every operand is valid. + pub(crate) valid: Mask, + /// Number of output rows. + pub(crate) len: usize, + /// Output nullability from the scalar function's return dtype. + pub(crate) nullability: Nullability, } diff --git a/vortex-geo/src/scalar_fn/execute/binary.rs b/vortex-geo/src/scalar_fn/execute/binary.rs new file mode 100644 index 00000000000..f2c03bd1beb --- /dev/null +++ b/vortex-geo/src/scalar_fn/execute/binary.rs @@ -0,0 +1,334 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Binary pairwise dispatch, plus an adapter for row-oriented `geo_types` kernels. + +use geo::BoundingRect; +use geo_types::Geometry; +use geo_types::Rect; +use vortex_array::ArrayRef; +use vortex_array::ExecutionCtx; +use vortex_array::IntoArray; +use vortex_array::arrays::Constant; +use vortex_array::arrays::ConstantArray; +use vortex_array::dtype::DType; +use vortex_array::dtype::Nullability; +use vortex_array::scalar::Scalar; +use vortex_error::VortexResult; +use vortex_mask::Mask; + +use super::Execution; +use super::Operand; +use super::geo_types::GeoTypesOutput; +use super::geo_types::eval_column; +use super::geo_types::eval_column_pair; +use crate::extension::single_geometry; + +/// Dispatch a binary strict geometry kernel over constants and columns. +/// +/// A null constant or an empty combined validity mask short-circuits to an all-null constant +/// output. Otherwise, `kernel` receives both operand shapes and the mask of rows where both are +/// valid. Two columns are always paired by row index. The kernel remains responsible for physical +/// input interpretation and Vortex output construction. +pub(crate) fn dispatch_binary( + left: &ArrayRef, + right: &ArrayRef, + output_dtype: DType, + kernel: K, + ctx: &mut ExecutionCtx, +) -> VortexResult +where + K: FnOnce(Execution<2>, &mut ExecutionCtx) -> VortexResult, +{ + let len = left.len(); + for operand in [left, right] { + if operand + .as_opt::() + .is_some_and(|constant| constant.scalar().is_null()) + { + return Ok(ConstantArray::new(Scalar::null(output_dtype), len).into_array()); + } + } + + let (left, right, valid) = match (left.as_opt::(), right.as_opt::()) { + (Some(left), Some(right)) => ( + Operand::Constant(left.scalar().clone()), + Operand::Constant(right.scalar().clone()), + Mask::new_true(len), + ), + (Some(left), None) => ( + Operand::Constant(left.scalar().clone()), + Operand::Column(right.clone()), + right.validity()?.execute_mask(len, ctx)?, + ), + (None, Some(right)) => ( + Operand::Column(left.clone()), + Operand::Constant(right.scalar().clone()), + left.validity()?.execute_mask(len, ctx)?, + ), + (None, None) => { + let left_valid = left.validity()?.execute_mask(len, ctx)?; + let right_valid = right.validity()?.execute_mask(len, ctx)?; + ( + Operand::Column(left.clone()), + Operand::Column(right.clone()), + &left_valid & &right_valid, + ) + } + }; + + if len != 0 && valid.all_false() { + return Ok(ConstantArray::new(Scalar::null(output_dtype), len).into_array()); + } + kernel( + Execution { + operands: [left, right], + valid, + len, + nullability: output_dtype.nullability(), + }, + ctx, + ) +} + +/// A bounding-rectangle pre-check for [`execute_binary_geo_types`]'s one-constant paths. +/// +/// Called per row with rectangles in operand order, it returns `Some(result)` when they prove the +/// result and `None` when the exact kernel must run. +pub(crate) type BboxPrecheck = fn(&Rect, &Rect) -> Option; + +/// Run a binary row-oriented kernel whose inputs are decoded to `geo_types::Geometry`. +/// +/// The `geo_types` name describes the values passed to `compute`, not the output. `T` is converted +/// into a Vortex array before this function returns. Nulls propagate from either operand. With +/// exactly one constant operand, `bbox_precheck` may prove a result from the fixed constant +/// bounding rectangle and the current row's rectangle before the exact kernel runs. +pub(crate) fn execute_binary_geo_types( + left: &ArrayRef, + right: &ArrayRef, + compute: F, + bbox_precheck: Option>, + ctx: &mut ExecutionCtx, +) -> VortexResult +where + T: GeoTypesOutput, + F: Fn(&Geometry, &Geometry) -> T + Copy, +{ + let nullability = Nullability::from(left.dtype().is_nullable() || right.dtype().is_nullable()); + dispatch_binary( + left, + right, + T::dtype(nullability), + |execution, ctx| match execution.operands { + [Operand::Constant(left), Operand::Constant(right)] => { + let left = single_geometry(&left, ctx)?; + let right = single_geometry(&right, ctx)?; + Ok(ConstantArray::new( + compute(&left, &right).into_scalar(execution.nullability), + execution.len, + ) + .into_array()) + } + [Operand::Constant(left), Operand::Column(right)] => { + let left = single_geometry(&left, ctx)?; + let prescreen = bbox_precheck.zip(left.bounding_rect()); + eval_column( + &right, + &execution.valid, + |right| { + prescreen + .and_then(|(precheck, fixed)| precheck(&fixed, &right.bounding_rect()?)) + .unwrap_or_else(|| compute(&left, right)) + }, + execution.nullability, + ctx, + ) + } + [Operand::Column(left), Operand::Constant(right)] => { + let right = single_geometry(&right, ctx)?; + let prescreen = bbox_precheck.zip(right.bounding_rect()); + eval_column( + &left, + &execution.valid, + |left| { + prescreen + .and_then(|(precheck, fixed)| precheck(&left.bounding_rect()?, &fixed)) + .unwrap_or_else(|| compute(left, &right)) + }, + execution.nullability, + ctx, + ) + } + [Operand::Column(left), Operand::Column(right)] => eval_column_pair( + &left, + &right, + &execution.valid, + compute, + execution.nullability, + ctx, + ), + }, + ctx, + ) +} + +#[cfg(test)] +mod tests { + use std::cell::Cell; + + use geo::Contains; + use geo::Intersects; + use geo_types::Geometry; + use vortex_array::ArrayRef; + use vortex_array::ExecutionCtx; + use vortex_array::IntoArray; + use vortex_array::VortexSessionExecute; + use vortex_array::arrays::BoolArray; + use vortex_array::arrays::ConstantArray; + use vortex_array::assert_arrays_eq; + use vortex_array::validity::Validity; + use vortex_buffer::BitBuffer; + use vortex_error::VortexResult; + + use super::BboxPrecheck; + use super::execute_binary_geo_types; + use crate::test_harness::linestring_column; + use crate::test_harness::nullable_point_column; + use crate::test_harness::point_column; + use crate::test_harness::polygon_column; + + const DISJOINT_PRECHECK: BboxPrecheck = + |left, right| (!left.intersects(right)).then_some(false); + + fn triangle_constant(len: usize, ctx: &mut ExecutionCtx) -> VortexResult { + let ring = vec![(0.0, 0.0), (10.0, 0.0), (0.0, 10.0), (0.0, 0.0)]; + let scalar = polygon_column(vec![vec![ring]])?.execute_scalar(0, ctx)?; + Ok(ConstantArray::new(scalar, len).into_array()) + } + + fn counting_intersects( + counter: &Cell, + ) -> impl Fn(&Geometry, &Geometry) -> bool + Copy { + move |left, right| { + counter.set(counter.get() + 1); + left.intersects(right) + } + } + + #[test] + fn bbox_precheck_skips_exact_test() -> VortexResult<()> { + let session = vortex_array::array_session(); + let mut ctx = session.create_execution_ctx(); + let triangle = triangle_constant(3, &mut ctx)?; + let probes = point_column(vec![50.0, 8.0, 2.0], vec![50.0, 8.0, 2.0])?; + let exact_runs = Cell::new(0); + + let result = execute_binary_geo_types( + &triangle, + &probes, + counting_intersects(&exact_runs), + Some(DISJOINT_PRECHECK), + &mut ctx, + )?; + + assert_arrays_eq!(result, BoolArray::from_iter([false, false, true]), &mut ctx); + assert_eq!(exact_runs.get(), 2); + Ok(()) + } + + #[test] + fn bbox_precheck_leaves_nulls_alone() -> VortexResult<()> { + let session = vortex_array::array_session(); + let mut ctx = session.create_execution_ctx(); + let triangle = triangle_constant(3, &mut ctx)?; + let probes = nullable_point_column(vec![Some((50.0, 50.0)), None, Some((2.0, 2.0))])?; + let exact_runs = Cell::new(0); + + let result = execute_binary_geo_types( + &triangle, + &probes, + counting_intersects(&exact_runs), + Some(DISJOINT_PRECHECK), + &mut ctx, + )?; + let expected = BoolArray::new( + BitBuffer::from_iter([false, false, true]), + Validity::from_iter([true, false, true]), + ) + .into_array(); + + assert_arrays_eq!(result, expected, &mut ctx); + assert_eq!(exact_runs.get(), 1); + Ok(()) + } + + #[test] + fn bbox_precheck_sees_rects_in_operand_order() -> VortexResult<()> { + let session = vortex_array::array_session(); + let mut ctx = session.create_execution_ctx(); + let probes = point_column(vec![2.0, 50.0], vec![2.0, 50.0])?; + let triangle = triangle_constant(2, &mut ctx)?; + let exact_runs = Cell::new(0); + let counted = |left: &Geometry, right: &Geometry| { + exact_runs.set(exact_runs.get() + 1); + left.contains(right) + }; + + let result = execute_binary_geo_types( + &probes, + &triangle, + counted, + Some(|left, right| (!left.contains(right)).then_some(false)), + &mut ctx, + )?; + + assert_arrays_eq!(result, BoolArray::from_iter([false, false]), &mut ctx); + assert_eq!(exact_runs.get(), 0); + Ok(()) + } + + #[test] + fn empty_constant_falls_through_to_exact() -> VortexResult<()> { + let session = vortex_array::array_session(); + let mut ctx = session.create_execution_ctx(); + let scalar = linestring_column(vec![vec![]])?.execute_scalar(0, &mut ctx)?; + let empty = ConstantArray::new(scalar, 2).into_array(); + let probes = point_column(vec![2.0, 50.0], vec![2.0, 50.0])?; + let exact_runs = Cell::new(0); + + let result = execute_binary_geo_types( + &empty, + &probes, + counting_intersects(&exact_runs), + Some(DISJOINT_PRECHECK), + &mut ctx, + )?; + + assert_arrays_eq!(result, BoolArray::from_iter([false, false]), &mut ctx); + assert_eq!(exact_runs.get(), 2); + Ok(()) + } + + #[test] + fn bbox_precheck_matches_exact_results() -> VortexResult<()> { + let session = vortex_array::array_session(); + let mut ctx = session.create_execution_ctx(); + let triangle = triangle_constant(6, &mut ctx)?; + let probes = nullable_point_column(vec![ + Some((50.0, 50.0)), + Some((8.0, 8.0)), + Some((2.0, 2.0)), + None, + Some((0.0, 0.0)), + Some((10.0, 0.0)), + ])?; + let exact = |left: &Geometry, right: &Geometry| left.intersects(right); + + let with_precheck = + execute_binary_geo_types(&triangle, &probes, exact, Some(DISJOINT_PRECHECK), &mut ctx)?; + let exact_only = execute_binary_geo_types(&triangle, &probes, exact, None, &mut ctx)?; + + assert_arrays_eq!(with_precheck, exact_only, &mut ctx); + Ok(()) + } +} diff --git a/vortex-geo/src/scalar_fn/execute/geo_types.rs b/vortex-geo/src/scalar_fn/execute/geo_types.rs new file mode 100644 index 00000000000..038aca46502 --- /dev/null +++ b/vortex-geo/src/scalar_fn/execute/geo_types.rs @@ -0,0 +1,144 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Shared input decoding and Vortex output construction for `geo_types` kernels. +//! +//! `geo_types` is the row representation consumed by the kernel. These helpers always construct +//! and return Vortex arrays; they do not expose `geo_types` values as scalar-function outputs. + +use geo_types::Geometry; +use vortex_array::ArrayRef; +use vortex_array::ExecutionCtx; +use vortex_array::IntoArray; +use vortex_array::arrays::BoolArray; +use vortex_array::arrays::PrimitiveArray; +use vortex_array::dtype::DType; +use vortex_array::dtype::Nullability; +use vortex_array::dtype::PType; +use vortex_array::scalar::Scalar; +use vortex_array::validity::Validity; +use vortex_buffer::BitBuffer; +use vortex_error::VortexResult; +use vortex_mask::AllOr; +use vortex_mask::Mask; + +use crate::extension::geometries; + +/// A primitive result produced after kernel inputs are decoded to `geo_types`. +pub(crate) trait GeoTypesOutput: Copy { + /// The Vortex dtype used to represent this output. + fn dtype(nullability: Nullability) -> DType; + + /// Convert one computed value into a Vortex scalar for constant output. + fn into_scalar(self, nullability: Nullability) -> Scalar; + + /// Scatter values computed for valid rows into a full-length output array. + fn build_array( + len: usize, + valid: &Mask, + values: Vec, + nullability: Nullability, + ) -> ArrayRef; +} + +impl GeoTypesOutput for f64 { + fn dtype(nullability: Nullability) -> DType { + DType::Primitive(PType::F64, nullability) + } + + fn into_scalar(self, nullability: Nullability) -> Scalar { + Scalar::primitive(self, nullability) + } + + fn build_array( + len: usize, + valid: &Mask, + values: Vec, + nullability: Nullability, + ) -> ArrayRef { + let validity = Validity::from_mask(valid.clone(), nullability); + match valid.indices() { + AllOr::All => PrimitiveArray::new(values, validity).into_array(), + AllOr::None => PrimitiveArray::new(vec![0.0f64; len], validity).into_array(), + AllOr::Some(rows) => { + let mut data = vec![0.0f64; len]; + for (&row, value) in rows.iter().zip(values) { + data[row] = value; + } + PrimitiveArray::new(data, validity).into_array() + } + } + } +} + +impl GeoTypesOutput for bool { + fn dtype(nullability: Nullability) -> DType { + DType::Bool(nullability) + } + + fn into_scalar(self, nullability: Nullability) -> Scalar { + Scalar::bool(self, nullability) + } + + fn build_array( + len: usize, + valid: &Mask, + values: Vec, + nullability: Nullability, + ) -> ArrayRef { + let validity = Validity::from_mask(valid.clone(), nullability); + match valid.indices() { + AllOr::All => BoolArray::new(BitBuffer::from_iter(values), validity).into_array(), + AllOr::None => BoolArray::new(BitBuffer::new_unset(len), validity).into_array(), + AllOr::Some(rows) => { + let mut data = vec![false; len]; + for (&row, value) in rows.iter().zip(values) { + data[row] = value; + } + BoolArray::new(BitBuffer::from_iter(data), validity).into_array() + } + } + } +} + +/// Evaluate a decoded kernel over each valid row of one geometry column. +pub(super) fn eval_column( + column: &ArrayRef, + valid: &Mask, + compute: F, + nullability: Nullability, + ctx: &mut ExecutionCtx, +) -> VortexResult +where + T: GeoTypesOutput, + F: Fn(&Geometry) -> T, +{ + let len = column.len(); + let decoded = geometries(&column.filter(valid.clone())?, ctx)?; + let values = decoded.iter().map(compute).collect(); + Ok(T::build_array(len, valid, values, nullability)) +} + +/// Evaluate a decoded kernel over rows where both geometry columns are valid. +pub(super) fn eval_column_pair( + left: &ArrayRef, + right: &ArrayRef, + valid: &Mask, + compute: F, + nullability: Nullability, + ctx: &mut ExecutionCtx, +) -> VortexResult +where + T: GeoTypesOutput, + F: Fn(&Geometry, &Geometry) -> T, +{ + let len = left.len(); + let left = geometries(&left.filter(valid.clone())?, ctx)?; + let right = geometries(&right.filter(valid.clone())?, ctx)?; + let values = left + .iter() + .zip(&right) + .map(|(left, right)| compute(left, right)) + .collect(); + Ok(T::build_array(len, valid, values, nullability)) +} diff --git a/vortex-geo/src/scalar_fn/execute/unary.rs b/vortex-geo/src/scalar_fn/execute/unary.rs new file mode 100644 index 00000000000..8af7474491d --- /dev/null +++ b/vortex-geo/src/scalar_fn/execute/unary.rs @@ -0,0 +1,62 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Unary operand dispatch for native geometry kernels. + +use vortex_array::ArrayRef; +use vortex_array::ExecutionCtx; +use vortex_array::IntoArray; +use vortex_array::arrays::Constant; +use vortex_array::arrays::ConstantArray; +use vortex_array::dtype::DType; +use vortex_array::scalar::Scalar; +use vortex_error::VortexResult; +use vortex_mask::Mask; + +use super::Execution; +use super::Operand; + +/// Dispatch a unary strict geometry kernel over a constant or column. +/// +/// A null constant or all-null column short-circuits to an all-null constant output. Otherwise, +/// `kernel` receives the operand shape and its valid-row mask. The kernel remains responsible for +/// interpreting the native input and constructing its Vortex output. +pub(crate) fn dispatch_unary( + array: &ArrayRef, + output_dtype: DType, + kernel: K, + ctx: &mut ExecutionCtx, +) -> VortexResult +where + K: FnOnce(Execution<1>, &mut ExecutionCtx) -> VortexResult, +{ + let len = array.len(); + if let Some(constant) = array.as_opt::() { + if constant.scalar().is_null() { + return Ok(ConstantArray::new(Scalar::null(output_dtype), len).into_array()); + } + return kernel( + Execution { + operands: [Operand::Constant(constant.scalar().clone())], + valid: Mask::new_true(len), + len, + nullability: output_dtype.nullability(), + }, + ctx, + ); + } + + let valid = array.validity()?.execute_mask(len, ctx)?; + if len != 0 && valid.all_false() { + return Ok(ConstantArray::new(Scalar::null(output_dtype), len).into_array()); + } + kernel( + Execution { + operands: [Operand::Column(array.clone())], + valid, + len, + nullability: output_dtype.nullability(), + }, + ctx, + ) +} diff --git a/vortex-geo/src/scalar_fn/intersects.rs b/vortex-geo/src/scalar_fn/intersects.rs index 3f3842e4e1c..8fa65aba34a 100644 --- a/vortex-geo/src/scalar_fn/intersects.rs +++ b/vortex-geo/src/scalar_fn/intersects.rs @@ -19,11 +19,28 @@ use vortex_array::scalar_fn::ScalarFnId; use vortex_array::scalar_fn::ScalarFnVTable; use vortex_array::scalar_fn::TypedScalarFnInstance; use vortex_error::VortexResult; +use vortex_error::vortex_ensure; use vortex_session::VortexSession; use vortex_session::registry::CachedId; -use crate::extension::validate_geometry_operands; -use crate::scalar_fn::execute::execute_null_propagating; +use crate::extension::is_native_geometry; +use crate::scalar_fn::execute::execute_binary_geo_types; + +/// Validate the two native geometry operands accepted by `ST_Intersects`. +fn validate_intersects_operands(dtypes: &[DType]) -> VortexResult<()> { + vortex_ensure!( + dtypes.len() == 2, + "geo: intersects requires exactly two geometry operands, got {}", + dtypes.len() + ); + for dtype in dtypes { + vortex_ensure!( + is_native_geometry(dtype), + "geo: intersects operand {dtype} is not a native geometry type" + ); + } + Ok(()) +} /// OGC `ST_Intersects` (not disjoint; boundary contact counts) between two native geometry /// operands, each a column or a constant literal. @@ -70,7 +87,7 @@ impl ScalarFnVTable for GeoIntersects { } fn return_dtype(&self, _: &Self::Options, dtypes: &[DType]) -> VortexResult { - validate_geometry_operands(dtypes)?; + validate_intersects_operands(dtypes)?; let nullability = Nullability::from(dtypes.iter().any(DType::is_nullable)); Ok(DType::Bool(nullability)) } @@ -85,7 +102,7 @@ impl ScalarFnVTable for GeoIntersects { let b = args.get(1)?; // Disjoint bounding rects prove the geometries disjoint; rect contact (closed test) // falls through to the exact test. - execute_null_propagating( + execute_binary_geo_types( &a, &b, |x, y| x.intersects(y),