[SPARK-59438][SQL][TESTS] Add planner-stage test coverage for ASOF JOIN sort-merge operator

sql/core/src/test/scala/org/apache/spark/sql/SortMergeAsOfJoinSuite.scala

@@ -602,6 +602,40 @@ class SortMergeAsOfJoinSuite extends QueryTest
     )
   }
 
+  test("null-safe equi-key (<=>) in ON matches null keys, unlike EqualTo") {
+    // <=> is a residual (empty equi-keys), so null keys match, unlike EqualTo.
+    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"), Row(null, 10, "c")).asJava, schema1)
+    val df2 = spark.createDataFrame(
+      List(Row(null, 3, "x"), Row(1, 4, "y"), Row(null, 8, "z")).asJava, schema2)
+    val joined = df1.joinAsOf(
+      df2, df1.col("ts"), df2.col("ts"),
+      joinExprs = df1.col("grp") <=> df2.col("grp"),
+      joinType = "inner", tolerance = null,
+      allowExactMatches = true, direction = "backward")
+    checkAnswer(
+      joined,
+      Seq(
+        Row(null, 5, "a", null, 3, "x"),
+        Row(1, 5, "b", 1, 4, "y"),
+        Row(null, 10, "c", null, 8, "z")
+      )
+    )
+    val plan = joined.queryExecution.executedPlan
+    val asOfExecs = collectWithSubqueries(plan) { case j: SortMergeAsOfJoinExec => j }
+    assert(asOfExecs.length == 1, s"expected one SortMergeAsOfJoinExec in:\n$plan")
+    assert(asOfExecs.head.leftKeys.isEmpty && asOfExecs.head.rightKeys.isEmpty,
+      s"<=> must be a residual, so equi-keys must be empty, got ${asOfExecs.head}")
+  }
+
   test("residual condition via joinExprs") {
     // Test that pair-correlated residual predicates are routed into the
     // scanner's residualCondition (not a post-join FilterExec).

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,56 @@ class PlannerSuite extends SharedSparkSession with AdaptiveSparkPlanHelper {
     }
   }
 
+  test("SPARK-59438: as-of join with no equi-keys requires a single partition") {
+    // No equi-keys: distribution is AllTuples, so both sides shuffle to one partition.
+    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)
+    assert(asOfExec.requiredChildDistribution == Seq(AllTuples, AllTuples))
+    assert(asOfExec.requiredChildOrdering == Seq(Seq(orderingA), Seq(orderingB)))
+    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-59438: as-of join with equi-keys hash-partitions on the keys") {
+    // With equi-keys: distribution is Clustered, so each side hash-partitions on its key.
+    val asOfExec = SortMergeAsOfJoinExec(
+      leftKeys = exprC :: Nil,
+      rightKeys = exprC :: Nil,
+      leftSortExprs = exprA :: Nil,
+      rightSortExprs = exprB :: Nil,
+      asOfCondition = GreaterThanOrEqual(exprA, exprB),
+      orderExpression = Subtract(exprA, exprB),
+      joinType = Inner,
+      condition = None,
+      left = planA,
+      right = planB)
+    assert(asOfExec.requiredChildDistribution ==
+      Seq(ClusteredDistribution(exprC :: Nil), ClusteredDistribution(exprC :: Nil)))
+    assert(asOfExec.requiredChildOrdering ==
+      Seq(Seq(orderingC, orderingA), Seq(orderingC, orderingB)))
+    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 match {
+      case h: HashPartitioning => h.expressions == Seq(exprC)
+      case _ => false
+    }), s"Both sides must hash-partition on the equi-key:\n$outputPlan")
+  }
+
   test("SPARK-24500: create union with stream of children") {
     withSQLConf(
       SQLConf.ANALYZER_SINGLE_PASS_RESOLVER_ENABLED.key -> "false",