[SPARK-53738][SQL] PlannedWrite should preserve custom sort order when query output contains literal

pan3793 · pan3793 · commit a0aa9f40ca02 · 2025-09-28T07:21:34.000+08:00
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
@@ -28,7 +28,7 @@ import org.apache.spark.sql.internal.SQLConf
  */
 trait AliasAwareOutputExpression extends SQLConfHelper {
   protected val aliasCandidateLimit = conf.getConf(SQLConf.EXPRESSION_PROJECTION_CANDIDATE_LIMIT)
-  protected def outputExpressions: Seq[NamedExpression]
+  def outputExpressions: Seq[NamedExpression]
   /**
    * This method can be used to strip expression which does not affect the result, for example:
    * strip the expression which is ordering agnostic for output ordering.
@@ -88,7 +88,7 @@ trait AliasAwareOutputExpression extends SQLConfHelper {
  */
 trait AliasAwareQueryOutputOrdering[T <: QueryPlan[T]]
   extends AliasAwareOutputExpression { self: QueryPlan[T] =>
-  protected def orderingExpressions: Seq[SortOrder]
+  def orderingExpressions: Seq[SortOrder]
 
   override protected def strip(expr: Expression): Expression = expr match {
     case e: Empty2Null => strip(e.child)
diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/logical/LogicalPlan.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/logical/LogicalPlan.scala
@@ -293,8 +293,16 @@ trait BinaryNode extends LogicalPlan with BinaryLike[LogicalPlan]
 
 trait OrderPreservingUnaryNode extends UnaryNode
   with AliasAwareQueryOutputOrdering[LogicalPlan] {
-  override protected def outputExpressions: Seq[NamedExpression] = child.output
-  override protected def orderingExpressions: Seq[SortOrder] = child.outputOrdering
+
+  override def outputExpressions: Seq[NamedExpression] = child match {
+    case o: OrderPreservingUnaryNode => o.outputExpressions
+    case _ => child.output
+  }
+
+  override def orderingExpressions: Seq[SortOrder] = child match {
+    case o: OrderPreservingUnaryNode => o.orderingExpressions
+    case _ => child.outputOrdering
+  }
 }
 
 object LogicalPlanIntegrity {
diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/logical/basicLogicalOperators.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/plans/logical/basicLogicalOperators.scala
@@ -73,7 +73,7 @@ object Subquery {
 case class Project(projectList: Seq[NamedExpression], child: LogicalPlan)
     extends OrderPreservingUnaryNode {
   override def output: Seq[Attribute] = projectList.map(_.toAttribute)
-  override protected def outputExpressions: Seq[NamedExpression] = projectList
+  override def outputExpressions: Seq[NamedExpression] = projectList
   override def maxRows: Option[Long] = child.maxRows
   override def maxRowsPerPartition: Option[Long] = child.maxRowsPerPartition
 
@@ -906,13 +906,13 @@ case class Sort(
     order: Seq[SortOrder],
     global: Boolean,
     child: LogicalPlan,
-    hint: Option[SortHint] = None) extends UnaryNode {
+    hint: Option[SortHint] = None) extends UnaryNode with OrderPreservingUnaryNode {
   override def output: Seq[Attribute] = child.output
+  override def orderingExpressions: Seq[SortOrder] = order
   override def maxRows: Option[Long] = child.maxRows
   override def maxRowsPerPartition: Option[Long] = {
     if (global) maxRows else child.maxRowsPerPartition
   }
-  override def outputOrdering: Seq[SortOrder] = order
   final override val nodePatterns: Seq[TreePattern] = Seq(SORT)
   override protected def withNewChildInternal(newChild: LogicalPlan): Sort = copy(child = newChild)
 }
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/SortExec.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/SortExec.scala
@@ -42,11 +42,16 @@ case class SortExec(
     global: Boolean,
     child: SparkPlan,
     testSpillFrequency: Int = 0)
-  extends UnaryExecNode with BlockingOperatorWithCodegen {
+  extends UnaryExecNode with BlockingOperatorWithCodegen with OrderPreservingUnaryExecNode {
 
   override def output: Seq[Attribute] = child.output
 
-  override def outputOrdering: Seq[SortOrder] = sortOrder
+  override def outputExpressions: Seq[NamedExpression] = child match {
+    case o: OrderPreservingUnaryExecNode => o.outputExpressions
+    case _ => child.output
+  }
+
+  override def orderingExpressions: Seq[SortOrder] = sortOrder
 
   // sort performed is local within a given partition so will retain
   // child operator's partitioning
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/WholeStageCodegenExec.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/WholeStageCodegenExec.scala
@@ -633,13 +633,21 @@ object WholeStageCodegenExec {
  * used to generated code for [[BoundReference]].
  */
 case class WholeStageCodegenExec(child: SparkPlan)(val codegenStageId: Int)
-    extends UnaryExecNode with CodegenSupport {
+    extends UnaryExecNode with CodegenSupport with OrderPreservingUnaryExecNode {
 
   override def output: Seq[Attribute] = child.output
 
   override def outputPartitioning: Partitioning = child.outputPartitioning
 
-  override def outputOrdering: Seq[SortOrder] = child.outputOrdering
+  override def outputExpressions: Seq[NamedExpression] = child match {
+    case o: OrderPreservingUnaryExecNode => o.outputExpressions
+    case _ => child.output
+  }
+
+  override def orderingExpressions: Seq[SortOrder] = child match {
+    case o: OrderPreservingUnaryExecNode => o.orderingExpressions
+    case _ => child.outputOrdering
+  }
 
   // This is not strictly needed because the codegen transformation happens after the columnar
   // transformation but just for consistency
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/aggregate/BaseAggregateExec.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/aggregate/BaseAggregateExec.scala
@@ -91,7 +91,7 @@ trait BaseAggregateExec extends UnaryExecNode with PartitioningPreservingUnaryEx
 
   override def output: Seq[Attribute] = resultExpressions.map(_.toAttribute)
 
-  override protected def outputExpressions: Seq[NamedExpression] = resultExpressions
+  override def outputExpressions: Seq[NamedExpression] = resultExpressions
 
   override def requiredChildDistribution: List[Distribution] = {
     requiredChildDistributionExpressions match {
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/aggregate/SortAggregateExec.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/aggregate/SortAggregateExec.scala
@@ -51,7 +51,7 @@ case class SortAggregateExec(
     groupingExpressions.map(SortOrder(_, Ascending)) :: Nil
   }
 
-  override protected def orderingExpressions: Seq[SortOrder] = {
+  override def orderingExpressions: Seq[SortOrder] = {
     groupingExpressions.map(SortOrder(_, Ascending))
   }
 
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/basicPhysicalOperators.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/basicPhysicalOperators.scala
@@ -102,9 +102,9 @@ case class ProjectExec(projectList: Seq[NamedExpression], child: SparkPlan)
     }
   }
 
-  override protected def outputExpressions: Seq[NamedExpression] = projectList
+  override def outputExpressions: Seq[NamedExpression] = projectList
 
-  override protected def orderingExpressions: Seq[SortOrder] = child.outputOrdering
+  override def orderingExpressions: Seq[SortOrder] = child.outputOrdering
 
   override def verboseStringWithOperatorId(): String = {
     s"""
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/FileFormatWriter.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/datasources/FileFormatWriter.scala
@@ -19,6 +19,8 @@ package org.apache.spark.sql.execution.datasources
 
 import java.util.{Date, UUID}
 
+import scala.annotation.tailrec
+
 import org.apache.hadoop.conf.Configuration
 import org.apache.hadoop.fs.Path
 import org.apache.hadoop.mapreduce._
@@ -37,7 +39,7 @@ import org.apache.spark.sql.catalyst.expressions.BindReferences.bindReferences
 import org.apache.spark.sql.catalyst.util.{CaseInsensitiveMap, DateTimeUtils}
 import org.apache.spark.sql.classic.SparkSession
 import org.apache.spark.sql.connector.write.WriterCommitMessage
-import org.apache.spark.sql.execution.{ProjectExec, SortExec, SparkPlan, SQLExecution, UnsafeExternalRowSorter}
+import org.apache.spark.sql.execution.{OrderPreservingUnaryExecNode, ProjectExec, SortExec, SparkPlan, SQLExecution, UnsafeExternalRowSorter}
 import org.apache.spark.sql.execution.adaptive.AdaptiveSparkPlanExec
 import org.apache.spark.util.{SerializableConfiguration, Utils}
 import org.apache.spark.util.ArrayImplicits._
@@ -138,10 +140,6 @@ object FileFormatWriter extends Logging {
       statsTrackers = statsTrackers
     )
 
-    // We should first sort by dynamic partition columns, then bucket id, and finally sorting
-    // columns.
-    val requiredOrdering = partitionColumns.drop(numStaticPartitionCols) ++
-        writerBucketSpec.map(_.bucketIdExpression) ++ sortColumns
     val writeFilesOpt = V1WritesUtils.getWriteFilesOpt(plan)
 
     // SPARK-40588: when planned writing is disabled and AQE is enabled,
@@ -153,10 +151,34 @@ object FileFormatWriter extends Logging {
       case p: SparkPlan => p.withNewChildren(p.children.map(materializeAdaptiveSparkPlan))
     }
 
+    val query = writeFilesOpt.map(_.child).getOrElse(materializeAdaptiveSparkPlan(plan))
+
     // the sort order doesn't matter
-    val actualOrdering = writeFilesOpt.map(_.child)
-      .getOrElse(materializeAdaptiveSparkPlan(plan))
-      .outputOrdering
+    val actualOrdering = query.outputOrdering
+
+    val queryOutput = query match {
+      case o: OrderPreservingUnaryExecNode => o.outputExpressions
+      case _ => query.output
+    }
+
+    @tailrec
+    def isLiteral(e: Expression, name: String): Option[String] =
+      e match {
+        case Alias(child, n) => isLiteral(child, n)
+        case _: Literal => Some(name)
+        case _ => None
+      }
+
+    val literalColumns = queryOutput.flatMap { ne => isLiteral(ne, ne.name) }
+
+    // We should first sort by dynamic partition columns, then bucket id, and finally sorting
+    // columns, then drop literal columns
+    val requiredOrdering = (partitionColumns.drop(numStaticPartitionCols) ++
+      writerBucketSpec.map(_.bucketIdExpression) ++ sortColumns).dropWhile {
+      case attr: Attribute => literalColumns.contains(attr.name)
+      case _ => false
+    }
+
     val orderingMatched = V1WritesUtils.isOrderingMatched(requiredOrdering, actualOrdering)
 
     SQLExecution.checkSQLExecutionId(sparkSession)
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/limit.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/limit.scala
@@ -385,9 +385,9 @@ case class TakeOrderedAndProjectExec(
     }
   }
 
-  override protected def outputExpressions: Seq[NamedExpression] = projectList
+  override def outputExpressions: Seq[NamedExpression] = projectList
 
-  override protected def orderingExpressions: Seq[SortOrder] = sortOrder
+  override def orderingExpressions: Seq[SortOrder] = sortOrder
 
   override def outputPartitioning: Partitioning = SinglePartition
 
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/V1WriteCommandSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/V1WriteCommandSuite.scala
@@ -63,10 +63,23 @@ trait V1WriteCommandSuiteBase extends SQLTestUtils with AdaptiveSparkPlanHelper
       hasLogicalSort: Boolean,
       orderingMatched: Boolean,
       hasEmpty2Null: Boolean = false)(query: => Unit): Unit = {
-    var optimizedPlan: LogicalPlan = null
+    executeAndCheckOrderingAndCustomValidate(
+      hasLogicalSort, orderingMatched, hasEmpty2Null)(query)(_ => ())
+  }
+
+  /**
+   * Execute a write query and check ordering of the plan, then do custom validation
+   */
+  protected def executeAndCheckOrderingAndCustomValidate(
+      hasLogicalSort: Boolean,
+      orderingMatched: Boolean,
+      hasEmpty2Null: Boolean = false)(query: => Unit)(
+      customValidate: LogicalPlan => Unit): Unit = {
+    @volatile var optimizedPlan: LogicalPlan = null
 
     val listener = new QueryExecutionListener {
       override def onSuccess(funcName: String, qe: QueryExecution, durationNs: Long): Unit = {
+        val conf = qe.sparkSession.sessionState.conf
         qe.optimizedPlan match {
           case w: V1WriteCommand =>
             if (hasLogicalSort && conf.getConf(SQLConf.PLANNED_WRITE_ENABLED)) {
@@ -103,6 +116,8 @@ trait V1WriteCommandSuiteBase extends SQLTestUtils with AdaptiveSparkPlanHelper
     assert(empty2nullExpr == hasEmpty2Null,
       s"Expect hasEmpty2Null: $hasEmpty2Null, Actual: $empty2nullExpr. Plan:\n$optimizedPlan")
 
+    customValidate(optimizedPlan)
+
     spark.listenerManager.unregister(listener)
   }
 }
@@ -391,4 +406,30 @@ class V1WriteCommandSuite extends QueryTest with SharedSparkSession with V1Write
       }
     }
   }
+
+  test("v1 write with sort by literal column preserve custom order") {
+    withPlannedWrite { _ =>
+      withTable("t") {
+        sql(
+          """
+            |CREATE TABLE t(i INT, j INT, k STRING) USING PARQUET
+            |PARTITIONED BY (k)
+            |""".stripMargin)
+        executeAndCheckOrderingAndCustomValidate(hasLogicalSort = true, orderingMatched = true) {
+          sql(
+            """
+              |INSERT OVERWRITE t
+              |SELECT i, j, '0' as k FROM t0 SORT BY k, i
+              |""".stripMargin)
+        } { optimizedPlan =>
+          assert {
+            optimizedPlan.outputOrdering.exists {
+              case SortOrder(attr: AttributeReference, _, _, _) => attr.name == "i"
+              case _ => false
+            }
+          }
+        }
+      }
+    }
+  }
 }
diff --git a/sql/hive/src/test/scala/org/apache/spark/sql/hive/execution/command/V1WriteHiveCommandSuite.scala b/sql/hive/src/test/scala/org/apache/spark/sql/hive/execution/command/V1WriteHiveCommandSuite.scala
@@ -18,6 +18,7 @@
 package org.apache.spark.sql.hive.execution.command
 
 import org.apache.spark.sql.QueryTest
+import org.apache.spark.sql.catalyst.expressions.{AttributeReference, SortOrder}
 import org.apache.spark.sql.execution.datasources.V1WriteCommandSuiteBase
 import org.apache.spark.sql.hive.HiveUtils._
 import org.apache.spark.sql.hive.test.TestHiveSingleton
@@ -126,4 +127,35 @@ class V1WriteHiveCommandSuite
       }
     }
   }
+
+  test("v1 write to hive table with sort by literal column preserve custom order") {
+    withCovnertMetastore { _ =>
+      withPlannedWrite { _ =>
+        withSQLConf("hive.exec.dynamic.partition.mode" -> "nonstrict") {
+          withTable("t") {
+            sql(
+              """
+                |CREATE TABLE t(i INT, j INT, k STRING) STORED AS PARQUET
+                |PARTITIONED BY (k)
+                |""".stripMargin)
+            executeAndCheckOrderingAndCustomValidate(
+              hasLogicalSort = true, orderingMatched = true) {
+              sql(
+                """
+                  |INSERT OVERWRITE t
+                  |SELECT i, j, '0' as k FROM t0 SORT BY k, i
+                  |""".stripMargin)
+            } { optimizedPlan =>
+              assert {
+                optimizedPlan.outputOrdering.exists {
+                  case SortOrder(attr: AttributeReference, _, _, _) => attr.name == "i"
+                  case _ => false
+                }
+              }
+            }
+          }
+        }
+      }
+    }
+  }
 }

Original file line number	Diff line number	Diff line change
`@@ -51,7 +51,7 @@ case class SortAggregateExec(`
`51`	`51`	`groupingExpressions.map(SortOrder(_, Ascending)) :: Nil`
`52`	`52`	`}`
`53`	`53`
`54`		`- override protected def orderingExpressions: Seq[SortOrder] = {`
	`54`	`+ override def orderingExpressions: Seq[SortOrder] = {`
`55`	`55`	`groupingExpressions.map(SortOrder(_, Ascending))`
`56`	`56`	`}`
`57`	`57`
Original file line number	Diff line number	Diff line change
`@@ -102,9 +102,9 @@ case class ProjectExec(projectList: Seq[NamedExpression], child: SparkPlan)`
`102`	`102`	`}`
`103`	`103`	`}`
`104`	`104`
`105`		`- override protected def outputExpressions: Seq[NamedExpression] = projectList`
	`105`	`+ override def outputExpressions: Seq[NamedExpression] = projectList`
`106`	`106`
`107`		`- override protected def orderingExpressions: Seq[SortOrder] = child.outputOrdering`
	`107`	`+ override def orderingExpressions: Seq[SortOrder] = child.outputOrdering`
`108`	`108`
`109`	`109`	`override def verboseStringWithOperatorId(): String = {`
`110`	`110`	`s"""`
Original file line number	Diff line number	Diff line change
`@@ -385,9 +385,9 @@ case class TakeOrderedAndProjectExec(`
`385`	`385`	`}`
`386`	`386`	`}`
`387`	`387`
`388`		`- override protected def outputExpressions: Seq[NamedExpression] = projectList`
	`388`	`+ override def outputExpressions: Seq[NamedExpression] = projectList`
`389`	`389`
`390`		`- override protected def orderingExpressions: Seq[SortOrder] = sortOrder`
	`390`	`+ override def orderingExpressions: Seq[SortOrder] = sortOrder`
`391`	`391`
`392`	`392`	`override def outputPartitioning: Partitioning = SinglePartition`
`393`	`393`