Implement deduplication for join key pairs and partition expressions; add inferred predicate alias handling

Author: kosiew

datafusion/core/src/physical_planner.rs

@@ -18,7 +18,7 @@
 //! Planner for [`LogicalPlan`] to [`ExecutionPlan`]
 
 use std::borrow::Cow;
-use std::collections::HashMap;
+use std::collections::{HashMap, HashSet};
 use std::sync::Arc;
 
 use crate::datasource::file_format::file_type_to_format;
@@ -151,6 +151,46 @@ pub trait ExtensionPlanner {
     ) -> Result<Option<Arc<dyn ExecutionPlan>>>;
 }
 
+fn expr_canonical_string(expr: &Arc<dyn PhysicalExpr>) -> String {
+    format!("{}", expr)
+}
+
+fn dedupe_join_key_pairs(
+    left: Vec<Arc<dyn PhysicalExpr>>,
+    right: Vec<Arc<dyn PhysicalExpr>>,
+) -> (Vec<Arc<dyn PhysicalExpr>>, Vec<Arc<dyn PhysicalExpr>>) {
+    assert_eq!(left.len(), right.len());
+    let mut seen = HashSet::new();
+    let mut out_left = Vec::new();
+    let mut out_right = Vec::new();
+    for (l, r) in left.into_iter().zip(right.into_iter()) {
+        let l_s = expr_canonical_string(&l);
+        let r_s = expr_canonical_string(&r);
+        let key = if l_s <= r_s {
+            format!("{}={}", l_s, r_s)
+        } else {
+            format!("{}={}", r_s, l_s)
+        };
+        if seen.insert(key) {
+            out_left.push(l);
+            out_right.push(r);
+        }
+    }
+    (out_left, out_right)
+}
+
+fn dedupe_partition_exprs(exprs: Vec<Arc<dyn PhysicalExpr>>) -> Vec<Arc<dyn PhysicalExpr>> {
+    let mut seen = HashSet::new();
+    let mut out = Vec::new();
+    for e in exprs.into_iter() {
+        let key = expr_canonical_string(&e);
+        if seen.insert(key) {
+            out.push(e);
+        }
+    }
+    out
+}
+
 /// Default single node physical query planner that converts a
 /// `LogicalPlan` to an `ExecutionPlan` suitable for execution.
 ///
@@ -858,6 +898,7 @@ impl DefaultPhysicalPlanner {
                                 )
                             })
                             .collect::<Result<Vec<_>>>()?;
+                        let runtime_expr = dedupe_partition_exprs(runtime_expr);
                         Partitioning::Hash(runtime_expr, *n)
                     }
                     LogicalPartitioning::DistributeBy(_) => {
@@ -1088,6 +1129,10 @@ impl DefaultPhysicalPlanner {
                         Ok((l, r))
                     })
                     .collect::<Result<join_utils::JoinOn>>()?;
+                let (left_keys, right_keys): (Vec<_>, Vec<_>) = join_on.into_iter().unzip();
+                let (left_keys, right_keys) = dedupe_join_key_pairs(left_keys, right_keys);
+                let join_on: join_utils::JoinOn =
+                    left_keys.into_iter().zip(right_keys).collect();
 
                 let join_filter = match filter {
                     Some(expr) => {
@@ -2345,6 +2390,7 @@ impl<'a> OptimizationInvariantChecker<'a> {
 
         Ok(())
     }
+
 }
 
 impl<'n> TreeNodeVisitor<'n> for OptimizationInvariantChecker<'_> {
@@ -3520,4 +3566,29 @@ digraph {
 
         Ok(())
     }
+
+    #[test]
+    fn test_dedupe_join_key_pairs() {
+        let l1: Arc<dyn PhysicalExpr> = Arc::new(Column::new("a", 0));
+        let r1: Arc<dyn PhysicalExpr> = Arc::new(Column::new("b", 0));
+        let l2: Arc<dyn PhysicalExpr> = Arc::new(Column::new("b", 0));
+        let r2: Arc<dyn PhysicalExpr> = Arc::new(Column::new("a", 0));
+        let (left, right) = dedupe_join_key_pairs(vec![l1.clone(), l2], vec![r1.clone(), r2]);
+        assert_eq!(left.len(), 1);
+        assert_eq!(format!("{}", left[0]), "a@0");
+        assert_eq!(format!("{}", right[0]), "b@0");
+    }
+
+    #[test]
+    fn test_dedupe_partition_exprs() {
+        let exprs: Vec<Arc<dyn PhysicalExpr>> = vec![
+            Arc::new(Column::new("a", 0)),
+            Arc::new(Column::new("a", 0)),
+            Arc::new(Column::new("b", 1)),
+        ];
+        let result = dedupe_partition_exprs(exprs);
+        assert_eq!(result.len(), 2);
+        assert_eq!(format!("{}", result[0]), "a@0");
+        assert_eq!(format!("{}", result[1]), "b@1");
+    }
 }

datafusion/expr/src/expr.rs

@@ -38,6 +38,7 @@ use datafusion_common::tree_node::{
 use datafusion_common::{
     Column, DFSchema, HashMap, Result, ScalarValue, Spans, TableReference,
 };
+use crate::utils::INFERRED_PREDICATE_ALIAS;
 use datafusion_functions_window_common::field::WindowUDFFieldArgs;
 use sqlparser::ast::{
     display_comma_separated, ExceptSelectItem, ExcludeSelectItem, IlikeSelectItem,
@@ -3183,7 +3184,13 @@ pub const UNNEST_COLUMN_PREFIX: &str = "UNNEST";
 impl Display for Expr {
     fn fmt(&self, f: &mut Formatter) -> fmt::Result {
         match self {
-            Expr::Alias(Alias { expr, name, .. }) => write!(f, "{expr} AS {name}"),
+            Expr::Alias(Alias { expr, name, .. }) => {
+                if name == INFERRED_PREDICATE_ALIAS {
+                    write!(f, "{expr}")
+                } else {
+                    write!(f, "{expr} AS {name}")
+                }
+            }
             Expr::Column(c) => write!(f, "{c}"),
             Expr::OuterReferenceColumn(_, c) => {
                 write!(f, "{OUTER_REFERENCE_COLUMN_PREFIX}({c})")

datafusion/expr/src/utils.rs

@@ -28,6 +28,15 @@ use crate::{
 };
 use datafusion_expr_common::signature::{Signature, TypeSignature};
 
+/// Alias used to mark inferred join predicates that should remain join filters
+/// and not be converted into additional equijoin keys during optimization.
+pub const INFERRED_PREDICATE_ALIAS: &str = "__datafusion_inferred_join_predicate";
+
+/// Returns true if expr is an alias used to mark inferred join predicates.
+pub fn is_inferred_alias(expr: &Expr) -> bool {
+    matches!(expr, Expr::Alias(Alias { name, .. }) if name == INFERRED_PREDICATE_ALIAS)
+}
+
 use arrow::datatypes::{DataType, Field, Schema};
 use datafusion_common::tree_node::{
     Transformed, TransformedResult, TreeNode, TreeNodeRecursion,
@@ -946,7 +955,9 @@ fn split_conjunction_impl<'a>(expr: &'a Expr, mut exprs: Vec<&'a Expr>) -> Vec<&
             let exprs = split_conjunction_impl(left, exprs);
             split_conjunction_impl(right, exprs)
         }
-        Expr::Alias(Alias { expr, .. }) => split_conjunction_impl(expr, exprs),
+        Expr::Alias(alias) if !is_inferred_alias(expr) => {
+            split_conjunction_impl(&alias.expr, exprs)
+        }
         other => {
             exprs.push(other);
             exprs
@@ -970,7 +981,11 @@ pub fn iter_conjunction(expr: &Expr) -> impl Iterator<Item = &Expr> {
                     stack.push(right);
                     stack.push(left);
                 }
-                Expr::Alias(Alias { expr, .. }) => stack.push(expr),
+                Expr::Alias(Alias { expr, name, .. })
+                    if name != INFERRED_PREDICATE_ALIAS =>
+                {
+                    stack.push(expr);
+                }
                 other => return Some(other),
             }
         }
@@ -994,7 +1009,11 @@ pub fn iter_conjunction_owned(expr: Expr) -> impl Iterator<Item = Expr> {
                     stack.push(*right);
                     stack.push(*left);
                 }
-                Expr::Alias(Alias { expr, .. }) => stack.push(*expr),
+                Expr::Alias(Alias { expr, name, .. })
+                    if name != INFERRED_PREDICATE_ALIAS =>
+                {
+                    stack.push(*expr);
+                }
                 other => return Some(other),
             }
         }
@@ -1063,8 +1082,13 @@ fn split_binary_owned_impl(
             let exprs = split_binary_owned_impl(*left, operator, exprs);
             split_binary_owned_impl(*right, operator, exprs)
         }
-        Expr::Alias(Alias { expr, .. }) => {
-            split_binary_owned_impl(*expr, operator, exprs)
+        Expr::Alias(alias) => {
+            if is_inferred_alias(&Expr::Alias(alias.clone())) {
+                exprs.push(Expr::Alias(alias));
+                exprs
+            } else {
+                split_binary_owned_impl(*alias.expr, operator, exprs)
+            }
         }
         other => {
             exprs.push(other);
@@ -1090,7 +1114,9 @@ fn split_binary_impl<'a>(
             let exprs = split_binary_impl(left, operator, exprs);
             split_binary_impl(right, operator, exprs)
         }
-        Expr::Alias(Alias { expr, .. }) => split_binary_impl(expr, operator, exprs),
+        Expr::Alias(alias) if !is_inferred_alias(expr) => {
+            split_binary_impl(&alias.expr, operator, exprs)
+        }
         other => {
             exprs.push(other);
             exprs
@@ -1648,6 +1674,17 @@ mod tests {
         assert_eq!(result, vec![&expr1, &expr2]);
     }
 
+    #[test]
+    fn test_split_conjunction_inferred_alias() {
+        let expr = col("a").eq(col("b")).alias(INFERRED_PREDICATE_ALIAS);
+        let result = split_conjunction(&expr);
+        assert_eq!(result.len(), 1);
+        match result[0] {
+            Expr::Alias(alias) => assert_eq!(alias.name, INFERRED_PREDICATE_ALIAS),
+            _ => panic!("expected alias"),
+        }
+    }
+
     #[test]
     fn test_split_conjunction_or() {
         let expr = col("a").eq(lit(5)).or(col("b"));
@@ -1705,6 +1742,17 @@ mod tests {
         );
     }
 
+    #[test]
+    fn test_split_conjunction_owned_inferred_alias() {
+        let expr = col("a").eq(col("b")).alias(INFERRED_PREDICATE_ALIAS);
+        let result = split_conjunction_owned(expr);
+        assert_eq!(result.len(), 1);
+        match &result[0] {
+            Expr::Alias(alias) => assert_eq!(alias.name, INFERRED_PREDICATE_ALIAS),
+            _ => panic!("expected alias"),
+        }
+    }
+
     #[test]
     fn test_conjunction_empty() {
         assert_eq!(conjunction(vec![]), None);