diff --git a/vortex-array/src/expr/transform/match_between.rs b/vortex-array/src/expr/transform/match_between.rs index 56f03dbac4d..7cb427596f4 100644 --- a/vortex-array/src/expr/transform/match_between.rs +++ b/vortex-array/src/expr/transform/match_between.rs @@ -125,7 +125,18 @@ fn is_strict_comparison(op: Operator) -> Option { #[cfg(test)] mod tests { + use vortex_buffer::buffer; + use vortex_error::VortexResult; + use super::find_between; + use crate::IntoArray; + use crate::VortexSessionExecute; + use crate::array_session; + use crate::arrays::BoolArray; + use crate::arrays::StructArray; + use crate::dtype::DType; + use crate::dtype::Nullability; + use crate::dtype::PType; use crate::expr::and; use crate::expr::between; use crate::expr::col; @@ -134,9 +145,42 @@ mod tests { use crate::expr::lit; use crate::expr::lt; use crate::expr::lt_eq; + use crate::scalar::Scalar; use crate::scalar_fn::fns::between::BetweenOptions; use crate::scalar_fn::fns::between::StrictComparison; + /// A null literal bound must not change the values of the rewritten expression. Kleene `AND` + /// keeps a row false when the surviving comparison is false, so the rewrite cannot null it. + #[test] + fn test_null_literal_bound_is_value_preserving() -> VortexResult<()> { + let session = array_session(); + let ctx = &mut session.create_execution_ctx(); + let data = StructArray::from_fields(&[("x", buffer![10, 1].into_array())])?.into_array(); + + let null_lit = lit(Scalar::null(DType::Primitive( + PType::I32, + Nullability::Nullable, + ))); + let expr = and(gt_eq(col("x"), null_lit), lt_eq(col("x"), lit(5i32))); + + let before = data + .clone() + .apply(&expr)? + .execute::(ctx)? + .opt_bool_vec(ctx); + + let after = data + .apply(&find_between(expr))? + .execute::(ctx)? + .opt_bool_vec(ctx); + + // Row 0 is false rather than null because `$.x <= 5` falsifies it on its own. + assert_eq!(before, [Some(false), None]); + assert_eq!(before, after); + + Ok(()) + } + #[test] fn test_bad_match() { // An impossible expression diff --git a/vortex-array/src/scalar_fn/fns/between/kernel.rs b/vortex-array/src/scalar_fn/fns/between/kernel.rs index ee4fb688982..a4e03af623f 100644 --- a/vortex-array/src/scalar_fn/fns/between/kernel.rs +++ b/vortex-array/src/scalar_fn/fns/between/kernel.rs @@ -70,7 +70,7 @@ where let lower = &children[1]; let upper = &children[2]; let arr = array.array().clone(); - if let Some(result) = precondition(&arr, lower, upper)? { + if let Some(result) = precondition(&arr, lower, upper, parent.options)? { return Ok(Some(result)); } ::between(array, lower, upper, parent.options) @@ -105,8 +105,11 @@ where let lower = &children[1]; let upper = &children[2]; let arr = array.array().clone(); - if let Some(result) = precondition(&arr, lower, upper)? { - return Ok(Some(result)); + if let Some(result) = precondition(&arr, lower, upper, parent.options)? { + // TODO(joe): return the lazy array directly, blocked on the same executor support as + // the fallback in `between_canonical`. The reduce adaptor above already passes it + // through unexecuted, since a reduce rule can return a lazy array. + return result.execute::(ctx).map(Some); } ::between(array, lower, upper, parent.options, ctx) } diff --git a/vortex-array/src/scalar_fn/fns/between/mod.rs b/vortex-array/src/scalar_fn/fns/between/mod.rs index 9e8c4649b5d..0039d9e91ec 100644 --- a/vortex-array/src/scalar_fn/fns/between/mod.rs +++ b/vortex-array/src/scalar_fn/fns/between/mod.rs @@ -8,7 +8,6 @@ use std::fmt::Formatter; pub use kernel::*; use prost::Message; -use vortex_array::expr::and; use vortex_error::VortexResult; use vortex_error::vortex_bail; use vortex_proto::expr as pb; @@ -33,7 +32,6 @@ use crate::scalar_fn::ChildName; use crate::scalar_fn::ExecutionArgs; use crate::scalar_fn::ScalarFnId; use crate::scalar_fn::ScalarFnVTable; -use crate::scalar_fn::fns::binary::execute_boolean; use crate::scalar_fn::fns::operators::CompareOperator; use crate::scalar_fn::fns::operators::Operator; @@ -91,12 +89,19 @@ impl StrictComparison { /// Common preconditions for between operations that apply to all arrays. /// /// Returns `Some(result)` if the precondition short-circuits the between operation -/// (empty array, null bounds), or `None` if between should proceed with the -/// encoding-specific implementation. +/// (empty array, null bounds), or `None` if between must proceed with the +/// encoding-specific implementation. Kernels can therefore rely on both bounds being +/// non-null. +/// +/// The result can be a lazy [`ScalarFn`] array, so a caller that needs a computed array +/// **must** execute it. +/// +/// [`ScalarFn`]: crate::arrays::ScalarFn pub(super) fn precondition( arr: &ArrayRef, lower: &ArrayRef, upper: &ArrayRef, + options: &BetweenOptions, ) -> VortexResult> { let return_dtype = Bool(arr.dtype().nullability() | lower.dtype().nullability() | upper.dtype().nullability()); @@ -106,20 +111,43 @@ pub(super) fn precondition( return Ok(Some(Canonical::empty(&return_dtype).into_array())); } - if lower.as_constant().is_some_and(|v| v.is_null()) - || upper.as_constant().is_some_and(|v| v.is_null()) - { + let lower_is_null = lower.as_constant().is_some_and(|v| v.is_null()); + let upper_is_null = upper.as_constant().is_some_and(|v| v.is_null()); + + // `Between` is not strict, and Kleene `AND` gives `null AND false = false`, so a null bound + // cannot falsify a row on its own. Every row is null only when both bounds are null. + if lower_is_null && upper_is_null { return Ok(Some( ConstantArray::new(Scalar::null(return_dtype), arr.len()).into_array(), )); } + // Every kernel requires non-null constant bounds, so a single null bound leaves nothing to + // dispatch to. The two compares keep the surviving bound, which can still falsify rows. + if lower_is_null || upper_is_null { + return as_two_compares(arr, lower, upper, options).map(Some); + } + Ok(None) } +/// The two compares that `Between` stands for, combined with Kleene `AND`. +/// +/// The returned array is lazy, so a reduce rule can call this function. +fn as_two_compares( + arr: &ArrayRef, + lower: &ArrayRef, + upper: &ArrayRef, + options: &BetweenOptions, +) -> VortexResult { + let lower_cmp = lower.binary(arr.clone(), options.lower_strict.to_operator())?; + let upper_cmp = arr.binary(upper.clone(), options.upper_strict.to_operator())?; + lower_cmp.binary(upper_cmp, Operator::And) +} + /// Between on a canonical array by directly dispatching to the appropriate kernel. /// -/// Falls back to compare + boolean and if no kernel handles the input. +/// Falls back to [`as_two_compares`] if no kernel handles the input. fn between_canonical( arr: &ArrayRef, lower: &ArrayRef, @@ -127,8 +155,10 @@ fn between_canonical( options: &BetweenOptions, ctx: &mut ExecutionCtx, ) -> VortexResult { - if let Some(result) = precondition(arr, lower, upper)? { - return Ok(result); + if let Some(result) = precondition(arr, lower, upper, options)? { + // TODO(joe): return the lazy array directly, blocked on the same executor support as the + // fallback below. Only the single-null-bound case is lazy, so this forces it for now. + return result.execute::(ctx); } // Try type-specific kernels @@ -146,15 +176,7 @@ fn between_canonical( // TODO(joe): return lazy compare once the executor supports this // Fall back to compare + boolean and - let lower_cmp = lower.clone().binary( - arr.clone(), - Operator::from(options.lower_strict.to_compare_operator()), - )?; - let upper_cmp = arr.clone().binary( - upper.clone(), - Operator::from(options.upper_strict.to_compare_operator()), - )?; - execute_boolean(lower_cmp, upper_cmp, Operator::And, ctx) + as_two_compares(arr, lower, upper, options)?.execute::(ctx) } /// An optimized scalar expression to compute whether values fall between two bounds. @@ -299,15 +321,17 @@ impl ScalarFnVTable for Between { fn validity( &self, _options: &Self::Options, - expression: &Expression, + _expression: &Expression, ) -> VortexResult> { - let arr = expression.child(0).validity()?; - let lower = expression.child(1).validity()?; - let upper = expression.child(2).validity()?; - Ok(Some(and(and(arr, lower), upper))) + // `Between` stands for two compares under Kleene `AND`, and `null AND false` is `false`, + // so a null bound does not make a row null. There is no validity expression to derive, + // which is also why `Binary` returns `None` for `Operator::And`. + Ok(None) } fn is_strict(&self, _options: &Self::Options) -> bool { + // Not strict for the same reason `validity` returns `None` above: under Kleene `AND` a + // null bound does not force a null row. false } @@ -328,12 +352,15 @@ mod tests { use crate::VortexSessionExecute; use crate::arrays::BoolArray; use crate::arrays::DecimalArray; + use crate::arrays::PrimitiveArray; + use crate::arrays::StructArray; use crate::assert_arrays_eq; use crate::dtype::DType; use crate::dtype::DecimalDType; use crate::dtype::Nullability; use crate::dtype::PType; use crate::expr::between; + use crate::expr::col; use crate::expr::get_item; use crate::expr::lit; use crate::expr::root; @@ -344,6 +371,56 @@ mod tests { static SESSION: LazyLock = LazyLock::new(crate::array_session); + const NON_STRICT: BetweenOptions = BetweenOptions { + lower_strict: StrictComparison::NonStrict, + upper_strict: StrictComparison::NonStrict, + }; + + /// `len` null `i32` values held as a [`ConstantArray`], which is what `as_constant` sees. + fn null_i32s(len: usize) -> ArrayRef { + let null = Scalar::null(DType::Primitive(PType::I32, Nullability::Nullable)); + ConstantArray::new(null, len).into_array() + } + + /// A declared validity expression must agree with the mask of the executed result. + /// + /// The bounds are columns rather than literals so that a null bound reaches execution + /// instead of being intercepted as a constant. + #[test] + fn validity_agrees_with_execution() -> VortexResult<()> { + let ctx = &mut SESSION.create_execution_ctx(); + + // x lo hi expected + // 10 null 5 false, since the upper bound alone falsifies the row + // 10 null 50 null, since neither bound falsifies the row + // 1 0 5 true + let x = PrimitiveArray::from_option_iter([Some(10), Some(10), Some(1)]).into_array(); + let lo = PrimitiveArray::from_option_iter([None, None, Some(0)]).into_array(); + let hi = PrimitiveArray::from_option_iter([Some(5), Some(50), Some(5)]).into_array(); + let data = StructArray::from_fields(&[("x", x), ("lo", lo), ("hi", hi)])?.into_array(); + + let expr = between(col("x"), col("lo"), col("hi"), NON_STRICT); + + let executed = data + .clone() + .apply(&expr)? + .execute::(ctx)? + .opt_bool_vec(ctx); + + let declared = data + .apply(&expr.validity()?)? + .execute::(ctx)? + .bool_vec(ctx); + + assert_eq!(executed, [Some(false), None, Some(true)]); + assert_eq!( + executed.iter().map(Option::is_some).collect::>(), + declared + ); + + Ok(()) + } + #[test] fn is_not_strict() { let expr = between( @@ -444,8 +521,11 @@ mod tests { .execute::(ctx) .unwrap(); - let indices = to_int_indices(matches, ctx).unwrap(); - assert!(indices.is_empty()); + // The rows the lower bound already falsified stay false rather than becoming null. + assert_eq!( + matches.opt_bool_vec(ctx), + [None, None, Some(false), None, Some(false)] + ); // upper is a fixed constant let upper = ConstantArray::new(Scalar::from(2), 5).into_array(); @@ -485,6 +565,41 @@ mod tests { assert_eq!(indices, vec![0, 1, 2, 3, 4]); } + /// `Between` is not strict, so a null bound only makes a row null when the surviving + /// comparison is not already false. This must not depend on how the bound is encoded, and + /// compression stores an all-null chunk as a [`ConstantArray`]. + #[rstest] + #[case::primitive_nulls(PrimitiveArray::from_option_iter([None::, None]).into_array())] + #[case::constant_null(null_i32s(2))] + fn null_lower_bound(#[case] lower: ArrayRef) -> VortexResult<()> { + let ctx = &mut SESSION.create_execution_ctx(); + let array = buffer![10, 10].into_array(); + let upper = buffer![5, 50].into_array(); + + let result = between_canonical(&array, &lower, &upper, &NON_STRICT, ctx)? + .execute::(ctx)?; + + // Row 0 stays false because the upper bound falsifies it on its own. + assert_eq!(result.opt_bool_vec(ctx), [Some(false), None]); + + Ok(()) + } + + /// With both bounds null no comparison can falsify a row, so every row is null. + #[test] + fn both_bounds_null() -> VortexResult<()> { + let ctx = &mut SESSION.create_execution_ctx(); + let array = buffer![10, 10].into_array(); + let bound = null_i32s(2); + + let result = between_canonical(&array, &bound, &bound, &NON_STRICT, ctx)? + .execute::(ctx)?; + + assert_eq!(result.opt_bool_vec(ctx), [None, None]); + + Ok(()) + } + #[test] fn test_between_decimal() { let ctx = &mut SESSION.create_execution_ctx();