(spark) branch master updated: [SPARK-56395][SQL][FOLLOWUP] Fix RewriteNearestByJoin nullability for LEFT OUTER

[we...@apache.org]

This is an automated email from the ASF dual-hosted git repository.

cloud-fan pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/spark.git


The following commit(s) were added to refs/heads/master by this push:
     new 62e4e166698f [SPARK-56395][SQL][FOLLOWUP] Fix RewriteNearestByJoin nullability for LEFT OUTER
62e4e166698f is described below

commit 62e4e166698f95bcbeb16687ceaa61a3a8f14dbc
Author: Wenchen Fan <we...@databricks.com>
AuthorDate: Sun Jun 14 17:43:49 2026 -0700

    [SPARK-56395][SQL][FOLLOWUP] Fix RewriteNearestByJoin nullability for LEFT OUTER
    
    ### What changes were proposed in this pull request?
    
    Followup to https://github.com/apache/spark/pull/55682.
    
    In `RewriteNearestByJoin`, when the `NEAREST BY` join type is `LEFT OUTER`, the synthesized `Join` widens the right-side columns to nullable. However, the synthesized `Aggregate` (and the optional `__ranking__` `Project`) built on top of that join still referenced the right-side columns via `right.output` and `rankingExpression` with their original (non-nullable) nullability. As a result the rewritten plan can declare a right-side column as non-nullable while its child -- the join --  [...]
    
    This PR maps the right-side attributes to their widened (nullable) form for `LEFT OUTER` and rewrites both the `CreateStruct(right.*)` and the ranking expression to use that widened nullability, so the rewritten plan's schema is consistent with its child. For `INNER` joins the right side is not widened, so this is a no-op.
    
    ### Why are the changes needed?
    
    Without this fix the rewritten plan for a `LEFT OUTER NEAREST BY` declares right-side columns non-nullable while its join child produces them nullable -- an inconsistency that nullability/plan-integrity validation flags as a regression.
    
    ### Does this PR introduce any user-facing change?
    
    No.
    
    ### How was this patch tested?
    
    Added a regression test to `RewriteNearestByJoinSuite` that, for both `INNER` and `LEFT OUTER`, asserts every right-side attribute the synthesized `Aggregate` references agrees on nullability with its join child. The test uses **non-nullable** right-side columns so that `LEFT OUTER`'s widening is observable -- it fails without this fix (`x#.. declared nullable=false but its child produces nullable=true`) and passes with it, while `INNER` stays a no-op. The suite's expected-plan helper [...]
    
    ### Was this patch authored or co-authored using generative AI tooling?
    
    No.
    
    Closes #56484 from cloud-fan/SPARK-56395-followup.
    
    Authored-by: Wenchen Fan <we...@databricks.com>
    Signed-off-by: Wenchen Fan <we...@databricks.com>
---
 .../catalyst/optimizer/RewriteNearestByJoin.scala  | 23 +++++--
 .../optimizer/RewriteNearestByJoinSuite.scala      | 75 +++++++++++++++++++++-
 2 files changed, 91 insertions(+), 7 deletions(-)

diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/RewriteNearestByJoin.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/RewriteNearestByJoin.scala
index ee8e61b457fa..974c5191ad99 100644
--- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/RewriteNearestByJoin.scala
+++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/RewriteNearestByJoin.scala
@@ -92,11 +92,26 @@ object RewriteNearestByJoin extends Rule[LogicalPlan] {
       //    cross-product today -- should not silently bypass that choice.
       val join = Join(taggedLeft, right, joinType, None, JoinHint.NONE)
 
-      val (aggInput, rankingForAgg) = if (!rankingExpression.deterministic) {
-        val rankingAlias = Alias(rankingExpression, "__ranking__")()
+      // A LEFT OUTER join widens the right-side columns to nullable. The synthesized Aggregate
+      // (and the optional `__ranking__` Project) below sit directly on top of this join, so every
+      // reference to a right-side column must carry that widened nullability. Otherwise the
+      // rewritten plan would declare a right column non-nullable while its child -- the join --
+      // produces it as nullable, which plan-integrity validation flags as a nullability
+      // regression. INNER joins do not widen the right side, so this is a no-op there.
+      val rightAttrs = joinType match {
+        case LeftOuter => right.output.map(_.withNullability(true))
+        case _ => right.output
+      }
+      val rightNullabilityMap = AttributeMap(right.output.zip(rightAttrs))
+      val rankingInJoin = rankingExpression.transform {
+        case a: Attribute => rightNullabilityMap.getOrElse(a, a)
+      }
+
+      val (aggInput, rankingForAgg) = if (!rankingInJoin.deterministic) {
+        val rankingAlias = Alias(rankingInJoin, "__ranking__")()
         Project(join.output :+ rankingAlias, join) -> rankingAlias.toAttribute
       } else {
-        join -> rankingExpression
+        join -> rankingInJoin
       }
 
       // 4. Aggregate grouped by `__qid`:
@@ -104,7 +119,7 @@ object RewriteNearestByJoin extends Rule[LogicalPlan] {
       //      - max_by/min_by(struct(right.*), ranking, k) as `_matches`.
       //    The ranking expression references left and right columns directly; no outer
       //    reference is needed because both sides are present in the joined input.
-      val rightStruct = CreateStruct(right.output)
+      val rightStruct = CreateStruct(rightAttrs)
       // reverse = true  -> MIN_BY (smallest ranking value first, for DISTANCE)
       // reverse = false -> MAX_BY (largest ranking value first, for SIMILARITY)
       val reverse = direction match {
diff --git a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/optimizer/RewriteNearestByJoinSuite.scala b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/optimizer/RewriteNearestByJoinSuite.scala
index 729b58394d4b..90d6ad806587 100644
--- a/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/optimizer/RewriteNearestByJoinSuite.scala
+++ b/sql/catalyst/src/test/scala/org/apache/spark/sql/catalyst/optimizer/RewriteNearestByJoinSuite.scala
@@ -19,7 +19,7 @@ package org.apache.spark.sql.catalyst.optimizer
 
 import org.apache.spark.sql.catalyst.dsl.expressions._
 import org.apache.spark.sql.catalyst.dsl.plans._
-import org.apache.spark.sql.catalyst.expressions.{Alias, AttributeReference, CreateStruct, Inline, Literal, Rand, Uuid}
+import org.apache.spark.sql.catalyst.expressions.{Alias, Attribute, AttributeMap, AttributeReference, CreateStruct, Inline, Literal, Rand, Uuid}
 import org.apache.spark.sql.catalyst.expressions.aggregate.{AggregateExpression, First, MaxMinByK}
 import org.apache.spark.sql.catalyst.plans.{Inner, JoinType, LeftOuter, NearestByDistance, NearestBySimilarity, PlanTest}
 import org.apache.spark.sql.catalyst.plans.logical.{Aggregate, Generate, Join, JoinHint, LocalRelation, NearestByJoin, Project}
@@ -46,9 +46,19 @@ class RewriteNearestByJoinSuite extends PlanTest {
     val taggedLeft = Project(left.output :+ qidAlias, left)
     val join = Join(taggedLeft, right, joinType, None, JoinHint.NONE)
 
-    val rightStruct = CreateStruct(right.output)
+    // Mirror the rewrite: a LEFT OUTER join widens right-side columns to nullable, so the
+    // struct and ranking that sit on top of the join must reference them with that nullability.
+    val rightAttrs = joinType match {
+      case LeftOuter => right.output.map(_.withNullability(true))
+      case _ => right.output
+    }
+    val rightNullabilityMap = AttributeMap(right.output.zip(rightAttrs))
+    val rankingInJoin = ranking.transform {
+      case a: Attribute => rightNullabilityMap.getOrElse(a, a)
+    }
+    val rightStruct = CreateStruct(rightAttrs)
     val topKAgg = MaxMinByK(
-      rightStruct, ranking, Literal(numResults), reverse = reverse)
+      rightStruct, rankingInJoin, Literal(numResults), reverse = reverse)
       .toAggregateExpression()
     val matchesAlias = Alias(topKAgg, "__nearest_matches__")()
     val firstLeftAggs = left.output.map { attr =>
@@ -145,6 +155,65 @@ class RewriteNearestByJoinSuite extends PlanTest {
     comparePlans(normalizeUuidSeed(rewritten), expected, checkAnalysis = false)
   }
 
+  test("SPARK-56395: LEFT OUTER rewrite keeps right-side nullability consistent with its child") {
+    // A LEFT OUTER NEAREST BY widens the synthetic join's right-side columns to nullable. Every
+    // operator built on top of that join that references those columns (the `_matches` struct,
+    // the ranking) must carry the widened nullability -- otherwise the rewritten plan declares a
+    // column non-nullable while its child produces it as nullable, an internal inconsistency that
+    // no framework check catches (`LogicalPlanIntegrity` compares types `asNullable` and schemas
+    // `equalsIgnoreNullability`), so this assertion is the only guard. INNER does not widen the
+    // right side, so it stays a no-op.
+    //
+    // The right-side columns are declared non-nullable here: that is what makes LEFT OUTER's
+    // widening observable (with nullable columns the widening is a no-op and the bug is hidden).
+    // The ranking is exercised both deterministic (the reference lands directly in the Aggregate)
+    // and nondeterministic (the rule pre-materializes it into a `__ranking__` Project above the
+    // Join), so the widening is checked wherever the reference ends up.
+    val left = LocalRelation($"a".int, $"b".int)
+    val right = LocalRelation(
+      AttributeReference("x", IntegerType, nullable = false)(),
+      AttributeReference("y", IntegerType, nullable = false)())
+    val rankings = Seq(
+      "deterministic" -> (left.output(0) + right.output(0)),
+      "nondeterministic" -> (Rand(Literal(0L)) + right.output(0)))
+    for ((joinType, rightNullable) <- Seq(Inner -> false, LeftOuter -> true);
+         (label, ranking) <- rankings) {
+      val query = NearestByJoin(
+        left, right, joinType, approx = true, numResults = 1,
+        rankingExpression = ranking,
+        direction = NearestBySimilarity)
+
+      val rewritten = RewriteNearestByJoin(query.analyze)
+      val join = rewritten.collect { case j: Join => j }.head
+
+      // Sanity-check the fixture: the synthetic join widens its right-side output to nullable
+      // iff it is LEFT OUTER. (`join.right` is the right relation as it appears in the rewritten
+      // plan, so its ExprIds line up with the join's output.)
+      val rightExprIds = join.right.output.map(_.exprId).toSet
+      val joinRightOutput = join.output.filter(a => rightExprIds.contains(a.exprId))
+      assert(joinRightOutput.nonEmpty)
+      assert(joinRightOutput.forall(_.nullable == rightNullable))
+
+      // Whole-plan integrity: at every operator, an attribute reference whose ExprId is produced
+      // by one of that operator's children must agree with the child on nullability -- this is
+      // exactly what the fix corrects for LEFT OUTER. Walking the whole plan (rather than just
+      // the Aggregate) also covers the `__ranking__` Project that the nondeterministic path
+      // inserts above the Join, where the widened ranking reference lands.
+      rewritten.foreach { node =>
+        val childNullability =
+          node.children.flatMap(_.output).map(a => a.exprId -> a.nullable).toMap
+        node.expressions.foreach(_.foreach {
+          case ref: AttributeReference if childNullability.contains(ref.exprId) =>
+            assert(ref.nullable == childNullability(ref.exprId),
+              s"$joinType/$label: ${ref.name}#${ref.exprId.id} declared " +
+                s"nullable=${ref.nullable} but its child produces " +
+                s"nullable=${childNullability(ref.exprId)}")
+          case _ =>
+        })
+      }
+    }
+  }
+
   test("synthetic Join uses the user's joinType") {
     // Locks in that the rewrite's synthetic Join carries the user's `joinType`
     // (Inner or LeftOuter).


---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@spark.apache.org
For additional commands, e-mail: commits-help@spark.apache.org