diff --git a/vortex-array/src/scalar_fn/fns/is_not_null.rs b/vortex-array/src/scalar_fn/fns/is_not_null.rs index e9a106156d6..a2a229a5bbd 100644 --- a/vortex-array/src/scalar_fn/fns/is_not_null.rs +++ b/vortex-array/src/scalar_fn/fns/is_not_null.rs @@ -9,6 +9,7 @@ use vortex_error::VortexResult; use vortex_session::VortexSession; use vortex_session::registry::CachedId; +use crate::AnyColumnar; use crate::ArrayRef; use crate::ExecutionCtx; use crate::arrays::ScalarFnArray; @@ -16,6 +17,7 @@ use crate::dtype::DType; use crate::dtype::Nullability; use crate::expr::display::ExprDisplay; use crate::scalar_fn::Arity; +use crate::scalar_fn::ArrayReduceNode; use crate::scalar_fn::ChildName; use crate::scalar_fn::EmptyOptions; use crate::scalar_fn::ExecutionArgs; @@ -35,6 +37,9 @@ use crate::scalar_fn::is_not_null_node; /// IsNull(x) -> if !x.nullable lit(false) else not(x.validity()) /// IsNotNull(x) -> if !x.nullable lit(true) or x.validity() /// +/// These rewrites never evaluate x, so they require that x cannot fail +/// ([`ReduceNode::contains_fallible`]). +/// /// In expression and array contexts for x, y where x is nullable but y is not, /// reduce /// @@ -46,19 +51,22 @@ use crate::scalar_fn::is_not_null_node; /// Latter optimizations make sense because calculating IsNull(x)/IsNotNull(x) /// is at most expensive as calculating x, but usually much cheaper. Although /// in two cases you exchange 4 computations to 4 computations, the latter -/// four are cheaper. +/// four are cheaper. They still evaluate x and y, so they need no fallibility +/// check. pub(crate) fn reduce_null(is_null: bool, node: &T) -> VortexResult> { let child = node.child(0); - if !child.node_dtype()?.is_nullable() { - return Ok(Some(node.new_constant((!is_null).into()))); - } + if !child.contains_fallible() { + if !child.node_dtype()?.is_nullable() { + return Ok(Some(node.new_constant((!is_null).into()))); + } - if let ReduceNodeValidity::Reduced(validity) = child.validity()? { - return Ok(Some(if is_null { - validity.new_node(Not.bind(EmptyOptions), std::slice::from_ref(&validity))? - } else { - validity - })); + if let ReduceNodeValidity::Reduced(validity) = child.validity()? { + return Ok(Some(if is_null { + validity.new_node(Not.bind(EmptyOptions), std::slice::from_ref(&validity))? + } else { + validity + })); + } } let Some(child_fn) = child.scalar_fn() else { @@ -103,6 +111,22 @@ pub(crate) fn reduce_null(is_null: bool, node: &T) -> VortexResul Ok(Some(node.new_node(Binary.bind(combine), &[left, right])?)) } +/// Executes the input of a null check to columnar form if evaluating it can fail. +/// +/// [`ArrayRef::validity`] alone does not evaluate the array, so it misses the errors of a fallible +/// input. Columnar form keeps a constant input constant. The null check discards the computed +/// values, so a plan that also reads them evaluates the input twice. +pub(crate) fn execute_if_fallible( + input: ArrayRef, + ctx: &mut ExecutionCtx, +) -> VortexResult { + if !ArrayReduceNode::new(&input).contains_fallible() { + return Ok(input); + } + + input.execute_until::(ctx) +} + /// Expression that checks for non-null values. #[derive(Clone)] pub struct IsNotNull; @@ -166,9 +190,11 @@ impl ScalarFnVTable for IsNotNull { &self, _data: &Self::Options, args: &dyn ExecutionArgs, - _ctx: &mut ExecutionCtx, + ctx: &mut ExecutionCtx, ) -> VortexResult { - Ok(args.get(0)?.validity()?.to_array(args.row_count())) + let input = execute_if_fallible(args.get(0)?, ctx)?; + + Ok(input.validity()?.to_array(args.row_count())) } fn reduce(&self, _options: &Self::Options, node: &T) -> VortexResult> { @@ -189,18 +215,28 @@ impl ScalarFnVTable for IsNotNull { mod tests { use std::sync::LazyLock; + use rstest::rstest; use vortex_buffer::buffer; + use vortex_error::VortexError; use vortex_error::VortexExpect as _; use vortex_error::VortexResult; use vortex_session::VortexSession; + use crate::ArrayRef; + use crate::Canonical; + use crate::ExecutionCtx; use crate::IntoArray; use crate::VortexSessionExecute; use crate::array_session; + use crate::arrays::BoolArray; + use crate::arrays::DictArray; use crate::arrays::PrimitiveArray; use crate::arrays::StructArray; + use crate::assert_arrays_eq; + use crate::builtins::ArrayBuiltins; use crate::dtype::DType; use crate::dtype::Nullability; + use crate::dtype::PType; use crate::dtype::StructFields; use crate::expr::BoundExpression; use crate::expr::Expression; @@ -218,10 +254,12 @@ mod tests { use crate::scalar::Scalar; use crate::scalar_fn::EmptyOptions; use crate::scalar_fn::ScalarFnVTableExt; + use crate::scalar_fn::fns::operators::Operator; use crate::scalar_fn::internal::row_count::RowCount; use crate::stats::StatsSession; use crate::stats::all_null; use crate::stats::null_count; + use crate::validity::Validity; static STATS_SESSION: LazyLock = LazyLock::new(|| VortexSession::empty().with::()); @@ -413,4 +451,99 @@ mod tests { ); Ok(()) } + + #[derive(Clone, Copy, Debug)] + enum NullCheck { + IsNull, + IsNotNull, + } + + /// Runs the check through the array builtins, so the reduction rules apply before the kernel. + fn run_null_check( + input: &ArrayRef, + check: NullCheck, + ctx: &mut ExecutionCtx, + ) -> VortexResult { + let result = match check { + NullCheck::IsNull => input.is_null()?, + NullCheck::IsNotNull => input.is_not_null()?, + }; + + result.execute::(ctx) + } + + /// Returns the error text without the backtrace that `RUST_BACKTRACE` appends. + fn without_backtrace(error: &VortexError) -> String { + let text = error.to_string(); + match text.split_once("\nBacktrace:") { + Some((message, _)) => message.to_string(), + None => text, + } + } + + /// Covers `reduce_null` on a non-nullable input. + fn null_into_non_nullable_cast() -> VortexResult { + PrimitiveArray::from_option_iter([Some(1i64), None]) + .into_array() + .cast(PType::I64.into()) + } + + /// Covers a fallible input below an infallible parent. + fn comparison_with_failing_cast() -> VortexResult { + PrimitiveArray::from_option_iter([Some(1i64), None]) + .into_array() + .binary(null_into_non_nullable_cast()?, Operator::Lt) + } + + /// Covers a fallible input below an encoding. + fn dictionary_of_failing_cast() -> VortexResult { + let codes = PrimitiveArray::from_option_iter([Some(0u8), None]).into_array(); + Ok(DictArray::try_new(codes, null_into_non_nullable_cast()?)?.into_array()) + } + + /// The out-of-range value is null, so the cast succeeds. + fn successful_cast() -> VortexResult { + PrimitiveArray::new(buffer![300i64, 1], Validity::from_iter([false, true])) + .into_array() + .cast(DType::Primitive(PType::I8, Nullability::Nullable)) + } + + /// A Kleene `and` over two nullable inputs has irreducible validity. + fn kleene_and() -> VortexResult { + let lhs = BoolArray::from_iter([Some(false), None, None, Some(true)]).into_array(); + let rhs = BoolArray::from_iter([None, Some(false), None, Some(true)]).into_array(); + lhs.binary(rhs, Operator::And) + } + + /// A null check over a lazy input must match the check over the evaluated input, including any + /// error. + #[rstest] + #[case::null_into_non_nullable_cast(null_into_non_nullable_cast)] + #[case::comparison_with_failing_cast(comparison_with_failing_cast)] + #[case::dictionary_of_failing_cast(dictionary_of_failing_cast)] + #[case::successful_cast(successful_cast)] + #[case::kleene_and(kleene_and)] + fn null_checks_match_eager_evaluation( + #[case] input: fn() -> VortexResult, + #[values(NullCheck::IsNull, NullCheck::IsNotNull)] check: NullCheck, + ) -> VortexResult<()> { + let mut ctx = array_session().create_execution_ctx(); + let input = input()?; + + let expected = input + .clone() + .execute::(&mut ctx) + .and_then(|evaluated| run_null_check(&evaluated.into_array(), check, &mut ctx)); + let actual = run_null_check(&input, check, &mut ctx); + + match (actual, expected) { + (Ok(actual), Ok(expected)) => assert_arrays_eq!(actual, expected, &mut ctx), + (Err(actual), Err(expected)) => { + assert_eq!(without_backtrace(&actual), without_backtrace(&expected)); + } + (actual, expected) => panic!("expected {expected:?}, got {actual:?}"), + } + + Ok(()) + } } diff --git a/vortex-array/src/scalar_fn/fns/is_null.rs b/vortex-array/src/scalar_fn/fns/is_null.rs index b3e1d6bdd51..ebbf99f5e34 100644 --- a/vortex-array/src/scalar_fn/fns/is_null.rs +++ b/vortex-array/src/scalar_fn/fns/is_null.rs @@ -22,6 +22,7 @@ use crate::scalar_fn::ReduceNode; use crate::scalar_fn::ScalarFnId; use crate::scalar_fn::ScalarFnVTable; use crate::scalar_fn::ScalarFnVTableExt; +use crate::scalar_fn::fns::is_not_null::execute_if_fallible; use crate::scalar_fn::fns::is_not_null::reduce_null; use crate::validity::Validity; @@ -77,9 +78,9 @@ impl ScalarFnVTable for IsNull { &self, _data: &Self::Options, args: &dyn ExecutionArgs, - _ctx: &mut ExecutionCtx, + ctx: &mut ExecutionCtx, ) -> VortexResult { - let child = args.get(0)?; + let child = execute_if_fallible(args.get(0)?, ctx)?; if let Some(scalar) = child.as_constant() { return Ok(ConstantArray::new(scalar.is_null(), args.row_count()).into_array()); } diff --git a/vortex-array/src/scalar_fn/vtable.rs b/vortex-array/src/scalar_fn/vtable.rs index cf13ba98bcf..02e2e75781a 100644 --- a/vortex-array/src/scalar_fn/vtable.rs +++ b/vortex-array/src/scalar_fn/vtable.rs @@ -45,6 +45,8 @@ pub enum ReduceNodeValidity { /// expression reduce node, one example is byte_length(x). /// validity(byte_length(x)) can be symbolically reduced to validity(x) /// since byte_length doesn't change validity. + /// + /// Reducing does not evaluate T, so the reduced node can skip the semantic errors of T. Reduced(T), /// Validity of T can't be symbolically reduced to anything, and all /// further reductions require evaluating T first. For an expression reduce @@ -269,7 +271,8 @@ pub trait ScalarFnVTable: 'static + Sized + Clone + Send + Sync { /// /// Returning `true` permits optimizations that evaluate the function over values that no input /// row references. Dictionary push-down, for example, evaluates every dictionary value, so a - /// fallible function could error on a value that row-wise evaluation would never reach. + /// fallible function could error on a value that row-wise evaluation would never reach. It also + /// permits null checks to skip evaluating the function. /// /// This applies only to the scalar function, not its child expressions, and only to inputs /// accepted by [`ScalarFnVTable::return_dtype`]. The default is conservatively `false`. @@ -311,6 +314,28 @@ pub trait ReduceNode: Clone { /// Produce a new constant node in the same scope as "self" fn new_constant(&self, value: Scalar) -> Self; + /// Returns whether evaluating this subtree can raise a semantic error. + /// + /// This is `true` if any scalar function in the subtree is not + /// [infallible](ScalarFnVTable::is_infallible), including one below an infallible parent, as + /// in `cast(x, i8) < 5`. Array encodings and expression roots raise no errors of their own. + fn contains_fallible(&self) -> bool { + // A recursive search can overflow the call stack on deep expression trees. + let mut pending = vec![self.clone()]; + while let Some(node) = pending.pop() { + if node + .scalar_fn() + .is_some_and(|scalar_fn| !scalar_fn.signature().is_infallible()) + { + return true; + } + + pending.extend((0..node.child_count()).map(|idx| node.child(idx))); + } + + false + } + /// Symbolic validity of this node. Reduced() if you can get from node's /// validity to validity of its children or a constant without evaluating /// node.