test: audit CI failures from enabling the Iceberg write split-operator plan by default

.github/workflows/pr_build_linux.yml

@@ -314,6 +314,7 @@ jobs:
               org.apache.comet.CometIcebergNativeSuite
               org.apache.comet.CometIcebergEncryptionSuite
               org.apache.comet.CometIcebergRewriteActionSuite
+              org.apache.comet.CometIcebergWriteActionSuite
               org.apache.comet.iceberg.IcebergReflectionSuite
               org.apache.comet.csv.CometCsvNativeReadSuite
               org.apache.comet.CometFuzzTestSuite

.github/workflows/pr_build_macos.yml

@@ -130,6 +130,7 @@ jobs:
               org.apache.comet.CometIcebergNativeSuite
               org.apache.comet.CometIcebergEncryptionSuite
               org.apache.comet.CometIcebergRewriteActionSuite
+              org.apache.comet.CometIcebergWriteActionSuite
               org.apache.comet.iceberg.IcebergReflectionSuite
               org.apache.comet.csv.CometCsvNativeReadSuite
               org.apache.comet.CometFuzzTestSuite

docs/source/user-guide/latest/iceberg-writes.md

@@ -0,0 +1,92 @@
+<!---
+  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.
+-->
+
+# Iceberg Writes: Comet's Split-Operator Plan (Experimental)
+
+**This feature is experimental and enabled by default.** Set
+`spark.comet.write.iceberg.splitOperator.enabled=false` to restore Spark's stock combined
+write operator.
+
+## Overview
+
+Spark writes an Iceberg table through a single physical operator that combines data-file
+writing with metadata writing, committing, and catalog validation. Because that operator sits
+outside Spark's Adaptive Query Execution (AQE), the sub-query feeding the write — the scans,
+projects, sorts, and exchanges producing the rows — cannot be re-planned at runtime.
+
+When `spark.comet.write.iceberg.splitOperator.enabled=true`, Comet rewrites eligible Iceberg
+writes into two operators:
+
+1. **`IcebergWrite`** — writes the data files on the executors, exactly as iceberg-java does
+   today, and returns each task's serialized commit message. This operator and the sub-query
+   feeding it run inside AQE.
+2. **`IcebergCommit`** — collects the commit messages on the driver and performs the normal
+   Iceberg commit (including commit-time validation), outside AQE, exactly once.
+
+Data files are still written by iceberg-java; only the plan shape changes. The split makes the
+write's input visible to AQE and to Comet's columnar rules, and it is the groundwork for a
+planned follow-up in which Comet writes the data files natively via
+[iceberg-rust](https://github.com/apache/iceberg-rust).
+
+## Configuration
+
+Standard Comet + Iceberg setup (see [`iceberg.md`](iceberg.md)) is all that is required; the
+split-operator plan is applied automatically. To turn it off:
+
+```
+# Standard Comet / Iceberg wiring
+spark.plugins=org.apache.spark.CometPlugin
+spark.sql.extensions=org.apache.iceberg.spark.extensions.IcebergSparkSessionExtensions
+spark.sql.catalog.<name>=org.apache.iceberg.spark.SparkCatalog
+spark.sql.catalog.<name>.type=hadoop                          # or hive / glue / rest / ...
+spark.sql.catalog.<name>.warehouse=...
+
+# Split-operator plan (experimental, on by default); set to false to opt out
+spark.comet.write.iceberg.splitOperator.enabled=false
+```
+
+## Supported operations
+
+The split-operator plan supports the following operations on every Spark version Comet
+supports:
+
+- `INSERT INTO` / DataFrame `append` (`AppendData`)
+- `INSERT OVERWRITE`, static and dynamic (`OverwriteByExpression`,
+  `OverwritePartitionsDynamic`)
+- Copy-on-write `DELETE` / `UPDATE` / `MERGE` (`ReplaceData`)
+
+The mechanism behind row-level DML differs by Spark version: on Spark 4.0+ the analyzer emits
+operation-coded rows that Comet's writer dispatches through `ReplaceData`'s projections, while
+on Spark 3.4/3.5 the rewritten rows are written as a plain row stream. The supported set of
+operations is the same either way.
+
+## When Comet falls back to Spark's write operator
+
+The rewrite is skipped — and the write runs through Spark's stock combined operator — when:
+
+- `spark.comet.write.iceberg.splitOperator.enabled` is `false`;
+- the write is not an Iceberg `SparkWrite` (any other V2 data source);
+- the table uses merge-on-read: delta writes (Iceberg `WriteDelta`) are not intercepted;
+- the write requires Spark's commit coordinator, which Comet's per-task commit protocol does
+  not use;
+- Comet cannot reflect the Iceberg internals needed to build the two-operator plan (for
+  example an unrecognised write class or a `ReplaceData` projection it cannot map).
+
+In every fallback case the write is planned as if Comet were absent; there is no correctness
+trade-off, only no plan change.

docs/source/user-guide/latest/index.rst

@@ -82,6 +82,7 @@ to read more.
    :hidden:
 
    Iceberg Guide <iceberg>
+   Iceberg Writes <iceberg-writes>
    S3 Credential Providers <s3-credential-providers>
    Kubernetes Guide <kubernetes>
 

spark/src/main/scala/org/apache/comet/CometConf.scala

@@ -121,6 +121,16 @@ object CometConf extends ShimCometConf {
       .booleanConf
       .createWithDefault(true)
 
+  val COMET_ICEBERG_WRITE_SPLIT_OPERATOR_ENABLED: ConfigEntry[Boolean] =
+    conf("spark.comet.write.iceberg.splitOperator.enabled")
+      .category(CATEGORY_TESTING)
+      .doc(
+        "Whether to rewrite Iceberg V2 writes from Spark's combined V2 write/commit operator " +
+          "into Comet's two-operator shape: a file writer exec (inside AQE) and a committer " +
+          "(outside AQE).")
+      .booleanConf
+      .createWithDefault(true)
+
   val COMET_ICEBERG_DATA_FILE_CONCURRENCY_LIMIT: ConfigEntry[Int] =
     conf("spark.comet.scan.icebergNative.dataFileConcurrencyLimit")
       .category(CATEGORY_SCAN)

spark/src/main/scala/org/apache/comet/CometSparkSessionExtensions.scala

@@ -32,6 +32,7 @@ import org.apache.spark.sql.execution._
 import org.apache.spark.sql.internal.SQLConf
 
 import org.apache.comet.CometConf._
+import org.apache.comet.iceberg.IcebergWriteStrategy
 import org.apache.comet.rules.{CometExecRule, CometPlanAdaptiveDynamicPruningFilters, CometReuseSubquery, CometScanRule, CometSpark34AqeDppFallbackRule, EliminateRedundantTransitions, RevertNativeForTransitionHeavyStages}
 import org.apache.comet.shims.ShimCometSparkSessionExtensions
 
@@ -99,6 +100,7 @@ class CometSparkSessionExtensions
     extensions.injectQueryStagePrepRule { session => CometExecRule(session) }
     injectQueryStageOptimizerRuleShim(extensions, CometPlanAdaptiveDynamicPruningFilters)
     injectQueryStageOptimizerRuleShim(extensions, CometReuseSubquery)
+    extensions.injectPlannerStrategy { session => IcebergWriteStrategy(session) }
   }
 
   case class CometScanColumnar(session: SparkSession) extends ColumnarRule {

spark/src/main/scala/org/apache/comet/iceberg/IcebergReflection.scala

@@ -54,6 +54,10 @@ object IcebergReflection extends Logging {
     val SPARK_SCHEMA_UTIL = "org.apache.iceberg.spark.SparkSchemaUtil"
     val TABLE = "org.apache.iceberg.Table"
     val PARTITIONING = "org.apache.iceberg.Partitioning"
+    val SPARK_WRITE = "org.apache.iceberg.spark.source.SparkWrite"
+
+    // Iceberg 1.5.2 uses its own `ReplaceIcebergData` due to lack of `ReplaceData` in Spark 3.4.
+    val REPLACE_ICEBERG_DATA = "org.apache.spark.sql.catalyst.plans.logical.ReplaceIcebergData"
   }
 
   /**
@@ -128,6 +132,45 @@ object IcebergReflection extends Logging {
     val UNKNOWN = "unknown"
   }
 
+  /** Loads a class, returning `None` when it's absent (e.g. Iceberg not on the classpath). */
+  private def tryLoadClass(name: String): Option[Class[_]] =
+    try Some(loadClass(name))
+    catch { case _: ClassNotFoundException => None }
+
+  private lazy val sparkWriteClassOpt: Option[Class[_]] = tryLoadClass(ClassNames.SPARK_WRITE)
+
+  /** Whether `write` is an Iceberg `SparkWrite` (false if Iceberg isn't on the classpath). */
+  def isIcebergSparkWrite(write: Any): Boolean =
+    sparkWriteClassOpt.exists(_.isInstance(write))
+
+  def isReplaceIcebergData(plan: Any): Boolean =
+    plan != null && plan.getClass.getName == ClassNames.REPLACE_ICEBERG_DATA
+
+  private def reflectField(plan: Any, fieldName: String): Option[AnyRef] =
+    try {
+      val field = plan.getClass.getDeclaredField(fieldName)
+      field.setAccessible(true)
+      Option(field.get(plan))
+    } catch {
+      case e: Exception =>
+        logError(
+          s"Iceberg reflection failure: $fieldName on ${plan.getClass.getName}: ${e.getMessage}")
+        None
+    }
+
+  def extractReplaceIcebergDataFields(plan: Any): Option[(AnyRef, AnyRef, AnyRef, AnyRef)] = {
+    if (!isReplaceIcebergData(plan)) return None
+    for {
+      table <- reflectField(plan, "table")
+      query <- reflectField(plan, "query")
+      originalTable <- reflectField(plan, "originalTable")
+      write <- reflectField(
+        plan,
+        "write"
+      ) // Option[Write]; field can be Some(null) so kept AnyRef
+    } yield (table, query, originalTable, write)
+  }
+
   /**
    * Loads a class using the thread context classloader first, then falls back to the system
    * classloader.

spark/src/main/scala/org/apache/comet/iceberg/IcebergWriteLogical.scala

@@ -0,0 +1,44 @@
+/*
+ * 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.comet.iceberg
+
+import org.apache.spark.sql.catalyst.expressions.{Attribute, AttributeReference}
+import org.apache.spark.sql.catalyst.plans.logical.{LogicalPlan, UnaryNode}
+import org.apache.spark.sql.comet.IcebergWriteExec
+import org.apache.spark.sql.connector.write.{BatchWrite, Write}
+import org.apache.spark.sql.types.BinaryType
+
+/** Logical anchor for the writer. See `IcebergWriteStrategy` for the rationale. */
+case class IcebergWriteLogical(
+    child: LogicalPlan,
+    // Driver-side only: AQE re-planning is driver-local and write commands aren't cached.
+    @transient batchWrite: BatchWrite,
+    @transient write: Write,
+    replaceDataDispatch: Option[ReplaceDataDispatchInfo] = None)
+    extends UnaryNode {
+
+  // Owns the commit-message attribute so the physical writer keeps the same exprId across
+  // AQE re-plans.
+  override val output: Seq[Attribute] = Seq(
+    AttributeReference(IcebergWriteExec.CommitMessageColumn, BinaryType, nullable = false)())
+
+  override protected def withNewChildInternal(newChild: LogicalPlan): IcebergWriteLogical =
+    copy(child = newChild)
+}

spark/src/main/scala/org/apache/comet/iceberg/IcebergWriteStrategy.scala

@@ -0,0 +1,122 @@
+/*
+ * 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.comet.iceberg
+
+import org.apache.spark.sql.SparkSession
+import org.apache.spark.sql.catalyst.plans.logical.{AppendData, LogicalPlan, OverwriteByExpression, OverwritePartitionsDynamic, ReplaceData}
+import org.apache.spark.sql.comet.{IcebergCommitExec, IcebergWriteExec}
+import org.apache.spark.sql.connector.write.Write
+import org.apache.spark.sql.execution.{SparkPlan, SparkStrategy}
+import org.apache.spark.sql.execution.datasources.v2.DataSourceV2Relation
+
+import org.apache.comet.CometConf
+
+/**
+ * Spark Strategy that intercepts Iceberg V2 copy-on-write logical writes and emits Comet's
+ * two-operator physical tree.
+ */
+case class IcebergWriteStrategy(session: SparkSession) extends SparkStrategy {
+
+  override def apply(plan: LogicalPlan): Seq[SparkPlan] = {
+    if (!CometConf.COMET_ICEBERG_WRITE_SPLIT_OPERATOR_ENABLED.get(session.sessionState.conf)) {
+      return Nil
+    }
+
+    plan match {
+      case ad: AppendData =>
+        matchedSparkWrite(ad.table, ad.write, ad.query, replaceDataDispatch = None).toList
+      case obe: OverwriteByExpression =>
+        matchedSparkWrite(obe.table, obe.write, obe.query, replaceDataDispatch = None).toList
+      case opd: OverwritePartitionsDynamic =>
+        matchedSparkWrite(opd.table, opd.write, opd.query, replaceDataDispatch = None).toList
+      case rd: ReplaceData =>
+        matchedSparkWrite(
+          rd.originalTable,
+          rd.write,
+          rd.query,
+          replaceDataDispatch = IcebergReplaceDataShim.extractProjections(rd)).toList
+      case plan if IcebergReflection.isReplaceIcebergData(plan) =>
+        IcebergReflection
+          .extractReplaceIcebergDataFields(plan)
+          .flatMap { case (_, query, originalTable, write) =>
+            matchedSparkWrite(
+              originalTable.asInstanceOf[org.apache.spark.sql.catalyst.analysis.NamedRelation],
+              write.asInstanceOf[Option[Write]],
+              query.asInstanceOf[LogicalPlan],
+              replaceDataDispatch = None)
+          }
+          .toList
+      // Hit by AQE.
+      case l @ IcebergWriteLogical(child, batchWrite, write, replaceDataDispatch) =>
+        Seq(IcebergWriteExec(batchWrite, write, l.output, planLater(child), replaceDataDispatch))
+      case _ => Nil
+    }
+  }
+
+  private def matchedSparkWrite(
+      table: org.apache.spark.sql.catalyst.analysis.NamedRelation,
+      write: Option[Write],
+      query: LogicalPlan,
+      replaceDataDispatch: Option[ReplaceDataDispatchInfo]): Option[SparkPlan] = {
+    table match {
+      case rel: DataSourceV2Relation =>
+        write.flatMap { w =>
+          if (IcebergReflection.isIcebergSparkWrite(w)) {
+            buildTwoOp(w, rel, query, replaceDataDispatch)
+          } else {
+            None
+          }
+        }
+      case _ => None
+    }
+  }
+
+  /**
+   * Builds the two-op tree. The committer and writer share one `BatchWrite` (also reused across
+   * AQE re-plans): `toBatch()` returns a fresh instance per call, but the committer's commit-time
+   * validation must see the same instance the writer wrote through, hence we store it. The
+   * writer's child is wrapped in [[IcebergWriteLogical]] so AQE re-emits only the data-writing
+   * operator on each re-plan as opposed to multiple new commit operators.
+   *
+   * Iceberg's `SparkWrite` never asks for Spark's commit coordinator, so the
+   * `useCommitCoordinator` fallback below is defensive coverage in case a future Iceberg version
+   * changes that; the split writer's per-task commit protocol does not use it.
+   */
+  private def buildTwoOp(
+      write: Write,
+      rel: DataSourceV2Relation,
+      query: LogicalPlan,
+      replaceDataDispatch: Option[ReplaceDataDispatchInfo]): Option[SparkPlan] = {
+    val batchWrite = write.toBatch
+    if (batchWrite.useCommitCoordinator()) {
+      return None
+    }
+    // To mirror Spark ReplaceData semantics we invalidate our cache of the state of
+    // `originalTable`.
+    val refresh: () => Unit = () => IcebergRefreshCacheShim.recacheByPlan(session, rel)
+    Some(
+      IcebergCommitExec(
+        batchWrite,
+        write,
+        refresh,
+        // `replaceDataDispatch` may project the data into the format the writer expects.
+        planLater(IcebergWriteLogical(query, batchWrite, write, replaceDataDispatch))))
+  }
+}