chore: refactor usage of reassign_predicate_columns (#17703)

* chore: refactor usage of `reassign_predicate_columns`

* chore: Address PR comments

---------

Co-authored-by: Andrew Lamb <andrew@nerdnetworks.org>

Author: rkrishn7

datafusion/datasource-parquet/src/row_filter.rs

@@ -76,7 +76,7 @@ use datafusion_common::tree_node::{TreeNode, TreeNodeRecursion, TreeNodeVisitor}
 use datafusion_common::Result;
 use datafusion_datasource::schema_adapter::{SchemaAdapterFactory, SchemaMapper};
 use datafusion_physical_expr::expressions::Column;
-use datafusion_physical_expr::utils::reassign_predicate_columns;
+use datafusion_physical_expr::utils::reassign_expr_columns;
 use datafusion_physical_expr::{split_conjunction, PhysicalExpr};
 
 use datafusion_physical_plan::metrics;
@@ -119,9 +119,8 @@ impl DatafusionArrowPredicate {
         rows_matched: metrics::Count,
         time: metrics::Time,
     ) -> Result<Self> {
-        let projected_schema = Arc::clone(&candidate.filter_schema);
         let physical_expr =
-            reassign_predicate_columns(candidate.expr, &projected_schema, true)?;
+            reassign_expr_columns(candidate.expr, &candidate.filter_schema)?;
 
         Ok(Self {
             physical_expr,

datafusion/datasource/src/file_scan_config.rs

@@ -48,21 +48,21 @@ use datafusion_common::{
 use datafusion_execution::{
     object_store::ObjectStoreUrl, SendableRecordBatchStream, TaskContext,
 };
-use datafusion_physical_expr::{expressions::Column, utils::reassign_predicate_columns};
-use datafusion_physical_expr::{EquivalenceProperties, Partitioning};
+use datafusion_expr::Operator;
+use datafusion_physical_expr::expressions::BinaryExpr;
+use datafusion_physical_expr::{expressions::Column, utils::reassign_expr_columns};
+use datafusion_physical_expr::{split_conjunction, EquivalenceProperties, Partitioning};
 use datafusion_physical_expr_adapter::PhysicalExprAdapterFactory;
 use datafusion_physical_expr_common::physical_expr::PhysicalExpr;
 use datafusion_physical_expr_common::sort_expr::{LexOrdering, PhysicalSortExpr};
+use datafusion_physical_plan::filter_pushdown::FilterPushdownPropagation;
 use datafusion_physical_plan::projection::ProjectionExpr;
 use datafusion_physical_plan::{
     display::{display_orderings, ProjectSchemaDisplay},
     metrics::ExecutionPlanMetricsSet,
     projection::{all_alias_free_columns, new_projections_for_columns},
     DisplayAs, DisplayFormatType,
 };
-use datafusion_physical_plan::{
-    filter::collect_columns_from_predicate, filter_pushdown::FilterPushdownPropagation,
-};
 
 use datafusion_physical_plan::coop::cooperative;
 use datafusion_physical_plan::execution_plan::SchedulingType;
@@ -588,27 +588,14 @@ impl DataSource for FileScanConfig {
         if let Some(filter) = self.file_source.filter() {
             // We need to remap column indexes to match the projected schema since that's what the equivalence properties deal with.
             // Note that this will *ignore* any non-projected columns: these don't factor into ordering / equivalence.
-            match reassign_predicate_columns(filter, &schema, true) {
-                Ok(filter) => {
-                    match Self::add_filter_equivalence_info(
-                        filter,
-                        &mut eq_properties,
-                        &schema,
-                    ) {
-                        Ok(()) => {}
-                        Err(e) => {
-                            warn!("Failed to add filter equivalence info: {e}");
-                            #[cfg(debug_assertions)]
-                            panic!("Failed to add filter equivalence info: {e}");
-                        }
-                    }
-                }
+            match Self::add_filter_equivalence_info(filter, &mut eq_properties, &schema) {
+                Ok(()) => {}
                 Err(e) => {
-                    warn!("Failed to reassign predicate columns: {e}");
+                    warn!("Failed to add filter equivalence info: {e}");
                     #[cfg(debug_assertions)]
-                    panic!("Failed to reassign predicate columns: {e}");
+                    panic!("Failed to add filter equivalence info: {e}");
                 }
-            };
+            }
         }
         eq_properties
     }
@@ -764,24 +751,24 @@ impl FileScanConfig {
         eq_properties: &mut EquivalenceProperties,
         schema: &Schema,
     ) -> Result<()> {
-        macro_rules! ignore_dangling_col {
-            ($col:expr) => {
-                if let Some(col) = $col.as_any().downcast_ref::<Column>() {
-                    if schema.index_of(col.name()).is_err() {
-                        continue;
+        // Gather valid equality pairs from the filter expression
+        let equal_pairs = split_conjunction(&filter).into_iter().filter_map(|expr| {
+            // Ignore any binary expressions that reference non-existent columns in the current schema
+            // (e.g. due to unnecessary projections being removed)
+            reassign_expr_columns(Arc::clone(expr), schema)
+                .ok()
+                .and_then(|expr| match expr.as_any().downcast_ref::<BinaryExpr>() {
+                    Some(expr) if expr.op() == &Operator::Eq => {
+                        Some((Arc::clone(expr.left()), Arc::clone(expr.right())))
                     }
-                }
-            };
-        }
+                    _ => None,
+                })
+        });
 
-        let (equal_pairs, _) = collect_columns_from_predicate(&filter);
         for (lhs, rhs) in equal_pairs {
-            // Ignore any binary expressions that reference non-existent columns in the current schema
-            // (e.g. due to unnecessary projections being removed)
-            ignore_dangling_col!(lhs);
-            ignore_dangling_col!(rhs);
-            eq_properties.add_equal_conditions(Arc::clone(lhs), Arc::clone(rhs))?
+            eq_properties.add_equal_conditions(lhs, rhs)?
         }
+
         Ok(())
     }
 

datafusion/physical-expr/src/expressions/dynamic_filters.rs

@@ -337,7 +337,7 @@ impl PhysicalExpr for DynamicFilterPhysicalExpr {
 mod test {
     use crate::{
         expressions::{col, lit, BinaryExpr},
-        utils::reassign_predicate_columns,
+        utils::reassign_expr_columns,
     };
     use arrow::{
         array::RecordBatch,
@@ -375,18 +375,16 @@ mod test {
         ]));
         // Each ParquetExec calls `with_new_children` on the DynamicFilterPhysicalExpr
         // and remaps the children to the file schema.
-        let dynamic_filter_1 = reassign_predicate_columns(
+        let dynamic_filter_1 = reassign_expr_columns(
             Arc::clone(&dynamic_filter) as Arc<dyn PhysicalExpr>,
             &filter_schema_1,
-            false,
         )
         .unwrap();
         let snap = dynamic_filter_1.snapshot().unwrap().unwrap();
         insta::assert_snapshot!(format!("{snap:?}"), @r#"BinaryExpr { left: Column { name: "a", index: 0 }, op: Eq, right: Literal { value: Int32(42), field: Field { name: "lit", data_type: Int32, nullable: false, dict_id: 0, dict_is_ordered: false, metadata: {} } }, fail_on_overflow: false }"#);
-        let dynamic_filter_2 = reassign_predicate_columns(
+        let dynamic_filter_2 = reassign_expr_columns(
             Arc::clone(&dynamic_filter) as Arc<dyn PhysicalExpr>,
             &filter_schema_2,
-            false,
         )
         .unwrap();
         let snap = dynamic_filter_2.snapshot().unwrap().unwrap();

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

@@ -238,22 +238,23 @@ pub fn collect_columns(expr: &Arc<dyn PhysicalExpr>) -> HashSet<Column> {
     columns
 }
 
-/// Re-assign column indices referenced in predicate according to given schema.
-/// This may be helpful when dealing with projections.
-pub fn reassign_predicate_columns(
-    pred: Arc<dyn PhysicalExpr>,
+/// Re-assign indices of [`Column`]s within the given [`PhysicalExpr`] according to
+/// the provided [`Schema`].
+///
+/// This can be useful when attempting to map an expression onto a different schema.
+///
+/// # Errors
+///
+/// This function will return an error if any column in the expression cannot be found
+/// in the provided schema.
+pub fn reassign_expr_columns(
+    expr: Arc<dyn PhysicalExpr>,
     schema: &Schema,
-    ignore_not_found: bool,
 ) -> Result<Arc<dyn PhysicalExpr>> {
-    pred.transform_down(|expr| {
-        let expr_any = expr.as_any();
-
-        if let Some(column) = expr_any.downcast_ref::<Column>() {
-            let index = match schema.index_of(column.name()) {
-                Ok(idx) => idx,
-                Err(_) if ignore_not_found => usize::MAX,
-                Err(e) => return Err(e.into()),
-            };
+    expr.transform_down(|expr| {
+        if let Some(column) = expr.as_any().downcast_ref::<Column>() {
+            let index = schema.index_of(column.name())?;
+
             return Ok(Transformed::yes(Arc::new(Column::new(
                 column.name(),
                 index,
@@ -506,7 +507,7 @@ pub(crate) mod tests {
     }
 
     #[test]
-    fn test_reassign_predicate_columns_in_list() {
+    fn test_reassign_expr_columns_in_list() {
         let int_field = Field::new("should_not_matter", DataType::Int64, true);
         let dict_field = Field::new(
             "id",
@@ -526,7 +527,7 @@ pub(crate) mod tests {
         )
         .unwrap();
 
-        let actual = reassign_predicate_columns(pred, &schema_small, false).unwrap();
+        let actual = reassign_expr_columns(pred, &schema_small).unwrap();
 
         let expected = in_list(
             Arc::new(Column::new_with_schema("id", &schema_small).unwrap()),

datafusion/physical-plan/src/filter.rs

@@ -265,7 +265,7 @@ impl FilterExec {
             default_selectivity,
         )?;
         let mut eq_properties = input.equivalence_properties().clone();
-        let (equal_pairs, _) = collect_columns_from_predicate(predicate);
+        let (equal_pairs, _) = collect_columns_from_predicate_inner(predicate);
         for (lhs, rhs) in equal_pairs {
             eq_properties.add_equal_conditions(Arc::clone(lhs), Arc::clone(rhs))?
         }
@@ -716,8 +716,18 @@ impl RecordBatchStream for FilterExecStream {
 }
 
 /// Return the equals Column-Pairs and Non-equals Column-Pairs
+#[deprecated(
+    since = "51.0.0",
+    note = "This function will be internal in the future"
+)]
 pub fn collect_columns_from_predicate(
     predicate: &'_ Arc<dyn PhysicalExpr>,
+) -> EqualAndNonEqual<'_> {
+    collect_columns_from_predicate_inner(predicate)
+}
+
+fn collect_columns_from_predicate_inner(
+    predicate: &'_ Arc<dyn PhysicalExpr>,
 ) -> EqualAndNonEqual<'_> {
     let mut eq_predicate_columns = Vec::<PhysicalExprPairRef>::new();
     let mut ne_predicate_columns = Vec::<PhysicalExprPairRef>::new();
@@ -787,7 +797,7 @@ mod tests {
             &schema,
         )?;
 
-        let (equal_pairs, ne_pairs) = collect_columns_from_predicate(&predicate);
+        let (equal_pairs, ne_pairs) = collect_columns_from_predicate_inner(&predicate);
         assert_eq!(2, equal_pairs.len());
         assert!(equal_pairs[0].0.eq(&col("c2", &schema)?));
         assert!(equal_pairs[0].1.eq(&lit(4u32)));

datafusion/physical-plan/src/filter_pushdown.rs

@@ -38,7 +38,7 @@ use std::collections::HashSet;
 use std::sync::Arc;
 
 use datafusion_common::Result;
-use datafusion_physical_expr::utils::{collect_columns, reassign_predicate_columns};
+use datafusion_physical_expr::utils::{collect_columns, reassign_expr_columns};
 use datafusion_physical_expr_common::physical_expr::PhysicalExpr;
 use itertools::Itertools;
 
@@ -343,7 +343,7 @@ impl ChildFilterDescription {
                 // All columns exist in child - we can push down
                 // Need to reassign column indices to match child schema
                 let reassigned_filter =
-                    reassign_predicate_columns(Arc::clone(filter), &child_schema, false)?;
+                    reassign_expr_columns(Arc::clone(filter), &child_schema)?;
                 child_parent_filters
                     .push(PushedDownPredicate::supported(reassigned_filter));
             } else {