Support datafusion dynamic filters

vortex-array/src/expr/exprs.rs

@@ -33,6 +33,9 @@ use crate::scalar_fn::fns::cast::Cast;
 use crate::scalar_fn::fns::dynamic::DynamicComparison;
 use crate::scalar_fn::fns::dynamic::DynamicComparisonExpr;
 use crate::scalar_fn::fns::dynamic::Rhs;
+use crate::scalar_fn::fns::dynamic_expression::DynamicExpression;
+use crate::scalar_fn::fns::dynamic_expression::DynamicExpressionOptions;
+use crate::scalar_fn::fns::dynamic_expression::DynamicExpressionSource;
 use crate::scalar_fn::fns::ext_storage::ExtStorage;
 use crate::scalar_fn::fns::fill_null::FillNull;
 use crate::scalar_fn::fns::get_item::GetItem;
@@ -58,6 +61,8 @@ use crate::scalar_fn::fns::replace_time_zone::ReplaceTimeZone;
 use crate::scalar_fn::fns::replace_time_zone::ReplaceTimeZoneOptions;
 use crate::scalar_fn::fns::select::FieldSelection;
 use crate::scalar_fn::fns::select::Select;
+use crate::scalar_fn::fns::switch::Switch;
+use crate::scalar_fn::fns::switch::SwitchState;
 use crate::scalar_fn::fns::variant_get::VariantGet;
 use crate::scalar_fn::fns::variant_get::VariantGetOptions;
 use crate::scalar_fn::fns::variant_get::VariantPath;
@@ -1107,6 +1112,45 @@ pub fn bound_dynamic(
     )
 }
 
+/// Creates a dynamic expression that evaluates the current predicate of `source`.
+///
+/// The `inputs` are the expressions that every predicate of the source can read. A predicate
+/// reads input `i` as the field named by
+/// [`input_field_name(i)`](crate::scalar_fn::fns::dynamic_expression::input_field_name) of its
+/// scope.
+pub fn dynamic_expression(
+    source: Arc<dyn DynamicExpressionSource>,
+    inputs: impl IntoIterator<Item = Expression>,
+) -> Expression {
+    DynamicExpression.new_expr(DynamicExpressionOptions::new(source), inputs)
+}
+
+/// Creates a bound dynamic expression that evaluates the current predicate of `source`.
+pub fn bound_dynamic_expression(
+    source: Arc<dyn DynamicExpressionSource>,
+    inputs: impl IntoIterator<Item = BoundExpression>,
+) -> BoundExpression {
+    DynamicExpression
+        .try_new_bound_expr(DynamicExpressionOptions::new(source), inputs)
+        .vortex_expect("dynamic expressions accept any inputs")
+}
+
+// ---- Switch ----
+
+/// Creates a switch, initially on, over a predicate that a scan may turn off: `on ⇒ predicate`.
+///
+/// See [`Switch`] for when a predicate may be turned off.
+pub fn switch(predicate: Expression) -> Expression {
+    Switch.new_expr(SwitchState::new(), [predicate])
+}
+
+/// Creates a bound switch, initially on, over a predicate that a scan may turn off.
+pub fn bound_switch(predicate: BoundExpression) -> BoundExpression {
+    Switch
+        .try_new_bound_expr(SwitchState::new(), [predicate])
+        .vortex_expect("switch expressions require a boolean predicate")
+}
+
 // ---- ListContains ----
 
 /// Creates an expression that checks whether a list contains a value.
@@ -1403,6 +1447,7 @@ pub mod bound {
     pub use super::bound_checked_add as checked_add;
     pub use super::bound_col as col;
     pub use super::bound_dynamic as dynamic;
+    pub use super::bound_dynamic_expression as dynamic_expression;
     pub use super::bound_dynamic_with_options as dynamic_with_options;
     pub use super::bound_eq as eq;
     pub use super::bound_ext_storage as ext_storage;
@@ -1438,6 +1483,7 @@ pub mod bound {
     pub use super::bound_root as root;
     pub use super::bound_select as select;
     pub use super::bound_select_exclude as select_exclude;
+    pub use super::bound_switch as switch;
     pub use super::bound_variant_get as variant_get;
     pub use super::bound_zip_expr as zip_expr;
 }

vortex-array/src/expr/mod.rs

@@ -71,6 +71,7 @@ pub use exprs::cast;
 pub use exprs::checked_add;
 pub use exprs::col;
 pub use exprs::dynamic;
+pub use exprs::dynamic_expression;
 pub use exprs::dynamic_with_options;
 pub use exprs::eq;
 pub use exprs::ext_storage;
@@ -109,6 +110,7 @@ pub use exprs::replace_time_zone;
 pub use exprs::root;
 pub use exprs::select;
 pub use exprs::select_exclude;
+pub use exprs::switch;
 pub use exprs::variant_get;
 pub use exprs::zip_expr;
 pub use scope::*;

vortex-array/src/scalar_fn/demand.rs

@@ -0,0 +1,105 @@
+// SPDX-License-Identifier: Apache-2.0
+// SPDX-FileCopyrightText: Copyright the Vortex contributors
+
+//! Execution of an array for the rows a caller demands.
+//!
+//! A demand is a mask of the rows whose values the caller reads. The values of the other rows are
+//! unspecified, so a scalar function may ignore the demand and compute every row. A function can
+//! also use it to skip work: `AND` and `OR` demand of their right operand only the rows their left
+//! operand does not decide.
+
+use vortex_error::VortexResult;
+use vortex_mask::Mask;
+
+use crate::ArrayRef;
+use crate::Canonical;
+use crate::ExecutionCtx;
+use crate::IntoArray;
+use crate::arrays::ConstantArray;
+use crate::arrays::ScalarFn;
+use crate::arrays::scalar_fn::ScalarFnArrayExt;
+use crate::scalar::Scalar;
+use crate::scalar_fn::VecExecutionArgs;
+
+/// Execute `array` to its canonical form for the rows of `demand`, or for every row without one.
+///
+/// The values of rows outside the demand are unspecified. A computation that is not evaluated yet,
+/// such as a comparison, receives the demand through its [`ExecutionArgs`] and decides how to use
+/// it. An encoded or canonical array is executed whole.
+///
+/// [`ExecutionArgs`]: crate::scalar_fn::ExecutionArgs
+pub fn execute_demanded(
+    array: ArrayRef,
+    demand: Option<&Mask>,
+    ctx: &mut ExecutionCtx,
+) -> VortexResult<ArrayRef> {
+    let Some(demand) = demand.filter(|demand| !demand.all_true()) else {
+        return Ok(array.execute::<Canonical>(ctx)?.into_array());
+    };
+    if demand.all_false() {
+        return Ok(
+            ConstantArray::new(Scalar::default_value(array.dtype()), array.len()).into_array(),
+        );
+    }
+    let Some(scalar_fn_array) = array.as_opt::<ScalarFn>() else {
+        return Ok(array.execute::<Canonical>(ctx)?.into_array());
+    };
+
+    let args =
+        VecExecutionArgs::new(scalar_fn_array.children(), array.len()).with_demand(demand.clone());
+    Ok(scalar_fn_array
+        .scalar_fn()
+        .execute(&args, ctx)?
+        .execute::<Canonical>(ctx)?
+        .into_array())
+}
+
+#[cfg(test)]
+mod tests {
+    use vortex_error::VortexResult;
+    use vortex_mask::Mask;
+
+    use super::*;
+    use crate::VortexSessionExecute;
+    use crate::array_session;
+    use crate::arrays::Constant;
+    use crate::arrays::PrimitiveArray;
+    use crate::builtins::ArrayBuiltins;
+    use crate::scalar_fn::fns::operators::Operator;
+
+    /// The demanded rows hold the values of a full evaluation.
+    #[test]
+    fn demanded_rows_match_full_evaluation() -> VortexResult<()> {
+        let mut ctx = array_session().create_execution_ctx();
+        let len = 4096;
+        let step = 64;
+        let values = PrimitiveArray::from_iter(0..len as u64).into_array();
+        let comparison = values.binary(
+            ConstantArray::new(Scalar::from(1000u64), len).into_array(),
+            Operator::Lt,
+        )?;
+        let demand = Mask::from_indices(len, (0..len).step_by(step).collect::<Vec<_>>());
+
+        let full = comparison
+            .clone()
+            .execute::<Canonical>(&mut ctx)?
+            .into_array();
+        let demanded = execute_demanded(comparison, Some(&demand), &mut ctx)?;
+        for idx in (0..len).step_by(step) {
+            assert_eq!(
+                demanded.execute_scalar(idx, &mut ctx)?,
+                full.execute_scalar(idx, &mut ctx)?
+            );
+        }
+        Ok(())
+    }
+
+    #[test]
+    fn nothing_demanded_is_not_evaluated() -> VortexResult<()> {
+        let mut ctx = array_session().create_execution_ctx();
+        let values = PrimitiveArray::from_iter([1u64, 2, 3]).into_array();
+        let demanded = execute_demanded(values, Some(&Mask::new_false(3)), &mut ctx)?;
+        assert!(demanded.is::<Constant>());
+        Ok(())
+    }
+}

vortex-array/src/scalar_fn/fns/binary/boolean.rs

@@ -24,6 +24,7 @@ use crate::arrays::BoolArray;
 use crate::arrays::Constant;
 use crate::arrays::ConstantArray;
 use crate::arrays::ScalarFn;
+use crate::arrays::bool::BoolArrayExt;
 use crate::arrays::scalar_fn::ExactScalarFn;
 use crate::arrays::scalar_fn::ScalarFnArrayExt;
 use crate::arrays::scalar_fn::ScalarFnArrayView;
@@ -33,6 +34,7 @@ use crate::dtype::Nullability;
 use crate::kernel::ExecuteParentKernel;
 use crate::scalar::BoolScalar;
 use crate::scalar::Scalar;
+use crate::scalar_fn::demand::execute_demanded;
 use crate::scalar_fn::fns::binary::Binary;
 use crate::scalar_fn::fns::operators::Operator;
 use crate::validity::Validity;
@@ -114,10 +116,14 @@ pub fn or_kleene(lhs: &ArrayRef, rhs: &ArrayRef) -> VortexResult<ArrayRef> {
 ///
 /// This is the entry point for boolean operations from the binary expression.
 /// Handles constants and canonical boolean arrays directly, otherwise falls back to Arrow.
+///
+/// The right operand is evaluated only for the demanded rows that the left operand does not
+/// decide: a valid `true` decides `OR` and a valid `false` decides `AND`.
 pub(crate) fn execute_boolean(
     lhs: ArrayRef,
     rhs: ArrayRef,
     op: Operator,
+    demand: Option<&Mask>,
     ctx: &mut ExecutionCtx,
 ) -> VortexResult<ArrayRef> {
     let nullable = boolean_nullability(&lhs, &rhs);
@@ -130,19 +136,45 @@ pub(crate) fn execute_boolean(
         return Ok(result);
     }
 
-    let lhs = lhs.execute::<BoolArray>(ctx)?;
+    let lhs = execute_demanded(lhs, demand, ctx)?.execute::<BoolArray>(ctx)?;
     if let Some(result) = <Bool as BooleanKernel>::boolean(lhs.as_view(), &rhs, op, ctx)? {
         return Ok(result);
     }
 
-    let rhs = rhs.execute::<BoolArray>(ctx)?;
+    let rhs_demand = undecided_rows(&lhs, op, demand, ctx)?;
+    if rhs_demand.all_false() {
+        return lhs.into_array().cast(DType::Bool(nullable));
+    }
+    let rhs = execute_demanded(rhs, Some(&rhs_demand), ctx)?.execute::<BoolArray>(ctx)?;
     let Some(result) = <Bool as BooleanKernel>::boolean(rhs.as_view(), &lhs.into_array(), op, ctx)?
     else {
         vortex_bail!("No boolean kernel for two BoolArrays");
     };
     Ok(result)
 }
 
+/// The demanded rows whose result `lhs` does not decide: every row but those holding a valid `true`
+/// for `OR` or a valid `false` for `AND`.
+fn undecided_rows(
+    lhs: &BoolArray,
+    op: Operator,
+    demand: Option<&Mask>,
+    ctx: &mut ExecutionCtx,
+) -> VortexResult<Mask> {
+    let values = Mask::from_buffer(lhs.to_bit_buffer());
+    let valid = lhs.validity()?.execute_mask(lhs.len(), ctx)?;
+    let decided = if op == Operator::Or {
+        &values & &valid
+    } else {
+        valid.bitand_not(&values)
+    };
+    let undecided = !decided;
+    Ok(match demand {
+        Some(demand) => &undecided & demand,
+        None => undecided,
+    })
+}
+
 /// Handles boolean operations where at least one operand is a constant array.
 fn constant_boolean(
     lhs: &ArrayRef,
@@ -757,13 +789,47 @@ mod tests {
     use crate::array_session;
     use crate::arrays::BoolArray;
     use crate::arrays::ConstantArray;
+    use crate::arrays::PrimitiveArray;
+    use crate::arrays::ScalarFn;
     use crate::assert_arrays_eq;
     use crate::builtins::ArrayBuiltins;
     use crate::dtype::DType;
     use crate::dtype::Nullability;
     use crate::scalar::Scalar;
     use crate::scalar_fn::fns::operators::Operator;
 
+    /// The right operand is a lazy comparison, so the rows the left operand decides are not
+    /// demanded of it.
+    #[rstest]
+    #[case::or(Operator::Or)]
+    #[case::and(Operator::And)]
+    fn test_short_circuit_matches_full_evaluation(#[case] op: Operator) -> VortexResult<()> {
+        let mut ctx = array_session().create_execution_ctx();
+        let len = 4096;
+        // Every eighth row is null, the rest mostly decide the result.
+        let decides = op == Operator::Or;
+        let lhs =
+            BoolArray::from_iter((0..len).map(|idx| {
+                (idx % 8 != 0).then_some(if idx % 16 == 1 { !decides } else { decides })
+            }))
+            .into_array();
+        let values = PrimitiveArray::from_option_iter(
+            (0..len).map(|idx| (idx % 3 != 0).then_some(idx as u64)),
+        )
+        .into_array();
+        let rhs = values.binary(
+            ConstantArray::new(Scalar::from(2000u64), len).into_array(),
+            Operator::Lt,
+        )?;
+        assert!(rhs.is::<ScalarFn>());
+
+        let expected = lhs
+            .clone()
+            .binary(rhs.clone().execute::<BoolArray>(&mut ctx)?.into_array(), op)?;
+        assert_arrays_eq!(lhs.binary(rhs, op)?, expected, &mut ctx);
+        Ok(())
+    }
+
     #[test]
     fn test_kleene_truth_table() -> VortexResult<()> {
         let mut ctx = array_session().create_execution_ctx();

vortex-array/src/scalar_fn/fns/binary/mod.rs

@@ -192,8 +192,8 @@ impl ScalarFnVTable for Binary {
             Operator::Lte => execute_compare(&lhs, &rhs, CompareOperator::Lte, ctx),
             Operator::Gt => execute_compare(&lhs, &rhs, CompareOperator::Gt, ctx),
             Operator::Gte => execute_compare(&lhs, &rhs, CompareOperator::Gte, ctx),
-            Operator::And => execute_boolean(lhs, rhs, Operator::And, ctx),
-            Operator::Or => execute_boolean(lhs, rhs, Operator::Or, ctx),
+            Operator::And => execute_boolean(lhs, rhs, Operator::And, args.demand(), ctx),
+            Operator::Or => execute_boolean(lhs, rhs, Operator::Or, args.demand(), ctx),
             Operator::Add => execute_numeric(&lhs, &rhs, NumericOperator::Add, ctx),
             Operator::Sub => execute_numeric(&lhs, &rhs, NumericOperator::Sub, ctx),
             Operator::Mul => execute_numeric(&lhs, &rhs, NumericOperator::Mul, ctx),