Backport ExtendedV2Writes rule for Spark 3.1

spark/v3.1/spark-extensions/src/main/scala/org/apache/iceberg/spark/extensions/IcebergSparkSessionExtensions.scala

@@ -31,6 +31,7 @@ import org.apache.spark.sql.catalyst.optimizer.RewriteMergeInto
 import org.apache.spark.sql.catalyst.optimizer.RewriteUpdate
 import org.apache.spark.sql.catalyst.parser.extensions.IcebergSparkSqlExtensionsParser
 import org.apache.spark.sql.execution.datasources.v2.ExtendedDataSourceV2Strategy
+import org.apache.spark.sql.execution.datasources.v2.ExtendedV2Writes
 
 class IcebergSparkSessionExtensions extends (SparkSessionExtensions => Unit) {
 
@@ -51,6 +52,11 @@ class IcebergSparkSessionExtensions extends (SparkSessionExtensions => Unit) {
     extensions.injectOptimizerRule { spark => RewriteUpdate(spark) }
     extensions.injectOptimizerRule { spark => RewriteMergeInto(spark) }
 
+    // pre-CBO extensions
+    // attach the Iceberg table's required distribution and ordering to V2 writes after
+    // the optimizer has resolved the write target but before physical planning
+    extensions.injectPreCBORule { _ => ExtendedV2Writes }
+
     // planner extensions
     extensions.injectPlannerStrategy { spark => ExtendedDataSourceV2Strategy(spark) }
   }

spark/v3.1/spark-extensions/src/main/scala/org/apache/spark/sql/execution/datasources/v2/ExtendedV2Writes.scala

@@ -0,0 +1,97 @@
+/*
+ * 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.spark.sql.execution.datasources.v2
+
+import org.apache.iceberg.spark.Spark3Util
+import org.apache.spark.sql.catalyst.plans.logical.AppendData
+import org.apache.spark.sql.catalyst.plans.logical.LogicalPlan
+import org.apache.spark.sql.catalyst.plans.logical.OverwriteByExpression
+import org.apache.spark.sql.catalyst.plans.logical.OverwritePartitionsDynamic
+import org.apache.spark.sql.catalyst.plans.logical.RepartitionByExpression
+import org.apache.spark.sql.catalyst.plans.logical.Sort
+import org.apache.spark.sql.catalyst.rules.Rule
+import org.apache.spark.sql.catalyst.utils.DistributionAndOrderingUtils
+import org.apache.spark.sql.catalyst.utils.PlanUtils.isIcebergRelation
+import org.apache.spark.sql.connector.iceberg.distributions.Distribution
+import org.apache.spark.sql.connector.iceberg.distributions.Distributions
+import org.apache.spark.sql.connector.iceberg.distributions.OrderedDistribution
+import org.apache.spark.sql.connector.iceberg.expressions.SortOrder
+
+/**
+ * Backport of Spark 3.2's V2Writes for v3.1 AppendData/OverwriteByExpression/
+ * OverwritePartitionsDynamic. Attaches a local Sort only when the table has an explicit sort
+ * order or RANGE distribution; never attaches an Exchange. Unsorted partitioned writes thus
+ * require write.spark.fanout.enabled=true or pre-clustered input.
+ *
+ * No distribution: Spark 3.1 only has strict RepartitionByExpression. Spark 3.4+'s
+ * RebalancePartitions (used by Spark 3.5's V2Writes) doesn't exist here, and a strict
+ * repartition would turn skewed partition keys into stragglers.
+ *
+ * MERGE/UPDATE/DELETE are skipped — RewriteRowLevelOperationHelper.buildWritePlan already
+ * prepares those queries (using the unwrapped Spark3Util.buildRequiredOrdering, which still
+ * synthesizes the partition prefix); alreadyPrepared() detects its shape to avoid double-wrap.
+ */
+object ExtendedV2Writes extends Rule[LogicalPlan] {
+
+  override def apply(plan: LogicalPlan): LogicalPlan = plan transformDown {
+    case a @ AppendData(r: DataSourceV2Relation, query, _, _)
+      if isIcebergRelation(r) && !alreadyPrepared(query) =>
+      a.withNewQuery(prepareQuery(r, query))
+
+    case o @ OverwriteByExpression(r: DataSourceV2Relation, _, query, _, _)
+      if isIcebergRelation(r) && !alreadyPrepared(query) =>
+      o.withNewQuery(prepareQuery(r, query))
+
+    case o @ OverwritePartitionsDynamic(r: DataSourceV2Relation, query, _, _)
+      if isIcebergRelation(r) && !alreadyPrepared(query) =>
+      o.withNewQuery(prepareQuery(r, query))
+  }
+
+  // Matches the shapes RewriteRowLevelOperationHelper.buildWritePlan produces. Bare
+  // Sort(_, false, _) is intentionally NOT matched — it would swallow a user's
+  // sortWithinPartitions on the wrong columns and skip the table-required ordering.
+  private def alreadyPrepared(query: LogicalPlan): Boolean = query match {
+    case Sort(_, false, RepartitionByExpression(_, _, None)) => true
+    case RepartitionByExpression(_, _, None) => true
+    case _ => false
+  }
+
+  private def prepareQuery(r: DataSourceV2Relation, query: LogicalPlan): LogicalPlan = {
+    val icebergTable = Spark3Util.toIcebergTable(r.table)
+    // Distribution is computed only to surface OrderedDistribution.ordering; we always pass
+    // unspecified() to prepareQuery so no Exchange is attached.
+    val tableDistribution = Spark3Util.buildRequiredDistribution(icebergTable)
+    val ordering = requiredOrdering(tableDistribution, icebergTable)
+    DistributionAndOrderingUtils.prepareQuery(
+      Distributions.unspecified(), ordering, query, conf)
+  }
+
+  // Delegate to Spark3Util.buildRequiredOrdering only for OrderedDistribution or sorted tables.
+  // Unsorted tables get an empty ordering — fanout or pre-clustered input is required.
+  private def requiredOrdering(
+      distribution: Distribution,
+      icebergTable: org.apache.iceberg.Table): Array[SortOrder] = {
+    if (distribution.isInstanceOf[OrderedDistribution] || !icebergTable.sortOrder().isUnsorted) {
+      Spark3Util.buildRequiredOrdering(distribution, icebergTable)
+    } else {
+      Array.empty[SortOrder]
+    }
+  }
+}