From 3535a81ca47cc639486f58cc226cd23514395cc9 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Luka=20Zdravi=C4=87?= Date: Fri, 11 Sep 2026 13:49:27 +0000 Subject: [PATCH] [SPARK-59438][SQL][TESTS] Add planner-stage test coverage for ASOF JOIN sort-merge operator Adds two planner-stage tests for the ASOF JOIN sort-merge path, both mirroring existing sibling-join test patterns: - `PlannerSuite`: a no-equi-key `AsOfJoin` requires a single partition on both sides (the `AllTuples` branch of `SortMergeAsOfJoinExec`'s `requiredChildDistribution`), verified by running `EnsureRequirements` and asserting a single-partition shuffle on each side. Follows the existing `SortMergeJoinExec` distribution tests (SPARK-24495 / SPARK-27485). - `SortMergeAsOfJoinSuite`: a null-safe (`<=>`) equi-key in `ON` is routed to the residual condition rather than treated as an equi-key, so null keys on both sides do match -- the counterpart to the existing `EqualTo` case where null keys never match. Mirrors `InnerJoinSuite`'s "inner join, null safe". Test-only change; no production code is modified. --- .../spark/sql/SortMergeAsOfJoinSuite.scala | 31 +++++++++++++++++++ .../spark/sql/execution/PlannerSuite.scala | 25 ++++++++++++++- 2 files changed, 55 insertions(+), 1 deletion(-) diff --git a/sql/core/src/test/scala/org/apache/spark/sql/SortMergeAsOfJoinSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/SortMergeAsOfJoinSuite.scala index 547b2fa0dd237..40c16183f287e 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/SortMergeAsOfJoinSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/SortMergeAsOfJoinSuite.scala @@ -602,6 +602,37 @@ class SortMergeAsOfJoinSuite extends QueryTest ) } + test("null-safe equi-key (<=>) in ON matches null keys, unlike EqualTo") { + // The AsOfJoinSelection strategy excludes EqualNullSafe from the equi-keys and + // routes it to the residual condition, so - unlike the EqualTo case above where + // null equi-keys never match - null keys on both sides DO match under <=>. + val schema1 = StructType( + StructField("grp", IntegerType, nullable = true) :: + StructField("ts", IntegerType) :: + StructField("val", StringType) :: Nil) + val schema2 = StructType( + StructField("grp", IntegerType, nullable = true) :: + StructField("ts", IntegerType) :: + StructField("val", StringType) :: Nil) + val df1 = spark.createDataFrame( + List(Row(null, 5, "a"), Row(1, 5, "b")).asJava, schema1) + val df2 = spark.createDataFrame( + List(Row(null, 3, "x"), Row(1, 4, "y")).asJava, schema2) + checkAnswer( + df1.joinAsOf( + df2, df1.col("ts"), df2.col("ts"), + joinExprs = df1.col("grp") <=> df2.col("grp"), + joinType = "inner", tolerance = null, + allowExactMatches = true, direction = "backward"), + Seq( + // grp=null <=> grp=null is true, so this left row matches (EqualTo would drop it) + Row(null, 5, "a", null, 3, "x"), + // grp=1: right.ts=4 <= left.ts=5 -> match + Row(1, 5, "b", 1, 4, "y") + ) + ) + } + test("residual condition via joinExprs") { // Test that pair-correlated residual predicates are routed into the // scanner's residualCondition (not a post-join FilterExec). diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/PlannerSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/PlannerSuite.scala index ab2ad6e9df055..75026390d32a6 100644 --- a/sql/core/src/test/scala/org/apache/spark/sql/execution/PlannerSuite.scala +++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/PlannerSuite.scala @@ -31,7 +31,7 @@ import org.apache.spark.sql.execution.adaptive.{AdaptiveSparkPlanHelper, Disable import org.apache.spark.sql.execution.aggregate.{HashAggregateExec, ObjectHashAggregateExec, SortAggregateExec} import org.apache.spark.sql.execution.columnar.{InMemoryRelation, InMemoryTableScanExec} import org.apache.spark.sql.execution.exchange.{BroadcastExchangeLike, EnsureRequirements, REPARTITION_BY_COL, ReusedExchangeExec, ShuffleExchangeExec, ShuffleExchangeLike} -import org.apache.spark.sql.execution.joins.{BroadcastHashJoinExec, SortMergeJoinExec} +import org.apache.spark.sql.execution.joins.{BroadcastHashJoinExec, SortMergeAsOfJoinExec, SortMergeJoinExec} import org.apache.spark.sql.execution.reuse.ReuseExchangeAndSubquery import org.apache.spark.sql.functions._ import org.apache.spark.sql.internal.SQLConf @@ -746,6 +746,29 @@ class PlannerSuite extends SharedSparkSession with AdaptiveSparkPlanHelper { } } + test("SPARK-59438: as-of join with no equi-keys requires a single partition") { + // SortMergeAsOfJoinExec overrides requiredChildDistribution to AllTuples on both + // sides when there are no equi-keys, so EnsureRequirements must shuffle each side + // to a single partition (rather than hash-partition on join keys). + val asOfExec = SortMergeAsOfJoinExec( + leftKeys = Nil, + rightKeys = Nil, + leftSortExprs = exprA :: Nil, + rightSortExprs = exprB :: Nil, + asOfCondition = GreaterThanOrEqual(exprA, exprB), + orderExpression = Subtract(exprA, exprB), + joinType = Inner, + condition = None, + left = planA, + right = planB) + val outputPlan = EnsureRequirements.apply(asOfExec) + assertDistributionRequirementsAreSatisfied(outputPlan) + val exchanges = outputPlan.collect { case e: ShuffleExchangeExec => e } + assert(exchanges.length == 2, s"Expected a shuffle on each side:\n$outputPlan") + assert(exchanges.forall(_.outputPartitioning == SinglePartition), + s"Both sides must be shuffled to a single partition:\n$outputPlan") + } + test("SPARK-24500: create union with stream of children") { withSQLConf( SQLConf.ANALYZER_SINGLE_PASS_RESOLVER_ENABLED.key -> "false",