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 547b2fa0dd23..40c16183f287 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 ab2ad6e9df05..75026390d32a 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",