Refactor: expose predicate constant inference from physical-expr (#44)

Co-authored-by: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

Author: xudong963

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,66 @@ pub fn split_conjunction(
     split_impl(Operator::And, predicate, vec![])
 }
 
+impl ConstExpr {
+    /// 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.
+    ///
+    /// For example, given predicate `a = 5 AND b = c` where `c` is already
+    /// known constant, this returns constants for both `a` (Uniform with value
+    /// 5) and `b` (propagating `c`'s across-partitions 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 +623,31 @@ 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 =
+            ConstExpr::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

@@ -55,12 +55,12 @@ use datafusion_common::{
 use datafusion_execution::TaskContext;
 use datafusion_expr::Operator;
 use datafusion_physical_expr::equivalence::ProjectionMapping;
-use datafusion_physical_expr::expressions::{BinaryExpr, Column, Literal, lit};
+use datafusion_physical_expr::expressions::{BinaryExpr, Column, 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,
+    conjunction, split_conjunction,
 };
 
 use datafusion_physical_expr_common::physical_expr::fmt_sql;
@@ -243,55 +243,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>,
@@ -329,7 +280,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(ConstExpr::collect_predicate_constants(
+            input.equivalence_properties(),
+            predicate,
+        ))?;
 
         let mut output_partitioning = input.output_partitioning().clone();
         // If contains projection, update the PlanProperties.