[AURON #2437] Support inner residual join conditions in native broadcast hash join

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 {

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) => {

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<JoinFilter>,
     ) -> Result<Self> {
-        // 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,

native-engine/datafusion-ext-plans/src/joins/bhj/full_join.rs

@@ -167,7 +167,7 @@ impl<const P: JoinerParams> FullJoiner<P> {
 
         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

native-engine/datafusion-ext-plans/src/joins/test.rs

@@ -29,15 +29,19 @@ 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},
     };
 
     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<dyn ExecutionPlan>,
+        right: Arc<dyn ExecutionPlan>,
+        on: JoinOn,
+        join_filter: JoinFilter,
+    ) -> Result<(Vec<String>, Vec<RecordBatch>)> {
+        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<dyn ExecutionPlan>,
+            Arc<dyn ExecutionPlan>,
+            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 {

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)
 

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 {

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") {

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);
 

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 {

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