From ca3b8da7e8aeeba183dfc8e3adb5d084673e5dd1 Mon Sep 17 00:00:00 2001 From: zhangyizhong Date: Sun, 20 Sep 2026 17:01:13 +0800 Subject: [PATCH] Support residual conditions in native broadcast nested loop joins --- .../src/broadcast_join_build_hash_map_exec.rs | 226 +++++++++++++++++- .../src/broadcast_join_exec.rs | 18 +- .../auron/plan/NativeBroadcastJoinExec.scala | 12 + .../AuronBroadcastNestedLoopJoinSuite.scala | 201 ++++++++++++++++ .../SparkAuronConfiguration.java | 2 +- .../spark/sql/auron/AuronConverters.scala | 4 +- 6 files changed, 448 insertions(+), 15 deletions(-) create mode 100644 spark-extension-shims-spark/src/test/scala/org/apache/auron/AuronBroadcastNestedLoopJoinSuite.scala diff --git a/native-engine/datafusion-ext-plans/src/broadcast_join_build_hash_map_exec.rs b/native-engine/datafusion-ext-plans/src/broadcast_join_build_hash_map_exec.rs index 07d9dc81b..666248501 100644 --- a/native-engine/datafusion-ext-plans/src/broadcast_join_build_hash_map_exec.rs +++ b/native-engine/datafusion-ext-plans/src/broadcast_join_build_hash_map_exec.rs @@ -31,7 +31,7 @@ use auron_jni_bridge::{ use datafusion::{ common::Result, execution::{SendableRecordBatchStream, TaskContext}, - physical_expr::{EquivalenceProperties, Partitioning, PhysicalExprRef}, + physical_expr::{EquivalenceProperties, Partitioning, PhysicalExprRef, expressions::lit}, physical_plan::{ DisplayAs, DisplayFormatType, ExecutionPlan, ExecutionPlanProperties, PlanProperties, execution_plan::{Boundedness, EmissionType}, @@ -135,6 +135,16 @@ impl ExecutionPlan for BroadcastJoinBuildHashMapExec { } } +/// Use a constant key for keyless joins so SMJ fallback forms one Cartesian +/// group. +pub(crate) fn smj_fallback_keys(keys: &[PhysicalExprRef]) -> Vec { + if keys.is_empty() { + vec![lit(0i32)] + } else { + keys.to_vec() + } +} + pub fn execute_build_hash_map( mut input: SendableRecordBatchStream, keys: Vec, @@ -205,7 +215,7 @@ pub fn execute_build_hash_map( let input_exec = create_record_batch_stream_exec(input, exec_ctx.partition_id())?; let sort_exec = create_default_ascending_sort_exec( input_exec, - &keys, + &smj_fallback_keys(&keys), Some(exec_ctx.execution_plan_metrics().clone()), false, // do not record output metric ); @@ -234,3 +244,215 @@ pub fn execute_build_hash_map( Ok(()) })) } + +#[cfg(test)] +mod tests { + use arrow::array::Int32Array; + use arrow_schema::{Field, Schema}; + use auron_memmgr::MemManager; + use datafusion::{ + common::{JoinSide, ScalarValue}, + logical_expr::Operator, + physical_expr::expressions::{BinaryExpr, Column}, + physical_plan::{common::collect, joins::utils::build_join_schema, test::TestMemoryExec}, + prelude::SessionContext, + }; + + use super::*; + use crate::{ + broadcast_join_exec::BroadcastJoinExec, + common::column_pruning::ExecuteWithColumnPruning, + joins::{ColumnIndex, JoinFilter, join_utils::JoinType}, + }; + + fn memory(batch: RecordBatch) -> Result> { + Ok(Arc::new(TestMemoryExec::try_new( + &[vec![batch.clone()]], + batch.schema(), + None, + )?)) + } + + #[tokio::test] + async fn nested_loop_sorted_fallback_preserves_condition_and_projection() -> Result<()> { + MemManager::init(1000000); + let values = [ + vec![None, Some(1), Some(3), Some(8)], + vec![None, Some(2), Some(4)], + ]; + let batches = values + .iter() + .enumerate() + .map(|(side, values)| { + RecordBatch::try_new( + Arc::new(Schema::new(vec![Field::new( + if side == 0 { "l" } else { "r" }, + DataType::Int32, + true, + )])), + vec![Arc::new(Int32Array::from(values.clone()))], + ) + .expect("valid test batch") + }) + .collect::>(); + for (jt, build_left) in [ + (JoinType::Inner, false), + (JoinType::Inner, true), + (JoinType::Left, false), + (JoinType::Right, true), + (JoinType::LeftSemi, false), + (JoinType::LeftAnti, false), + (JoinType::Existence, false), + ] { + let schema = if jt == JoinType::Existence { + Arc::new(Schema::new(vec![ + batches[0].schema().field(0).clone(), + Field::new("exists", DataType::Boolean, false), + ])) + } else { + Arc::new( + build_join_schema(&batches[0].schema(), &batches[1].schema(), &jt.try_into()?) + .0, + ) + }; + let mut expected = vec![]; + let mut right_matched = vec![false; values[1].len()]; + for &l in &values[0] { + let mut matched = false; + for (ri, &r) in values[1].iter().enumerate() { + if l.zip(r).is_some_and(|(l, r)| l < r) { + matched = true; + right_matched[ri] = true; + if matches!(jt, JoinType::Inner | JoinType::Left | JoinType::Right) { + expected.push(vec![ScalarValue::Int32(l), ScalarValue::Int32(r)]); + } + } + } + match jt { + JoinType::Left if !matched => { + expected.push(vec![ScalarValue::Int32(l), ScalarValue::Int32(None)]) + } + JoinType::LeftSemi if matched => expected.push(vec![ScalarValue::Int32(l)]), + JoinType::LeftAnti if !matched => expected.push(vec![ScalarValue::Int32(l)]), + JoinType::Existence => expected.push(vec![ + ScalarValue::Int32(l), + ScalarValue::Boolean(Some(matched)), + ]), + _ => {} + } + } + if jt == JoinType::Right { + for (&r, matched) in values[1].iter().zip(right_matched) { + if !matched { + expected.push(vec![ScalarValue::Int32(None), ScalarValue::Int32(r)]); + } + } + } + let ctx = SessionContext::new().task_ctx(); + let build_batch = batches[usize::from(!build_left)].clone(); + // A null table column marks the build stream as sorted for SMJ fallback. + let sorted = create_default_ascending_sort_exec( + memory(build_batch)?, + &smj_fallback_keys(&[]), + None, + false, + ); + let sorted_batches = collect(sorted.execute(0, ctx.clone())?).await?; + let mut spill_batches = vec![]; + for batch in sorted_batches { + let schema = join_hash_map_schema(&batch.schema()); + let cols = [ + batch.columns().to_vec(), + vec![new_null_array(&DataType::Binary, batch.num_rows())], + ] + .concat(); + spill_batches.push(RecordBatch::try_new(schema, cols)?); + } + let spill_schema = spill_batches[0].schema(); + let built = Arc::new(TestMemoryExec::try_new( + &[spill_batches], + spill_schema, + None, + )?) as Arc; + let (left, right) = if build_left { + (built, memory(batches[1].clone())?) + } else { + (memory(batches[0].clone())?, built) + }; + let filter = JoinFilter { + expression: Arc::new(BinaryExpr::new( + Arc::new(Column::new("l", 0)), + Operator::Lt, + Arc::new(Column::new("r", 1)), + )), + column_indices: vec![ + ColumnIndex { + side: JoinSide::Left, + index: 0, + }, + ColumnIndex { + side: JoinSide::Right, + index: 0, + }, + ], + schema: Arc::new(Schema::new(vec![ + batches[0].schema().field(0).clone(), + batches[1].schema().field(0).clone(), + ])), + }; + let join = BroadcastJoinExec::try_new( + schema.clone(), + left, + right, + vec![], + jt, + if build_left { + JoinSide::Left + } else { + JoinSide::Right + }, + true, + None, + false, + Some(filter), + )?; + for projection in [ + (0..schema.fields().len()).collect::>(), + vec![schema.fields().len() - 1], + vec![], + ] { + let output = collect(join.execute_projected(0, ctx.clone(), &projection)?).await?; + let mut actual = vec![]; + for batch in output { + for row in 0..batch.num_rows() { + actual.push( + batch + .columns() + .iter() + .map(|col| { + ScalarValue::try_from_array(col, row).map(|v| v.to_string()) + }) + .collect::>>()?, + ); + } + } + let mut wanted = expected + .iter() + .map(|row| { + projection + .iter() + .map(|&i| row[i].to_string()) + .collect::>() + }) + .collect::>(); + actual.sort(); + wanted.sort(); + assert_eq!( + actual, wanted, + "join={jt:?} build_left={build_left} projection={projection:?}" + ); + } + } + Ok(()) + } +} diff --git a/native-engine/datafusion-ext-plans/src/broadcast_join_exec.rs b/native-engine/datafusion-ext-plans/src/broadcast_join_exec.rs index 1b59d57a8..485f420a9 100644 --- a/native-engine/datafusion-ext-plans/src/broadcast_join_exec.rs +++ b/native-engine/datafusion-ext-plans/src/broadcast_join_exec.rs @@ -50,7 +50,7 @@ use once_cell::sync::OnceCell; use parking_lot::Mutex; use crate::{ - broadcast_join_build_hash_map_exec::execute_build_hash_map, + broadcast_join_build_hash_map_exec::{execute_build_hash_map, smj_fallback_keys}, common::{ column_pruning::ExecuteWithColumnPruning, execution_context::{ExecutionContext, WrappedRecordBatchSender}, @@ -458,13 +458,16 @@ async fn execute_join_with_smj_fallback( create_record_batch_stream_exec(remoted_stream, exec_ctx.partition_id())? }; + let left_keys = smj_fallback_keys(&join_params.left_keys); + let right_keys = smj_fallback_keys(&join_params.right_keys); + // create sorted streams, build side is already sorted let (left_exec, right_exec) = match broadcast_side { JoinSide::Left => ( built_sorted, create_default_ascending_sort_exec( probed_plan, - &join_params.right_keys, + &right_keys, Some(exec_ctx.execution_plan_metrics().clone()), false, // do not record output metric ), @@ -472,7 +475,7 @@ async fn execute_join_with_smj_fallback( JoinSide::Right => ( create_default_ascending_sort_exec( probed_plan, - &join_params.left_keys, + &left_keys, Some(exec_ctx.execution_plan_metrics().clone()), false, // do not record output metric ), @@ -486,15 +489,10 @@ async fn execute_join_with_smj_fallback( join_params.output_schema, left_exec.clone(), right_exec.clone(), - join_params - .left_keys - .to_vec() - .into_iter() - .zip(join_params.right_keys.to_vec()) - .collect(), + left_keys.iter().cloned().zip(right_keys).collect(), join_params.join_type, join_params.join_filter.clone(), - vec![SortOptions::default(); join_params.left_keys.len()], + vec![SortOptions::default(); left_keys.len()], )?); let mut projection = match join_params.join_type { Inner | Left | Right | Full => join_params diff --git a/spark-extension-shims-spark/src/main/scala/org/apache/spark/sql/execution/joins/auron/plan/NativeBroadcastJoinExec.scala b/spark-extension-shims-spark/src/main/scala/org/apache/spark/sql/execution/joins/auron/plan/NativeBroadcastJoinExec.scala index de2a7d2c8..b7b1699c7 100644 --- a/spark-extension-shims-spark/src/main/scala/org/apache/spark/sql/execution/joins/auron/plan/NativeBroadcastJoinExec.scala +++ b/spark-extension-shims-spark/src/main/scala/org/apache/spark/sql/execution/joins/auron/plan/NativeBroadcastJoinExec.scala @@ -18,12 +18,14 @@ package org.apache.spark.sql.execution.joins.auron.plan import org.apache.spark.sql.auron.join.JoinBuildSides.{JoinBuildLeft, JoinBuildRight, JoinBuildSide} import org.apache.spark.sql.catalyst.expressions.Expression +import org.apache.spark.sql.catalyst.expressions.SortOrder import org.apache.spark.sql.catalyst.plans.JoinType import org.apache.spark.sql.catalyst.plans.physical.Partitioning import org.apache.spark.sql.execution.SparkPlan import org.apache.spark.sql.execution.auron.plan.NativeBroadcastJoinBase import org.apache.spark.sql.execution.joins.HashJoin +import org.apache.auron.spark.configuration.SparkAuronConfiguration import org.apache.auron.sparkver case class NativeBroadcastJoinExec( @@ -48,6 +50,16 @@ case class NativeBroadcastJoinExec( isNullAwareAntiJoin) with HashJoin { + // Keyless joins and SMJ fallback cannot guarantee the original probe ordering. + override def outputOrdering: Seq[SortOrder] = { + if ((leftKeys.isEmpty && rightKeys.isEmpty) || SparkAuronConfiguration.SMJ_FALLBACK_ENABLE + .get()) { + Nil + } else { + super.outputOrdering + } + } + @sparkver("3.1 / 3.2 / 3.3 / 3.4 / 3.5 / 4.0 / 4.1 / 4.2") override def buildSide: org.apache.spark.sql.catalyst.optimizer.BuildSide = broadcastSide match { diff --git a/spark-extension-shims-spark/src/test/scala/org/apache/auron/AuronBroadcastNestedLoopJoinSuite.scala b/spark-extension-shims-spark/src/test/scala/org/apache/auron/AuronBroadcastNestedLoopJoinSuite.scala new file mode 100644 index 000000000..0bfeab317 --- /dev/null +++ b/spark-extension-shims-spark/src/test/scala/org/apache/auron/AuronBroadcastNestedLoopJoinSuite.scala @@ -0,0 +1,201 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package org.apache.auron + +import org.apache.spark.SparkEnv +import org.apache.spark.sql.auron.NativeHelper +import org.apache.spark.sql.auron.Shims +import org.apache.spark.sql.auron.join.JoinBuildSides.{JoinBuildLeft, JoinBuildRight} +import org.apache.spark.sql.catalyst.expressions.{Ascending, SortOrder} +import org.apache.spark.sql.execution.SortExec +import org.apache.spark.sql.execution.joins.BroadcastNestedLoopJoinExec +import org.apache.spark.sql.execution.joins.auron.plan.NativeBroadcastJoinExec + +import org.apache.auron.util.AuronTestUtils + +class AuronBroadcastNestedLoopJoinSuite extends AuronJoinConditionTestBase { + import testImplicits._ + + for ((joinType, buildAlias, native) <- Seq( + ("inner", "l", true), + ("inner", "r", true), + ("left outer", "r", true), + ("right outer", "l", true), + ("left semi", "r", true), + ("left anti", "r", true), + ("left outer", "l", false), + ("right outer", "r", false), + ("left semi", "l", false), + ("left anti", "l", false), + ("full outer", "l", false), + ("full outer", "r", false))) { + test(s"conditional BNLJ $joinType building $buildAlias with multiple probe partitions") { + withSQLConf( + "spark.sql.adaptive.enabled" -> "false", + "spark.sql.shuffle.partitions" -> "2", + // Keep the explicit probe repartition below semi/anti joins. + "spark.sql.optimizer.excludedRules" -> + "org.apache.spark.sql.catalyst.optimizer.PushDownLeftSemiAntiJoin", + "spark.auron.enable.bnlj" -> "true", + "spark.auron.enable.broadcastExchange" -> "true") { + withTempView("bnlj_left", "bnlj_right") { + spark.sparkContext + .parallelize(Seq[Integer](null, 0, 1, 3), 2) + .toDF("k") + .repartition(2) + .createOrReplaceTempView("bnlj_left") + spark.sparkContext + .parallelize(Seq[Integer](null, 1, 2), 2) + .toDF("k") + .repartition(2) + .createOrReplaceTempView("bnlj_right") + + val df = checkSparkAnswer(s""" + |SELECT /*+ BROADCAST($buildAlias) */ * + |FROM bnlj_left l $joinType JOIN bnlj_right r ON l.k < r.k + |""".stripMargin) + val plan = df.queryExecution.executedPlan + val nativeJoin = collectFirst(plan) { case join: NativeBroadcastJoinExec => join } + val sparkJoin = collectFirst(plan) { case join: BroadcastNestedLoopJoinExec => join } + val expectedBuildSide = if (buildAlias == "l") JoinBuildLeft else JoinBuildRight + if (native) { + assert(nativeJoin.isDefined, plan.toString) + assert(nativeJoin.get.leftKeys.isEmpty && nativeJoin.get.rightKeys.isEmpty) + assert(nativeJoin.get.condition.isDefined) + assert(nativeJoin.get.broadcastSide == expectedBuildSide) + val probe = if (buildAlias == "l") nativeJoin.get.right else nativeJoin.get.left + assert(NativeHelper.executeNative(probe).getNumPartitions == 2) + val orderedProbe = + SortExec(Seq(SortOrder(probe.output.head, Ascending)), global = false, probe) + val orderedJoin = if (buildAlias == "l") { + nativeJoin.get.copy(right = orderedProbe) + } else { + nativeJoin.get.copy(left = orderedProbe) + } + assert(orderedProbe.outputOrdering.nonEmpty) + assert(orderedJoin.outputOrdering.isEmpty) + assert(sparkJoin.isEmpty, plan.toString) + } else { + assert(nativeJoin.isEmpty, plan.toString) + assert(sparkJoin.isDefined, plan.toString) + assert(Shims.get.getJoinBuildSide(sparkJoin.get) == expectedBuildSide) + } + } + } + } + } + + test("BNLJ count-only projection retains condition columns") { + withSQLConf("spark.sql.adaptive.enabled" -> "false") { + withJoinInputs { + checkJoinPlan( + """ + |SELECT /*+ BROADCAST(r) */ COUNT(*) + |FROM condition_left l LEFT OUTER JOIN condition_right r ON l.v < r.v + |""".stripMargin, + _.isInstanceOf[NativeBroadcastJoinExec]) + } + } + } + + test("BNLJ residual condition respects the configuration switch") { + withSQLConf( + "spark.sql.adaptive.enabled" -> "false", + "spark.auron.enable.native.join.condition" -> "false") { + withJoinInputs { + checkJoinPlan( + """ + |SELECT /*+ BROADCAST(r) */ * + |FROM condition_left l LEFT OUTER JOIN condition_right r ON l.v < r.v + |""".stripMargin, + _.isInstanceOf[NativeBroadcastJoinExec], + native = false) + } + } + } + + for ((joinType, buildAlias) <- Seq(("inner", "l"), ("left outer", "r"))) { + test(s"BNLJ $joinType building $buildAlias preserves residuals in forced SMJ fallback") { + withSparkConf( + "spark.auron.smjfallback.enable" -> "true", + "spark.auron.smjfallback.rows.threshold" -> "0") { + withSQLConf( + "spark.sql.adaptive.enabled" -> "false", + "spark.sql.shuffle.partitions" -> "2") { + withJoinInputs { + val df = checkSparkAnswer(s""" + |SELECT /*+ BROADCAST($buildAlias) */ * + |FROM condition_left l $joinType JOIN condition_right r ON l.v < r.v + |""".stripMargin) + val plan = df.queryExecution.executedPlan + val join = collectFirst(plan) { case native: NativeBroadcastJoinExec => native } + assert(join.isDefined, plan.toString) + assert(join.get.leftKeys.isEmpty && join.get.rightKeys.isEmpty) + assert(join.get.metrics("fallback_sort_merge_join_time").value > 0) + assert(join.get.outputOrdering.isEmpty) + } + } + } + } + } + + test("broadcast hash join ordering accounts for enabled SMJ fallback") { + withSQLConf("spark.sql.adaptive.enabled" -> "false") { + withTempView("ordering_left", "ordering_right") { + Seq((1, 0), (1, 1)) + .toDF("k", "v") + .repartition(2) + .createOrReplaceTempView("ordering_left") + Seq((1, 0)) + .toDF("k", "v") + .repartition(2) + .createOrReplaceTempView("ordering_right") + val df = sql(""" + |SELECT /*+ BROADCAST(r) */ * + |FROM ordering_left l LEFT OUTER JOIN ordering_right r + |ON l.k = r.k AND l.v > r.v + |""".stripMargin) + val plan = df.queryExecution.executedPlan + val join = collectFirst(plan) { case native: NativeBroadcastJoinExec => native } + assert(join.isDefined, plan.toString) + assert(join.get.leftKeys.nonEmpty) + val orderedProbe = SortExec( + Seq(SortOrder(join.get.left.output.last, Ascending)), + global = false, + join.get.left) + val orderedJoin = join.get.copy(left = orderedProbe) + val conf = SparkEnv.get.conf + val key = "spark.auron.smjfallback.enable" + val previous = conf.getOption(key) + try { + conf.set(key, "false") + // Spark 3.0 HashJoin does not expose probe-side ordering. + val expectedOrdering = + if (AuronTestUtils.isSparkV31OrGreater) orderedProbe.outputOrdering else Nil + assert(orderedJoin.outputOrdering == expectedOrdering) + conf.set(key, "true") + assert(orderedJoin.outputOrdering.isEmpty) + } finally { + previous match { + case Some(value) => conf.set(key, value) + case None => conf.remove(key) + } + } + } + } + } +} diff --git a/spark-extension/src/main/java/org/apache/auron/spark/configuration/SparkAuronConfiguration.java b/spark-extension/src/main/java/org/apache/auron/spark/configuration/SparkAuronConfiguration.java index a8a9d43b6..86c52a7f1 100644 --- a/spark-extension/src/main/java/org/apache/auron/spark/configuration/SparkAuronConfiguration.java +++ b/spark-extension/src/main/java/org/apache/auron/spark/configuration/SparkAuronConfiguration.java @@ -405,7 +405,7 @@ public class SparkAuronConfiguration extends AuronConfiguration { .withKey("auron.enable.native.join.condition") .withCategory("Operator Supports") .withDescription( - "Enable native SMJ/SHJ/BHJ residual join condition evaluation. Disable this to fall back to Spark " + "Enable native SMJ/SHJ/BHJ/BNLJ residual join condition evaluation. Disable this to fall back to Spark " + "for joins with residual conditions when native evaluation is not desirable.") .withDefaultValue(true); diff --git a/spark-extension/src/main/scala/org/apache/spark/sql/auron/AuronConverters.scala b/spark-extension/src/main/scala/org/apache/spark/sql/auron/AuronConverters.scala index e6cc43f40..93ab58bdf 100644 --- a/spark-extension/src/main/scala/org/apache/spark/sql/auron/AuronConverters.scala +++ b/spark-extension/src/main/scala/org/apache/spark/sql/auron/AuronConverters.scala @@ -794,7 +794,7 @@ object AuronConverters extends Logging { require( JoinBuildSides.supportsBroadcastJoin(joinType, buildSide), s"native broadcast join does not support $joinType with $buildSide") - assert(condition.isEmpty, "join condition is not supported") + validateNativeJoinCondition(condition) // verify build side is native buildSide match { @@ -812,7 +812,7 @@ object AuronConverters extends Logging { Nil, Nil, joinType, - None, + condition, buildSide, isNullAwareAntiJoin = false) } catch {