Add non-column expression equality tracking to filter exec (apache#9819)

* Add non-column expression equality tracking to filter exec

* Minor changes

datafusion/physical-plan/src/filter.rs

@@ -29,7 +29,7 @@ use super::{
 };
 use crate::{
  metrics::{BaselineMetrics, ExecutionPlanMetricsSet, MetricsSet},
- Column, DisplayFormatType, ExecutionPlan,
+ DisplayFormatType, ExecutionPlan,
 };
 
 use arrow::compute::filter_record_batch;
@@ -192,9 +192,7 @@ impl FilterExec {
  let mut eq_properties = input.equivalence_properties().clone();
  let (equal_pairs, _) = collect_columns_from_predicate(predicate);
  for (lhs, rhs) in equal_pairs {
- let lhs_expr = Arc::new(lhs.clone()) as _;
- let rhs_expr = Arc::new(rhs.clone()) as _;
- eq_properties.add_equal_conditions(&lhs_expr, &rhs_expr)
+ eq_properties.add_equal_conditions(lhs, rhs)
  }
  // Add the columns that have only one viable value (singleton) after
  // filtering to constants.
@@ -405,34 +403,33 @@ impl RecordBatchStream for FilterExecStream {
 
 /// Return the equals Column-Pairs and Non-equals Column-Pairs
 fn collect_columns_from_predicate(predicate: &Arc<dyn PhysicalExpr>) -> EqualAndNonEqual {
- let mut eq_predicate_columns = Vec::<(&Column, &Column)>::new();
- let mut ne_predicate_columns = Vec::<(&Column, &Column)>::new();
+ let mut eq_predicate_columns = Vec::<PhysicalExprPairRef>::new();
+ let mut ne_predicate_columns = Vec::<PhysicalExprPairRef>::new();
 
  let predicates = split_conjunction(predicate);
  predicates.into_iter().for_each(|p| {
  if let Some(binary) = p.as_any().downcast_ref::<BinaryExpr>() {
- if let (Some(left_column), Some(right_column)) = (
- binary.left().as_any().downcast_ref::<Column>(),
- binary.right().as_any().downcast_ref::<Column>(),
- ) {
- match binary.op() {
- Operator::Eq => {
- eq_predicate_columns.push((left_column, right_column))
- }
- Operator::NotEq => {
- ne_predicate_columns.push((left_column, right_column))
- }
- _ => {}
+ match binary.op() {
+ Operator::Eq => {
+ eq_predicate_columns.push((binary.left(), binary.right()))
+ }
+ Operator::NotEq => {
+ ne_predicate_columns.push((binary.left(), binary.right()))
  }
+ _ => {}
  }
  }
  });
 
  (eq_predicate_columns, ne_predicate_columns)
 }
+
+/// Pair of `Arc<dyn PhysicalExpr>`s
+pub type PhysicalExprPairRef<'a> = (&'a Arc<dyn PhysicalExpr>, &'a Arc<dyn PhysicalExpr>);
+
 /// The equals Column-Pairs and Non-equals Column-Pairs in the Predicates
 pub type EqualAndNonEqual<'a> =
- (Vec<(&'a Column, &'a Column)>, Vec<(&'a Column, &'a Column)>);
+ (Vec<PhysicalExprPairRef<'a>>, Vec<PhysicalExprPairRef<'a>>);
 
 #[cfg(test)]
 mod tests {
@@ -482,14 +479,16 @@ mod tests {
  )?;
 
  let (equal_pairs, ne_pairs) = collect_columns_from_predicate(&predicate);
+ assert_eq!(2, equal_pairs.len());
+ assert!(equal_pairs[0].0.eq(&col("c2", &schema)?));
+ assert!(equal_pairs[0].1.eq(&lit(4u32)));
 
- assert_eq!(1, equal_pairs.len());
- assert_eq!(equal_pairs[0].0.name(), "c2");
- assert_eq!(equal_pairs[0].1.name(), "c9");
+ assert!(equal_pairs[1].0.eq(&col("c2", &schema)?));
+ assert!(equal_pairs[1].1.eq(&col("c9", &schema)?));
 
  assert_eq!(1, ne_pairs.len());
- assert_eq!(ne_pairs[0].0.name(), "c1");
- assert_eq!(ne_pairs[0].1.name(), "c13");
+ assert!(ne_pairs[0].0.eq(&col("c1", &schema)?));
+ assert!(ne_pairs[0].1.eq(&col("c13", &schema)?));
 
  Ok(())
  }

datafusion/physical-plan/src/lib.rs

@@ -33,7 +33,6 @@ use datafusion_common::config::ConfigOptions;
 use datafusion_common::utils::DataPtr;
 use datafusion_common::Result;
 use datafusion_execution::TaskContext;
-use datafusion_physical_expr::expressions::Column;
 use datafusion_physical_expr::{
  EquivalenceProperties, LexOrdering, PhysicalSortExpr, PhysicalSortRequirement,
 };

datafusion/sqllogictest/test_files/select.slt

@@ -1386,6 +1386,27 @@ AggregateExec: mode=FinalPartitioned, gby=[c2@0 as c2], aggr=[COUNT(*)]
 --------RepartitionExec: partitioning=RoundRobinBatch(2), input_partitions=1
 ----------CsvExec: file_groups={1 group: [[WORKSPACE_ROOT/testing/data/csv/aggregate_test_100.csv]]}, projection=[c2], has_header=true
 
+# FilterExec can track equality of non-column expressions.
+# plan below shouldn't have a SortExec because given column 'a' is ordered.
+# 'CAST(ROUND(b) as INT)' is also ordered. After filter is applied.
+query TT
+EXPLAIN SELECT *
+FROM annotated_data_finite2
+WHERE CAST(ROUND(b) as INT) = a
+ORDER BY CAST(ROUND(b) as INT);
+----
+logical_plan
+Sort: CAST(round(CAST(annotated_data_finite2.b AS Float64)) AS Int32) ASC NULLS LAST
+--Filter: CAST(round(CAST(annotated_data_finite2.b AS Float64)) AS Int32) = annotated_data_finite2.a
+----TableScan: annotated_data_finite2 projection=[a0, a, b, c, d], partial_filters=[CAST(round(CAST(annotated_data_finite2.b AS Float64)) AS Int32) = annotated_data_finite2.a]
+physical_plan
+SortPreservingMergeExec: [CAST(round(CAST(b@2 AS Float64)) AS Int32) ASC NULLS LAST]
+--CoalesceBatchesExec: target_batch_size=8192
+----FilterExec: CAST(round(CAST(b@2 AS Float64)) AS Int32) = a@1
+------RepartitionExec: partitioning=RoundRobinBatch(2), input_partitions=1
+--------CsvExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/core/tests/data/window_2.csv]]}, projection=[a0, a, b, c, d], output_ordering=[a@1 ASC NULLS LAST, b@2 ASC NULLS LAST, c@3 ASC NULLS LAST], has_header=true
+
+
 statement ok
 drop table annotated_data_finite2;