[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.

@cloud-fan if I revert changes in FileFormatWriter.scala, this is not required.

Do you want me to keep it or revert it?

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.

cloud-fan · 2025-10-14T01:49:51Z

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

+      // columns.
+      val ordering = partitionColumns.drop(numStaticPartitionCols) ++
+        writerBucketSpec.map(_.bucketIdExpression) ++ sortColumns
+      plan.logicalLink match {


I'm a bit worried about this. In AQE we have a fallback to find logical link in the children, so that it's more reliable. Now we have the risk of perf regression if the logical link is not present and we add an extra sort.

Shall we remove the adding sort here completly if planned write is enabled (WriteFiles is present)?

I'm a bit worried about this. In AQE we have a fallback to find logical link in the children, so that it's more reliable.

@cloud-fan do you suggest

- plan.logicalLink match { + plan.logicalLink.orElse { + plan.collectFirst { case p if p.logicalLink.isDefined => p.logicalLink.get } + } match {

Shall we remove the adding sort here completly if planned write is enabled (WriteFiles is present)?

I think the current code has already satisfied your expectation, when planned write is enabled:

if concurrent writer is disabled, the calculated required ordering won't be used.

if concurrent writer is enabled, the calculated required ordering is only used in the concurrent writer step 2.

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

Lines 393 to 406 in 29434ea

/**

* Dynamic partition writer with concurrent writers, meaning multiple concurrent writers are opened

* for writing.

*

* The process has the following steps:

* - Step 1: Maintain a map of output writers per each partition and/or bucket columns. Keep all

* writers opened and write rows one by one.

* - Step 2: If number of concurrent writers exceeds limit, sort rest of rows on partition and/or

* bucket column(s). Write rows one by one, and eagerly close the writer when finishing

* each partition and/or bucket.

*

* Caller is expected to call `writeWithIterator()` instead of `write()` to write records.

*/

class DynamicPartitionDataConcurrentWriter(

@cloud-fan I have updated the code to fallback to find logical link in the children, then setLogicalLink for WholeStageCodegenExec is unnecessary for this PR, please let me know if you want me to keep it or restore it.

cloud-fan · 2025-10-15T18:18:22Z

what happens if we don't modify FileFormatWriter.scala at all? I think it only affects non-planned-write and planned-write with concurrent write, and we can improve them later. (ideally the sort should be handled at the logical plan phase)

I think the only issue is for test: if (Utils.isTesting) outputOrderingMatched = orderingMatched. We can fix the affected tests and specify the required ordering explicitly.

pan3793 · 2025-10-15T18:44:34Z

@cloud-fan I agree with your summary. Only the newly added tests are affected if I don't touch FileFormatWriter.scala, so the simplest way is to skip checking orderingMatched temporarily for the new tests.

Have updated the code, please take another look.

pan3793 · 2025-10-17T08:30:31Z

Kindly ping @cloud-fan, do you have further concerns with this PR?

[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 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

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

cloud-fan reviewed Oct 14, 2025

View reviewed changes

fallback get logicalLink from children

7129fb2

pan3793 requested a review from cloud-fan October 15, 2025 03:15

revert change of FileFormatWriter

f1df0ec

	/**
	* Dynamic partition writer with concurrent writers, meaning multiple concurrent writers are opened
	* for writing.
	*
	* The process has the following steps:
	* - Step 1: Maintain a map of output writers per each partition and/or bucket columns. Keep all
	* writers opened and write rows one by one.
	* - Step 2: If number of concurrent writers exceeds limit, sort rest of rows on partition and/or
	* bucket column(s). Write rows one by one, and eagerly close the writer when finishing
	* each partition and/or bucket.
	*
	* Caller is expected to call `writeWithIterator()` instead of `write()` to write records.
	*/
	class DynamicPartitionDataConcurrentWriter(

[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!

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!

Choose a reason for hiding this comment

Uh oh!

pan3793 Oct 14, 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!

Choose a reason for hiding this comment

Uh oh!

cloud-fan commented Oct 15, 2025

Uh oh!

pan3793 commented Oct 15, 2025

Uh oh!

pan3793 commented Oct 17, 2025

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

pan3793 Oct 13, 2025 •

edited

Loading

pan3793 Oct 14, 2025 •

edited

Loading