(spark) branch branch-4.0 updated: [SPARK-58323][4.0][SQL] Materialize AliasAware output ordering and partitioning to avoid StackOverflowError

[pt...@apache.org]

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

peter-toth pushed a commit to branch branch-4.0
in repository https://gitbox.apache.org/repos/asf/spark.git


The following commit(s) were added to refs/heads/branch-4.0 by this push:
     new 59b72972926c [SPARK-58323][4.0][SQL] Materialize AliasAware output ordering and partitioning to avoid StackOverflowError
59b72972926c is described below

commit 59b72972926caa0fcf645a60318ab7a006223d32
Author: Peter Toth <pe...@gmail.com>
AuthorDate: Sun Jul 26 10:02:52 2026 +0200

    [SPARK-58323][4.0][SQL] Materialize AliasAware output ordering and partitioning to avoid StackOverflowError
    
    ### What changes were proposed in this pull request?
    
    Backport of SPARK-58323 (#57502) to branch-4.0.
    
    `AliasAwareQueryOutputOrdering.outputOrdering` and `BroadcastHashJoinExec`'s output-partitioning expansion build their results with `multiTransform`, which returns a `LazyList`, and store the bounded result unforced. This forces them into strict collections with `.toList`, right after the existing `.take(...)` so the lazy short-circuit during candidate generation is preserved.
    
    Tailored for branch-4.0 (differs from the master PR):
    - The `PartitioningPreservingUnaryExecNode.outputPartitioning` change is omitted: on branch-4.0 that method already materializes a strict collection (`...iterator.flatMap(...).take(...).toSeq` over an `Iterator`), so there is no unforced `LazyList` to fix.
    - The `BroadcastJoinSuite` single-element assertion change is omitted: on branch-4.0 `expandOutputPartitioning` always returns a `PartitioningCollection` (no `case p :: Nil` unwrap), so the single-element output shape is unchanged.
    
    ### Why are the changes needed?
    
    Each plan node re-wraps the child ordering/partitioning's unforced `LazyList`, so across a deep projection chain the deferred nesting grows unbounded and overflows the stack when the ordering/partitioning is later serialized or deeply traversed. On master this was observed as a driver `StackOverflowError` at `DAGScheduler.submitMissingTasks` during task serialization of a `SortedMergeCoalescedRDD` (SPARK-55715) on a large bucketed `MERGE`. That RDD and code path do not exist on branch [...]
    
    ### Does this PR introduce _any_ user-facing change?
    
    No.
    
    ### How was this patch tested?
    
    New test in `ProjectedOrderingAndPartitioningSuite` asserting `sameOrderExpressions` is strict (not a `LazyList`); it fails before this change and passes after. Existing `ProjectedOrderingAndPartitioningSuite` / `BroadcastJoinSuite` / `BroadcastJoinSuiteAE` pass.
    
    ### Was this patch authored or co-authored using generative AI tooling?
    
    Generated-by: Claude Opus 4.8
    
    Closes #57532 from peter-toth/SPARK-58323-aliasaware-strict-ordering-partitioning-4.0.
    
    Authored-by: Peter Toth <pe...@gmail.com>
    Signed-off-by: Peter Toth <pe...@gmail.com>
---
 .../sql/catalyst/plans/AliasAwareOutputExpression.scala |  6 ++++++
 .../sql/execution/joins/BroadcastHashJoinExec.scala     |  5 ++++-
 .../ProjectedOrderingAndPartitioningSuite.scala         | 17 +++++++++++++++++
 3 files changed, 27 insertions(+), 1 deletion(-)

diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/AliasAwareOutputExpression.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/AliasAwareOutputExpression.scala
index e1a9e8b5ea81..e7d2ab8538a4 100644
--- a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/AliasAwareOutputExpression.scala
+++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/AliasAwareOutputExpression.scala
@@ -107,6 +107,12 @@ trait AliasAwareQueryOutputOrdering[T <: QueryPlan[T]]
           .flatMap(projectExpression)
           .filter(e => orderingSet.add(e.canonicalized))
           .take(aliasCandidateLimit)
+          // Materialize the bounded result into a strict collection. The `LazyList` above is only
+          // needed so `take` can short-circuit candidate generation; storing `sameOrderExpressions`
+          // as an unforced `LazyList` lets each plan node re-wrap the child ordering's lazy list,
+          // and across a deep projection chain that nesting overflows the stack when the ordering
+          // is later serialized or deeply traversed.
+          .toList
         if (sameOrderings.nonEmpty) {
           Some(sortOrder.copy(child = sameOrderings.head,
             sameOrderExpressions = sameOrderings.tail))
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/joins/BroadcastHashJoinExec.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/joins/BroadcastHashJoinExec.scala
index 368534d05b1f..0d1a35fd89e9 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/joins/BroadcastHashJoinExec.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/joins/BroadcastHashJoinExec.scala
@@ -117,7 +117,10 @@ case class BroadcastHashJoinExec(
       case e: Expression if streamedKeyToBuildKeyMapping.contains(e.canonicalized) =>
         e +: streamedKeyToBuildKeyMapping(e.canonicalized)
     }.asInstanceOf[LazyList[HashPartitioningLike]]
-      .take(conf.broadcastHashJoinOutputPartitioningExpandLimit))
+      .take(conf.broadcastHashJoinOutputPartitioningExpandLimit)
+      // Materialize the bounded expansion into a strict `List` -- never an unforced `LazyList`,
+      // which chained across nodes can overflow the stack when serialized or deeply traversed.
+      .toList)
   }
 
   protected override def doExecute(): RDD[InternalRow] = {
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/ProjectedOrderingAndPartitioningSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/ProjectedOrderingAndPartitioningSuite.scala
index ec13d48d45f8..d55e5bd68187 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/execution/ProjectedOrderingAndPartitioningSuite.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/ProjectedOrderingAndPartitioningSuite.scala
@@ -84,6 +84,23 @@ class ProjectedOrderingAndPartitioningSuite
     }
   }
 
+  test("SPARK-58323: AliasAware output ordering is strict, not lazy") {
+    // A multi-alias projection makes `sameOrderExpressions` non-trivial. It must be a strict
+    // collection: the underlying `multiTransform` produces a `LazyList`, and storing it unforced
+    // lets each plan node re-wrap the child ordering's lazy list. Across a deep projection chain
+    // that nesting overflows the stack when the ordering is later serialized or deeply traversed.
+    withSQLConf(SQLConf.EXPRESSION_PROJECTION_CANDIDATE_LIMIT.key -> "5") {
+      // id -> {x, y, z} gives sameOrderExpressions.
+      val df = spark.range(2).orderBy($"id").selectExpr("id as x", "id as y", "id as z")
+      val outputOrdering = df.queryExecution.optimizedPlan.outputOrdering
+      assert(outputOrdering.head.sameOrderExpressions.nonEmpty)
+      outputOrdering.foreach { so =>
+        assert(!so.sameOrderExpressions.isInstanceOf[LazyList[_]],
+          s"sameOrderExpressions must be strict, was ${so.sameOrderExpressions.getClass.getName}")
+      }
+    }
+  }
+
   test("SPARK-42049: Improve AliasAwareOutputExpression - ordering - multi-references") {
     val df = spark.range(2).selectExpr("id as a", "id as b")
       .orderBy($"a" + $"b").selectExpr("a as x", "b as y")


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