[SPARK-53738][SQL] Fix planned write when query output contains foldable orderings #52584

pan3793 · 2025-10-13T05:48:08Z

What changes were proposed in this pull request?

This is the second try of #52474, following the suggestion from cloud-fan

This PR fixes a bug in plannedWrite, where the query has foldable orderings in the partition columns.

CREATE TABLE t (i INT, j INT, k STRING) USING PARQUET PARTITIONED BY (k);

INSERT OVERWRITE t SELECT j AS i, i AS j, '0' as k FROM t0 SORT BY k, i;

The evaluation of FileFormatWriter.orderingMatched fails because SortOrder(Literal) is eliminated by EliminateSorts.

Why are the changes needed?

V1Writes will override the custom sort order when the query output ordering does not satisfy the required ordering. Before SPARK-53707, when the query's output contains literals in partition columns, the judgment produces a false-negative result, thus causing the sort order not to take effect.

SPARK-53707 partially fixes the issue on the logical plan by adding a Project of query in V1Writes.

Before SPARK-53707

Sort [0 ASC NULLS FIRST, i#280 ASC NULLS FIRST], false
+- Project [j#287 AS i#280, i#286 AS j#281, 0 AS k#282]
   +- Relation spark_catalog.default.t0[i#286,j#287,k#288] parquet

After SPARK-53707

Project [i#284, j#285, 0 AS k#290]
+- Sort [0 ASC NULLS FIRST, i#284 ASC NULLS FIRST], false
   +- Project [i#284, j#285]
      +- Relation spark_catalog.default.t0[i#284,j#285,k#286] parquet

Note, note the issue still exists because there is another place to check the ordering match again in FileFormatWriter.

This PR fixes the issue thoroughly, with new UTs added.

Does this PR introduce any user-facing change?

Yes, it's a bug fix.

How was this patch tested?

New UTs are added.

Was this patch authored or co-authored using generative AI tooling?

No.

…n query output contains literal

pan3793 · 2025-10-13T05:50:04Z

sql/core/src/main/scala/org/apache/spark/sql/execution/WholeStageCodegenExec.scala

-        WholeStageCodegenExec(insertInputAdapter(plan))(codegenStageCounter.incrementAndGet())
+        val newId = codegenStageCounter.incrementAndGet()
+        val newPlan = WholeStageCodegenExec(insertInputAdapter(plan))(newId)
+        plan.logicalLink.foreach(newPlan.setLogicalLink)


It appears that WholeStageCodegenExec misses setting logicalLink, is it by design?

interesting, and it never caused issue with AQE before?

Haven't seen the real issues in both production and existing UT.

pan3793 · 2025-10-13T05:51:26Z

sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/FileFormatWriter.scala

+      plan.logicalLink match {
+        case Some(WriteFiles(query, _, _, _, _, _)) =>
+          V1WritesUtils.eliminateFoldableOrdering(ordering, query).outputOrdering
+        case Some(query) =>


the query can be WholeStageCodegenExec, that's why I set logicalLink on WholeStageCodegenExec

pan3793 · 2025-10-13T05:52:34Z

sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/V1WriteCommandSuite.scala


    val listener = new QueryExecutionListener {
      override def onSuccess(funcName: String, qe: QueryExecution, durationNs: Long): Unit = {
+        val conf = qe.sparkSession.sessionState.conf


this is a bugfix, the listener runs in another thread, without this change, conf.getConf actually gets conf from the thread local, thus may cause issues on concurrency running tests

pan3793 · 2025-10-13T06:03:15Z

sql/core/src/main/scala/org/apache/spark/sql/execution/columnar/InMemoryRelation.scala

+  def withOutput(newOutput: Seq[Attribute]): InMemoryRelation = {
+    val map = AttributeMap(output.zip(newOutput))
+    val newOutputOrdering = outputOrdering
+      .map(_.transform { case a: Attribute => map(a) })
+      .asInstanceOf[Seq[SortOrder]]
+    InMemoryRelation(newOutput, cacheBuilder, newOutputOrdering, statsOfPlanToCache)


issue was identified in previous try, see #52474 (comment)

pan3793 · 2025-10-13T06:03:51Z

sql/core/src/main/scala/org/apache/spark/sql/execution/columnar/InMemoryRelation.scala


+  override def makeCopy(newArgs: Array[AnyRef]): LogicalPlan = {
+    val copied = super.makeCopy(newArgs).asInstanceOf[InMemoryRelation]
+    copied.statsOfPlanToCache = this.statsOfPlanToCache


ditto, issue was identified in previous try, see #52474 (comment)

pan3793 · 2025-10-13T06:52:59Z

cc @cloud-fan @peter-toth @ulysses-you

pan3793 · 2025-10-13T06:59:15Z

@cloud-fan BTW, the "planned write" switch (an internal config) was added since 3.4, do we have a plan to remove it to simplify code, or tend to preserve it forever?

sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/FileFormatWriter.scala

sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/V1Writes.scala

cloud-fan · 2025-10-13T12:12:53Z

sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/V1Writes.scala

    expressions.exists(_.exists(_.isInstanceOf[Empty2Null]))
  }

+  def eliminateFoldableOrdering(ordering: Seq[SortOrder], query: LogicalPlan): LogicalPlan =


let's add comments to explain the reason behind it.

cloud-fan · 2025-10-13T12:13:46Z

sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/FileFormatWriter.scala

      .getOrElse(materializeAdaptiveSparkPlan(plan))
      .outputOrdering
+
+    val requiredOrdering = {


is this the code path when planned write is disabled?

I think we can leave it unfixed, as this code path is rarely reached and this fix is kind of an optimization: it's only about perf.

it's a necessary change for the "planned write" to make UT happy

if (Utils.isTesting) outputOrderingMatched = orderingMatched

OK this is a necessary for the current codebase, but do we really need to do it in theory? The planned write should have added the sort already, ideally we don't need to try to add sort again here.

The planned write should have added the sort already, ideally we don't need to try to add sort again here.

yes, exactly

peter-toth

LGTM, pending CI.

dongjoon-hyun

+1, LGTM.

[SPARK-53738][SQL] PlannedWrite should preserve custom sort order whe…

2c66ba9

…n query output contains literal

github-actions bot added the SQL label Oct 13, 2025

pan3793 commented Oct 13, 2025

View reviewed changes

fix codegenStageCounter

a1bf09e

peter-toth reviewed Oct 13, 2025

View reviewed changes

sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/FileFormatWriter.scala Show resolved Hide resolved

peter-toth reviewed Oct 13, 2025

View reviewed changes

sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/V1Writes.scala Outdated Show resolved Hide resolved

peter-toth reviewed Oct 13, 2025

View reviewed changes

sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/V1Writes.scala Outdated Show resolved Hide resolved

cloud-fan reviewed Oct 13, 2025

View reviewed changes

pan3793 added 2 commits October 13, 2025 20:45

address comments

2384ce8

nit

fd401c0

peter-toth approved these changes Oct 13, 2025

View reviewed changes

dongjoon-hyun approved these changes Oct 13, 2025

View reviewed changes

[SPARK-53738][SQL] Fix planned write when query output contains foldable orderings #52584

Are you sure you want to change the base?

[SPARK-53738][SQL] Fix planned write when query output contains foldable orderings #52584

Conversation

pan3793 commented Oct 13, 2025

What changes were proposed in this pull request?

Why are the changes needed?

Does this PR introduce any user-facing change?

How was this patch tested?

Was this patch authored or co-authored using generative AI tooling?

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

pan3793 commented Oct 13, 2025

Uh oh!

pan3793 commented Oct 13, 2025

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

pan3793 Oct 13, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

peter-toth left a comment

Choose a reason for hiding this comment

Uh oh!

dongjoon-hyun left a comment

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

pan3793 Oct 13, 2025 •

edited

Loading