NestedLoopJoin Projection Pushdown (apache#14120)

* nlj proj pushdown Signed-off-by: Jay Zhan <[email protected]> * fmt Signed-off-by: Jay Zhan <[email protected]> * move swap proj to util Signed-off-by: Jay Zhan <[email protected]> * fmt Signed-off-by: Jay Zhan <[email protected]> * fix proto Signed-off-by: Jay Zhan <[email protected]> * fmt Signed-off-by: Jay Zhan <[email protected]> * use none Signed-off-by: Jay Zhan <[email protected]> * proto fix Signed-off-by: Jay Zhan <[email protected]> * fix slt Signed-off-by: Jay Zhan <[email protected]> * Update projection_pushdown.rs * refactor: streamline projection pushdown logic for join operations * minor * fmt Signed-off-by: Jay Zhan <[email protected]> --------- Signed-off-by: Jay Zhan <[email protected]> Co-authored-by: berkaysynnada <[email protected]>
jonahgao · Jan 16, 2025 · 05f4e5a · 05f4e5a
1 parent 3cd31af
commit 05f4e5a
Show file tree

Hide file tree

Showing 14 changed files with 395 additions and 177 deletions.
diff --git a/datafusion/core/src/physical_optimizer/projection_pushdown.rs b/datafusion/core/src/physical_optimizer/projection_pushdown.rs
diff --git a/datafusion/core/src/physical_planner.rs b/datafusion/core/src/physical_planner.rs
@@ -1110,6 +1110,7 @@ impl DefaultPhysicalPlanner {
                             physical_right,
                             join_filter,
                             join_type,
+                            None,
                         )?)
                     }
                 } else if session_state.config().target_partitions() > 1

diff --git a/datafusion/core/tests/fuzz_cases/join_fuzz.rs b/datafusion/core/tests/fuzz_cases/join_fuzz.rs
@@ -486,7 +486,7 @@ impl JoinFuzzTestCase {
         let filter = JoinFilter::new(expression, column_indices, intermediate_schema);
 
         Arc::new(
-            NestedLoopJoinExec::try_new(left, right, Some(filter), &self.join_type)
+            NestedLoopJoinExec::try_new(left, right, Some(filter), &self.join_type, None)
                 .unwrap(),
         )
     }

diff --git a/datafusion/physical-optimizer/src/join_selection.rs b/datafusion/physical-optimizer/src/join_selection.rs
@@ -1055,6 +1055,7 @@ mod tests_statistical {
                 Arc::clone(&small),
                 nl_join_filter(),
                 &join_type,
+                None,
             )
             .unwrap(),
         );
@@ -1123,6 +1124,7 @@ mod tests_statistical {
                 Arc::clone(&small),
                 nl_join_filter(),
                 &join_type,
+                None,
             )
             .unwrap(),
         );

diff --git a/datafusion/physical-plan/src/joins/hash_join.rs b/datafusion/physical-plan/src/joins/hash_join.rs
@@ -26,7 +26,7 @@ use std::{any::Any, vec};
 
 use super::utils::{
     asymmetric_join_output_partitioning, get_final_indices_from_shared_bitmap,
-    reorder_output_after_swap,
+    reorder_output_after_swap, swap_join_projection,
 };
 use super::{
     utils::{OnceAsync, OnceFut},
@@ -477,7 +477,7 @@ impl HashJoinExec {
     }
 
     /// Return whether the join contains a projection
-    pub fn contain_projection(&self) -> bool {
+    pub fn contains_projection(&self) -> bool {
         self.projection.is_some()
     }
 
@@ -626,38 +626,6 @@ impl HashJoinExec {
     }
 }
 
-/// This function swaps the given join's projection.
-fn swap_join_projection(
-    left_schema_len: usize,
-    right_schema_len: usize,
-    projection: Option<&Vec<usize>>,
-    join_type: &JoinType,
-) -> Option<Vec<usize>> {
-    match join_type {
-        // For Anti/Semi join types, projection should remain unmodified,
-        // since these joins output schema remains the same after swap
-        JoinType::LeftAnti
-        | JoinType::LeftSemi
-        | JoinType::RightAnti
-        | JoinType::RightSemi => projection.cloned(),
-
-        _ => projection.map(|p| {
-            p.iter()
-                .map(|i| {
-                    // If the index is less than the left schema length, it is from
-                    // the left schema, so we add the right schema length to it.
-                    // Otherwise, it is from the right schema, so we subtract the left
-                    // schema length from it.
-                    if *i < left_schema_len {
-                        *i + right_schema_len
-                    } else {
-                        *i - left_schema_len
-                    }
-                })
-                .collect()
-        }),
-    }
-}
 impl DisplayAs for HashJoinExec {
     fn fmt_as(&self, t: DisplayFormatType, f: &mut fmt::Formatter) -> fmt::Result {
         match t {
@@ -666,7 +634,7 @@ impl DisplayAs for HashJoinExec {
                     || "".to_string(),
                     |f| format!(", filter={}", f.expression()),
                 );
-                let display_projections = if self.contain_projection() {
+                let display_projections = if self.contains_projection() {
                     format!(
                         ", projection=[{}]",
                         self.projection