diff --git a/backends-velox/src/main/scala/org/apache/gluten/backendsapi/velox/VeloxRuleApi.scala b/backends-velox/src/main/scala/org/apache/gluten/backendsapi/velox/VeloxRuleApi.scala index d63928527df..802a233a7f8 100644 --- a/backends-velox/src/main/scala/org/apache/gluten/backendsapi/velox/VeloxRuleApi.scala +++ b/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. diff --git a/gluten-substrait/src/main/scala/org/apache/gluten/extension/columnar/EnsureRowBasedWriteFilesOrdering.scala b/gluten-substrait/src/main/scala/org/apache/gluten/extension/columnar/EnsureRowBasedWriteFilesOrdering.scala new file mode 100644 index 00000000000..795fde6306e --- /dev/null +++ b/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) + } +} diff --git a/shims/common/src/main/scala/org/apache/gluten/sql/shims/SparkShims.scala b/shims/common/src/main/scala/org/apache/gluten/sql/shims/SparkShims.scala index 750bea0c41a..bd42c951669 100644 --- a/shims/common/src/main/scala/org/apache/gluten/sql/shims/SparkShims.scala +++ b/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. diff --git a/shims/spark34/src/main/scala/org/apache/gluten/sql/shims/spark34/Spark34Shims.scala b/shims/spark34/src/main/scala/org/apache/gluten/sql/shims/spark34/Spark34Shims.scala index 97ff19a84a1..061318441e1 100644 --- a/shims/spark34/src/main/scala/org/apache/gluten/sql/shims/spark34/Spark34Shims.scala +++ b/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) } diff --git a/shims/spark35/src/main/scala/org/apache/gluten/sql/shims/spark35/Spark35Shims.scala b/shims/spark35/src/main/scala/org/apache/gluten/sql/shims/spark35/Spark35Shims.scala index d62fdfea193..dd467e6067a 100644 --- a/shims/spark35/src/main/scala/org/apache/gluten/sql/shims/spark35/Spark35Shims.scala +++ b/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) } diff --git a/shims/spark40/src/main/scala/org/apache/gluten/sql/shims/spark40/Spark40Shims.scala b/shims/spark40/src/main/scala/org/apache/gluten/sql/shims/spark40/Spark40Shims.scala index 5847e62c106..c18eb89caa8 100644 --- a/shims/spark40/src/main/scala/org/apache/gluten/sql/shims/spark40/Spark40Shims.scala +++ b/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) } diff --git a/shims/spark41/src/main/scala/org/apache/gluten/sql/shims/spark41/Spark41Shims.scala b/shims/spark41/src/main/scala/org/apache/gluten/sql/shims/spark41/Spark41Shims.scala index bb0b94cc022..23b40a9a3cd 100644 --- a/shims/spark41/src/main/scala/org/apache/gluten/sql/shims/spark41/Spark41Shims.scala +++ b/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) }