From 59dc5820da4ac68d2970c10e3b7af5c0a7ac59d3 Mon Sep 17 00:00:00 2001 From: weimingdiit Date: Tue, 28 Jul 2026 14:40:06 +0800 Subject: [PATCH] [AURON #2437] Support inner residual join conditions in native broadcast hash join --- native-engine/auron-planner/proto/auron.proto | 1 + native-engine/auron-planner/src/planner.rs | 15 ++- .../src/broadcast_join_exec.rs | 8 +- .../src/joins/bhj/full_join.rs | 2 +- .../datafusion-ext-plans/src/joins/test.rs | 118 +++++++++++++++++- .../apache/spark/sql/auron/ShimsImpl.scala | 2 + .../auron/plan/NativeBroadcastJoinExec.scala | 4 +- .../org/apache/auron/AuronQuerySuite.scala | 46 +++++++ .../SparkAuronConfiguration.java | 2 +- .../spark/sql/auron/AuronConverters.scala | 4 +- .../org/apache/spark/sql/auron/Shims.scala | 1 + .../auron/plan/NativeBroadcastJoinBase.scala | 8 ++ 12 files changed, 194 insertions(+), 17 deletions(-) diff --git a/native-engine/auron-planner/proto/auron.proto b/native-engine/auron-planner/proto/auron.proto index 74b0bebb0..f922a7d74 100644 --- a/native-engine/auron-planner/proto/auron.proto +++ b/native-engine/auron-planner/proto/auron.proto @@ -500,6 +500,7 @@ message BroadcastJoinExecNode { JoinSide broadcast_side = 6; string cached_build_hash_map_id = 7; bool is_null_aware_anti_join = 8; + JoinFilter filter = 9; } message RenameColumnsExecNode { diff --git a/native-engine/auron-planner/src/planner.rs b/native-engine/auron-planner/src/planner.rs index b31e53636..9ce75d643 100644 --- a/native-engine/auron-planner/src/planner.rs +++ b/native-engine/auron-planner/src/planner.rs @@ -426,6 +426,15 @@ impl PhysicalPlanner { let join_type = protobuf::JoinType::try_from(broadcast_join.join_type) .expect("invalid JoinType"); + let join_type = join_type + .try_into() + .map_err(|_| proto_error("invalid JoinType"))?; + let join_filter = self.parse_join_filter(broadcast_join.filter.as_ref())?; + if join_filter.is_some() && join_type != JoinType::Inner { + return Err(proto_error( + "broadcast join filter is only supported for inner join", + )); + } let broadcast_side = protobuf::JoinSide::try_from(broadcast_join.broadcast_side) .expect("invalid BroadcastSide"); @@ -438,16 +447,14 @@ impl PhysicalPlanner { left, right, on, - join_type - .try_into() - .map_err(|_| proto_error("invalid JoinType"))?, + join_type, broadcast_side .try_into() .map_err(|_| proto_error("invalid BroadcastSide"))?, true, Some(cached_build_hash_map_id), is_null_aware_anti_join, - None, + join_filter, )?)) } PhysicalPlanType::Union(union) => { 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 001b74451..0df0959bd 100644 --- a/native-engine/datafusion-ext-plans/src/broadcast_join_exec.rs +++ b/native-engine/datafusion-ext-plans/src/broadcast_join_exec.rs @@ -107,12 +107,8 @@ impl BroadcastJoinExec { is_null_aware_anti_join: bool, join_filter: Option, ) -> Result { - // A broadcast hash join may reuse a prebuilt hash map whose stored - // rows are not laid out for evaluating Spark's residual join condition. - // Only the shuffled hash join path builds both sides in-process and - // can filter candidate pairs before projection today. - if join_filter.is_some() && (is_built || join_type != JoinType::Inner) { - df_execution_err!("join filter is only supported for inner shuffled hash join")?; + if join_filter.is_some() && join_type != JoinType::Inner { + df_execution_err!("join filter is only supported for inner hash join")?; } Ok(Self { left, diff --git a/native-engine/datafusion-ext-plans/src/joins/bhj/full_join.rs b/native-engine/datafusion-ext-plans/src/joins/bhj/full_join.rs index 14cb34229..2974595e2 100644 --- a/native-engine/datafusion-ext-plans/src/joins/bhj/full_join.rs +++ b/native-engine/datafusion-ext-plans/src/joins/bhj/full_join.rs @@ -167,7 +167,7 @@ impl FullJoiner

{ if let Some(join_filter) = &self.join_params.join_filter { if P.probe_side_outer || P.build_side_outer { - df_execution_err!("join filter is only supported for inner shuffled hash join")?; + df_execution_err!("join filter is only supported for inner hash join")?; } // Materialize candidate pairs from the hash lookup before // evaluating the residual condition. This keeps the filter inside diff --git a/native-engine/datafusion-ext-plans/src/joins/test.rs b/native-engine/datafusion-ext-plans/src/joins/test.rs index e82766273..70d0d2fca 100644 --- a/native-engine/datafusion-ext-plans/src/joins/test.rs +++ b/native-engine/datafusion-ext-plans/src/joins/test.rs @@ -29,7 +29,8 @@ mod tests { use datafusion::{ assert_batches_sorted_eq, common::{JoinSide, Result}, - physical_expr::expressions::Column, + logical_expr::Operator, + physical_expr::expressions::{BinaryExpr, Column}, physical_plan::{ExecutionPlan, common, joins::utils::*, test::TestMemoryExec}, prelude::{SessionConfig, SessionContext}, }; @@ -37,7 +38,10 @@ mod tests { use crate::{ broadcast_join_build_hash_map_exec::BroadcastJoinBuildHashMapExec, broadcast_join_exec::BroadcastJoinExec, - joins::join_utils::{JoinType, JoinType::*}, + joins::{ + ColumnIndex, JoinFilter, + join_utils::{JoinType, JoinType::*}, + }, sort_merge_join_exec::SortMergeJoinExec, }; @@ -382,6 +386,60 @@ mod tests { Ok((columns, batches)) } + async fn bhj_collect_with_filter( + test_type: TestType, + left: Arc, + right: Arc, + on: JoinOn, + join_filter: JoinFilter, + ) -> Result<(Vec, Vec)> { + MemManager::init(1000000); + let session_ctx = SessionContext::new(); + let task_ctx = session_ctx.task_ctx(); + let schema = build_join_schema_for_test(&left.schema(), &right.schema(), Inner)?; + + let (left, right, broadcast_side): ( + Arc, + Arc, + JoinSide, + ) = match test_type { + BHJLeftProbed => ( + left, + Arc::new(BroadcastJoinBuildHashMapExec::new( + right, + on.iter().map(|(_, right_key)| right_key.clone()).collect(), + )), + JoinSide::Right, + ), + BHJRightProbed => ( + Arc::new(BroadcastJoinBuildHashMapExec::new( + left, + on.iter().map(|(left_key, _)| left_key.clone()).collect(), + )), + right, + JoinSide::Left, + ), + _ => unreachable!("expected broadcast hash join test type"), + }; + + let join = BroadcastJoinExec::try_new( + schema, + left, + right, + on, + Inner, + broadcast_side, + true, + None, + false, + Some(join_filter), + )?; + let columns = columns(&join.schema()); + let stream = join.execute(0, task_ctx)?; + let batches = common::collect(stream).await?; + Ok((columns, batches)) + } + const ALL_TEST_TYPE: [TestType; 5] = [ SMJ, BHJLeftProbed, @@ -546,6 +604,62 @@ mod tests { Ok(()) } + #[tokio::test(flavor = "multi_thread", worker_threads = 1)] + async fn broadcast_join_inner_with_filter() -> Result<()> { + for test_type in [BHJLeftProbed, BHJRightProbed] { + let left = build_table( + ("a1", &vec![10, 20, 30]), + ("b1", &vec![1, 1, 2]), + ("c1", &vec![1, 5, 3]), + )?; + let right = build_table( + ("a2", &vec![100, 200, 300]), + ("b1", &vec![1, 1, 2]), + ("c2", &vec![2, 4, 2]), + )?; + let on: JoinOn = vec![( + Arc::new(Column::new_with_schema("b1", &left.schema())?), + Arc::new(Column::new_with_schema("b1", &right.schema())?), + )]; + + let filter_schema = Arc::new(Schema::new(vec![ + Field::new("left_c1", DataType::Int32, false), + Field::new("right_c2", DataType::Int32, false), + ])); + let join_filter = JoinFilter { + expression: Arc::new(BinaryExpr::new( + Arc::new(Column::new("left_c1", 0)), + Operator::Lt, + Arc::new(Column::new("right_c2", 1)), + )), + column_indices: vec![ + ColumnIndex { + side: JoinSide::Left, + index: 2, + }, + ColumnIndex { + side: JoinSide::Right, + index: 2, + }, + ], + schema: filter_schema, + }; + + let (_, batches) = + bhj_collect_with_filter(test_type, left, right, on, join_filter).await?; + let expected = vec![ + "+----+----+----+-----+----+----+", + "| a1 | b1 | c1 | a2 | b1 | c2 |", + "+----+----+----+-----+----+----+", + "| 10 | 1 | 1 | 100 | 1 | 2 |", + "| 10 | 1 | 1 | 200 | 1 | 4 |", + "+----+----+----+-----+----+----+", + ]; + assert_batches_sorted_eq!(expected, &batches); + } + Ok(()) + } + #[tokio::test(flavor = "multi_thread", worker_threads = 1)] async fn join_inner_batchsize() -> Result<()> { for test_type in ALL_TEST_TYPE { diff --git a/spark-extension-shims-spark/src/main/scala/org/apache/spark/sql/auron/ShimsImpl.scala b/spark-extension-shims-spark/src/main/scala/org/apache/spark/sql/auron/ShimsImpl.scala index 99a9e088b..a4cf6ea16 100644 --- a/spark-extension-shims-spark/src/main/scala/org/apache/spark/sql/auron/ShimsImpl.scala +++ b/spark-extension-shims-spark/src/main/scala/org/apache/spark/sql/auron/ShimsImpl.scala @@ -237,6 +237,7 @@ class ShimsImpl extends Shims with Logging { leftKeys: Seq[Expression], rightKeys: Seq[Expression], joinType: JoinType, + condition: Option[Expression], broadcastSide: JoinBuildSide, isNullAwareAntiJoin: Boolean): NativeBroadcastJoinBase = NativeBroadcastJoinExec( @@ -246,6 +247,7 @@ class ShimsImpl extends Shims with Logging { leftKeys, rightKeys, joinType, + condition, broadcastSide, isNullAwareAntiJoin) 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 f04770f08..efb78386d 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 @@ -33,6 +33,7 @@ case class NativeBroadcastJoinExec( override val leftKeys: Seq[Expression], override val rightKeys: Seq[Expression], override val joinType: JoinType, + override val condition: Option[Expression], broadcastSide: JoinBuildSide, isNullAwareAntiJoin: Boolean) extends NativeBroadcastJoinBase( @@ -42,12 +43,11 @@ case class NativeBroadcastJoinExec( leftKeys, rightKeys, joinType, + condition, broadcastSide, isNullAwareAntiJoin) with HashJoin { - override val condition: Option[Expression] = None - @sparkver("3.1 / 3.2 / 3.3 / 3.4 / 3.5 / 4.0 / 4.1") 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/AuronQuerySuite.scala b/spark-extension-shims-spark/src/test/scala/org/apache/auron/AuronQuerySuite.scala index 56119a7cb..f91ad270a 100644 --- a/spark-extension-shims-spark/src/test/scala/org/apache/auron/AuronQuerySuite.scala +++ b/spark-extension-shims-spark/src/test/scala/org/apache/auron/AuronQuerySuite.scala @@ -17,6 +17,7 @@ package org.apache.auron import org.apache.spark.sql.{AuronQueryTest, Row} +import org.apache.spark.sql.auron.join.JoinBuildSides.{JoinBuildLeft, JoinBuildRight} import org.apache.spark.sql.execution.auron.plan.NativeFilterBase import org.apache.spark.sql.execution.auron.plan.NativeShuffledHashJoinBase import org.apache.spark.sql.execution.auron.plan.NativeSortMergeJoinBase @@ -883,6 +884,51 @@ class AuronQuerySuite extends AuronQueryTest with BaseAuronSQLSuite with AuronSQ } } + test("native broadcast hash join supports inner residual condition") { + withSQLConf("spark.sql.adaptive.enabled" -> "false") { + withTable("bhj_left", "bhj_right") { + sql(""" + |CREATE TABLE bhj_left USING parquet AS + |SELECT * FROM VALUES + | (1, 1), + | (1, 5), + | (2, null), + | (3, 7) + |AS t(id, lv) + |""".stripMargin) + + sql(""" + |CREATE TABLE bhj_right USING parquet AS + |SELECT * FROM VALUES + | (1, 2), + | (1, 4), + | (2, 3), + | (3, 8) + |AS t(id, rv) + |""".stripMargin) + + Seq(("BROADCAST(r)", JoinBuildRight), ("BROADCAST(l)", JoinBuildLeft)).foreach { + case (hint, expectedBuildSide) => + val df = checkSparkAnswerAndOperator(s""" + |SELECT /*+ $hint */ l.id + |FROM bhj_left l + |JOIN bhj_right r + | ON l.id = r.id AND l.lv < r.rv + |ORDER BY l.id + |""".stripMargin) + + val plan = stripAQEPlan(df.queryExecution.executedPlan) + val nativeBhj = plan + .collectFirst { case join: NativeBroadcastJoinExec => join } + .getOrElse( + fail(s"expected NativeBroadcastJoinExec in executed plan, but got:\n$plan")) + assert(nativeBhj.condition.nonEmpty) + assert(nativeBhj.broadcastSide == expectedBuildSide) + } + } + } + } + test("left join with NOT IN subquery should filter NULL values") { // This test verifies the fix for the NULL handling issue in Anti join. withSQLConf("spark.sql.autoBroadcastJoinThreshold" -> "-1") { 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 366f32c13..9bf46bbe1 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 @@ -373,7 +373,7 @@ public class SparkAuronConfiguration extends AuronConfiguration { .withKey("auron.enable.native.join.condition") .withCategory("Operator Supports") .withDescription( - "Enable native SMJ/SHJ residual join condition evaluation. Disable this to fall back to Spark " + "Enable native SMJ/SHJ/BHJ 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 979aafaf6..bcd05235c 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 @@ -739,7 +739,7 @@ object AuronConverters extends Logging { "joinType" -> joinType, "condition" -> condition, "buildSide" -> buildSide)) - assert(condition.isEmpty, "join condition is not supported") + validateNativeInnerJoinCondition(joinType, condition) // verify build side is native buildSide match { @@ -756,6 +756,7 @@ object AuronConverters extends Logging { leftKeys, rightKeys, joinType, + condition, buildSide, naaj) @@ -797,6 +798,7 @@ object AuronConverters extends Logging { Nil, Nil, joinType, + None, buildSide, isNullAwareAntiJoin = false) } catch { diff --git a/spark-extension/src/main/scala/org/apache/spark/sql/auron/Shims.scala b/spark-extension/src/main/scala/org/apache/spark/sql/auron/Shims.scala index 9dcdfcf1c..a3853f52d 100644 --- a/spark-extension/src/main/scala/org/apache/spark/sql/auron/Shims.scala +++ b/spark-extension/src/main/scala/org/apache/spark/sql/auron/Shims.scala @@ -90,6 +90,7 @@ abstract class Shims { leftKeys: Seq[Expression], rightKeys: Seq[Expression], joinType: JoinType, + condition: Option[Expression], broadcastSide: JoinBuildSide, isNullAwareAntiJoin: Boolean): NativeBroadcastJoinBase diff --git a/spark-extension/src/main/scala/org/apache/spark/sql/execution/auron/plan/NativeBroadcastJoinBase.scala b/spark-extension/src/main/scala/org/apache/spark/sql/execution/auron/plan/NativeBroadcastJoinBase.scala index 9540a8b3a..82a50ca16 100644 --- a/spark-extension/src/main/scala/org/apache/spark/sql/execution/auron/plan/NativeBroadcastJoinBase.scala +++ b/spark-extension/src/main/scala/org/apache/spark/sql/execution/auron/plan/NativeBroadcastJoinBase.scala @@ -29,6 +29,7 @@ import org.apache.spark.sql.auron.Shims 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.plans.FullOuter +import org.apache.spark.sql.catalyst.plans.InnerLike import org.apache.spark.sql.catalyst.plans.JoinType import org.apache.spark.sql.catalyst.plans.LeftAnti import org.apache.spark.sql.catalyst.plans.LeftOuter @@ -53,6 +54,7 @@ abstract class NativeBroadcastJoinBase( leftKeys: Seq[Expression], rightKeys: Seq[Expression], joinType: JoinType, + condition: Option[Expression], broadcastSide: JoinBuildSide, isNullAwareAntiJoin: Boolean) extends BinaryExecNode @@ -113,6 +115,9 @@ abstract class NativeBroadcastJoinBase( private def nativeJoinType = NativeConverters.convertJoinType(joinType) + private def nativeJoinFilter = + condition.map(NativeConverters.convertJoinFilter(_, left.output, right.output)) + private def nativeBroadcastSide = broadcastSide match { case JoinBuildLeft => pb.JoinSide.LEFT_SIDE case JoinBuildRight => pb.JoinSide.RIGHT_SIDE @@ -121,6 +126,7 @@ abstract class NativeBroadcastJoinBase( protected def rewriteKeyExprToLong(exprs: Seq[Expression]): Seq[Expression] // check whether native converting is supported + assert(condition.isEmpty || joinType.isInstanceOf[InnerLike], "join condition is not supported") nativeSchema nativeJoinType nativeJoinOn @@ -133,6 +139,7 @@ abstract class NativeBroadcastJoinBase( val nativeSchema = this.nativeSchema val nativeJoinType = this.nativeJoinType val nativeJoinOn = this.nativeJoinOn + val nativeJoinFilter = this.nativeJoinFilter val (probedRDD, builtRDD) = broadcastSide match { case JoinBuildLeft => (rightRDD, leftRDD) @@ -179,6 +186,7 @@ abstract class NativeBroadcastJoinBase( .setCachedBuildHashMapId(cachedBuildHashMapId) .setIsNullAwareAntiJoin(isNullAwareAntiJoin) .addAllOn(nativeJoinOn.asJava) + nativeJoinFilter.foreach(broadcastJoinExec.setFilter) pb.PhysicalPlanNode.newBuilder().setBroadcastJoin(broadcastJoinExec).build() },