From 90e45480fd7bbbaf27b6c0e52a40de16814e879b Mon Sep 17 00:00:00 2001 From: Nemo Yu Date: Thu, 16 Jul 2026 15:26:56 -0400 Subject: [PATCH 1/4] feat(vortex-geo): null-propagating scalar functions Make the binary geo scalar functions (ST_Distance, ST_Intersects, ST_Contains) null-propagating, matching SQL/OGC semantics and the other Vortex binary kernels: a nullable geometry operand is now allowed, and any row whose geometry input is null yields a null result (output nullable iff an operand is nullable). 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. The shared execute_null_propagating dispatch in scalar_fn/mod.rs instead filters the null rows out up front, computes over the rows valid in both operands, and scatters the results back under the combined null mask. - validate_geometry_operands no longer rejects nullable operands - return_dtype mirrors operand nullability - tests for nullable columns, constant-null operands, and column/column nulls Signed-off-by: Nemo Yu --- Cargo.lock | 2 + vortex-geo/Cargo.toml | 2 + vortex-geo/src/extension/mod.rs | 9 +- vortex-geo/src/scalar_fn/contains.rs | 116 ++++++------ vortex-geo/src/scalar_fn/distance.rs | 124 +++++++------ vortex-geo/src/scalar_fn/intersects.rs | 105 +++++------ vortex-geo/src/scalar_fn/mod.rs | 243 +++++++++++++++++++++++++ vortex-geo/src/test_harness.rs | 22 +++ 8 files changed, 449 insertions(+), 174 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index e23ac0d40f8..847e14a2013 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -10264,8 +10264,10 @@ dependencies = [ "rstest", "vortex-array", "vortex-arrow", + "vortex-buffer", "vortex-error", "vortex-layout", + "vortex-mask", "vortex-session", "wkb", ] diff --git a/vortex-geo/Cargo.toml b/vortex-geo/Cargo.toml index 9a7980bc631..351161a6ec4 100644 --- a/vortex-geo/Cargo.toml +++ b/vortex-geo/Cargo.toml @@ -24,7 +24,9 @@ geoarrow-cast = { workspace = true } prost = { workspace = true } vortex-array = { workspace = true } vortex-arrow = { workspace = true } +vortex-buffer = { workspace = true } vortex-error = { workspace = true } +vortex-mask = { workspace = true } vortex-session = { workspace = true } wkb = { workspace = true } diff --git a/vortex-geo/src/extension/mod.rs b/vortex-geo/src/extension/mod.rs index 4bd8ae3800f..a9ff7ff66de 100644 --- a/vortex-geo/src/extension/mod.rs +++ b/vortex-geo/src/extension/mod.rs @@ -72,18 +72,15 @@ 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) and non-nullable (geometry arrays never carry nulls). +/// Validate the operands of a geo scalar function: each must be a native geometry type so the +/// kernel can decode it. 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" ); - vortex_ensure!( - !dtype.is_nullable(), - "geo: nullable operand {dtype} is unsupported" - ); } Ok(()) } diff --git a/vortex-geo/src/scalar_fn/contains.rs b/vortex-geo/src/scalar_fn/contains.rs index 5638c224e20..531e723a9f9 100644 --- a/vortex-geo/src/scalar_fn/contains.rs +++ b/vortex-geo/src/scalar_fn/contains.rs @@ -6,14 +6,9 @@ use geo::Contains; 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::ScalarFnArray; use vortex_array::dtype::DType; use vortex_array::dtype::Nullability; -use vortex_array::scalar::Scalar; use vortex_array::scalar_fn::Arity; use vortex_array::scalar_fn::ChildName; use vortex_array::scalar_fn::EmptyOptions; @@ -22,13 +17,11 @@ 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_eq; use vortex_session::VortexSession; use vortex_session::registry::CachedId; -use crate::extension::geometries; -use crate::extension::single_geometry; use crate::extension::validate_geometry_operands; +use crate::scalar_fn::execute_null_propagating; /// 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 @@ -77,7 +70,8 @@ impl ScalarFnVTable for GeoContains { fn return_dtype(&self, _: &Self::Options, dtypes: &[DType]) -> VortexResult { validate_geometry_operands(dtypes)?; - Ok(DType::Bool(Nullability::NonNullable)) + let nullability = Nullability::from(dtypes.iter().any(DType::is_nullable)); + Ok(DType::Bool(nullability)) } fn execute( @@ -89,59 +83,10 @@ impl ScalarFnVTable for GeoContains { let a = args.get(0)?; let b = args.get(1)?; // Containment is not symmetric: `a` is always the container and `b` the contained. - match (a.as_opt::(), b.as_opt::()) { - (Some(qa), Some(qb)) => { - let ga = single_geometry(qa.scalar(), ctx)?; - let gb = single_geometry(qb.scalar(), ctx)?; - Ok(ConstantArray::new( - Scalar::bool(ga.contains(&gb), Nullability::NonNullable), - a.len(), - ) - .into_array()) - } - (Some(qa), None) => constant_contains_column(qa.scalar(), &b, ctx), - (None, Some(qb)) => column_contains_constant(&a, qb.scalar(), ctx), - (None, None) => { - vortex_ensure_eq!( - a.len(), - b.len(), - "geo contains: operand length mismatch {} vs {}", - a.len(), - b.len() - ); - let ag = geometries(&a, ctx)?; - let bg = geometries(&b, ctx)?; - let hits = ag.iter().zip(&bg).map(|(x, y)| x.contains(y)); - Ok(BoolArray::from_iter(hits).into_array()) - } - } + execute_null_propagating(&a, &b, |a, b| a.contains(b), ctx) } } -/// Whether the constant `container` contains each row of `contained`. -fn constant_contains_column( - container: &Scalar, - contained: &ArrayRef, - ctx: &mut ExecutionCtx, -) -> VortexResult { - let container = single_geometry(container, ctx)?; - let geoms = geometries(contained, ctx)?; - let hits = geoms.iter().map(|g| container.contains(g)); - Ok(BoolArray::from_iter(hits).into_array()) -} - -/// Whether each row of `container` contains the constant `contained`. -fn column_contains_constant( - container: &ArrayRef, - contained: &Scalar, - ctx: &mut ExecutionCtx, -) -> VortexResult { - let contained = single_geometry(contained, ctx)?; - let geoms = geometries(container, ctx)?; - let hits = geoms.iter().map(|g| g.contains(&contained)); - Ok(BoolArray::from_iter(hits).into_array()) -} - #[cfg(test)] mod tests { use geo_types::Geometry; @@ -160,13 +105,17 @@ mod tests { use vortex_array::dtype::DType; use vortex_array::dtype::Nullability; use vortex_array::dtype::PType; + use vortex_array::scalar::Scalar; use vortex_array::scalar_fn::EmptyOptions; use vortex_array::scalar_fn::ScalarFnVTable; + use vortex_array::validity::Validity; + use vortex_buffer::BitBuffer; use vortex_error::VortexResult; use vortex_error::vortex_err; use wkb::writer::WriteOptions; use super::GeoContains; + use crate::test_harness::nullable_point_column; use crate::test_harness::point_column; /// A rectangle polygon with corners `(x0, y0)` and `(x1, y1)`, no holes. @@ -286,12 +235,53 @@ mod tests { assert_contains(polygons, points, [true, false]) } - /// Geometry arrays are never nullable, so a nullable operand dtype is rejected. + /// Output nullability mirrors the operands: nullable if any operand is nullable, otherwise + /// non-nullable. #[test] - fn nullable_operand_is_rejected() -> VortexResult<()> { + fn output_nullability_mirrors_operands() -> VortexResult<()> { let dtype = point_column(vec![0.0], vec![0.0])?.dtype().clone(); - let result = GeoContains.return_dtype(&EmptyOptions, &[dtype.as_nullable(), dtype]); - assert!(result.is_err()); + let non_nullable = + GeoContains.return_dtype(&EmptyOptions, &[dtype.clone(), dtype.clone()])?; + assert!(!non_nullable.is_nullable()); + let nullable = GeoContains.return_dtype(&EmptyOptions, &[dtype.as_nullable(), dtype])?; + assert!(nullable.is_nullable()); + Ok(()) + } + + /// A null row in the contained operand yields a null verdict; valid rows keep their verdict + /// (a strictly interior point is contained, an outside point is not). + #[test] + fn contains_propagates_null_rows() -> VortexResult<()> { + let session = vortex_array::array_session(); + let mut ctx = session.create_execution_ctx(); + + let container = geometry_constant(&Geometry::Polygon(rect_polygon(0.0, 0.0, 4.0, 4.0)), 3)?; + let points = nullable_point_column(vec![Some((2.0, 2.0)), None, Some((10.0, 10.0))])?; + let contains = GeoContains::try_new_array(container, points)?.into_array(); + + let expected = BoolArray::new( + BitBuffer::from_iter([true, false, false]), + Validity::from_iter([true, false, true]), + ) + .into_array(); + assert_arrays_eq!(contains, expected, &mut ctx); + Ok(()) + } + + /// A constant-null operand produces an all-null output. + #[test] + fn contains_constant_null_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 null_const = ConstantArray::new(Scalar::null(point_dtype), 2).into_array(); + let points = point_column(vec![2.0, 10.0], vec![2.0, 10.0])?; + let contains = GeoContains::try_new_array(null_const, points)?.into_array(); + + let expected = + BoolArray::new(BitBuffer::from_iter([false, false]), Validity::AllInvalid).into_array(); + assert_arrays_eq!(contains, expected, &mut ctx); Ok(()) } diff --git a/vortex-geo/src/scalar_fn/distance.rs b/vortex-geo/src/scalar_fn/distance.rs index 2e196706ec2..e9e739e663b 100644 --- a/vortex-geo/src/scalar_fn/distance.rs +++ b/vortex-geo/src/scalar_fn/distance.rs @@ -7,15 +7,10 @@ use geo::Distance; use geo::Euclidean; 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::arrays::PrimitiveArray; use vortex_array::arrays::ScalarFnArray; use vortex_array::dtype::DType; use vortex_array::dtype::Nullability; use vortex_array::dtype::PType; -use vortex_array::scalar::Scalar; use vortex_array::scalar_fn::Arity; use vortex_array::scalar_fn::ChildName; use vortex_array::scalar_fn::EmptyOptions; @@ -24,13 +19,11 @@ 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_eq; use vortex_session::VortexSession; use vortex_session::registry::CachedId; -use crate::extension::geometries; -use crate::extension::single_geometry; use crate::extension::validate_geometry_operands; +use crate::scalar_fn::execute_null_propagating; /// Planar (Euclidean) `ST_Distance` (no geodesic correction) between two native geometry /// operands, each a column or a constant literal. @@ -78,7 +71,8 @@ impl ScalarFnVTable for GeoDistance { fn return_dtype(&self, _: &Self::Options, dtypes: &[DType]) -> VortexResult { validate_geometry_operands(dtypes)?; - Ok(DType::Primitive(PType::F64, Nullability::NonNullable)) + let nullability = Nullability::from(dtypes.iter().any(DType::is_nullable)); + Ok(DType::Primitive(PType::F64, nullability)) } fn execute( @@ -89,49 +83,10 @@ impl ScalarFnVTable for GeoDistance { ) -> VortexResult { let a = args.get(0)?; let b = args.get(1)?; - match (a.as_opt::(), b.as_opt::()) { - (Some(qa), Some(qb)) => { - let ga = single_geometry(qa.scalar(), ctx)?; - let gb = single_geometry(qb.scalar(), ctx)?; - let distance = Euclidean.distance(&ga, &gb); - Ok(ConstantArray::new( - Scalar::primitive(distance, Nullability::NonNullable), - a.len(), - ) - .into_array()) - } - (Some(query), None) => distances_to_constant(&b, query.scalar(), ctx), - (None, Some(query)) => distances_to_constant(&a, query.scalar(), ctx), - (None, None) => { - vortex_ensure_eq!( - a.len(), - b.len(), - "geo distance: operand length mismatch {} vs {}", - a.len(), - b.len() - ); - let ag = geometries(&a, ctx)?; - let bg = geometries(&b, ctx)?; - let distances = ag.iter().zip(&bg).map(|(x, y)| Euclidean.distance(x, y)); - Ok(PrimitiveArray::from_iter(distances).into_array()) - } - } + execute_null_propagating(&a, &b, |x, y| Euclidean.distance(x, y), ctx) } } -/// Distance from each row of `operand` to the constant `query` geometry. Distance is symmetric, -/// so this serves a constant on either side. -fn distances_to_constant( - operand: &ArrayRef, - query: &Scalar, - ctx: &mut ExecutionCtx, -) -> VortexResult { - let query = single_geometry(query, ctx)?; - let geoms = geometries(operand, ctx)?; - let distances = geoms.iter().map(|g| Euclidean.distance(g, &query)); - Ok(PrimitiveArray::from_iter(distances).into_array()) -} - #[cfg(test)] mod tests { use vortex_array::ArrayRef; @@ -140,14 +95,19 @@ mod tests { use vortex_array::IntoArray; use vortex_array::VortexSessionExecute; use vortex_array::arrays::ConstantArray; + use vortex_array::arrays::PrimitiveArray; + 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::Scalar; use vortex_array::scalar_fn::EmptyOptions; use vortex_array::scalar_fn::ScalarFnVTable; + use vortex_array::validity::Validity; use vortex_error::VortexResult; use super::GeoDistance; + use crate::test_harness::nullable_point_column; use crate::test_harness::point_column; /// A constant `Point` column of length `len`, every row at `(x, y)`. @@ -227,12 +187,70 @@ mod tests { Ok(()) } - /// Geometry arrays are never nullable, so a nullable operand dtype is rejected. + /// Output nullability mirrors the operands: nullable if any operand is nullable, otherwise + /// non-nullable. #[test] - fn nullable_operand_is_rejected() -> VortexResult<()> { + fn output_nullability_mirrors_operands() -> VortexResult<()> { let dtype = point_column(vec![0.0], vec![0.0])?.dtype().clone(); - let result = GeoDistance.return_dtype(&EmptyOptions, &[dtype.as_nullable(), dtype]); - assert!(result.is_err()); + let non_nullable = + GeoDistance.return_dtype(&EmptyOptions, &[dtype.clone(), dtype.clone()])?; + assert!(!non_nullable.is_nullable()); + let nullable = GeoDistance.return_dtype(&EmptyOptions, &[dtype.as_nullable(), dtype])?; + assert!(nullable.is_nullable()); + Ok(()) + } + + /// A null row in a geometry operand yields a null result; valid rows are unaffected. + #[test] + fn distance_propagates_null_rows() -> VortexResult<()> { + let session = vortex_array::array_session(); + let mut ctx = session.create_execution_ctx(); + + let a = nullable_point_column(vec![Some((0.0, 0.0)), None, Some((3.0, 4.0))])?; + let b = point_constant(0.0, 0.0, 3, &mut ctx)?; + let distance = GeoDistance::try_new_array(a, b)?.into_array(); + + let expected = PrimitiveArray::new( + vec![0.0f64, 0.0, 5.0], + Validity::from_iter([true, false, true]), + ) + .into_array(); + assert_arrays_eq!(distance, expected, &mut ctx); + Ok(()) + } + + /// Both operands nullable: a row is null if either operand is null there. + #[test] + fn distance_propagates_column_pair_nulls() -> VortexResult<()> { + let session = vortex_array::array_session(); + let mut ctx = session.create_execution_ctx(); + + let a = nullable_point_column(vec![Some((0.0, 0.0)), None, Some((0.0, 0.0))])?; + let b = nullable_point_column(vec![Some((3.0, 4.0)), Some((1.0, 1.0)), None])?; + let distance = GeoDistance::try_new_array(a, b)?.into_array(); + + let expected = PrimitiveArray::new( + vec![5.0f64, 0.0, 0.0], + Validity::from_iter([true, false, false]), + ) + .into_array(); + assert_arrays_eq!(distance, expected, &mut ctx); + Ok(()) + } + + /// A constant-null operand produces an all-null output. + #[test] + fn distance_constant_null_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 null_const = ConstantArray::new(Scalar::null(point_dtype), 3).into_array(); + let b = point_column(vec![0.0, 3.0, 0.0], vec![0.0, 0.0, 4.0])?; + let distance = GeoDistance::try_new_array(null_const, b)?.into_array(); + + let expected = PrimitiveArray::new(vec![0.0f64; 3], Validity::AllInvalid).into_array(); + assert_arrays_eq!(distance, expected, &mut ctx); Ok(()) } diff --git a/vortex-geo/src/scalar_fn/intersects.rs b/vortex-geo/src/scalar_fn/intersects.rs index 801745dc277..ee6dca8028c 100644 --- a/vortex-geo/src/scalar_fn/intersects.rs +++ b/vortex-geo/src/scalar_fn/intersects.rs @@ -6,14 +6,9 @@ use geo::Intersects; 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::ScalarFnArray; use vortex_array::dtype::DType; use vortex_array::dtype::Nullability; -use vortex_array::scalar::Scalar; use vortex_array::scalar_fn::Arity; use vortex_array::scalar_fn::ChildName; use vortex_array::scalar_fn::EmptyOptions; @@ -22,13 +17,11 @@ 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_eq; use vortex_session::VortexSession; use vortex_session::registry::CachedId; -use crate::extension::geometries; -use crate::extension::single_geometry; use crate::extension::validate_geometry_operands; +use crate::scalar_fn::execute_null_propagating; /// OGC `ST_Intersects` (not disjoint; boundary contact counts) between two native geometry /// operands, each a column or a constant literal. @@ -76,7 +69,8 @@ impl ScalarFnVTable for GeoIntersects { fn return_dtype(&self, _: &Self::Options, dtypes: &[DType]) -> VortexResult { validate_geometry_operands(dtypes)?; - Ok(DType::Bool(Nullability::NonNullable)) + let nullability = Nullability::from(dtypes.iter().any(DType::is_nullable)); + Ok(DType::Bool(nullability)) } fn execute( @@ -87,48 +81,10 @@ impl ScalarFnVTable for GeoIntersects { ) -> VortexResult { let a = args.get(0)?; let b = args.get(1)?; - match (a.as_opt::(), b.as_opt::()) { - (Some(qa), Some(qb)) => { - let ga = single_geometry(qa.scalar(), ctx)?; - let gb = single_geometry(qb.scalar(), ctx)?; - Ok(ConstantArray::new( - Scalar::bool(ga.intersects(&gb), Nullability::NonNullable), - a.len(), - ) - .into_array()) - } - (Some(query), None) => intersects_constant(&b, query.scalar(), ctx), - (None, Some(query)) => intersects_constant(&a, query.scalar(), ctx), - (None, None) => { - vortex_ensure_eq!( - a.len(), - b.len(), - "geo intersects: operand length mismatch {} vs {}", - a.len(), - b.len() - ); - let ag = geometries(&a, ctx)?; - let bg = geometries(&b, ctx)?; - let hits = ag.iter().zip(&bg).map(|(x, y)| x.intersects(y)); - Ok(BoolArray::from_iter(hits).into_array()) - } - } + execute_null_propagating(&a, &b, |x, y| x.intersects(y), ctx) } } -/// Whether each row of `operand` intersects the constant `query` geometry. Intersection is -/// symmetric, so this serves a constant on either side. -fn intersects_constant( - operand: &ArrayRef, - query: &Scalar, - ctx: &mut ExecutionCtx, -) -> VortexResult { - let query = single_geometry(query, ctx)?; - let geoms = geometries(operand, ctx)?; - let hits = geoms.iter().map(|g| g.intersects(&query)); - Ok(BoolArray::from_iter(hits).into_array()) -} - #[cfg(test)] mod tests { use geo_types::Coord; @@ -148,13 +104,17 @@ mod tests { use vortex_array::dtype::DType; use vortex_array::dtype::Nullability; use vortex_array::dtype::PType; + use vortex_array::scalar::Scalar; use vortex_array::scalar_fn::EmptyOptions; use vortex_array::scalar_fn::ScalarFnVTable; + use vortex_array::validity::Validity; + use vortex_buffer::BitBuffer; use vortex_error::VortexResult; use vortex_error::vortex_err; use wkb::writer::WriteOptions; use super::GeoIntersects; + use crate::test_harness::nullable_point_column; use crate::test_harness::point_column; /// A rectangle polygon with corners `(x0, y0)` and `(x1, y1)`, no holes. @@ -313,12 +273,53 @@ mod tests { assert_intersects(points, geometry_constant(&empty, 2)?, [false, false]) } - /// Geometry arrays are never nullable, so a nullable operand dtype is rejected. + /// Output nullability mirrors the operands: nullable if any operand is nullable, otherwise + /// non-nullable. #[test] - fn nullable_operand_is_rejected() -> VortexResult<()> { + fn output_nullability_mirrors_operands() -> VortexResult<()> { let dtype = point_column(vec![0.0], vec![0.0])?.dtype().clone(); - let result = GeoIntersects.return_dtype(&EmptyOptions, &[dtype.as_nullable(), dtype]); - assert!(result.is_err()); + let non_nullable = + GeoIntersects.return_dtype(&EmptyOptions, &[dtype.clone(), dtype.clone()])?; + assert!(!non_nullable.is_nullable()); + let nullable = GeoIntersects.return_dtype(&EmptyOptions, &[dtype.as_nullable(), dtype])?; + assert!(nullable.is_nullable()); + Ok(()) + } + + /// A null row in a geometry operand yields a null verdict; valid rows keep their verdict + /// (interior intersects, exterior does not). + #[test] + fn intersects_propagates_null_rows() -> VortexResult<()> { + let session = vortex_array::array_session(); + let mut ctx = session.create_execution_ctx(); + + let points = nullable_point_column(vec![Some((2.0, 2.0)), None, Some((20.0, 20.0))])?; + let query = geometry_constant(&donut(), 3)?; + let intersects = GeoIntersects::try_new_array(points, query)?.into_array(); + + let expected = BoolArray::new( + BitBuffer::from_iter([true, false, false]), + Validity::from_iter([true, false, true]), + ) + .into_array(); + assert_arrays_eq!(intersects, expected, &mut ctx); + Ok(()) + } + + /// A constant-null operand produces an all-null output. + #[test] + fn intersects_constant_null_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 null_const = ConstantArray::new(Scalar::null(point_dtype), 2).into_array(); + let points = point_column(vec![2.0, 20.0], vec![2.0, 20.0])?; + let intersects = GeoIntersects::try_new_array(null_const, points)?.into_array(); + + let expected = + BoolArray::new(BitBuffer::from_iter([false, false]), Validity::AllInvalid).into_array(); + assert_arrays_eq!(intersects, expected, &mut ctx); Ok(()) } diff --git a/vortex-geo/src/scalar_fn/mod.rs b/vortex-geo/src/scalar_fn/mod.rs index 2246f82621b..2293ca7140f 100644 --- a/vortex-geo/src/scalar_fn/mod.rs +++ b/vortex-geo/src/scalar_fn/mod.rs @@ -6,3 +6,246 @@ pub mod contains; pub mod distance; pub mod intersects; + +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::Constant; +use vortex_array::arrays::ConstantArray; +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_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; + + /// The nullable [`DType`] of this result (`f64` → `Primitive(F64, Nullable)`, `bool` → + /// `Bool(Nullable)`). + fn null_dtype() -> DType; + + /// Assemble the `len`-row output. `valid` marks which rows are non-null, and `values` holds + /// the results for those rows in order. + 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 null_dtype() -> DType { + DType::Primitive(PType::F64, Nullability::Nullable) + } + + 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 null_dtype() -> DType { + DType::Bool(Nullability::Nullable) + } + + 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::from_iter(vec![false; 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() + } + } + } +} + +/// An all-null output of length `len`. +fn all_null_array(len: usize) -> ArrayRef { + ConstantArray::new(Scalar::null(T::null_dtype()), len).into_array() +} + +/// 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. +pub(crate) fn execute_null_propagating( + a: &ArrayRef, + b: &ArrayRef, + compute: F, + 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()); + + match (a.as_opt::(), b.as_opt::()) { + // Both constant: compute once and broadcast across every row. + (Some(qa), Some(qb)) => { + // A null constant makes every row null. + if qa.scalar().is_null() || qb.scalar().is_null() { + return Ok(all_null_array::(len)); + } + 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. + (Some(qa), None) => { + if qa.scalar().is_null() { + return Ok(all_null_array::(len)); + } + let ga = single_geometry(qa.scalar(), ctx)?; + eval_column(b, |g| compute(&ga, g), nullability, ctx) + } + (None, Some(qb)) => { + if qb.scalar().is_null() { + return Ok(all_null_array::(len)); + } + let gb = single_geometry(qb.scalar(), ctx)?; + eval_column(a, |g| 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) + } + } +} + +/// 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)?; + // Every row null: nothing to decode. + if valid.all_false() { + return Ok(all_null_array::(len)); + } + // Drop the null rows before decoding, since a null row has no geometry to decode. The common + // all-valid case decodes the column directly and skips the filter. + let decoded = if valid.all_true() { + geometries(column, ctx)? + } else { + 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; + if valid.all_false() { + return Ok(all_null_array::(len)); + } + // Keep only the rows valid in both columns, so decoding never sees a null geometry. The + // common all-valid case decodes the columns directly and skips the filter. + let (a, b) = if valid.all_true() { + (a.clone(), b.clone()) + } else { + (a.filter(valid.clone())?, b.filter(valid.clone())?) + }; + let ag = geometries(&a, ctx)?; + let bg = geometries(&b, ctx)?; + let values = ag.iter().zip(&bg).map(|(x, y)| compute(x, y)).collect(); + Ok(T::build_array(len, &valid, values, nullability)) +} diff --git a/vortex-geo/src/test_harness.rs b/vortex-geo/src/test_harness.rs index 79dc2d3361d..f023b727089 100644 --- a/vortex-geo/src/test_harness.rs +++ b/vortex-geo/src/test_harness.rs @@ -10,6 +10,7 @@ use vortex_array::arrays::ListArray; use vortex_array::arrays::PrimitiveArray; use vortex_array::arrays::StructArray; 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::ExtVTable; @@ -111,6 +112,27 @@ pub(crate) fn point_column(xs: Vec, ys: Vec) -> VortexResult geo_column::(storage, storage_dtype) } +/// A nullable `Point` column: `None` rows are null. Null rows carry placeholder coordinates in +/// storage that the geo kernels must never decode (they filter nulls before decoding). +pub(crate) fn nullable_point_column(points: Vec>) -> VortexResult { + let len = points.len(); + let valid = points.iter().map(Option::is_some); + let xs = points.iter().map(|p| p.map_or(0.0, |(x, _)| x)); + let ys = points.iter().map(|p| p.map_or(0.0, |(_, y)| y)); + let storage = StructArray::try_new( + FieldNames::from(["x", "y"]), + vec![ + PrimitiveArray::from_iter(xs).into_array(), + PrimitiveArray::from_iter(ys).into_array(), + ], + len, + Validity::from_iter(valid), + )? + .into_array(); + let storage_dtype = storage.dtype().clone(); + geo_column::(storage, storage_dtype) +} + /// A `LineString` column: each line a list of `(x, y)` vertices, stored as `List>`. pub(crate) fn linestring_column(lines: Vec>) -> VortexResult { geo_column::( From 0716c2cbf48b1faeacc3b5ed22222178f6180806 Mon Sep 17 00:00:00 2001 From: Nemo Yu Date: Thu, 16 Jul 2026 16:08:12 -0400 Subject: [PATCH 2/4] refactor(vortex-geo): address null-propagation review - Add validity() + is_null_sensitive()=false vtable overrides to the three geo kernels so the planner can derive the output null mask symbolically (mask pushdown), matching vortex-tensor's binary functions. Shared via a binary_result_validity helper. is_fallible is left at the default since geometry decoding can genuinely fail. - Cover the two-column paths for the asymmetric ST_Contains: nulls on either side, plus the all-null-column and empty-combined-mask early returns. - Drop the unreachable AllOr::None arm in build_array; callers short-circuit the fully-null case to all_null_array. - Document that operands may differ in geometry type, and that the nullable test column's validity mirrors its storage. Signed-off-by: Nemo Yu --- vortex-geo/src/extension/mod.rs | 6 +- vortex-geo/src/scalar_fn/contains.rs | 79 ++++++++++++++++++++++++++ vortex-geo/src/scalar_fn/distance.rs | 14 +++++ vortex-geo/src/scalar_fn/intersects.rs | 14 +++++ vortex-geo/src/scalar_fn/mod.rs | 28 +++++++-- 5 files changed, 133 insertions(+), 8 deletions(-) diff --git a/vortex-geo/src/extension/mod.rs b/vortex-geo/src/extension/mod.rs index a9ff7ff66de..1a3db5a7b3e 100644 --- a/vortex-geo/src/extension/mod.rs +++ b/vortex-geo/src/extension/mod.rs @@ -73,8 +73,10 @@ 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. Nullable operands are allowed; the kernels propagate nulls (a null -/// geometry input yields a null result) rather than decoding null rows. +/// 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!( diff --git a/vortex-geo/src/scalar_fn/contains.rs b/vortex-geo/src/scalar_fn/contains.rs index 531e723a9f9..c1c8f4d8a22 100644 --- a/vortex-geo/src/scalar_fn/contains.rs +++ b/vortex-geo/src/scalar_fn/contains.rs @@ -9,6 +9,7 @@ use vortex_array::ExecutionCtx; use vortex_array::arrays::ScalarFnArray; use vortex_array::dtype::DType; use vortex_array::dtype::Nullability; +use vortex_array::expr::Expression; use vortex_array::scalar_fn::Arity; use vortex_array::scalar_fn::ChildName; use vortex_array::scalar_fn::EmptyOptions; @@ -21,6 +22,7 @@ use vortex_session::VortexSession; use vortex_session::registry::CachedId; use crate::extension::validate_geometry_operands; +use crate::scalar_fn::binary_result_validity; use crate::scalar_fn::execute_null_propagating; /// OGC `ST_Contains` between two native geometry operands, each a column or a constant @@ -85,6 +87,18 @@ impl ScalarFnVTable for GeoContains { // Containment is not symmetric: `a` is always the container and `b` the contained. execute_null_propagating(&a, &b, |a, b| a.contains(b), ctx) } + + fn validity( + &self, + _: &Self::Options, + expression: &Expression, + ) -> VortexResult> { + binary_result_validity(expression) + } + + fn is_null_sensitive(&self, _: &Self::Options) -> bool { + false + } } #[cfg(test)] @@ -285,6 +299,71 @@ mod tests { Ok(()) } + /// Both operands nullable columns: containment (asymmetric) is null wherever either the + /// container or the contained row is null, and computed on the rows valid in both. + #[test] + fn contains_propagates_column_pair_nulls() -> VortexResult<()> { + let session = vortex_array::array_session(); + let mut ctx = session.create_execution_ctx(); + + // A point contains another point only when they are equal. + let container = nullable_point_column(vec![ + Some((1.0, 1.0)), + None, + Some((2.0, 2.0)), + Some((3.0, 3.0)), + ])?; + let contained = nullable_point_column(vec![ + Some((1.0, 1.0)), + Some((5.0, 5.0)), + None, + Some((4.0, 4.0)), + ])?; + let contains = GeoContains::try_new_array(container, contained)?.into_array(); + + let expected = BoolArray::new( + BitBuffer::from_iter([true, false, false, false]), + Validity::from_iter([true, false, false, true]), + ) + .into_array(); + assert_arrays_eq!(contains, expected, &mut ctx); + Ok(()) + } + + /// An entirely-null geometry column yields an all-null output (the fully-null fast path in + /// `eval_column`). + #[test] + fn contains_all_null_column_is_all_null() -> VortexResult<()> { + let session = vortex_array::array_session(); + let mut ctx = session.create_execution_ctx(); + + let container = geometry_constant(&Geometry::Polygon(rect_polygon(0.0, 0.0, 4.0, 4.0)), 2)?; + let points = nullable_point_column(vec![None, None])?; + let contains = GeoContains::try_new_array(container, points)?.into_array(); + + let expected = + BoolArray::new(BitBuffer::from_iter([false, false]), Validity::AllInvalid).into_array(); + assert_arrays_eq!(contains, expected, &mut ctx); + Ok(()) + } + + /// Two nullable columns whose nulls never line up: the combined mask is empty, so the output + /// is all null (the `valid.all_false()` early return in `eval_column_pair`). + #[test] + fn contains_column_pair_all_null() -> VortexResult<()> { + let session = vortex_array::array_session(); + let mut ctx = session.create_execution_ctx(); + + let container = nullable_point_column(vec![Some((1.0, 1.0)), None])?; + let contained = nullable_point_column(vec![None, Some((2.0, 2.0))])?; + let contains = GeoContains::try_new_array(container, contained)?.into_array(); + + let expected = + BoolArray::new(BitBuffer::from_iter([false, false]), Validity::AllInvalid).into_array(); + assert_arrays_eq!(contains, expected, &mut ctx); + Ok(()) + } + /// A non-geometry operand dtype is rejected up front, before execution. #[test] fn non_geometry_operand_is_rejected() -> VortexResult<()> { diff --git a/vortex-geo/src/scalar_fn/distance.rs b/vortex-geo/src/scalar_fn/distance.rs index e9e739e663b..bc6923a03fd 100644 --- a/vortex-geo/src/scalar_fn/distance.rs +++ b/vortex-geo/src/scalar_fn/distance.rs @@ -11,6 +11,7 @@ use vortex_array::arrays::ScalarFnArray; use vortex_array::dtype::DType; use vortex_array::dtype::Nullability; use vortex_array::dtype::PType; +use vortex_array::expr::Expression; use vortex_array::scalar_fn::Arity; use vortex_array::scalar_fn::ChildName; use vortex_array::scalar_fn::EmptyOptions; @@ -23,6 +24,7 @@ use vortex_session::VortexSession; use vortex_session::registry::CachedId; use crate::extension::validate_geometry_operands; +use crate::scalar_fn::binary_result_validity; use crate::scalar_fn::execute_null_propagating; /// Planar (Euclidean) `ST_Distance` (no geodesic correction) between two native geometry @@ -85,6 +87,18 @@ impl ScalarFnVTable for GeoDistance { let b = args.get(1)?; execute_null_propagating(&a, &b, |x, y| Euclidean.distance(x, y), ctx) } + + fn validity( + &self, + _: &Self::Options, + expression: &Expression, + ) -> VortexResult> { + binary_result_validity(expression) + } + + fn is_null_sensitive(&self, _: &Self::Options) -> bool { + false + } } #[cfg(test)] diff --git a/vortex-geo/src/scalar_fn/intersects.rs b/vortex-geo/src/scalar_fn/intersects.rs index ee6dca8028c..b1d4e4e4567 100644 --- a/vortex-geo/src/scalar_fn/intersects.rs +++ b/vortex-geo/src/scalar_fn/intersects.rs @@ -9,6 +9,7 @@ use vortex_array::ExecutionCtx; use vortex_array::arrays::ScalarFnArray; use vortex_array::dtype::DType; use vortex_array::dtype::Nullability; +use vortex_array::expr::Expression; use vortex_array::scalar_fn::Arity; use vortex_array::scalar_fn::ChildName; use vortex_array::scalar_fn::EmptyOptions; @@ -21,6 +22,7 @@ use vortex_session::VortexSession; use vortex_session::registry::CachedId; use crate::extension::validate_geometry_operands; +use crate::scalar_fn::binary_result_validity; use crate::scalar_fn::execute_null_propagating; /// OGC `ST_Intersects` (not disjoint; boundary contact counts) between two native geometry @@ -83,6 +85,18 @@ impl ScalarFnVTable for GeoIntersects { let b = args.get(1)?; execute_null_propagating(&a, &b, |x, y| x.intersects(y), ctx) } + + fn validity( + &self, + _: &Self::Options, + expression: &Expression, + ) -> VortexResult> { + binary_result_validity(expression) + } + + fn is_null_sensitive(&self, _: &Self::Options) -> bool { + false + } } #[cfg(test)] diff --git a/vortex-geo/src/scalar_fn/mod.rs b/vortex-geo/src/scalar_fn/mod.rs index 2293ca7140f..fcc9cb0136d 100644 --- a/vortex-geo/src/scalar_fn/mod.rs +++ b/vortex-geo/src/scalar_fn/mod.rs @@ -18,6 +18,8 @@ use vortex_array::arrays::PrimitiveArray; use vortex_array::dtype::DType; use vortex_array::dtype::Nullability; use vortex_array::dtype::PType; +use vortex_array::expr::Expression; +use vortex_array::expr::and; use vortex_array::scalar::Scalar; use vortex_array::validity::Validity; use vortex_buffer::BitBuffer; @@ -71,8 +73,6 @@ impl GeoOutput for f64 { 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]; @@ -81,6 +81,11 @@ impl GeoOutput for f64 { } PrimitiveArray::new(data, validity).into_array() } + // The all-null case never reaches here, `eval_column` / `eval_column_pair` return + // `all_null_array` when the combined mask is empty. + AllOr::None => { + unreachable!("empty masks are handled by all_null_array in eval_column(_pair)") + } } } } @@ -104,10 +109,6 @@ impl GeoOutput for bool { 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::from_iter(vec![false; 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]; @@ -116,6 +117,11 @@ impl GeoOutput for bool { } BoolArray::new(BitBuffer::from_iter(data), validity).into_array() } + // The all-null case never reaches here — `eval_column` / `eval_column_pair` return + // `all_null_array` when the combined mask is empty. + AllOr::None => { + unreachable!("empty masks are handled by all_null_array in eval_column(_pair)") + } } } } @@ -125,6 +131,16 @@ fn all_null_array(len: usize) -> ArrayRef { ConstantArray::new(Scalar::null(T::null_dtype()), len).into_array() } +/// The validity expression for a binary geo kernel: the result is null iff either operand is +/// null. Returning this lets the planner derive the output's null mask symbolically (mask +/// pushdown) instead of executing the whole kernel. +pub(crate) fn binary_result_validity(expression: &Expression) -> VortexResult> { + Ok(Some(and( + expression.child(0).validity()?, + expression.child(1).validity()?, + ))) +} + /// 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 From 9ad937625ff6cc55a0cf27590688d376382d68ae Mon Sep 17 00:00:00 2001 From: Nemo Yu Date: Thu, 16 Jul 2026 16:42:01 -0400 Subject: [PATCH 3/4] refactor(vortex-geo): move null-propagating dispatch to its own module Address PR review round 2: - Move the shared null-propagating dispatch out of scalar_fn/mod.rs into scalar_fn/null_propagate.rs, keeping mod.rs to module-tree wiring. - Bring ST_Distance and ST_Intersects null-path coverage to parity with ST_Contains: all-null-column and empty-combined-mask early returns, plus the two-column path for ST_Intersects. Signed-off-by: Nemo Yu --- vortex-geo/src/scalar_fn/contains.rs | 4 +- vortex-geo/src/scalar_fn/distance.rs | 36 ++- vortex-geo/src/scalar_fn/intersects.rs | 68 +++++- vortex-geo/src/scalar_fn/mod.rs | 260 +------------------- vortex-geo/src/scalar_fn/null_propagate.rs | 268 +++++++++++++++++++++ 5 files changed, 371 insertions(+), 265 deletions(-) create mode 100644 vortex-geo/src/scalar_fn/null_propagate.rs diff --git a/vortex-geo/src/scalar_fn/contains.rs b/vortex-geo/src/scalar_fn/contains.rs index c1c8f4d8a22..83216e58833 100644 --- a/vortex-geo/src/scalar_fn/contains.rs +++ b/vortex-geo/src/scalar_fn/contains.rs @@ -22,8 +22,8 @@ use vortex_session::VortexSession; use vortex_session::registry::CachedId; use crate::extension::validate_geometry_operands; -use crate::scalar_fn::binary_result_validity; -use crate::scalar_fn::execute_null_propagating; +use crate::scalar_fn::null_propagate::binary_result_validity; +use crate::scalar_fn::null_propagate::execute_null_propagating; /// 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 diff --git a/vortex-geo/src/scalar_fn/distance.rs b/vortex-geo/src/scalar_fn/distance.rs index bc6923a03fd..fc667be7439 100644 --- a/vortex-geo/src/scalar_fn/distance.rs +++ b/vortex-geo/src/scalar_fn/distance.rs @@ -24,8 +24,8 @@ use vortex_session::VortexSession; use vortex_session::registry::CachedId; use crate::extension::validate_geometry_operands; -use crate::scalar_fn::binary_result_validity; -use crate::scalar_fn::execute_null_propagating; +use crate::scalar_fn::null_propagate::binary_result_validity; +use crate::scalar_fn::null_propagate::execute_null_propagating; /// Planar (Euclidean) `ST_Distance` (no geodesic correction) between two native geometry /// operands, each a column or a constant literal. @@ -268,6 +268,38 @@ mod tests { Ok(()) } + /// An entirely-null geometry column yields an all-null output (the fully-null fast path in + /// `eval_column`). + #[test] + fn distance_all_null_column_is_all_null() -> VortexResult<()> { + let session = vortex_array::array_session(); + let mut ctx = session.create_execution_ctx(); + + let a = nullable_point_column(vec![None, None])?; + let b = point_constant(0.0, 0.0, 2, &mut ctx)?; + let distance = GeoDistance::try_new_array(a, b)?.into_array(); + + let expected = PrimitiveArray::new(vec![0.0f64; 2], Validity::AllInvalid).into_array(); + assert_arrays_eq!(distance, expected, &mut ctx); + Ok(()) + } + + /// Two nullable columns whose nulls never line up: the combined mask is empty, so the output + /// is all null (the `valid.all_false()` early return in `eval_column_pair`). + #[test] + fn distance_column_pair_all_null() -> VortexResult<()> { + let session = vortex_array::array_session(); + let mut ctx = session.create_execution_ctx(); + + let a = nullable_point_column(vec![Some((0.0, 0.0)), None])?; + let b = nullable_point_column(vec![None, Some((1.0, 1.0))])?; + let distance = GeoDistance::try_new_array(a, b)?.into_array(); + + let expected = PrimitiveArray::new(vec![0.0f64; 2], Validity::AllInvalid).into_array(); + assert_arrays_eq!(distance, expected, &mut ctx); + Ok(()) + } + /// A non-geometry operand dtype is rejected up front, before execution. #[test] fn non_geometry_operand_is_rejected() -> VortexResult<()> { diff --git a/vortex-geo/src/scalar_fn/intersects.rs b/vortex-geo/src/scalar_fn/intersects.rs index b1d4e4e4567..301394e15da 100644 --- a/vortex-geo/src/scalar_fn/intersects.rs +++ b/vortex-geo/src/scalar_fn/intersects.rs @@ -22,8 +22,8 @@ use vortex_session::VortexSession; use vortex_session::registry::CachedId; use crate::extension::validate_geometry_operands; -use crate::scalar_fn::binary_result_validity; -use crate::scalar_fn::execute_null_propagating; +use crate::scalar_fn::null_propagate::binary_result_validity; +use crate::scalar_fn::null_propagate::execute_null_propagating; /// OGC `ST_Intersects` (not disjoint; boundary contact counts) between two native geometry /// operands, each a column or a constant literal. @@ -337,6 +337,70 @@ mod tests { Ok(()) } + /// Both operands nullable columns: the verdict is null wherever either row is null, and + /// computed on the rows valid in both (points intersect exactly when equal). + #[test] + fn intersects_propagates_column_pair_nulls() -> VortexResult<()> { + let session = vortex_array::array_session(); + let mut ctx = session.create_execution_ctx(); + + let a = nullable_point_column(vec![ + Some((0.0, 0.0)), + None, + Some((2.0, 2.0)), + Some((3.0, 3.0)), + ])?; + let b = nullable_point_column(vec![ + Some((0.0, 0.0)), + Some((5.0, 5.0)), + None, + Some((9.0, 9.0)), + ])?; + let intersects = GeoIntersects::try_new_array(a, b)?.into_array(); + + let expected = BoolArray::new( + BitBuffer::from_iter([true, false, false, false]), + Validity::from_iter([true, false, false, true]), + ) + .into_array(); + assert_arrays_eq!(intersects, expected, &mut ctx); + Ok(()) + } + + /// An entirely-null geometry column yields an all-null output (the fully-null fast path in + /// `eval_column`). + #[test] + fn intersects_all_null_column_is_all_null() -> VortexResult<()> { + let session = vortex_array::array_session(); + let mut ctx = session.create_execution_ctx(); + + let points = nullable_point_column(vec![None, None])?; + let query = geometry_constant(&donut(), 2)?; + let intersects = GeoIntersects::try_new_array(points, query)?.into_array(); + + let expected = + BoolArray::new(BitBuffer::from_iter([false, false]), Validity::AllInvalid).into_array(); + assert_arrays_eq!(intersects, expected, &mut ctx); + Ok(()) + } + + /// Two nullable columns whose nulls never line up: the combined mask is empty, so the output + /// is all null (the `valid.all_false()` early return in `eval_column_pair`). + #[test] + fn intersects_column_pair_all_null() -> VortexResult<()> { + let session = vortex_array::array_session(); + let mut ctx = session.create_execution_ctx(); + + let a = nullable_point_column(vec![Some((0.0, 0.0)), None])?; + let b = nullable_point_column(vec![None, Some((1.0, 1.0))])?; + let intersects = GeoIntersects::try_new_array(a, b)?.into_array(); + + let expected = + BoolArray::new(BitBuffer::from_iter([false, false]), Validity::AllInvalid).into_array(); + assert_arrays_eq!(intersects, expected, &mut ctx); + Ok(()) + } + /// A non-geometry operand dtype is rejected up front, before execution. #[test] fn non_geometry_operand_is_rejected() -> VortexResult<()> { diff --git a/vortex-geo/src/scalar_fn/mod.rs b/vortex-geo/src/scalar_fn/mod.rs index fcc9cb0136d..4bf2b719629 100644 --- a/vortex-geo/src/scalar_fn/mod.rs +++ b/vortex-geo/src/scalar_fn/mod.rs @@ -6,262 +6,4 @@ pub mod contains; pub mod distance; pub mod intersects; - -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::Constant; -use vortex_array::arrays::ConstantArray; -use vortex_array::arrays::PrimitiveArray; -use vortex_array::dtype::DType; -use vortex_array::dtype::Nullability; -use vortex_array::dtype::PType; -use vortex_array::expr::Expression; -use vortex_array::expr::and; -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; - - /// The nullable [`DType`] of this result (`f64` → `Primitive(F64, Nullable)`, `bool` → - /// `Bool(Nullable)`). - fn null_dtype() -> DType; - - /// Assemble the `len`-row output. `valid` marks which rows are non-null, and `values` holds - /// the results for those rows in order. - 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 null_dtype() -> DType { - DType::Primitive(PType::F64, Nullability::Nullable) - } - - 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(), - // 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() - } - // The all-null case never reaches here, `eval_column` / `eval_column_pair` return - // `all_null_array` when the combined mask is empty. - AllOr::None => { - unreachable!("empty masks are handled by all_null_array in eval_column(_pair)") - } - } - } -} - -impl GeoOutput for bool { - fn into_scalar(self, nullability: Nullability) -> Scalar { - Scalar::bool(self, nullability) - } - - fn null_dtype() -> DType { - DType::Bool(Nullability::Nullable) - } - - 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(), - // 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() - } - // The all-null case never reaches here — `eval_column` / `eval_column_pair` return - // `all_null_array` when the combined mask is empty. - AllOr::None => { - unreachable!("empty masks are handled by all_null_array in eval_column(_pair)") - } - } - } -} - -/// An all-null output of length `len`. -fn all_null_array(len: usize) -> ArrayRef { - ConstantArray::new(Scalar::null(T::null_dtype()), len).into_array() -} - -/// The validity expression for a binary geo kernel: the result is null iff either operand is -/// null. Returning this lets the planner derive the output's null mask symbolically (mask -/// pushdown) instead of executing the whole kernel. -pub(crate) fn binary_result_validity(expression: &Expression) -> VortexResult> { - Ok(Some(and( - expression.child(0).validity()?, - expression.child(1).validity()?, - ))) -} - -/// 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. -pub(crate) fn execute_null_propagating( - a: &ArrayRef, - b: &ArrayRef, - compute: F, - 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()); - - match (a.as_opt::(), b.as_opt::()) { - // Both constant: compute once and broadcast across every row. - (Some(qa), Some(qb)) => { - // A null constant makes every row null. - if qa.scalar().is_null() || qb.scalar().is_null() { - return Ok(all_null_array::(len)); - } - 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. - (Some(qa), None) => { - if qa.scalar().is_null() { - return Ok(all_null_array::(len)); - } - let ga = single_geometry(qa.scalar(), ctx)?; - eval_column(b, |g| compute(&ga, g), nullability, ctx) - } - (None, Some(qb)) => { - if qb.scalar().is_null() { - return Ok(all_null_array::(len)); - } - let gb = single_geometry(qb.scalar(), ctx)?; - eval_column(a, |g| 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) - } - } -} - -/// 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)?; - // Every row null: nothing to decode. - if valid.all_false() { - return Ok(all_null_array::(len)); - } - // Drop the null rows before decoding, since a null row has no geometry to decode. The common - // all-valid case decodes the column directly and skips the filter. - let decoded = if valid.all_true() { - geometries(column, ctx)? - } else { - 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; - if valid.all_false() { - return Ok(all_null_array::(len)); - } - // Keep only the rows valid in both columns, so decoding never sees a null geometry. The - // common all-valid case decodes the columns directly and skips the filter. - let (a, b) = if valid.all_true() { - (a.clone(), b.clone()) - } else { - (a.filter(valid.clone())?, b.filter(valid.clone())?) - }; - let ag = geometries(&a, ctx)?; - let bg = geometries(&b, ctx)?; - let values = ag.iter().zip(&bg).map(|(x, y)| compute(x, y)).collect(); - Ok(T::build_array(len, &valid, values, nullability)) -} +mod null_propagate; diff --git a/vortex-geo/src/scalar_fn/null_propagate.rs b/vortex-geo/src/scalar_fn/null_propagate.rs new file mode 100644 index 00000000000..3b0c4ac8139 --- /dev/null +++ b/vortex-geo/src/scalar_fn/null_propagate.rs @@ -0,0 +1,268 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright the Vortex contributors + +//! Null-propagating execution shared by the binary geo scalar functions. +//! +//! [`execute_null_propagating`] runs a binary geo kernel (`ST_Distance`, `ST_Intersects`, +//! `ST_Contains`) over its two operands with SQL null semantics: the result is null wherever +//! either operand is null. [`binary_result_validity`] is the matching validity expression the +//! planner can use for mask pushdown. + +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::Constant; +use vortex_array::arrays::ConstantArray; +use vortex_array::arrays::PrimitiveArray; +use vortex_array::dtype::DType; +use vortex_array::dtype::Nullability; +use vortex_array::dtype::PType; +use vortex_array::expr::Expression; +use vortex_array::expr::and; +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; + + /// The nullable [`DType`] of this result (`f64` → `Primitive(F64, Nullable)`, `bool` → + /// `Bool(Nullable)`). + fn null_dtype() -> DType; + + /// Assemble the `len`-row output. `valid` marks which rows are non-null, and `values` holds + /// the results for those rows in order. + 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 null_dtype() -> DType { + DType::Primitive(PType::F64, Nullability::Nullable) + } + + 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(), + // 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() + } + // The all-null case never reaches here, `eval_column` / `eval_column_pair` return + // `all_null_array` when the combined mask is empty. + AllOr::None => { + unreachable!("empty masks are handled by all_null_array in eval_column(_pair)") + } + } + } +} + +impl GeoOutput for bool { + fn into_scalar(self, nullability: Nullability) -> Scalar { + Scalar::bool(self, nullability) + } + + fn null_dtype() -> DType { + DType::Bool(Nullability::Nullable) + } + + 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(), + // 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() + } + // The all-null case never reaches here — `eval_column` / `eval_column_pair` return + // `all_null_array` when the combined mask is empty. + AllOr::None => { + unreachable!("empty masks are handled by all_null_array in eval_column(_pair)") + } + } + } +} + +/// An all-null output of length `len`. +fn all_null_array(len: usize) -> ArrayRef { + ConstantArray::new(Scalar::null(T::null_dtype()), len).into_array() +} + +/// The validity expression for a binary geo kernel: the result is null iff either operand is +/// null. Returning this lets the planner derive the output's null mask symbolically (mask +/// pushdown) instead of executing the whole kernel. +pub(crate) fn binary_result_validity(expression: &Expression) -> VortexResult> { + Ok(Some(and( + expression.child(0).validity()?, + expression.child(1).validity()?, + ))) +} + +/// 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. +pub(crate) fn execute_null_propagating( + a: &ArrayRef, + b: &ArrayRef, + compute: F, + 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()); + + match (a.as_opt::(), b.as_opt::()) { + // Both constant: compute once and broadcast across every row. + (Some(qa), Some(qb)) => { + // A null constant makes every row null. + if qa.scalar().is_null() || qb.scalar().is_null() { + return Ok(all_null_array::(len)); + } + 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. + (Some(qa), None) => { + if qa.scalar().is_null() { + return Ok(all_null_array::(len)); + } + let ga = single_geometry(qa.scalar(), ctx)?; + eval_column(b, |g| compute(&ga, g), nullability, ctx) + } + (None, Some(qb)) => { + if qb.scalar().is_null() { + return Ok(all_null_array::(len)); + } + let gb = single_geometry(qb.scalar(), ctx)?; + eval_column(a, |g| 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) + } + } +} + +/// 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)?; + // Every row null: nothing to decode. + if valid.all_false() { + return Ok(all_null_array::(len)); + } + // Drop the null rows before decoding, since a null row has no geometry to decode. The common + // all-valid case decodes the column directly and skips the filter. + let decoded = if valid.all_true() { + geometries(column, ctx)? + } else { + 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; + if valid.all_false() { + return Ok(all_null_array::(len)); + } + // Keep only the rows valid in both columns, so decoding never sees a null geometry. The + // common all-valid case decodes the columns directly and skips the filter. + let (a, b) = if valid.all_true() { + (a.clone(), b.clone()) + } else { + (a.filter(valid.clone())?, b.filter(valid.clone())?) + }; + let ag = geometries(&a, ctx)?; + let bg = geometries(&b, ctx)?; + let values = ag.iter().zip(&bg).map(|(x, y)| compute(x, y)).collect(); + Ok(T::build_array(len, &valid, values, nullability)) +} From 610a8c3c43322396954a8c3ce3891654e50d923e Mon Sep 17 00:00:00 2001 From: Nemo Yu Date: Thu, 16 Jul 2026 17:04:15 -0400 Subject: [PATCH 4/4] refactor: hoist null-propagating validity into a vortex-array helper MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Promote the "result is null iff any operand is null" validity — previously geo's binary_result_validity plus four inline copies in vortex-tensor — into a shared helper, vortex_array::expr::null_propagating_validity. It conjoins every child's validity, so it covers unary (l2_norm) and binary kernels alike. - vortex-array: add expr::null_propagating_validity - vortex-geo: GeoDistance / GeoContains / GeoIntersects use it; drop the local binary_result_validity - vortex-tensor: InnerProduct / CosineSimilarity / L2Denorm / L2Norm use it instead of inlining and(child validities) Signed-off-by: Nemo Yu --- vortex-array/src/expr/exprs.rs | 17 +++++++++++++++++ vortex-geo/src/scalar_fn/contains.rs | 4 ++-- vortex-geo/src/scalar_fn/distance.rs | 4 ++-- vortex-geo/src/scalar_fn/intersects.rs | 4 ++-- vortex-geo/src/scalar_fn/null_propagate.rs | 12 ------------ .../src/scalar_fns/cosine_similarity.rs | 7 ++----- vortex-tensor/src/scalar_fns/inner_product.rs | 7 ++----- vortex-tensor/src/scalar_fns/l2_denorm.rs | 7 ++----- vortex-tensor/src/scalar_fns/l2_norm.rs | 3 ++- 9 files changed, 31 insertions(+), 34 deletions(-) diff --git a/vortex-array/src/expr/exprs.rs b/vortex-array/src/expr/exprs.rs index da70427a0f7..aef254b144a 100644 --- a/vortex-array/src/expr/exprs.rs +++ b/vortex-array/src/expr/exprs.rs @@ -7,6 +7,7 @@ use std::sync::Arc; use std::sync::LazyLock; use vortex_error::VortexExpect; +use vortex_error::VortexResult; use vortex_error::vortex_panic; use vortex_utils::iter::ReduceBalancedIterExt; @@ -434,6 +435,22 @@ where iter.into_iter().reduce_balanced(and) } +/// The validity of a scalar function whose result is null exactly when any operand is null: the +/// conjunction of its children's validities. +/// +/// This is the `ScalarFnVTable::validity` for kernels that propagate nulls and never produce a +/// null from non-null inputs (comparisons, arithmetic, most geo and tensor ops). Returning it lets +/// the planner derive the output's null mask without executing the kernel. Yields `None` when the +/// expression has no children. +pub fn null_propagating_validity(expression: &Expression) -> VortexResult> { + let child_validities = expression + .children() + .iter() + .map(Expression::validity) + .collect::>>()?; + Ok(and_collect(child_validities)) +} + /// Create a new [`Binary`] using the [`Add`](Operator::Add) operator. /// /// ## Example usage diff --git a/vortex-geo/src/scalar_fn/contains.rs b/vortex-geo/src/scalar_fn/contains.rs index 83216e58833..57df3e72cd3 100644 --- a/vortex-geo/src/scalar_fn/contains.rs +++ b/vortex-geo/src/scalar_fn/contains.rs @@ -10,6 +10,7 @@ use vortex_array::arrays::ScalarFnArray; use vortex_array::dtype::DType; use vortex_array::dtype::Nullability; use vortex_array::expr::Expression; +use vortex_array::expr::null_propagating_validity; use vortex_array::scalar_fn::Arity; use vortex_array::scalar_fn::ChildName; use vortex_array::scalar_fn::EmptyOptions; @@ -22,7 +23,6 @@ use vortex_session::VortexSession; use vortex_session::registry::CachedId; use crate::extension::validate_geometry_operands; -use crate::scalar_fn::null_propagate::binary_result_validity; use crate::scalar_fn::null_propagate::execute_null_propagating; /// OGC `ST_Contains` between two native geometry operands, each a column or a constant @@ -93,7 +93,7 @@ impl ScalarFnVTable for GeoContains { _: &Self::Options, expression: &Expression, ) -> VortexResult> { - binary_result_validity(expression) + null_propagating_validity(expression) } fn is_null_sensitive(&self, _: &Self::Options) -> bool { diff --git a/vortex-geo/src/scalar_fn/distance.rs b/vortex-geo/src/scalar_fn/distance.rs index fc667be7439..fe56dc81635 100644 --- a/vortex-geo/src/scalar_fn/distance.rs +++ b/vortex-geo/src/scalar_fn/distance.rs @@ -12,6 +12,7 @@ use vortex_array::dtype::DType; use vortex_array::dtype::Nullability; use vortex_array::dtype::PType; use vortex_array::expr::Expression; +use vortex_array::expr::null_propagating_validity; use vortex_array::scalar_fn::Arity; use vortex_array::scalar_fn::ChildName; use vortex_array::scalar_fn::EmptyOptions; @@ -24,7 +25,6 @@ use vortex_session::VortexSession; use vortex_session::registry::CachedId; use crate::extension::validate_geometry_operands; -use crate::scalar_fn::null_propagate::binary_result_validity; use crate::scalar_fn::null_propagate::execute_null_propagating; /// Planar (Euclidean) `ST_Distance` (no geodesic correction) between two native geometry @@ -93,7 +93,7 @@ impl ScalarFnVTable for GeoDistance { _: &Self::Options, expression: &Expression, ) -> VortexResult> { - binary_result_validity(expression) + null_propagating_validity(expression) } fn is_null_sensitive(&self, _: &Self::Options) -> bool { diff --git a/vortex-geo/src/scalar_fn/intersects.rs b/vortex-geo/src/scalar_fn/intersects.rs index 301394e15da..a2f910f42c7 100644 --- a/vortex-geo/src/scalar_fn/intersects.rs +++ b/vortex-geo/src/scalar_fn/intersects.rs @@ -10,6 +10,7 @@ use vortex_array::arrays::ScalarFnArray; use vortex_array::dtype::DType; use vortex_array::dtype::Nullability; use vortex_array::expr::Expression; +use vortex_array::expr::null_propagating_validity; use vortex_array::scalar_fn::Arity; use vortex_array::scalar_fn::ChildName; use vortex_array::scalar_fn::EmptyOptions; @@ -22,7 +23,6 @@ use vortex_session::VortexSession; use vortex_session::registry::CachedId; use crate::extension::validate_geometry_operands; -use crate::scalar_fn::null_propagate::binary_result_validity; use crate::scalar_fn::null_propagate::execute_null_propagating; /// OGC `ST_Intersects` (not disjoint; boundary contact counts) between two native geometry @@ -91,7 +91,7 @@ impl ScalarFnVTable for GeoIntersects { _: &Self::Options, expression: &Expression, ) -> VortexResult> { - binary_result_validity(expression) + null_propagating_validity(expression) } fn is_null_sensitive(&self, _: &Self::Options) -> bool { diff --git a/vortex-geo/src/scalar_fn/null_propagate.rs b/vortex-geo/src/scalar_fn/null_propagate.rs index 3b0c4ac8139..24f4840ad57 100644 --- a/vortex-geo/src/scalar_fn/null_propagate.rs +++ b/vortex-geo/src/scalar_fn/null_propagate.rs @@ -19,8 +19,6 @@ use vortex_array::arrays::PrimitiveArray; use vortex_array::dtype::DType; use vortex_array::dtype::Nullability; use vortex_array::dtype::PType; -use vortex_array::expr::Expression; -use vortex_array::expr::and; use vortex_array::scalar::Scalar; use vortex_array::validity::Validity; use vortex_buffer::BitBuffer; @@ -132,16 +130,6 @@ fn all_null_array(len: usize) -> ArrayRef { ConstantArray::new(Scalar::null(T::null_dtype()), len).into_array() } -/// The validity expression for a binary geo kernel: the result is null iff either operand is -/// null. Returning this lets the planner derive the output's null mask symbolically (mask -/// pushdown) instead of executing the whole kernel. -pub(crate) fn binary_result_validity(expression: &Expression) -> VortexResult> { - Ok(Some(and( - expression.child(0).validity()?, - expression.child(1).validity()?, - ))) -} - /// 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 diff --git a/vortex-tensor/src/scalar_fns/cosine_similarity.rs b/vortex-tensor/src/scalar_fns/cosine_similarity.rs index 82328a9b2d9..4262b91c71b 100644 --- a/vortex-tensor/src/scalar_fns/cosine_similarity.rs +++ b/vortex-tensor/src/scalar_fns/cosine_similarity.rs @@ -15,7 +15,7 @@ use vortex_array::arrays::scalar_fn::plugin::ScalarFnArrayVTable; use vortex_array::dtype::DType; use vortex_array::dtype::Nullability; use vortex_array::expr::Expression; -use vortex_array::expr::and; +use vortex_array::expr::null_propagating_validity; use vortex_array::match_each_float_ptype; use vortex_array::scalar_fn::Arity; use vortex_array::scalar_fn::ChildName; @@ -180,10 +180,7 @@ impl ScalarFnVTable for CosineSimilarity { expression: &Expression, ) -> VortexResult> { // The result is null if either input tensor is null. - let lhs_validity = expression.child(0).validity()?; - let rhs_validity = expression.child(1).validity()?; - - Ok(Some(and(lhs_validity, rhs_validity))) + null_propagating_validity(expression) } fn is_null_sensitive(&self, _options: &Self::Options) -> bool { diff --git a/vortex-tensor/src/scalar_fns/inner_product.rs b/vortex-tensor/src/scalar_fns/inner_product.rs index cfea956f314..7fcf4b01261 100644 --- a/vortex-tensor/src/scalar_fns/inner_product.rs +++ b/vortex-tensor/src/scalar_fns/inner_product.rs @@ -18,7 +18,7 @@ use vortex_array::dtype::DType; use vortex_array::dtype::NativePType; use vortex_array::dtype::Nullability; use vortex_array::expr::Expression; -use vortex_array::expr::and; +use vortex_array::expr::null_propagating_validity; use vortex_array::match_each_float_ptype; use vortex_array::scalar_fn::Arity; use vortex_array::scalar_fn::ChildName; @@ -164,10 +164,7 @@ impl ScalarFnVTable for InnerProduct { expression: &Expression, ) -> VortexResult> { // The result is null if either input tensor is null. - let lhs_validity = expression.child(0).validity()?; - let rhs_validity = expression.child(1).validity()?; - - Ok(Some(and(lhs_validity, rhs_validity))) + null_propagating_validity(expression) } fn is_null_sensitive(&self, _options: &Self::Options) -> bool { diff --git a/vortex-tensor/src/scalar_fns/l2_denorm.rs b/vortex-tensor/src/scalar_fns/l2_denorm.rs index 4f4159b4a21..09d47a12bb1 100644 --- a/vortex-tensor/src/scalar_fns/l2_denorm.rs +++ b/vortex-tensor/src/scalar_fns/l2_denorm.rs @@ -31,7 +31,7 @@ use vortex_array::dtype::NativePType; use vortex_array::dtype::Nullability; use vortex_array::dtype::proto::dtype as pb; use vortex_array::expr::Expression; -use vortex_array::expr::and; +use vortex_array::expr::null_propagating_validity; use vortex_array::match_each_float_ptype; use vortex_array::scalar::Scalar; use vortex_array::scalar::ScalarValue; @@ -256,10 +256,7 @@ impl ScalarFnVTable for L2Denorm { _options: &Self::Options, expression: &Expression, ) -> VortexResult> { - let normalized_validity = expression.child(0).validity()?; - let norms_validity = expression.child(1).validity()?; - - Ok(Some(and(normalized_validity, norms_validity))) + null_propagating_validity(expression) } fn is_null_sensitive(&self, _options: &Self::Options) -> bool { diff --git a/vortex-tensor/src/scalar_fns/l2_norm.rs b/vortex-tensor/src/scalar_fns/l2_norm.rs index aa85ab4c78e..40dae03413b 100644 --- a/vortex-tensor/src/scalar_fns/l2_norm.rs +++ b/vortex-tensor/src/scalar_fns/l2_norm.rs @@ -25,6 +25,7 @@ use vortex_array::dtype::NativePType; use vortex_array::dtype::Nullability; use vortex_array::dtype::proto::dtype as pb; use vortex_array::expr::Expression; +use vortex_array::expr::null_propagating_validity; use vortex_array::match_each_float_ptype; use vortex_array::scalar::Scalar; use vortex_array::scalar_fn::Arity; @@ -184,7 +185,7 @@ impl ScalarFnVTable for L2Norm { expression: &Expression, ) -> VortexResult> { // The result is null if the input tensor is null. - Ok(Some(expression.child(0).validity()?)) + null_propagating_validity(expression) } fn is_null_sensitive(&self, _options: &Self::Options) -> bool {