feat: support input reordering for NestedLoopJoinExec

datafusion/core/src/physical_optimizer/join_selection.rs

@@ -30,8 +30,8 @@ use crate::error::Result;
 use crate::physical_optimizer::PhysicalOptimizerRule;
 use crate::physical_plan::joins::utils::{ColumnIndex, JoinFilter};
 use crate::physical_plan::joins::{
- CrossJoinExec, HashJoinExec, PartitionMode, StreamJoinPartitionMode,
- SymmetricHashJoinExec,
+ CrossJoinExec, HashJoinExec, NestedLoopJoinExec, PartitionMode,
+ StreamJoinPartitionMode, SymmetricHashJoinExec,
 };
 use crate::physical_plan::projection::ProjectionExec;
 use crate::physical_plan::{ExecutionPlan, ExecutionPlanProperties};
@@ -199,6 +199,38 @@ fn swap_hash_join(
  }
 }
 
+/// Swaps inputs of `NestedLoopJoinExec` and wraps it into `ProjectionExec` is required
+fn swap_nl_join(join: &NestedLoopJoinExec) -> Result<Arc<dyn ExecutionPlan>> {
+ let new_filter = swap_join_filter(join.filter());
+ let new_join_type = &swap_join_type(*join.join_type());
+
+ let new_join = NestedLoopJoinExec::try_new(
+ Arc::clone(join.right()),
+ Arc::clone(join.left()),
+ new_filter,
+ new_join_type,
+ )?;
+
+ // For Semi/Anti joins, swap result will produce same output schema,
+ // no need to wrap them into additional projection
+ let plan: Arc<dyn ExecutionPlan> = if matches!(
+ join.join_type(),
+ JoinType::LeftSemi
+ | JoinType::RightSemi
+ | JoinType::LeftAnti
+ | JoinType::RightAnti
+ ) {
+ Arc::new(new_join)
+ } else {
+ let projection =
+ swap_reverting_projection(&join.left().schema(), &join.right().schema());
+
+ Arc::new(ProjectionExec::try_new(projection, Arc::new(new_join))?)
+ };
+
+ Ok(plan)
+}
+
 /// When the order of the join is changed by the optimizer, the columns in
 /// the output should not be impacted. This function creates the expressions
 /// that will allow to swap back the values from the original left as the first
@@ -461,6 +493,14 @@ fn statistical_join_selection_subrule(
  } else {
  None
  }
+ } else if let Some(nl_join) = plan.as_any().downcast_ref::<NestedLoopJoinExec>() {
+ let left = nl_join.left();
+ let right = nl_join.right();
+ if should_swap_join_order(&**left, &**right)? {
+ swap_nl_join(nl_join).map(Some)?
+ } else {
+ None
+ }
  } else {
  None
  };
@@ -697,9 +737,12 @@ mod tests_statistical {
 
  use arrow::datatypes::{DataType, Field, Schema};
  use datafusion_common::{stats::Precision, JoinType, ScalarValue};
- use datafusion_physical_expr::expressions::Column;
+ use datafusion_expr::Operator;
+ use datafusion_physical_expr::expressions::{BinaryExpr, Column};
  use datafusion_physical_expr::{PhysicalExpr, PhysicalExprRef};
 
+ use rstest::rstest;
+
  /// Return statistcs for empty table
  fn empty_statistics() -> Statistics {
  Statistics {
@@ -785,6 +828,35 @@ mod tests_statistical {
  }]
  }
 
+ /// Create join filter for NLJoinExec with expression `big_col > small_col`
+ /// where both columns are 0-indexed and come from left and right inputs respectively
+ fn nl_join_filter() -> Option<JoinFilter> {
+ let column_indices = vec![
+ ColumnIndex {
+ index: 0,
+ side: JoinSide::Left,
+ },
+ ColumnIndex {
+ index: 0,
+ side: JoinSide::Right,
+ },
+ ];
+ let intermediate_schema = Schema::new(vec![
+ Field::new("big_col", DataType::Int32, false),
+ Field::new("small_col", DataType::Int32, false),
+ ]);
+ let expression = Arc::new(BinaryExpr::new(
+ Arc::new(Column::new_with_schema("big_col", &intermediate_schema).unwrap()),
+ Operator::Gt,
+ Arc::new(Column::new_with_schema("big_col", &intermediate_schema).unwrap()),
+ )) as _;
+ Some(JoinFilter::new(
+ expression,
+ column_indices,
+ intermediate_schema,
+ ))
+ }
+
  /// Returns three plans with statistics of (min, max, distinct_count)
  /// * big 100K rows @ (0, 50k, 50k)
  /// * medium 10K rows @ (1k, 5k, 1k)
@@ -1151,6 +1223,137 @@ mod tests_statistical {
  crosscheck_plans(join).unwrap();
  }
 
+ #[rstest(
+ join_type,
+ case::inner(JoinType::Inner),
+ case::left(JoinType::Left),
+ case::right(JoinType::Right),
+ case::full(JoinType::Full)
+ )]
+ #[tokio::test]
+ async fn test_nl_join_with_swap(join_type: JoinType) {
+ let (big, small) = create_big_and_small();
+
+ let join = Arc::new(
+ NestedLoopJoinExec::try_new(
+ Arc::clone(&big),
+ Arc::clone(&small),
+ nl_join_filter(),
+ &join_type,
+ )
+ .unwrap(),
+ );
+
+ let optimized_join = JoinSelection::new()
+ .optimize(join.clone(), &ConfigOptions::new())
+ .unwrap();
+
+ let swapping_projection = optimized_join
+ .as_any()
+ .downcast_ref::<ProjectionExec>()
+ .expect("A proj is required to swap columns back to their original order");
+
+ assert_eq!(swapping_projection.expr().len(), 2);
+ let (col, name) = &swapping_projection.expr()[0];
+ assert_eq!(name, "big_col");
+ assert_col_expr(col, "big_col", 1);
+ let (col, name) = &swapping_projection.expr()[1];
+ assert_eq!(name, "small_col");
+ assert_col_expr(col, "small_col", 0);
+
+ let swapped_join = swapping_projection
+ .input()
+ .as_any()
+ .downcast_ref::<NestedLoopJoinExec>()
+ .expect("The type of the plan should not be changed");
+
+ // Assert join side of big_col swapped in filter expression
+ let swapped_filter = swapped_join.filter().unwrap();
+ let swapped_big_col_idx = swapped_filter.schema().index_of("big_col").unwrap();
+ let swapped_big_col_side = swapped_filter
+ .column_indices()
+ .get(swapped_big_col_idx)
+ .unwrap()
+ .side;
+ assert_eq!(
+ swapped_big_col_side,
+ JoinSide::Right,
+ "Filter column side should be swapped"
+ );
+
+ assert_eq!(
+ swapped_join.left().statistics().unwrap().total_byte_size,
+ Precision::Inexact(8192)
+ );
+ assert_eq!(
+ swapped_join.right().statistics().unwrap().total_byte_size,
+ Precision::Inexact(2097152)
+ );
+ crosscheck_plans(join.clone()).unwrap();
+ }
+
+ #[rstest(
+ join_type,
+ case::left_semi(JoinType::LeftSemi),
+ case::left_anti(JoinType::LeftAnti),
+ case::right_semi(JoinType::RightSemi),
+ case::right_anti(JoinType::RightAnti)
+ )]
+ #[tokio::test]
+ async fn test_nl_join_with_swap_no_proj(join_type: JoinType) {
+ let (big, small) = create_big_and_small();
+
+ let join = Arc::new(
+ NestedLoopJoinExec::try_new(
+ Arc::clone(&big),
+ Arc::clone(&small),
+ nl_join_filter(),
+ &join_type,
+ )
+ .unwrap(),
+ );
+
+ let optimized_join = JoinSelection::new()
+ .optimize(join.clone(), &ConfigOptions::new())
+ .unwrap();
+
+ let swapped_join = optimized_join
+ .as_any()
+ .downcast_ref::<NestedLoopJoinExec>()
+ .expect("The type of the plan should not be changed");
+
+ // Assert before/after schemas are equal
+ assert_eq!(
+ join.schema(),
+ swapped_join.schema(),
+ "Join schema should not be modified while optimization"
+ );
+
+ // Assert join side of big_col swapped in filter expression
+ let swapped_filter = swapped_join.filter().unwrap();
+ let swapped_big_col_idx = swapped_filter.schema().index_of("big_col").unwrap();
+ let swapped_big_col_side = swapped_filter
+ .column_indices()
+ .get(swapped_big_col_idx)
+ .unwrap()
+ .side;
+ assert_eq!(
+ swapped_big_col_side,
+ JoinSide::Right,
+ "Filter column side should be swapped"
+ );
+
+ assert_eq!(
+ swapped_join.left().statistics().unwrap().total_byte_size,
+ Precision::Inexact(8192)
+ );
+ assert_eq!(
+ swapped_join.right().statistics().unwrap().total_byte_size,
+ Precision::Inexact(2097152)
+ );
+ crosscheck_plans(join.clone()).unwrap();
+ }
+
  #[tokio::test]
  async fn test_swap_reverting_projection() {
  let left_schema = Schema::new(vec![

datafusion/core/tests/fuzz_cases/join_fuzz.rs

@@ -21,13 +21,19 @@ use arrow::array::{ArrayRef, Int32Array};
 use arrow::compute::SortOptions;
 use arrow::record_batch::RecordBatch;
 use arrow::util::pretty::pretty_format_batches;
+use arrow_schema::Schema;
 use rand::Rng;
 
+use datafusion::common::JoinSide;
+use datafusion::logical_expr::{JoinType, Operator};
+use datafusion::physical_expr::expressions::BinaryExpr;
 use datafusion::physical_plan::collect;
 use datafusion::physical_plan::expressions::Column;
-use datafusion::physical_plan::joins::{HashJoinExec, PartitionMode, SortMergeJoinExec};
+use datafusion::physical_plan::joins::utils::{ColumnIndex, JoinFilter};
+use datafusion::physical_plan::joins::{
+ HashJoinExec, NestedLoopJoinExec, PartitionMode, SortMergeJoinExec,
+};
 use datafusion::physical_plan::memory::MemoryExec;
-use datafusion_expr::JoinType;
 
 use datafusion::prelude::{SessionConfig, SessionContext};
 use test_utils::stagger_batch_with_seed;
@@ -73,7 +79,7 @@ async fn test_full_join_1k() {
 }
 
 #[tokio::test]
-async fn test_semi_join_1k() {
+async fn test_semi_join_10k() {
  run_join_test(
  make_staggered_batches(10000),
  make_staggered_batches(10000),
@@ -83,7 +89,7 @@ async fn test_semi_join_1k() {
 }
 
 #[tokio::test]
-async fn test_anti_join_1k() {
+async fn test_anti_join_10k() {
  run_join_test(
  make_staggered_batches(10000),
  make_staggered_batches(10000),
@@ -118,6 +124,46 @@ async fn run_join_test(
  ),
  ];
 
+ // Nested loop join uses filter for joining records
+ let column_indices = vec![
+ ColumnIndex {
+ index: 0,
+ side: JoinSide::Left,
+ },
+ ColumnIndex {
+ index: 1,
+ side: JoinSide::Left,
+ },
+ ColumnIndex {
+ index: 0,
+ side: JoinSide::Right,
+ },
+ ColumnIndex {
+ index: 1,
+ side: JoinSide::Right,
+ },
+ ];
+ let intermediate_schema = Schema::new(vec![
+ schema1.field_with_name("a").unwrap().to_owned(),
+ schema1.field_with_name("b").unwrap().to_owned(),
+ schema2.field_with_name("a").unwrap().to_owned(),
+ schema2.field_with_name("b").unwrap().to_owned(),
+ ]);
+
+ let equal_a = Arc::new(BinaryExpr::new(
+ Arc::new(Column::new("a", 0)),
+ Operator::Eq,
+ Arc::new(Column::new("a", 2)),
+ )) as _;
+ let equal_b = Arc::new(BinaryExpr::new(
+ Arc::new(Column::new("b", 1)),
+ Operator::Eq,
+ Arc::new(Column::new("b", 3)),
+ )) as _;
+ let expression = Arc::new(BinaryExpr::new(equal_a, Operator::And, equal_b)) as _;
+
+ let on_filter = JoinFilter::new(expression, column_indices, intermediate_schema);
+
  // sort-merge join
  let left = Arc::new(
  MemoryExec::try_new(&[input1.clone()], schema1.clone(), None).unwrap(),
@@ -161,22 +207,55 @@ async fn run_join_test(
  );
  let hj_collected = collect(hj, task_ctx.clone()).await.unwrap();
 
+ // nested loop join
+ let left = Arc::new(
+ MemoryExec::try_new(&[input1.clone()], schema1.clone(), None).unwrap(),
+ );
+ let right = Arc::new(
+ MemoryExec::try_new(&[input2.clone()], schema2.clone(), None).unwrap(),
+ );
+ let nlj = Arc::new(
+ NestedLoopJoinExec::try_new(left, right, Some(on_filter), &join_type)
+ .unwrap(),
+ );
+ let nlj_collected = collect(nlj, task_ctx.clone()).await.unwrap();
+
  // compare
  let smj_formatted = pretty_format_batches(&smj_collected).unwrap().to_string();
  let hj_formatted = pretty_format_batches(&hj_collected).unwrap().to_string();
+ let nlj_formatted = pretty_format_batches(&nlj_collected).unwrap().to_string();
 
  let mut smj_formatted_sorted: Vec<&str> = smj_formatted.trim().lines().collect();
  smj_formatted_sorted.sort_unstable();
 
  let mut hj_formatted_sorted: Vec<&str> = hj_formatted.trim().lines().collect();
  hj_formatted_sorted.sort_unstable();
 
+ let mut nlj_formatted_sorted: Vec<&str> = nlj_formatted.trim().lines().collect();
+ nlj_formatted_sorted.sort_unstable();
+
  for (i, (smj_line, hj_line)) in smj_formatted_sorted
  .iter()
  .zip(&hj_formatted_sorted)
  .enumerate()
  {
- assert_eq!((i, smj_line), (i, hj_line));
+ assert_eq!(
+ (i, smj_line),
+ (i, hj_line),
+ "SortMergeJoinExec and HashJoinExec produced different results"
+ );
+ }
+
+ for (i, (nlj_line, hj_line)) in nlj_formatted_sorted
+ .iter()
+ .zip(&hj_formatted_sorted)
+ .enumerate()
+ {
+ assert_eq!(
+ (i, nlj_line),
+ (i, hj_line),
+ "NestedLoopJoinExec and HashJoinExec produced different results"
+ );
  }
  }
 }