From f252daf222e6c94ec69da1f7012bb21279726e05 Mon Sep 17 00:00:00 2001 From: Wenchen Fan Date: Sat, 13 Jun 2026 15:19:41 +0000 Subject: [PATCH 1/3] [SPARK-56395][SQL][FOLLOWUP] Fix RewriteNearestByJoin nullability for LEFT OUTER 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 -- produces it as nullable. This commit 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. --- .../optimizer/RewriteNearestByJoin.scala | 23 +++++++++++++++---- 1 file changed, 19 insertions(+), 4 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 ee8e61b457fa8..974c5191ad996 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 { From dd79e07e00d6928b7c1171ba7f8aa9fd8277ffd4 Mon Sep 17 00:00:00 2001 From: Wenchen Fan Date: Sat, 13 Jun 2026 15:36:19 +0000 Subject: [PATCH 2/3] [SPARK-56395][SQL][FOLLOWUP] Add regression test for LEFT OUTER nullability Update RewriteNearestByJoinSuite's expected-plan helper to mirror the widened right-side nullability for LEFT OUTER, and add a focused regression test that asserts every right-side attribute the synthesized Aggregate references agrees on nullability with its join child. The test uses non-nullable right columns so the LEFT OUTER widening is observable -- it fails without the production fix and passes with it. INNER stays a no-op. --- .../optimizer/RewriteNearestByJoinSuite.scala | 60 ++++++++++++++++++- 1 file changed, 57 insertions(+), 3 deletions(-) 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 729b58394d4bc..616aec511370c 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,50 @@ 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. The + // Aggregate built on top of that join references those columns (in the `_matches` struct and + // the ranking), so each reference must carry the widened nullability -- otherwise the + // rewritten plan declares a column non-nullable while its child produces it as nullable, an + // internal inconsistency. 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). + val left = LocalRelation($"a".int, $"b".int) + val right = LocalRelation( + AttributeReference("x", IntegerType, nullable = false)(), + AttributeReference("y", IntegerType, nullable = false)()) + Seq(Inner -> false, LeftOuter -> true).foreach { case (joinType, rightNullable) => + val query = NearestByJoin( + left, right, joinType, approx = true, numResults = 1, + rankingExpression = left.output(0) + right.output(0), + direction = NearestBySimilarity) + + val rewritten = RewriteNearestByJoin(query.analyze) + val join = rewritten.collect { case j: Join => j }.head + val aggregate = rewritten.collect { case a: Aggregate => a }.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)) + + // Every attribute the Aggregate references that its child (the join) produces must agree + // on nullability with the child -- this is exactly what the fix corrects for LEFT OUTER. + val childNullability = aggregate.child.output.map(a => a.exprId -> a.nullable).toMap + aggregate.aggregateExpressions.foreach(_.foreach { + case ref: AttributeReference if childNullability.contains(ref.exprId) => + assert(ref.nullable == childNullability(ref.exprId), + s"$joinType: ${ref.name}#${ref.exprId.id} declared nullable=${ref.nullable} " + + s"but its child produces 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). From ce977d245acdddcff76f4f213436f46944930646 Mon Sep 17 00:00:00 2001 From: Wenchen Fan Date: Sat, 13 Jun 2026 22:07:54 +0000 Subject: [PATCH 3/3] [SPARK-56395][SQL][FOLLOWUP] Generalize nullability test to cover the nondeterministic ranking path Generalize the LEFT OUTER nullability regression test in RewriteNearestByJoinSuite from an Aggregate-only assertion to a whole-plan integrity walk, and add a nondeterministic-ranking case alongside the deterministic one across both Inner and LeftOuter. The widened ranking reference lands in the Aggregate on the deterministic path but in a `__ranking__` Project (above the Join) on the nondeterministic path. The old assertion only inspected Aggregate expressions, so a regression isolated to the nondeterministic branch would not have been caught. No framework check guards this class of bug either -- LogicalPlanIntegrity compares types `asNullable` and schemas via `equalsIgnoreNullability` -- so this test is the only guard. Verified the new coverage has teeth: reverting just the nondeterministic alias to the un-widened expression fails the test specifically on LeftOuter/nondeterministic while the deterministic case still passes. Co-authored-by: Isaac --- .../optimizer/RewriteNearestByJoinSuite.scala | 51 ++++++++++++------- 1 file changed, 33 insertions(+), 18 deletions(-) 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 616aec511370c..90d6ad8065876 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 @@ -156,27 +156,35 @@ class RewriteNearestByJoinSuite extends PlanTest { } 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. The - // Aggregate built on top of that join references those columns (in the `_matches` struct and - // the ranking), so each reference must carry the widened nullability -- otherwise the - // rewritten plan declares a column non-nullable while its child produces it as nullable, an - // internal inconsistency. INNER does not widen the right side, so it stays a no-op. + // 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)()) - Seq(Inner -> false, LeftOuter -> true).foreach { case (joinType, rightNullable) => + 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 = left.output(0) + right.output(0), + rankingExpression = ranking, direction = NearestBySimilarity) val rewritten = RewriteNearestByJoin(query.analyze) val join = rewritten.collect { case j: Join => j }.head - val aggregate = rewritten.collect { case a: Aggregate => a }.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 @@ -186,16 +194,23 @@ class RewriteNearestByJoinSuite extends PlanTest { assert(joinRightOutput.nonEmpty) assert(joinRightOutput.forall(_.nullable == rightNullable)) - // Every attribute the Aggregate references that its child (the join) produces must agree - // on nullability with the child -- this is exactly what the fix corrects for LEFT OUTER. - val childNullability = aggregate.child.output.map(a => a.exprId -> a.nullable).toMap - aggregate.aggregateExpressions.foreach(_.foreach { - case ref: AttributeReference if childNullability.contains(ref.exprId) => - assert(ref.nullable == childNullability(ref.exprId), - s"$joinType: ${ref.name}#${ref.exprId.id} declared nullable=${ref.nullable} " + - s"but its child produces nullable=${childNullability(ref.exprId)}") - case _ => - }) + // 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 _ => + }) + } } }