Blaze Sort->MergeJoin reduces key row and column conversion (#1078)

Author: eden123456789

.idea/vcs.xml

@@ -409,20 +409,19 @@ async fn execute_join_with_smj_fallback(
     };
 
     // create sorted streams, build side is already sorted
-    let probed = exec_ctx.stat_input(exec_ctx.execute(&probed_plan)?);
     let (left_exec, right_exec) = match broadcast_side {
         JoinSide::Left => (
             built_sorted,
             create_default_ascending_sort_exec(
-                create_record_batch_stream_exec(probed, exec_ctx.partition_id())?,
+                probed_plan,
                 &join_params.right_keys,
                 Some(exec_ctx.execution_plan_metrics().clone()),
                 false, // do not record output metric
             ),
         ),
         JoinSide::Right => (
             create_default_ascending_sort_exec(
-                create_record_batch_stream_exec(probed, exec_ctx.partition_id())?,
+                probed_plan,
                 &join_params.left_keys,
                 Some(exec_ctx.execution_plan_metrics().clone()),
                 false, // do not record output metric
@@ -432,13 +431,18 @@ async fn execute_join_with_smj_fallback(
     };
 
     // run sort merge join
-    let mut smj_join_params = join_params.clone();
-    smj_join_params.sort_options = vec![SortOptions::default(); join_params.left_keys.len()];
-
-    let smj_exec = Arc::new(SortMergeJoinExec::try_new_with_join_params(
+    let smj_exec = Arc::new(SortMergeJoinExec::try_new(
+        join_params.output_schema,
         left_exec.clone(),
         right_exec.clone(),
-        smj_join_params,
+        join_params
+            .left_keys
+            .to_vec()
+            .into_iter()
+            .zip(join_params.right_keys.to_vec())
+            .collect(),
+        join_params.join_type,
+        vec![SortOptions::default(); join_params.left_keys.len()],
     )?);
     let mut join_output = smj_exec.execute(exec_ctx.partition_id(), exec_ctx.task_ctx())?;
 

native-engine/datafusion-ext-plans/src/broadcast_join_exec.rs

@@ -30,8 +30,10 @@ use datafusion::{
     physical_expr::{PhysicalExprRef, expressions::Column, utils::collect_columns},
     physical_plan::{ExecutionPlan, stream::RecordBatchStreamAdapter},
 };
+use datafusion_ext_commons::downcast_any;
 use futures::StreamExt;
 use itertools::Itertools;
+use parking_lot::Mutex;
 
 pub trait ExecuteWithColumnPruning {
     fn execute_projected(
@@ -78,24 +80,55 @@ pub fn prune_columns(exprs: &[PhysicalExprRef]) -> Result<(Vec<PhysicalExprRef>,
         .collect();
     let mapped_exprs: Vec<PhysicalExprRef> = exprs
         .iter()
-        .map(|expr| {
-            expr.clone()
-                .transform_down(&|node: PhysicalExprRef| {
-                    Ok(Transformed::yes(
-                        if let Some(column) = node.as_any().downcast_ref::<Column>() {
-                            let mapped_idx = required_columns_mapping[&column.index()];
-                            Arc::new(Column::new(column.name(), mapped_idx))
-                        } else {
-                            node
-                        },
-                    ))
-                })
-                .map(|r| r.data)
-        })
-        .collect::<Result<_>>()?;
+        .map(|expr| map_columns(expr, &required_columns_mapping))
+        .collect();
 
     Ok((
         mapped_exprs,
         required_columns.into_iter().map(|c| c.index()).collect(),
     ))
 }
+
+pub fn extend_projection_by_expr(
+    projection: &mut Vec<usize>,
+    expr: &PhysicalExprRef,
+) -> PhysicalExprRef {
+    let projection = Arc::new(Mutex::new(projection));
+    expr.clone()
+        .transform_down(&|node: PhysicalExprRef| {
+            Ok(Transformed::yes(
+                if let Ok(column) = downcast_any!(node, Column) {
+                    let mut projection = projection.lock();
+                    if let Some(existed_idx) =
+                        projection.iter().position(|&idx| idx == column.index())
+                    {
+                        Arc::new(Column::new(column.name(), existed_idx)) as PhysicalExprRef
+                    } else {
+                        let new_idx = projection.len();
+                        projection.push(column.index());
+                        Arc::new(Column::new(column.name(), new_idx)) as PhysicalExprRef
+                    }
+                } else {
+                    node
+                },
+            ))
+        })
+        .map(|r| r.data)
+        .unwrap()
+}
+
+pub fn map_columns(expr: &PhysicalExprRef, mapping: &HashMap<usize, usize>) -> PhysicalExprRef {
+    expr.clone()
+        .transform_down(&|node: PhysicalExprRef| {
+            Ok(Transformed::yes(
+                if let Ok(column) = downcast_any!(node, Column) {
+                    let mapped_idx = mapping[&column.index()];
+                    Arc::new(Column::new(column.name(), mapped_idx))
+                } else {
+                    node
+                },
+            ))
+        })
+        .map(|r| r.data)
+        .unwrap()
+}