Introduce Expr::is_volatile(), adjust TreeNode::exists() (apache#10191)

Author: peter-toth

datafusion/common/src/tree_node.rs

@@ -405,18 +405,17 @@ pub trait TreeNode: Sized {
     /// Returns true if `f` returns true for any node in the tree.
     ///
     /// Stops recursion as soon as a matching node is found
-    fn exists<F: FnMut(&Self) -> bool>(&self, mut f: F) -> bool {
+    fn exists<F: FnMut(&Self) -> Result<bool>>(&self, mut f: F) -> Result<bool> {
         let mut found = false;
         self.apply(|n| {
-            Ok(if f(n) {
+            Ok(if f(n)? {
                 found = true;
                 TreeNodeRecursion::Stop
             } else {
                 TreeNodeRecursion::Continue
             })
         })
-        .unwrap();
-        found
+        .map(|_| found)
     }
 
     /// Low-level API used to implement other APIs.

datafusion/expr/src/expr.rs

@@ -1271,7 +1271,16 @@ impl Expr {
 
     /// Return true when the expression contains out reference(correlated) expressions.
     pub fn contains_outer(&self) -> bool {
-        self.exists(|expr| matches!(expr, Expr::OuterReferenceColumn { .. }))
+        self.exists(|expr| Ok(matches!(expr, Expr::OuterReferenceColumn { .. })))
+            .unwrap()
+    }
+
+    /// Returns true if the expression is volatile, i.e. whether it can return different
+    /// results when evaluated multiple times with the same input.
+    pub fn is_volatile(&self) -> Result<bool> {
+        self.exists(|expr| {
+            Ok(matches!(expr, Expr::ScalarFunction(func) if func.func_def.is_volatile()?))
+        })
     }
 
     /// Recursively find all [`Expr::Placeholder`] expressions, and
@@ -1931,15 +1940,6 @@ fn create_names(exprs: &[Expr]) -> Result<String> {
         .join(", "))
 }
 
-/// Whether the given expression is volatile, i.e. whether it can return different results
-/// when evaluated multiple times with the same input.
-pub fn is_volatile(expr: &Expr) -> Result<bool> {
-    match expr {
-        Expr::ScalarFunction(func) => func.func_def.is_volatile(),
-        _ => Ok(false),
-    }
-}
-
 #[cfg(test)]
 mod test {
     use crate::expr::Cast;

datafusion/optimizer/src/common_subexpr_eliminate.rs

@@ -21,7 +21,6 @@ use std::collections::hash_map::Entry;
 use std::collections::{BTreeSet, HashMap};
 use std::sync::Arc;
 
-use crate::utils::is_volatile_expression;
 use crate::{utils, OptimizerConfig, OptimizerRule};
 
 use arrow::datatypes::{DataType, Field};
@@ -661,7 +660,7 @@ impl TreeNodeVisitor for ExprIdentifierVisitor<'_> {
     fn f_down(&mut self, expr: &Expr) -> Result<TreeNodeRecursion> {
         // related to https://github.com/apache/datafusion/issues/8814
         // If the expr contain volatile expression or is a short-circuit expression, skip it.
-        if expr.short_circuits() || is_volatile_expression(expr)? {
+        if expr.short_circuits() || expr.is_volatile()? {
             self.visit_stack
                 .push(VisitRecord::JumpMark(self.node_count));
             return Ok(TreeNodeRecursion::Jump); // go to f_up
@@ -717,7 +716,7 @@ impl TreeNodeRewriter for CommonSubexprRewriter<'_> {
         // The `CommonSubexprRewriter` relies on `ExprIdentifierVisitor` to generate
         // the `id_array`, which records the expr's identifier used to rewrite expr. So if we
         // skip an expr in `ExprIdentifierVisitor`, we should skip it here, too.
-        if expr.short_circuits() || is_volatile_expression(&expr)? {
+        if expr.short_circuits() || expr.is_volatile()? {
             return Ok(Transformed::new(expr, false, TreeNodeRecursion::Jump));
         }
 

datafusion/optimizer/src/push_down_filter.rs

@@ -18,7 +18,6 @@ use std::collections::{HashMap, HashSet};
 use std::sync::Arc;
 
 use crate::optimizer::ApplyOrder;
-use crate::utils::is_volatile_expression;
 use crate::{OptimizerConfig, OptimizerRule};
 
 use datafusion_common::tree_node::{
@@ -705,9 +704,7 @@ impl OptimizerRule for PushDownFilter {
 
                             (qualified_name(qualifier, field.name()), expr)
                         })
-                        .partition(|(_, value)| {
-                            is_volatile_expression(value).unwrap_or(true)
-                        });
+                        .partition(|(_, value)| value.is_volatile().unwrap_or(true));
 
                 let mut push_predicates = vec![];
                 let mut keep_predicates = vec![];

datafusion/optimizer/src/utils.rs

@@ -21,9 +21,7 @@ use std::collections::{BTreeSet, HashMap};
 
 use crate::{OptimizerConfig, OptimizerRule};
 
-use datafusion_common::tree_node::{TreeNode, TreeNodeRecursion};
 use datafusion_common::{Column, DFSchema, DFSchemaRef, Result};
-use datafusion_expr::expr::is_volatile;
 use datafusion_expr::expr_rewriter::replace_col;
 use datafusion_expr::utils as expr_utils;
 use datafusion_expr::{logical_plan::LogicalPlan, Expr, Operator};
@@ -97,20 +95,6 @@ pub fn log_plan(description: &str, plan: &LogicalPlan) {
     trace!("{description}::\n{}\n", plan.display_indent_schema());
 }
 
-/// check whether the expression is volatile predicates
-pub(crate) fn is_volatile_expression(e: &Expr) -> Result<bool> {
-    let mut is_volatile_expr = false;
-    e.apply(|expr| {
-        Ok(if is_volatile(expr)? {
-            is_volatile_expr = true;
-            TreeNodeRecursion::Stop
-        } else {
-            TreeNodeRecursion::Continue
-        })
-    })?;
-    Ok(is_volatile_expr)
-}
-
 /// Splits a conjunctive [`Expr`] such as `A AND B AND C` => `[A, B, C]`
 ///
 /// See [`split_conjunction_owned`] for more details and an example.