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-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/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..1a3db5a7b3e 100644 --- a/vortex-geo/src/extension/mod.rs +++ b/vortex-geo/src/extension/mod.rs @@ -72,18 +72,17 @@ 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. 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" ); - 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..57df3e72cd3 100644 --- a/vortex-geo/src/scalar_fn/contains.rs +++ b/vortex-geo/src/scalar_fn/contains.rs @@ -6,14 +6,11 @@ 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::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,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::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 @@ -77,7 +72,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,57 +85,20 @@ 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()) -} + fn validity( + &self, + _: &Self::Options, + expression: &Expression, + ) -> VortexResult> { + null_propagating_validity(expression) + } -/// 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()) + fn is_null_sensitive(&self, _: &Self::Options) -> bool { + false + } } #[cfg(test)] @@ -160,13 +119,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 +249,118 @@ 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(()) + } + + /// 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(()) } diff --git a/vortex-geo/src/scalar_fn/distance.rs b/vortex-geo/src/scalar_fn/distance.rs index 2e196706ec2..fe56dc81635 100644 --- a/vortex-geo/src/scalar_fn/distance.rs +++ b/vortex-geo/src/scalar_fn/distance.rs @@ -7,15 +7,12 @@ 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::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,13 +21,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::null_propagate::execute_null_propagating; /// Planar (Euclidean) `ST_Distance` (no geodesic correction) between two native geometry /// operands, each a column or a constant literal. @@ -78,7 +73,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,47 +85,20 @@ 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()) + fn validity( + &self, + _: &Self::Options, + expression: &Expression, + ) -> VortexResult> { + null_propagating_validity(expression) + } + + fn is_null_sensitive(&self, _: &Self::Options) -> bool { + false + } } #[cfg(test)] @@ -140,14 +109,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 +201,102 @@ 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(()) + } + + /// 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(()) } diff --git a/vortex-geo/src/scalar_fn/intersects.rs b/vortex-geo/src/scalar_fn/intersects.rs index 801745dc277..a2f910f42c7 100644 --- a/vortex-geo/src/scalar_fn/intersects.rs +++ b/vortex-geo/src/scalar_fn/intersects.rs @@ -6,14 +6,11 @@ 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::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,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::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. @@ -76,7 +71,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,46 +83,20 @@ 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()) + fn validity( + &self, + _: &Self::Options, + expression: &Expression, + ) -> VortexResult> { + null_propagating_validity(expression) + } + + fn is_null_sensitive(&self, _: &Self::Options) -> bool { + false + } } #[cfg(test)] @@ -148,13 +118,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 +287,117 @@ 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(()) + } + + /// 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(()) } diff --git a/vortex-geo/src/scalar_fn/mod.rs b/vortex-geo/src/scalar_fn/mod.rs index 2246f82621b..4bf2b719629 100644 --- a/vortex-geo/src/scalar_fn/mod.rs +++ b/vortex-geo/src/scalar_fn/mod.rs @@ -6,3 +6,4 @@ pub mod contains; pub mod distance; pub mod intersects; +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..24f4840ad57 --- /dev/null +++ b/vortex-geo/src/scalar_fn/null_propagate.rs @@ -0,0 +1,256 @@ +// 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::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() +} + +/// 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::( 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 {