Draft: [GLUTEN-12474][VL] Re-add partition-column sort for row-based fallback write with dynamic partitions.

backends-velox/src/main/scala/org/apache/gluten/backendsapi/velox/VeloxRuleApi.scala

@@ -122,6 +122,7 @@ object VeloxRuleApi {
     injector.injectPostTransform(_ => CollectLimitTransformerRule())
     injector.injectPostTransform(_ => CollectTailTransformerRule())
     injector.injectPostTransform(_ => V2WritePostRule())
+    injector.injectPostTransform(_ => EnsureRowBasedWriteFilesOrdering)
     injector.injectPostTransform(c => InsertTransitions.create(c.outputsColumnar, VeloxBatchType))
 
     // Gluten columnar: Fallback policies.

gluten-substrait/src/main/scala/org/apache/gluten/extension/columnar/EnsureRowBasedWriteFilesOrdering.scala

@@ -0,0 +1,110 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *    http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.gluten.extension.columnar
+
+import org.apache.gluten.execution.{SortExecTransformer, TransformSupport}
+import org.apache.gluten.sql.shims.SparkShimLoader
+
+import org.apache.spark.sql.catalyst.expressions.{Expression, SortOrder}
+import org.apache.spark.sql.catalyst.rules.Rule
+import org.apache.spark.sql.execution.{ColumnarWriteFilesExec, SortExec, SparkPlan}
+import org.apache.spark.sql.execution.datasources.WriteFilesExec
+
+/**
+ * Re-adds the local sort on dynamic partition (and bucket) columns that a row-based
+ * [[WriteFilesExec]] relies on, when Gluten has removed it while offloading the upstream
+ * `SortAggregate` to a native hash aggregate (see [[EliminateLocalSort]]).
+ *
+ * Background:
+ *   - With planned write enabled, `V1Writes` drops the write's required sort at logical
+ *     optimization time when the child already provides the ordering (e.g. a `SortAggregate` keyed
+ *     by the partition column produces output already ordered by that column). Hence the physical
+ *     `WriteFilesExec.requiredChildOrdering` is `Nil`.
+ *   - Gluten then replaces the `SortAggregate` with an (unordered) native hash aggregate and
+ *     [[EliminateLocalSort]] eagerly removes the now-redundant local sort feeding it.
+ *   - The ordering that the write implicitly depended on is thereby destroyed. When the native
+ *     Velox write is not applicable (e.g. ORC/Hive tables) the write falls back to the row-based
+ *     `WriteFilesExec`, whose `DynamicPartitionDataSingleWriter` requires the rows to be sorted by
+ *     the partition columns. Unsorted rows make it re-create an already-closed partition file and
+ *     fail with `FileAlreadyExistsException`.
+ *
+ * Native Velox writes ([[org.apache.gluten.execution.WriteFilesExecTransformer]]) do not need the
+ * ordering, so this rule only patches the row-based fallback write. The dummy no-op
+ * `WriteFilesExec` that [[ColumnarWriteFilesExec]] injects for the native path is skipped by
+ * checking for its [[ColumnarWriteFilesExec.NoopLeaf]] child.
+ *
+ * The required ordering is obtained through [[SparkShimLoader]] from `V1WritesUtils.getSortOrder`
+ * so that the behavior stays identical to vanilla `FileFormatWriter` (including the
+ * concurrent-writers and small-file-merge cases where `getSortOrder` returns empty). It is fetched
+ * via the shim because `V1WritesUtils` only exists since Spark 3.4; on Spark 3.2/3.3 the shim
+ * returns `Nil` and this rule is a no-op (the planned-write `WriteFilesExec` path is not used
+ * there). The "already satisfied" check mirrors `V1WritesUtils.isOrderingMatched`, which is pure
+ * catalyst and version-agnostic, so it is inlined here.
+ *
+ * This rule must run before [[org.apache.gluten.extension.columnar.transition.InsertTransitions]]:
+ * when the write's child is an offloaded transformer we insert a native [[SortExecTransformer]]
+ * (validated, falling back to a row-based [[SortExec]]) so the sort runs columnar in Velox, and
+ * InsertTransitions afterwards places the ColumnarToRow transition above the sort.
+ */
+object EnsureRowBasedWriteFilesOrdering extends Rule[SparkPlan] {
+  override def apply(plan: SparkPlan): SparkPlan = plan.transformUp {
+    case w: WriteFilesExec if !w.child.isInstanceOf[ColumnarWriteFilesExec.NoopLeaf] =>
+      val requiredOrdering: Seq[SortOrder] =
+        SparkShimLoader.getSparkShims.getWriteFilesRequiredOrdering(
+          w.child.output,
+          w.partitionColumns,
+          w.bucketSpec,
+          w.options,
+          w.staticPartitions.size)
+      if (
+        requiredOrdering.isEmpty ||
+        isOrderingMatched(requiredOrdering.map(_.child), w.child.outputOrdering)
+      ) {
+        w
+      } else {
+        w.withNewChildren(sortPlan(w.child, requiredOrdering) :: Nil)
+      }
+  }
+
+  // Mirrors `V1WritesUtils.isOrderingMatched`. Pure catalyst and identical across Spark versions,
+  // so it is inlined to avoid depending on `V1WritesUtils`, which is absent before Spark 3.4.
+  private def isOrderingMatched(
+      requiredOrdering: Seq[Expression],
+      outputOrdering: Seq[SortOrder]): Boolean = {
+    if (requiredOrdering.length > outputOrdering.length) {
+      false
+    } else {
+      requiredOrdering.zip(outputOrdering).forall {
+        case (requiredOrder, outputOrder) =>
+          outputOrder.satisfies(outputOrder.copy(child = requiredOrder))
+      }
+    }
+  }
+
+  private def sortPlan(child: SparkPlan, requiredOrdering: Seq[SortOrder]): SparkPlan =
+    child match {
+      case c: TransformSupport =>
+        val nativeSort = SortExecTransformer(requiredOrdering, global = false, c)
+        if (nativeSort.doValidate().ok()) {
+          nativeSort
+        } else {
+          SortExec(requiredOrdering, global = false, c)
+        }
+      case other =>
+        SortExec(requiredOrdering, global = false, other)
+    }
+}

shims/common/src/main/scala/org/apache/gluten/sql/shims/SparkShims.scala

@@ -24,7 +24,8 @@ import org.apache.spark.broadcast.Broadcast
 import org.apache.spark.internal.io.FileCommitProtocol
 import org.apache.spark.sql.{AnalysisException, SparkSession}
 import org.apache.spark.sql.catalyst.InternalRow
-import org.apache.spark.sql.catalyst.expressions.{Attribute, BinaryArithmetic, Expression, InputFileBlockLength, InputFileBlockStart, InputFileName, RaiseError, UnBase64}
+import org.apache.spark.sql.catalyst.catalog.BucketSpec
+import org.apache.spark.sql.catalyst.expressions.{Attribute, BinaryArithmetic, Expression, InputFileBlockLength, InputFileBlockStart, InputFileName, RaiseError, SortOrder, UnBase64}
 import org.apache.spark.sql.catalyst.plans.JoinType
 import org.apache.spark.sql.catalyst.plans.QueryPlan
 import org.apache.spark.sql.catalyst.plans.logical.LogicalPlan
@@ -129,6 +130,19 @@ trait SparkShims {
 
   def enableNativeWriteFilesByDefault(): Boolean = false
 
+  /**
+   * The local sort a row-based `WriteFilesExec` requires on its child for dynamic partition (and
+   * bucket) writes. Delegates to `V1WritesUtils.getSortOrder`, which only exists since Spark 3.4;
+   * on Spark 3.2/3.3 the planned-write `WriteFilesExec` path is not used at runtime, so the default
+   * returns `Nil`.
+   */
+  def getWriteFilesRequiredOrdering(
+      outputColumns: Seq[Attribute],
+      partitionColumns: Seq[Attribute],
+      bucketSpec: Option[BucketSpec],
+      options: Map[String, String],
+      numStaticPartitionCols: Int): Seq[SortOrder] = Nil
+
   def broadcastInternal[T: ClassTag](sc: SparkContext, value: T): Broadcast[T] = {
     // Since Spark 3.4, the `sc.broadcast` has been optimized to use `sc.broadcastInternal`.
     // More details see SPARK-39983.

shims/spark34/src/main/scala/org/apache/gluten/sql/shims/spark34/Spark34Shims.scala

@@ -27,6 +27,7 @@ import org.apache.spark.paths.SparkPath
 import org.apache.spark.sql.{AnalysisException, SparkSession}
 import org.apache.spark.sql.catalyst.InternalRow
 import org.apache.spark.sql.catalyst.analysis.DecimalPrecision
+import org.apache.spark.sql.catalyst.catalog.BucketSpec
 import org.apache.spark.sql.catalyst.expressions._
 import org.apache.spark.sql.catalyst.expressions.aggregate._
 import org.apache.spark.sql.catalyst.plans.QueryPlan
@@ -208,6 +209,19 @@ class Spark34Shims extends SparkShims {
 
   override def enableNativeWriteFilesByDefault(): Boolean = true
 
+  override def getWriteFilesRequiredOrdering(
+      outputColumns: Seq[Attribute],
+      partitionColumns: Seq[Attribute],
+      bucketSpec: Option[BucketSpec],
+      options: Map[String, String],
+      numStaticPartitionCols: Int): Seq[SortOrder] =
+    V1WritesUtils.getSortOrder(
+      outputColumns,
+      partitionColumns,
+      bucketSpec,
+      options,
+      numStaticPartitionCols)
+
   override def broadcastInternal[T: ClassTag](sc: SparkContext, value: T): Broadcast[T] = {
     SparkContextUtils.broadcastInternal(sc, value)
   }

shims/spark35/src/main/scala/org/apache/gluten/sql/shims/spark35/Spark35Shims.scala

@@ -27,6 +27,7 @@ import org.apache.spark.paths.SparkPath
 import org.apache.spark.sql.{AnalysisException, SparkSession}
 import org.apache.spark.sql.catalyst.{ExtendedAnalysisException, InternalRow}
 import org.apache.spark.sql.catalyst.analysis.DecimalPrecision
+import org.apache.spark.sql.catalyst.catalog.BucketSpec
 import org.apache.spark.sql.catalyst.expressions._
 import org.apache.spark.sql.catalyst.expressions.aggregate._
 import org.apache.spark.sql.catalyst.plans.QueryPlan
@@ -249,6 +250,19 @@ class Spark35Shims extends SparkShims {
 
   override def enableNativeWriteFilesByDefault(): Boolean = true
 
+  override def getWriteFilesRequiredOrdering(
+      outputColumns: Seq[Attribute],
+      partitionColumns: Seq[Attribute],
+      bucketSpec: Option[BucketSpec],
+      options: Map[String, String],
+      numStaticPartitionCols: Int): Seq[SortOrder] =
+    V1WritesUtils.getSortOrder(
+      outputColumns,
+      partitionColumns,
+      bucketSpec,
+      options,
+      numStaticPartitionCols)
+
   override def broadcastInternal[T: ClassTag](sc: SparkContext, value: T): Broadcast[T] = {
     SparkContextUtils.broadcastInternal(sc, value)
   }

shims/spark40/src/main/scala/org/apache/gluten/sql/shims/spark40/Spark40Shims.scala

@@ -27,6 +27,7 @@ import org.apache.spark.paths.SparkPath
 import org.apache.spark.sql.{AnalysisException, SparkSession}
 import org.apache.spark.sql.catalyst.{ExtendedAnalysisException, InternalRow}
 import org.apache.spark.sql.catalyst.analysis.DecimalPrecisionTypeCoercion
+import org.apache.spark.sql.catalyst.catalog.BucketSpec
 import org.apache.spark.sql.catalyst.expressions._
 import org.apache.spark.sql.catalyst.expressions.aggregate._
 import org.apache.spark.sql.catalyst.plans.{JoinType, LeftSingle}
@@ -254,6 +255,19 @@ class Spark40Shims extends SparkShims {
 
   override def enableNativeWriteFilesByDefault(): Boolean = true
 
+  override def getWriteFilesRequiredOrdering(
+      outputColumns: Seq[Attribute],
+      partitionColumns: Seq[Attribute],
+      bucketSpec: Option[BucketSpec],
+      options: Map[String, String],
+      numStaticPartitionCols: Int): Seq[SortOrder] =
+    V1WritesUtils.getSortOrder(
+      outputColumns,
+      partitionColumns,
+      bucketSpec,
+      options,
+      numStaticPartitionCols)
+
   override def broadcastInternal[T: ClassTag](sc: SparkContext, value: T): Broadcast[T] = {
     SparkContextUtils.broadcastInternal(sc, value)
   }

shims/spark41/src/main/scala/org/apache/gluten/sql/shims/spark41/Spark41Shims.scala

@@ -27,6 +27,7 @@ import org.apache.spark.paths.SparkPath
 import org.apache.spark.sql.{AnalysisException, SparkSession}
 import org.apache.spark.sql.catalyst.{ExtendedAnalysisException, InternalRow}
 import org.apache.spark.sql.catalyst.analysis.DecimalPrecisionTypeCoercion
+import org.apache.spark.sql.catalyst.catalog.BucketSpec
 import org.apache.spark.sql.catalyst.expressions._
 import org.apache.spark.sql.catalyst.expressions.aggregate._
 import org.apache.spark.sql.catalyst.plans.{JoinType, LeftSingle}
@@ -253,6 +254,19 @@ class Spark41Shims extends SparkShims {
 
   override def enableNativeWriteFilesByDefault(): Boolean = true
 
+  override def getWriteFilesRequiredOrdering(
+      outputColumns: Seq[Attribute],
+      partitionColumns: Seq[Attribute],
+      bucketSpec: Option[BucketSpec],
+      options: Map[String, String],
+      numStaticPartitionCols: Int): Seq[SortOrder] =
+    V1WritesUtils.getSortOrder(
+      outputColumns,
+      partitionColumns,
+      bucketSpec,
+      options,
+      numStaticPartitionCols)
+
   override def broadcastInternal[T: ClassTag](sc: SparkContext, value: T): Broadcast[T] = {
     SparkContextUtils.broadcastInternal(sc, value)
   }