Refactor: expose predicate constant inference from physical-expr

Author: xudong963

datafusion/physical-expr/src/lib.rs

@@ -72,7 +72,9 @@ pub use datafusion_physical_expr_common::sort_expr::{
 pub use planner::{create_physical_expr, create_physical_exprs};
 pub use scalar_function::ScalarFunctionExpr;
 pub use simplifier::PhysicalExprSimplifier;
-pub use utils::{conjunction, conjunction_opt, split_conjunction};
+pub use utils::{
+    collect_predicate_constants, conjunction, conjunction_opt, split_conjunction,
+};
 
 // For backwards compatibility
 pub mod tree_node {

datafusion/physical-expr/src/utils/mod.rs

@@ -21,10 +21,11 @@ pub use guarantee::{Guarantee, LiteralGuarantee};
 use std::borrow::Borrow;
 use std::sync::Arc;
 
-use crate::PhysicalExpr;
-use crate::PhysicalSortExpr;
-use crate::expressions::{BinaryExpr, Column};
+use crate::expressions::{BinaryExpr, Column, Literal};
 use crate::tree_node::ExprContext;
+use crate::{
+    AcrossPartitions, ConstExpr, EquivalenceProperties, PhysicalExpr, PhysicalSortExpr,
+};
 
 use arrow::datatypes::Schema;
 use datafusion_common::tree_node::{
@@ -45,6 +46,59 @@ pub fn split_conjunction(
     split_impl(Operator::And, predicate, vec![])
 }
 
+/// Collects predicate-derived constants from equality conjunctions.
+///
+/// For each equality predicate of the form `lhs = rhs`, if either side is
+/// already known constant according to `input_eqs`, or is a literal, then the
+/// other side is also constant and will be returned as a [`ConstExpr`].
+///
+/// Literals are treated as uniform constants across partitions, so
+/// `col = literal` produces a constant for `col` with the literal value.
+pub fn collect_predicate_constants(
+    input_eqs: &EquivalenceProperties,
+    predicate: &Arc<dyn PhysicalExpr>,
+) -> Vec<ConstExpr> {
+    /// Returns the `AcrossPartitions` value for `expr` if it is constant:
+    /// either already known constant in `input_eqs`, or a `Literal`
+    /// (which is inherently constant across all partitions).
+    fn expr_constant_or_literal(
+        expr: &Arc<dyn PhysicalExpr>,
+        input_eqs: &EquivalenceProperties,
+    ) -> Option<AcrossPartitions> {
+        input_eqs.is_expr_constant(expr).or_else(|| {
+            expr.as_any()
+                .downcast_ref::<Literal>()
+                .map(|l| AcrossPartitions::Uniform(Some(l.value().clone())))
+        })
+    }
+
+    let mut constants = Vec::new();
+    for conjunction in split_conjunction(predicate) {
+        if let Some(binary) = conjunction.as_any().downcast_ref::<BinaryExpr>()
+            && binary.op() == &Operator::Eq
+        {
+            // Check if either side is constant — either already known
+            // constant from the input equivalence properties, or a literal
+            // value (which is inherently constant across all partitions).
+            let left_const = expr_constant_or_literal(binary.left(), input_eqs);
+            let right_const = expr_constant_or_literal(binary.right(), input_eqs);
+
+            if let Some(left_across) = left_const {
+                // LEFT is constant, so RIGHT must also be constant.
+                // Use RIGHT's known across value if available, otherwise
+                // propagate LEFT's (e.g. Uniform from a literal).
+                let across = right_const.unwrap_or(left_across);
+                constants.push(ConstExpr::new(Arc::clone(binary.right()), across));
+            } else if let Some(right_across) = right_const {
+                // RIGHT is constant, so LEFT must also be constant.
+                constants.push(ConstExpr::new(Arc::clone(binary.left()), right_across));
+            }
+        }
+    }
+
+    constants
+}
+
 /// Create a conjunction of the given predicates.
 /// If the input is empty, return a literal true.
 /// If the input contains a single predicate, return the predicate.
@@ -562,4 +616,30 @@ pub(crate) mod tests {
         assert_eq!(collect_columns(&expr3), expected);
         Ok(())
     }
+
+    #[test]
+    fn test_collect_predicate_constants_propagates_uniform_literal_value() -> Result<()> {
+        let schema = Arc::new(Schema::new(vec![Field::new(
+            "ticker",
+            DataType::Utf8,
+            false,
+        )]));
+        let predicate = binary(
+            col("ticker", schema.as_ref())?,
+            Operator::Eq,
+            lit(ScalarValue::Utf8(Some("NGJ26".to_string()))),
+            schema.as_ref(),
+        )?;
+        let eq_properties = EquivalenceProperties::new(schema);
+
+        let constants = collect_predicate_constants(&eq_properties, &predicate);
+
+        assert_eq!(constants.len(), 1);
+        assert_eq!(
+            constants[0].across_partitions,
+            AcrossPartitions::Uniform(Some(ScalarValue::Utf8(Some("NGJ26".to_string()))))
+        );
+
+        Ok(())
+    }
 }

datafusion/physical-plan/src/filter.rs

@@ -63,8 +63,8 @@ use datafusion_physical_expr::expressions::{BinaryExpr, Column, Literal, lit};
 use datafusion_physical_expr::intervals::utils::check_support;
 use datafusion_physical_expr::utils::{collect_columns, reassign_expr_columns};
 use datafusion_physical_expr::{
-    AcrossPartitions, AnalysisContext, ConstExpr, EquivalenceProperties, ExprBoundaries,
-    PhysicalExpr, analyze, conjunction, split_conjunction,
+    AcrossPartitions, AnalysisContext, ConstExpr, ExprBoundaries, PhysicalExpr, analyze,
+    collect_predicate_constants, conjunction, split_conjunction,
 };
 
 use datafusion_physical_expr_common::physical_expr::fmt_sql;
@@ -350,55 +350,6 @@ impl FilterExec {
         })
     }
 
-    /// Returns the `AcrossPartitions` value for `expr` if it is constant:
-    /// either already known constant in `input_eqs`, or a `Literal`
-    /// (which is inherently constant across all partitions).
-    fn expr_constant_or_literal(
-        expr: &Arc<dyn PhysicalExpr>,
-        input_eqs: &EquivalenceProperties,
-    ) -> Option<AcrossPartitions> {
-        input_eqs.is_expr_constant(expr).or_else(|| {
-            expr.as_any()
-                .downcast_ref::<Literal>()
-                .map(|l| AcrossPartitions::Uniform(Some(l.value().clone())))
-        })
-    }
-
-    fn extend_constants(
-        input: &Arc<dyn ExecutionPlan>,
-        predicate: &Arc<dyn PhysicalExpr>,
-    ) -> Vec<ConstExpr> {
-        let mut res_constants = Vec::new();
-        let input_eqs = input.equivalence_properties();
-
-        let conjunctions = split_conjunction(predicate);
-        for conjunction in conjunctions {
-            if let Some(binary) = conjunction.as_any().downcast_ref::<BinaryExpr>()
-                && binary.op() == &Operator::Eq
-            {
-                // Check if either side is constant — either already known
-                // constant from the input equivalence properties, or a literal
-                // value (which is inherently constant across all partitions).
-                let left_const = Self::expr_constant_or_literal(binary.left(), input_eqs);
-                let right_const =
-                    Self::expr_constant_or_literal(binary.right(), input_eqs);
-
-                if let Some(left_across) = left_const {
-                    // LEFT is constant, so RIGHT must also be constant.
-                    // Use RIGHT's known across value if available, otherwise
-                    // propagate LEFT's (e.g. Uniform from a literal).
-                    let across = right_const.unwrap_or(left_across);
-                    res_constants
-                        .push(ConstExpr::new(Arc::clone(binary.right()), across));
-                } else if let Some(right_across) = right_const {
-                    // RIGHT is constant, so LEFT must also be constant.
-                    res_constants
-                        .push(ConstExpr::new(Arc::clone(binary.left()), right_across));
-                }
-            }
-        }
-        res_constants
-    }
     /// This function creates the cache object that stores the plan properties such as schema, equivalence properties, ordering, partitioning, etc.
     fn compute_properties(
         input: &Arc<dyn ExecutionPlan>,
@@ -436,7 +387,10 @@ impl FilterExec {
         eq_properties.add_constants(constants)?;
         // This is for logical constant (for example: a = '1', then a could be marked as a constant)
         // to do: how to deal with multiple situation to represent = (for example c1 between 0 and 0)
-        eq_properties.add_constants(Self::extend_constants(input, predicate))?;
+        eq_properties.add_constants(collect_predicate_constants(
+            input.equivalence_properties(),
+            predicate,
+        ))?;
 
         let mut output_partitioning = input.output_partitioning().clone();
         // If contains projection, update the PlanProperties.