This is an automated email from the ASF dual-hosted git repository.
andygrove pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/datafusion-comet.git
The following commit(s) were added to refs/heads/main by this push:
new 3598d55951 feat: Optionally split the Iceberg V2 write operator into
distinct writer and committer operations (#4658)
3598d55951 is described below
commit 3598d55951948e4d35227a6dce6562eb46751d14
Author: Jordan Epstein <[email protected]>
AuthorDate: Thu Aug 6 21:53:00 2026 -0500
feat: Optionally split the Iceberg V2 write operator into distinct writer
and committer operations (#4658)
* docs: add Iceberg V2 writes user guide
* feat: add split-operator plan for Iceberg V2 appends and overwrites
Iceberg's V2 writes run as a single `V2ExistingTableWriteExec` command that
both
writes files and commits. Spark's `InsertAdaptiveSparkPlan` treats it as a
leaf,
so AQE never sees the data sub-query inside it and Comet's columnar rules
can't
convert the scans / shuffles / sorts that feed the write.
This injects `IcebergWriteStrategy` ahead of Spark's `DataSourceV2Strategy`
(via
`experimentalMethods.extraStrategies`, which is prepended to the planner's
strategy list). For Iceberg `AppendData`, `OverwriteByExpression`, and
`OverwritePartitionsDynamic` it rewrites the single command into two
operators:
IcebergCommitExec(batchWrite, refreshCache) <- committer (driver-side)
└── AdaptiveSparkPlanExec <- AQE bubble
└── IcebergWriteExec(batchWrite) <- writer (UnaryExecNode)
└── <data sub-query> <- now visible to AQE /
Comet
The writer emits one row per Spark task carrying the Java-serialised
`WriterCommitMessage`; the committer collects them and calls
`BatchWrite.commit(messages)` -- the same call Iceberg-Java makes
internally --
then runs Spark's post-write cache refresh.
Two design points keep this stable under AQE, which re-runs the planner on
each
materialised stage's `logicalLink`:
- The writer's child is wrapped in `IcebergWriteLogical`, a stable logical
anchor, so each re-plan re-emits only the writer rather than vanishing it
or
re-firing the surrounding logical write and duplicating the commit.
- The `BatchWrite` is materialised once (`Write.toBatch()` mints a fresh
instance per call) and shared between committer and writer, so commit-time
validation sees the same scan / emitted-file state the writer used.
`WriteDelta` (merge-on-read) is intentionally not intercepted: its per-task
`DeltaWriter` is row-dispatched and the native writer can't emit
position-delete
files, so the split plan would add planning complexity for no acceleration.
Copy-on-write DELETE / UPDATE / MERGE (`ReplaceData`) is handled in a
follow-up
commit.
Guarded by `spark.comet.write.iceberg.splitOperator.enabled` (off by
default);
when off, writes go straight through Iceberg-Java unchanged. File writing
still
runs through Iceberg's JVM writer -- native Parquet write lands in a later
commit.
* feat: extend Iceberg split plan to copy-on-write DELETE/UPDATE/MERGE
Adds copy-on-write row-level DML to the split-operator plan. Spark lowers a
CoW
`DELETE` / `UPDATE` / `MERGE` on a V2 table into a `ReplaceData` logical
node;
`IcebergWriteStrategy` now matches it (refreshing `originalTable`'s cache)
and
routes it through the same committer + writer pair as appends/overwrites.
The writer gains a per-row dispatch path. On Spark 4.x a `ReplaceData`
rewrite
prefixes each row with an operation code and carries a
`ReplaceDataProjections`;
`IcebergWriteExec` applies the row / metadata projections per code (WRITE /
WRITE_WITH_METADATA) before handing rows to the `DataWriter`. The two-arg
`DataWriter.write(metadata, row)` is looked up reflectively since it is
4.x-only.
On Spark 3.4 / 3.5 the rewritten stream is already post-projection, so the
shim
returns `None` and the writer keeps its plain `write(row)` loop.
`ReplaceDataProjections` does not exist on 3.4 / 3.5, so
`ReplaceDataDispatchInfo`
mirrors it as a version-neutral carrier, populated by the per-version
`IcebergReplaceDataShim`.
Iceberg 1.5.2 (Spark 3.4) lacks native row-level operations and instead
routes
UPDATE / MERGE through its own `ReplaceIcebergData` logical node. It has
the same
field shape as Spark's `ReplaceData` and is matched by FQCN via reflection,
so
the main module takes no compile dependency on iceberg-spark-extensions.
* test: cover Iceberg V2 split-operator appends and overwrites
* test: cover Iceberg copy-on-write DELETE/UPDATE/MERGE writes
* fix: correct Iceberg split write metrics, abort guard, and dispatch
* test: cover review-requested Iceberg split write scenarios
* docs: note CTAS and RTAS planning in Iceberg split writes
---------
Co-authored-by: Jordan Epstein <[email protected]>
Co-authored-by: Andy Grove <[email protected]>
---
.github/workflows/pr_build_linux.yml | 1 +
.github/workflows/pr_build_macos.yml | 1 +
dev/diffs/iceberg/1.10.0.diff | 56 ++
dev/diffs/iceberg/1.11.0.diff | 28 +
dev/diffs/iceberg/1.8.1.diff | 56 ++
dev/diffs/iceberg/1.9.1.diff | 56 ++
docs/source/user-guide/latest/iceberg-writes.md | 97 ++
docs/source/user-guide/latest/index.rst | 1 +
.../main/scala/org/apache/comet/CometConf.scala | 10 +
.../apache/comet/CometSparkSessionExtensions.scala | 2 +
.../apache/comet/iceberg/IcebergReflection.scala | 43 +
.../apache/comet/iceberg/IcebergWriteLogical.scala | 43 +
.../comet/iceberg/IcebergWriteStrategy.scala | 122 +++
.../comet/iceberg/ReplaceDataDispatchInfo.scala | 31 +
.../apache/spark/sql/comet/IcebergCommitExec.scala | 125 +++
.../apache/spark/sql/comet/IcebergWriteExec.scala | 179 ++++
.../comet/iceberg/IcebergDriverMetricsShim.scala | 28 +
.../comet/iceberg/IcebergRefreshCacheShim.scala | 29 +
.../comet/iceberg/IcebergReplaceDataShim.scala | 26 +
.../spark/sql/comet/IcebergDataWriterShim.scala | 35 +
.../spark/sql/comet/IcebergWriteSummaryShim.scala | 33 +
.../comet/iceberg/IcebergRefreshCacheShim.scala | 33 +
.../spark/sql/comet/IcebergWriteSummaryShim.scala | 33 +
.../comet/iceberg/IcebergRefreshCacheShim.scala | 43 +
.../spark/sql/comet/IcebergWriteSummaryShim.scala | 55 ++
.../comet/iceberg/IcebergDriverMetricsShim.scala | 27 +
.../comet/iceberg/IcebergReplaceDataShim.scala | 34 +
.../spark/sql/comet/IcebergDataWriterShim.scala | 32 +
.../comet/CometIcebergRewriteActionSuite.scala | 13 +-
.../comet/CometIcebergWriteActionSuite.scala | 979 +++++++++++++++++++++
30 files changed, 2247 insertions(+), 4 deletions(-)
diff --git a/.github/workflows/pr_build_linux.yml
b/.github/workflows/pr_build_linux.yml
index 8fbf82d3f9..def2816607 100644
--- a/.github/workflows/pr_build_linux.yml
+++ b/.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
diff --git a/.github/workflows/pr_build_macos.yml
b/.github/workflows/pr_build_macos.yml
index 5b92c53c4a..bac1dcd15b 100644
--- a/.github/workflows/pr_build_macos.yml
+++ b/.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
diff --git a/dev/diffs/iceberg/1.10.0.diff b/dev/diffs/iceberg/1.10.0.diff
index fd5aad2a32..310b541499 100644
--- a/dev/diffs/iceberg/1.10.0.diff
+++ b/dev/diffs/iceberg/1.10.0.diff
@@ -3025,3 +3025,59 @@ index 9d2ce2b388..5e23368848 100644
} else {
assertThat(planAsString).as("Should be no post scan
filter").doesNotContain("Filter (");
}
+diff --git
a/spark/v3.4/spark-extensions/src/test/java/org/apache/iceberg/spark/extensions/TestSystemFunctionPushDownInRowLevelOperations.java
b/spark/v3.4/spark-extensions/src/test/java/org/apache/iceberg/spark/extensions/TestSystemFunctionPushDownInRowLevelOperations.java
+index 5a0a04edb..afe92a8b7 100644
+---
a/spark/v3.4/spark-extensions/src/test/java/org/apache/iceberg/spark/extensions/TestSystemFunctionPushDownInRowLevelOperations.java
++++
b/spark/v3.4/spark-extensions/src/test/java/org/apache/iceberg/spark/extensions/TestSystemFunctionPushDownInRowLevelOperations.java
+@@ -38,6 +38,7 @@ import
org.apache.spark.sql.catalyst.expressions.ApplyFunctionExpression;
+ import org.apache.spark.sql.catalyst.expressions.Expression;
+ import org.apache.spark.sql.catalyst.expressions.objects.StaticInvoke;
+ import org.apache.spark.sql.execution.CommandResultExec;
++import org.apache.spark.sql.execution.SparkPlan;
+ import org.apache.spark.sql.execution.datasources.v2.V2TableWriteExec;
+ import org.junit.jupiter.api.AfterEach;
+ import org.junit.jupiter.api.BeforeEach;
+@@ -313,9 +314,13 @@ public class
TestSystemFunctionPushDownInRowLevelOperations extends ExtensionsTe
+
+ private List<Expression> executeAndCollectFunctionCalls(String query,
Object... args) {
+ CommandResultExec command = (CommandResultExec) executeAndKeepPlan(query,
args);
+- V2TableWriteExec write = (V2TableWriteExec) command.commandPhysicalPlan();
++ // Comet's split-operator Iceberg write plans the command as
IcebergCommitExec, which is not
++ // a V2TableWriteExec; collect over the whole command plan in that case.
++ SparkPlan plan = command.commandPhysicalPlan();
++ SparkPlan queryPlan =
++ plan instanceof V2TableWriteExec ? ((V2TableWriteExec) plan).query()
: plan;
+ return SparkPlanUtil.collectExprs(
+- write.query(),
++ queryPlan,
+ expr -> expr instanceof StaticInvoke || expr instanceof
ApplyFunctionExpression);
+ }
+
+diff --git
a/spark/v3.5/spark-extensions/src/test/java/org/apache/iceberg/spark/extensions/TestSystemFunctionPushDownInRowLevelOperations.java
b/spark/v3.5/spark-extensions/src/test/java/org/apache/iceberg/spark/extensions/TestSystemFunctionPushDownInRowLevelOperations.java
+index 8cb92224f..777b9bb1d 100644
+---
a/spark/v3.5/spark-extensions/src/test/java/org/apache/iceberg/spark/extensions/TestSystemFunctionPushDownInRowLevelOperations.java
++++
b/spark/v3.5/spark-extensions/src/test/java/org/apache/iceberg/spark/extensions/TestSystemFunctionPushDownInRowLevelOperations.java
+@@ -38,6 +38,7 @@ import
org.apache.spark.sql.catalyst.expressions.ApplyFunctionExpression;
+ import org.apache.spark.sql.catalyst.expressions.Expression;
+ import org.apache.spark.sql.catalyst.expressions.objects.StaticInvoke;
+ import org.apache.spark.sql.execution.CommandResultExec;
++import org.apache.spark.sql.execution.SparkPlan;
+ import org.apache.spark.sql.execution.datasources.v2.V2TableWriteExec;
+ import org.junit.jupiter.api.AfterEach;
+ import org.junit.jupiter.api.BeforeEach;
+@@ -313,9 +314,13 @@ public class
TestSystemFunctionPushDownInRowLevelOperations extends ExtensionsTe
+
+ private List<Expression> executeAndCollectFunctionCalls(String query,
Object... args) {
+ CommandResultExec command = (CommandResultExec) executeAndKeepPlan(query,
args);
+- V2TableWriteExec write = (V2TableWriteExec) command.commandPhysicalPlan();
++ // Comet's split-operator Iceberg write plans the command as
IcebergCommitExec, which is not
++ // a V2TableWriteExec; collect over the whole command plan in that case.
++ SparkPlan plan = command.commandPhysicalPlan();
++ SparkPlan queryPlan =
++ plan instanceof V2TableWriteExec ? ((V2TableWriteExec) plan).query()
: plan;
+ return SparkPlanUtil.collectExprs(
+- write.query(),
++ queryPlan,
+ expr -> expr instanceof StaticInvoke || expr instanceof
ApplyFunctionExpression);
+ }
+
diff --git a/dev/diffs/iceberg/1.11.0.diff b/dev/diffs/iceberg/1.11.0.diff
index bddfe96160..050e2107d3 100644
--- a/dev/diffs/iceberg/1.11.0.diff
+++ b/dev/diffs/iceberg/1.11.0.diff
@@ -784,3 +784,31 @@ index 8cdcf22e58..820f9fc03b 100644
.enableHiveSupport()
.getOrCreate();
+diff --git
a/spark/v4.1/spark-extensions/src/test/java/org/apache/iceberg/spark/extensions/TestSystemFunctionPushDownInRowLevelOperations.java
b/spark/v4.1/spark-extensions/src/test/java/org/apache/iceberg/spark/extensions/TestSystemFunctionPushDownInRowLevelOperations.java
+index 934220e5d..92132d4b8 100644
+---
a/spark/v4.1/spark-extensions/src/test/java/org/apache/iceberg/spark/extensions/TestSystemFunctionPushDownInRowLevelOperations.java
++++
b/spark/v4.1/spark-extensions/src/test/java/org/apache/iceberg/spark/extensions/TestSystemFunctionPushDownInRowLevelOperations.java
+@@ -40,6 +40,7 @@ import
org.apache.spark.sql.catalyst.expressions.ApplyFunctionExpression;
+ import org.apache.spark.sql.catalyst.expressions.Expression;
+ import org.apache.spark.sql.catalyst.expressions.objects.StaticInvoke;
+ import org.apache.spark.sql.execution.CommandResultExec;
++import org.apache.spark.sql.execution.SparkPlan;
+ import org.apache.spark.sql.execution.datasources.v2.V2TableWriteExec;
+ import org.junit.jupiter.api.AfterEach;
+ import org.junit.jupiter.api.BeforeEach;
+@@ -330,9 +331,13 @@ public class
TestSystemFunctionPushDownInRowLevelOperations extends ExtensionsTe
+
+ private List<Expression> executeAndCollectFunctionCalls(String query,
Object... args) {
+ CommandResultExec command = (CommandResultExec) executeAndKeepPlan(query,
args);
+- V2TableWriteExec write = (V2TableWriteExec) command.commandPhysicalPlan();
++ // Comet's split-operator Iceberg write plans the command as
IcebergCommitExec, which is not
++ // a V2TableWriteExec; collect over the whole command plan in that case.
++ SparkPlan plan = command.commandPhysicalPlan();
++ SparkPlan queryPlan =
++ plan instanceof V2TableWriteExec ? ((V2TableWriteExec) plan).query()
: plan;
+ return SparkPlanUtil.collectExprs(
+- write.query(),
++ queryPlan,
+ expr -> expr instanceof StaticInvoke || expr instanceof
ApplyFunctionExpression);
+ }
+
diff --git a/dev/diffs/iceberg/1.8.1.diff b/dev/diffs/iceberg/1.8.1.diff
index aa89d3e83b..b674f43244 100644
--- a/dev/diffs/iceberg/1.8.1.diff
+++ b/dev/diffs/iceberg/1.8.1.diff
@@ -3261,3 +3261,59 @@ index 6719c45ca9..2515454401 100644
sourceColumnName,
tableName,
tableName(OTHER_TABLE_NAME),
+diff --git
a/spark/v3.4/spark-extensions/src/test/java/org/apache/iceberg/spark/extensions/TestSystemFunctionPushDownInRowLevelOperations.java
b/spark/v3.4/spark-extensions/src/test/java/org/apache/iceberg/spark/extensions/TestSystemFunctionPushDownInRowLevelOperations.java
+index 49119d319..04fd3d485 100644
+---
a/spark/v3.4/spark-extensions/src/test/java/org/apache/iceberg/spark/extensions/TestSystemFunctionPushDownInRowLevelOperations.java
++++
b/spark/v3.4/spark-extensions/src/test/java/org/apache/iceberg/spark/extensions/TestSystemFunctionPushDownInRowLevelOperations.java
+@@ -37,6 +37,7 @@ import
org.apache.spark.sql.catalyst.expressions.ApplyFunctionExpression;
+ import org.apache.spark.sql.catalyst.expressions.Expression;
+ import org.apache.spark.sql.catalyst.expressions.objects.StaticInvoke;
+ import org.apache.spark.sql.execution.CommandResultExec;
++import org.apache.spark.sql.execution.SparkPlan;
+ import org.apache.spark.sql.execution.datasources.v2.V2TableWriteExec;
+ import org.junit.After;
+ import org.junit.Before;
+@@ -316,9 +317,13 @@ public class
TestSystemFunctionPushDownInRowLevelOperations extends SparkExtensi
+
+ private List<Expression> executeAndCollectFunctionCalls(String query,
Object... args) {
+ CommandResultExec command = (CommandResultExec) executeAndKeepPlan(query,
args);
+- V2TableWriteExec write = (V2TableWriteExec) command.commandPhysicalPlan();
++ // Comet's split-operator Iceberg write plans the command as
IcebergCommitExec, which is not
++ // a V2TableWriteExec; collect over the whole command plan in that case.
++ SparkPlan plan = command.commandPhysicalPlan();
++ SparkPlan queryPlan =
++ plan instanceof V2TableWriteExec ? ((V2TableWriteExec) plan).query()
: plan;
+ return SparkPlanUtil.collectExprs(
+- write.query(),
++ queryPlan,
+ expr -> expr instanceof StaticInvoke || expr instanceof
ApplyFunctionExpression);
+ }
+
+diff --git
a/spark/v3.5/spark-extensions/src/test/java/org/apache/iceberg/spark/extensions/TestSystemFunctionPushDownInRowLevelOperations.java
b/spark/v3.5/spark-extensions/src/test/java/org/apache/iceberg/spark/extensions/TestSystemFunctionPushDownInRowLevelOperations.java
+index 8cb92224f..777b9bb1d 100644
+---
a/spark/v3.5/spark-extensions/src/test/java/org/apache/iceberg/spark/extensions/TestSystemFunctionPushDownInRowLevelOperations.java
++++
b/spark/v3.5/spark-extensions/src/test/java/org/apache/iceberg/spark/extensions/TestSystemFunctionPushDownInRowLevelOperations.java
+@@ -38,6 +38,7 @@ import
org.apache.spark.sql.catalyst.expressions.ApplyFunctionExpression;
+ import org.apache.spark.sql.catalyst.expressions.Expression;
+ import org.apache.spark.sql.catalyst.expressions.objects.StaticInvoke;
+ import org.apache.spark.sql.execution.CommandResultExec;
++import org.apache.spark.sql.execution.SparkPlan;
+ import org.apache.spark.sql.execution.datasources.v2.V2TableWriteExec;
+ import org.junit.jupiter.api.AfterEach;
+ import org.junit.jupiter.api.BeforeEach;
+@@ -313,9 +314,13 @@ public class
TestSystemFunctionPushDownInRowLevelOperations extends ExtensionsTe
+
+ private List<Expression> executeAndCollectFunctionCalls(String query,
Object... args) {
+ CommandResultExec command = (CommandResultExec) executeAndKeepPlan(query,
args);
+- V2TableWriteExec write = (V2TableWriteExec) command.commandPhysicalPlan();
++ // Comet's split-operator Iceberg write plans the command as
IcebergCommitExec, which is not
++ // a V2TableWriteExec; collect over the whole command plan in that case.
++ SparkPlan plan = command.commandPhysicalPlan();
++ SparkPlan queryPlan =
++ plan instanceof V2TableWriteExec ? ((V2TableWriteExec) plan).query()
: plan;
+ return SparkPlanUtil.collectExprs(
+- write.query(),
++ queryPlan,
+ expr -> expr instanceof StaticInvoke || expr instanceof
ApplyFunctionExpression);
+ }
+
diff --git a/dev/diffs/iceberg/1.9.1.diff b/dev/diffs/iceberg/1.9.1.diff
index b9153abb3e..636f7e958f 100644
--- a/dev/diffs/iceberg/1.9.1.diff
+++ b/dev/diffs/iceberg/1.9.1.diff
@@ -3219,3 +3219,59 @@ index 6719c45ca9..2515454401 100644
sourceColumnName,
tableName,
tableName(OTHER_TABLE_NAME),
+diff --git
a/spark/v3.4/spark-extensions/src/test/java/org/apache/iceberg/spark/extensions/TestSystemFunctionPushDownInRowLevelOperations.java
b/spark/v3.4/spark-extensions/src/test/java/org/apache/iceberg/spark/extensions/TestSystemFunctionPushDownInRowLevelOperations.java
+index 5a0a04edb..afe92a8b7 100644
+---
a/spark/v3.4/spark-extensions/src/test/java/org/apache/iceberg/spark/extensions/TestSystemFunctionPushDownInRowLevelOperations.java
++++
b/spark/v3.4/spark-extensions/src/test/java/org/apache/iceberg/spark/extensions/TestSystemFunctionPushDownInRowLevelOperations.java
+@@ -38,6 +38,7 @@ import
org.apache.spark.sql.catalyst.expressions.ApplyFunctionExpression;
+ import org.apache.spark.sql.catalyst.expressions.Expression;
+ import org.apache.spark.sql.catalyst.expressions.objects.StaticInvoke;
+ import org.apache.spark.sql.execution.CommandResultExec;
++import org.apache.spark.sql.execution.SparkPlan;
+ import org.apache.spark.sql.execution.datasources.v2.V2TableWriteExec;
+ import org.junit.jupiter.api.AfterEach;
+ import org.junit.jupiter.api.BeforeEach;
+@@ -313,9 +314,13 @@ public class
TestSystemFunctionPushDownInRowLevelOperations extends ExtensionsTe
+
+ private List<Expression> executeAndCollectFunctionCalls(String query,
Object... args) {
+ CommandResultExec command = (CommandResultExec) executeAndKeepPlan(query,
args);
+- V2TableWriteExec write = (V2TableWriteExec) command.commandPhysicalPlan();
++ // Comet's split-operator Iceberg write plans the command as
IcebergCommitExec, which is not
++ // a V2TableWriteExec; collect over the whole command plan in that case.
++ SparkPlan plan = command.commandPhysicalPlan();
++ SparkPlan queryPlan =
++ plan instanceof V2TableWriteExec ? ((V2TableWriteExec) plan).query()
: plan;
+ return SparkPlanUtil.collectExprs(
+- write.query(),
++ queryPlan,
+ expr -> expr instanceof StaticInvoke || expr instanceof
ApplyFunctionExpression);
+ }
+
+diff --git
a/spark/v3.5/spark-extensions/src/test/java/org/apache/iceberg/spark/extensions/TestSystemFunctionPushDownInRowLevelOperations.java
b/spark/v3.5/spark-extensions/src/test/java/org/apache/iceberg/spark/extensions/TestSystemFunctionPushDownInRowLevelOperations.java
+index 8cb92224f..777b9bb1d 100644
+---
a/spark/v3.5/spark-extensions/src/test/java/org/apache/iceberg/spark/extensions/TestSystemFunctionPushDownInRowLevelOperations.java
++++
b/spark/v3.5/spark-extensions/src/test/java/org/apache/iceberg/spark/extensions/TestSystemFunctionPushDownInRowLevelOperations.java
+@@ -38,6 +38,7 @@ import
org.apache.spark.sql.catalyst.expressions.ApplyFunctionExpression;
+ import org.apache.spark.sql.catalyst.expressions.Expression;
+ import org.apache.spark.sql.catalyst.expressions.objects.StaticInvoke;
+ import org.apache.spark.sql.execution.CommandResultExec;
++import org.apache.spark.sql.execution.SparkPlan;
+ import org.apache.spark.sql.execution.datasources.v2.V2TableWriteExec;
+ import org.junit.jupiter.api.AfterEach;
+ import org.junit.jupiter.api.BeforeEach;
+@@ -313,9 +314,13 @@ public class
TestSystemFunctionPushDownInRowLevelOperations extends ExtensionsTe
+
+ private List<Expression> executeAndCollectFunctionCalls(String query,
Object... args) {
+ CommandResultExec command = (CommandResultExec) executeAndKeepPlan(query,
args);
+- V2TableWriteExec write = (V2TableWriteExec) command.commandPhysicalPlan();
++ // Comet's split-operator Iceberg write plans the command as
IcebergCommitExec, which is not
++ // a V2TableWriteExec; collect over the whole command plan in that case.
++ SparkPlan plan = command.commandPhysicalPlan();
++ SparkPlan queryPlan =
++ plan instanceof V2TableWriteExec ? ((V2TableWriteExec) plan).query()
: plan;
+ return SparkPlanUtil.collectExprs(
+- write.query(),
++ queryPlan,
+ expr -> expr instanceof StaticInvoke || expr instanceof
ApplyFunctionExpression);
+ }
+
diff --git a/docs/source/user-guide/latest/iceberg-writes.md
b/docs/source/user-guide/latest/iceberg-writes.md
new file mode 100644
index 0000000000..8ef6b4cc16
--- /dev/null
+++ b/docs/source/user-guide/latest/iceberg-writes.md
@@ -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.
+-->
+
+# Iceberg Writes: Comet's Split-Operator Plan (Experimental)
+
+**This feature is experimental and disabled by default.** Enable it only after
validating it
+against your own workloads.
+
+## 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)) plus the
write-side toggle:
+
+```
+# 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, off by default)
+spark.comet.write.iceberg.splitOperator.enabled=true
+```
+
+## 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.
+
+On Spark 4.1+ the split plan matches two further stock-Spark behaviours: MERGE
metrics are
+forwarded to the writer's commit (Iceberg 1.11+ records them in the snapshot
summary), and
+cached catalog tables are recached by name after a write so cache entries
survive schema
+changes.
+
+## 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 default);
+- 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 statement is CTAS / RTAS on Spark 3.4, where the staged exec writes
inline; on Spark
+ 3.5+ those statements re-plan their inner append, which is intercepted
normally;
+- 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.
diff --git a/docs/source/user-guide/latest/index.rst
b/docs/source/user-guide/latest/index.rst
index 2cd1c7e16f..815e12289c 100644
--- a/docs/source/user-guide/latest/index.rst
+++ b/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>
diff --git a/spark/src/main/scala/org/apache/comet/CometConf.scala
b/spark/src/main/scala/org/apache/comet/CometConf.scala
index 1bb524349d..6e7bb4271b 100644
--- a/spark/src/main/scala/org/apache/comet/CometConf.scala
+++ b/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(false)
+
val COMET_ICEBERG_DATA_FILE_CONCURRENCY_LIMIT: ConfigEntry[Int] =
conf("spark.comet.scan.icebergNative.dataFileConcurrencyLimit")
.category(CATEGORY_SCAN)
diff --git
a/spark/src/main/scala/org/apache/comet/CometSparkSessionExtensions.scala
b/spark/src/main/scala/org/apache/comet/CometSparkSessionExtensions.scala
index 3bcb3f98e1..4dc3619433 100644
--- a/spark/src/main/scala/org/apache/comet/CometSparkSessionExtensions.scala
+++ b/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 {
diff --git
a/spark/src/main/scala/org/apache/comet/iceberg/IcebergReflection.scala
b/spark/src/main/scala/org/apache/comet/iceberg/IcebergReflection.scala
index c90d697bf7..edb321e8fc 100644
--- a/spark/src/main/scala/org/apache/comet/iceberg/IcebergReflection.scala
+++ b/spark/src/main/scala/org/apache/comet/iceberg/IcebergReflection.scala
@@ -57,6 +57,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"
}
/**
@@ -131,6 +135,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.
diff --git
a/spark/src/main/scala/org/apache/comet/iceberg/IcebergWriteLogical.scala
b/spark/src/main/scala/org/apache/comet/iceberg/IcebergWriteLogical.scala
new file mode 100644
index 0000000000..6ceb526d4c
--- /dev/null
+++ b/spark/src/main/scala/org/apache/comet/iceberg/IcebergWriteLogical.scala
@@ -0,0 +1,43 @@
+/*
+ * 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
+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,
+ 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)
+}
diff --git
a/spark/src/main/scala/org/apache/comet/iceberg/IcebergWriteStrategy.scala
b/spark/src/main/scala/org/apache/comet/iceberg/IcebergWriteStrategy.scala
new file mode 100644
index 0000000000..c501c90097
--- /dev/null
+++ b/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, replaceDataDispatch) =>
+ Seq(IcebergWriteExec(batchWrite, 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.refreshCache(session, rel)
+ Some(
+ IcebergCommitExec(
+ batchWrite,
+ write,
+ refresh,
+ // `replaceDataDispatch` may project the data into the format the
writer expects.
+ planLater(IcebergWriteLogical(query, batchWrite,
replaceDataDispatch))))
+ }
+}
diff --git
a/spark/src/main/scala/org/apache/comet/iceberg/ReplaceDataDispatchInfo.scala
b/spark/src/main/scala/org/apache/comet/iceberg/ReplaceDataDispatchInfo.scala
new file mode 100644
index 0000000000..8f9274a9e2
--- /dev/null
+++
b/spark/src/main/scala/org/apache/comet/iceberg/ReplaceDataDispatchInfo.scala
@@ -0,0 +1,31 @@
+/*
+ * 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.ProjectingInternalRow
+
+/**
+ * Mirror of Spark 4.x's `ReplaceDataProjections`, defined here so we can
compile against Spark
+ * 3.4 / 3.5 source trees (where the class doesn't exist). Populated by
`IcebergReplaceDataShim`
+ * on 4.x and left `None` on 3.x; consumed by `IcebergWriteExec`.
+ */
+case class ReplaceDataDispatchInfo(
+ rowProjection: ProjectingInternalRow,
+ metadataProjection: Option[ProjectingInternalRow])
diff --git
a/spark/src/main/scala/org/apache/spark/sql/comet/IcebergCommitExec.scala
b/spark/src/main/scala/org/apache/spark/sql/comet/IcebergCommitExec.scala
new file mode 100644
index 0000000000..5041c9fc7f
--- /dev/null
+++ b/spark/src/main/scala/org/apache/spark/sql/comet/IcebergCommitExec.scala
@@ -0,0 +1,125 @@
+/*
+ * 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.comet
+
+import org.apache.spark.internal.Logging
+import org.apache.spark.sql.catalyst.InternalRow
+import org.apache.spark.sql.catalyst.expressions.Attribute
+import org.apache.spark.sql.connector.write.{BatchWrite, Write,
WriterCommitMessage}
+import org.apache.spark.sql.execution.{SparkPlan, SQLExecution, UnaryExecNode}
+import org.apache.spark.sql.execution.datasources.v2.V2CommandExec
+import org.apache.spark.sql.execution.metric.{SQLMetric, SQLMetrics}
+
+import org.apache.comet.iceberg.IcebergDriverMetricsShim
+
+/**
+ * Driver-side committer for Comet's split-operator Iceberg V2 write.
+ */
+case class IcebergCommitExec(
+ // None of these fields are serialized, this is all run on the driver.
+ @transient batchWrite: BatchWrite,
+ @transient write: Write,
+ @transient refreshCache: IcebergCommitExec.RefreshCache,
+ child: SparkPlan)
+ extends V2CommandExec
+ with UnaryExecNode
+ with Logging {
+
+ override def output: Seq[Attribute] = Nil
+
+ override lazy val metrics: Map[String, SQLMetric] =
+ Map(
+ "numCommittedMessages" -> SQLMetrics
+ .createMetric(sparkContext, "number of task commit messages")) ++
+ write
+ .supportedCustomMetrics()
+ .map(m => m.name -> SQLMetrics.createV2CustomMetric(sparkContext, m))
+
+ // Exactly-once relies on V2CommandExec memoizing run() via its `result`
lazy val and on
+ // the writer executing once inside the AQE bubble anchored by
IcebergWriteLogical; pinned
+ // by the AQE re-plan test in CometIcebergWriteActionSuite.
+ override protected def run(): Seq[InternalRow] = {
+ try {
+ collectAndCommit()
+ } finally {
+ postDriverMetrics()
+ }
+ }
+
+ private def collectAndCommit(): Seq[InternalRow] = {
+ val messages: Array[WriterCommitMessage] =
+ try {
+ child.executeCollect().map { row =>
+ IcebergWriteExec.deserializeMessage(row.getBinary(0))
+ }
+ } catch {
+ case cause: Throwable =>
+ // The write job failed; the BatchWrite contract still expects a
job-level abort.
+ try batchWrite.abort(Array.empty[WriterCommitMessage])
+ catch {
+ case abortFailure: Throwable =>
+ cause.addSuppressed(abortFailure)
+ }
+ throw cause
+ }
+ longMetric("numCommittedMessages").add(messages.length)
+
+ try {
+ messages.foreach(batchWrite.onDataWriterCommit)
+ IcebergWriteSummaryShim.commit(batchWrite, messages, child)
+ logInfo(s"Iceberg commit succeeded with ${messages.length} task
message(s)")
+ } catch {
+ case cause: Throwable =>
+ logError(s"Iceberg commit failed; aborting ${messages.length} task
message(s)", cause)
+ try batchWrite.abort(messages)
+ catch {
+ case abortFailure: Throwable =>
+ cause.addSuppressed(abortFailure)
+ }
+ throw cause
+ }
+
+ refreshCache()
+ Nil
+ }
+
+ private def postDriverMetrics(): Unit = {
+ val driverMetrics = IcebergDriverMetricsShim.reportDriverMetrics(write)
+ if (driverMetrics.nonEmpty) {
+ val updated = driverMetrics.flatMap { taskMetric =>
+ metrics.get(taskMetric.name).map { sqlMetric =>
+ sqlMetric.set(taskMetric.value)
+ sqlMetric
+ }
+ }
+ val executionId =
sparkContext.getLocalProperty(SQLExecution.EXECUTION_ID_KEY)
+ SQLMetrics.postDriverMetricUpdates(sparkContext, executionId,
updated.toIndexedSeq)
+ }
+ }
+
+ override protected def withNewChildInternal(newChild: SparkPlan):
IcebergCommitExec =
+ copy(child = newChild)
+
+ override def nodeName: String = "IcebergCommit"
+}
+
+object IcebergCommitExec {
+ type RefreshCache = () => Unit
+}
diff --git
a/spark/src/main/scala/org/apache/spark/sql/comet/IcebergWriteExec.scala
b/spark/src/main/scala/org/apache/spark/sql/comet/IcebergWriteExec.scala
new file mode 100644
index 0000000000..90bfdaab04
--- /dev/null
+++ b/spark/src/main/scala/org/apache/spark/sql/comet/IcebergWriteExec.scala
@@ -0,0 +1,179 @@
+/*
+ * 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.comet
+
+import org.apache.spark.TaskContext
+import org.apache.spark.rdd.RDD
+import org.apache.spark.sql.catalyst.InternalRow
+import org.apache.spark.sql.catalyst.expressions.Attribute
+import org.apache.spark.sql.catalyst.expressions.UnsafeProjection
+import org.apache.spark.sql.catalyst.plans.physical.{Distribution,
UnspecifiedDistribution}
+import org.apache.spark.sql.connector.write.{BatchWrite, DataWriter,
PhysicalWriteInfoImpl, WriterCommitMessage}
+import org.apache.spark.sql.execution.{SparkPlan, UnaryExecNode}
+import org.apache.spark.sql.execution.metric.{SQLMetric, SQLMetrics}
+import org.apache.spark.sql.types.{BinaryType, StructField, StructType}
+import org.apache.spark.util.Utils
+
+import org.apache.comet.iceberg.ReplaceDataDispatchInfo
+
+/**
+ * Executor-side file writer for Comet's split-operator Iceberg V2 write.
+ */
+case class IcebergWriteExec(
+ // `batchWrite` only stored driver side, only the writer factory is
shipped to executors.
+ @transient batchWrite: BatchWrite,
+ override val output: Seq[Attribute],
+ child: SparkPlan,
+ replaceDataDispatch: Option[ReplaceDataDispatchInfo] = None)
+ extends UnaryExecNode {
+
+ // Spark already adds a distribution for the V2 write; adding another here
is redundant.
+ override def requiredChildDistribution: Seq[Distribution] =
Seq(UnspecifiedDistribution)
+
+ // The write's custom metrics are registered on IcebergCommitExec only:
Iceberg reports them
+ // through the driver-side commit report, and its DataWriters expose no
task-level values.
+ // Registering them here too would double-count them in the SQL UI.
+ override lazy val metrics: Map[String, SQLMetric] =
+ Map("numOutputRows" -> SQLMetrics.createMetric(sparkContext, "number of
output rows"))
+
+ override protected def doExecute(): RDD[InternalRow] = {
+ val rdd = {
+ val tempRdd = child.execute()
+ // SPARK-23271: run one write task even when the query RDD has zero
partitions.
+ if (tempRdd.getNumPartitions == 0) {
+ sparkContext.parallelize(Seq.empty[InternalRow], 1)
+ } else {
+ tempRdd
+ }
+ }
+ val factory =
batchWrite.createBatchWriterFactory(PhysicalWriteInfoImpl(rdd.getNumPartitions))
+ // Backstop only; IcebergWriteStrategy already falls back at planning time.
+ require(
+ !batchWrite.useCommitCoordinator(),
+ "Comet's Iceberg write path does not currently support BatchWrite
implementations that " +
+ "require Spark's commit coordinator; received: " +
batchWrite.getClass.getName)
+
+ val rowsMetric = longMetric("numOutputRows")
+ val schemaTypes = output.map(_.dataType).toArray
+ val capturedReplaceDataDispatch = replaceDataDispatch
+ rdd.mapPartitionsInternal { iter =>
+ val partId = TaskContext.getPartitionId()
+ val taskId = TaskContext.get().taskAttemptId()
+ val writer = factory.createWriter(partId, taskId)
+ val projection = UnsafeProjection.create(schemaTypes)
+ IcebergWriteExec.runWriter(
+ writer,
+ iter,
+ rowsMetric,
+ projection,
+ capturedReplaceDataDispatch)
+ }
+ }
+
+ override protected def withNewChildInternal(newChild: SparkPlan):
IcebergWriteExec =
+ copy(child = newChild)
+
+ override def nodeName: String = "IcebergWrite"
+}
+
+object IcebergWriteExec {
+
+ val CommitMessageColumn: String = "iceberg_commit_message"
+
+ val OutputSchema: StructType = StructType(
+ Seq(StructField(CommitMessageColumn, BinaryType, nullable = false)))
+
+ /** Writes data files and returns the serialised Iceberg commit message. */
+ def runWriter(
+ writer: DataWriter[InternalRow],
+ iter: Iterator[InternalRow],
+ rowsMetric: SQLMetric,
+ projection: UnsafeProjection,
+ replaceDataDispatch: Option[ReplaceDataDispatchInfo]):
Iterator[InternalRow] = {
+ val iterWithMetrics = new IteratorWithMetrics(iter, rowsMetric)
+ // Serialization happens inside the guarded block: if the commit message
cannot be
+ // serialised the task must abort so the already-finalised data files get
deleted.
+ val messageBytes = Utils.tryWithSafeFinallyAndFailureCallbacks(block = {
+ if (replaceDataDispatch.isDefined) {
+ runReplaceDataWriter(writer, iterWithMetrics, replaceDataDispatch.get)
+ } else {
+ while (iterWithMetrics.hasNext) {
+ writer.write(iterWithMetrics.next())
+ }
+ }
+ serializeMessage(writer.commit())
+ })(
+ catchBlock = {
+ writer.abort()
+ },
+ finallyBlock = {
+ writer.close()
+ })
+
+ Iterator.single(projection(InternalRow(messageBytes)).copy())
+ }
+
+ private class IteratorWithMetrics(iter: Iterator[InternalRow], rowsMetric:
SQLMetric)
+ extends Iterator[InternalRow] {
+
+ override def hasNext: Boolean = iter.hasNext
+
+ override def next(): InternalRow = {
+ rowsMetric.add(1L)
+ iter.next()
+ }
+ }
+
+ // Mirrors Spark RowDeltaUtils, which is private and changes location across
versions.
+ // These codes only exist on Spark 4.0+; on 3.x IcebergReplaceDataShim
returns None and rows
+ // take the plain write(row) loop instead, so this dispatch is never reached
there.
+ private val WRITE_OPERATION = 5
+ private val WRITE_WITH_METADATA_OPERATION = 6
+
+ def serializeMessage(message: WriterCommitMessage): Array[Byte] =
+ Utils.serialize(message)
+
+ def deserializeMessage(bytes: Array[Byte]): WriterCommitMessage =
+ Utils.deserialize[WriterCommitMessage](bytes,
Utils.getContextOrSparkClassLoader)
+
+ private def runReplaceDataWriter(
+ writer: DataWriter[InternalRow],
+ iter: Iterator[InternalRow],
+ dispatch: ReplaceDataDispatchInfo): Unit = {
+ val rowProjection = dispatch.rowProjection
+ val metadataProjection = dispatch.metadataProjection.orNull
+ while (iter.hasNext) {
+ val row = iter.next()
+ row.getInt(0) match {
+ case WRITE_OPERATION =>
+ rowProjection.project(row)
+ writer.write(rowProjection)
+ case WRITE_WITH_METADATA_OPERATION =>
+ rowProjection.project(row)
+ if (metadataProjection != null) metadataProjection.project(row)
+ IcebergDataWriterShim.writeWithMetadata(writer, metadataProjection,
rowProjection)
+ case other =>
+ throw new IllegalArgumentException(
+ s"Unexpected ReplaceData operation code $other; supported: " +
+ s"$WRITE_OPERATION (WRITE), $WRITE_WITH_METADATA_OPERATION
(WRITE_WITH_METADATA)")
+ }
+ }
+ }
+}
diff --git
a/spark/src/main/spark-3.x/org/apache/comet/iceberg/IcebergDriverMetricsShim.scala
b/spark/src/main/spark-3.x/org/apache/comet/iceberg/IcebergDriverMetricsShim.scala
new file mode 100644
index 0000000000..4f45f18cef
--- /dev/null
+++
b/spark/src/main/spark-3.x/org/apache/comet/iceberg/IcebergDriverMetricsShim.scala
@@ -0,0 +1,28 @@
+/*
+ * 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.connector.metric.CustomTaskMetric
+import org.apache.spark.sql.connector.write.Write
+
+/** `Write.reportDriverMetrics` (SPARK-50049) does not exist before Spark 4.0.
*/
+object IcebergDriverMetricsShim {
+ def reportDriverMetrics(write: Write): Array[CustomTaskMetric] = Array.empty
+}
diff --git
a/spark/src/main/spark-3.x/org/apache/comet/iceberg/IcebergRefreshCacheShim.scala
b/spark/src/main/spark-3.x/org/apache/comet/iceberg/IcebergRefreshCacheShim.scala
new file mode 100644
index 0000000000..a3dc459231
--- /dev/null
+++
b/spark/src/main/spark-3.x/org/apache/comet/iceberg/IcebergRefreshCacheShim.scala
@@ -0,0 +1,29 @@
+/*
+ * 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.execution.datasources.v2.DataSourceV2Relation
+
+private[iceberg] object IcebergRefreshCacheShim {
+ def refreshCache(session: SparkSession, rel: DataSourceV2Relation): Unit = {
+ session.sharedState.cacheManager.recacheByPlan(session, rel)
+ }
+}
diff --git
a/spark/src/main/spark-3.x/org/apache/comet/iceberg/IcebergReplaceDataShim.scala
b/spark/src/main/spark-3.x/org/apache/comet/iceberg/IcebergReplaceDataShim.scala
new file mode 100644
index 0000000000..5168e35012
--- /dev/null
+++
b/spark/src/main/spark-3.x/org/apache/comet/iceberg/IcebergReplaceDataShim.scala
@@ -0,0 +1,26 @@
+/*
+ * 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.plans.logical.ReplaceData
+
+private[iceberg] object IcebergReplaceDataShim {
+ def extractProjections(rd: ReplaceData): Option[ReplaceDataDispatchInfo] =
None
+}
diff --git
a/spark/src/main/spark-3.x/org/apache/spark/sql/comet/IcebergDataWriterShim.scala
b/spark/src/main/spark-3.x/org/apache/spark/sql/comet/IcebergDataWriterShim.scala
new file mode 100644
index 0000000000..a2ab14eacb
--- /dev/null
+++
b/spark/src/main/spark-3.x/org/apache/spark/sql/comet/IcebergDataWriterShim.scala
@@ -0,0 +1,35 @@
+/*
+ * 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.comet
+
+import org.apache.spark.sql.catalyst.InternalRow
+import org.apache.spark.sql.connector.write.DataWriter
+
+// Never reached on Spark 3.x: operation-coded ReplaceData rows only exist on
4.0+ and
+// IcebergReplaceDataShim returns no dispatch info there.
+private[comet] object IcebergDataWriterShim {
+
+ def writeWithMetadata(
+ writer: DataWriter[InternalRow],
+ metadata: InternalRow,
+ row: InternalRow): Unit =
+ throw new UnsupportedOperationException(
+ "DataWriter.write(metadata, row) requires Spark 4.0 or later")
+}
diff --git
a/spark/src/main/spark-3.x/org/apache/spark/sql/comet/IcebergWriteSummaryShim.scala
b/spark/src/main/spark-3.x/org/apache/spark/sql/comet/IcebergWriteSummaryShim.scala
new file mode 100644
index 0000000000..d6a78ca2ea
--- /dev/null
+++
b/spark/src/main/spark-3.x/org/apache/spark/sql/comet/IcebergWriteSummaryShim.scala
@@ -0,0 +1,33 @@
+/*
+ * 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.comet
+
+import org.apache.spark.sql.connector.write.{BatchWrite, WriterCommitMessage}
+import org.apache.spark.sql.execution.SparkPlan
+
+/** The `BatchWrite.commit(messages, summary)` overload only exists on Spark
4.1+. */
+private[comet] object IcebergWriteSummaryShim {
+ def commit(
+ batchWrite: BatchWrite,
+ messages: Array[WriterCommitMessage],
+ query: SparkPlan): Unit = {
+ batchWrite.commit(messages)
+ }
+}
diff --git
a/spark/src/main/spark-4.0/org/apache/comet/iceberg/IcebergRefreshCacheShim.scala
b/spark/src/main/spark-4.0/org/apache/comet/iceberg/IcebergRefreshCacheShim.scala
new file mode 100644
index 0000000000..2fa3172dd8
--- /dev/null
+++
b/spark/src/main/spark-4.0/org/apache/comet/iceberg/IcebergRefreshCacheShim.scala
@@ -0,0 +1,33 @@
+/*
+ * 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.classic.SparkSession
+import org.apache.spark.sql.execution.datasources.v2.DataSourceV2Relation
+
+/** Spark 4.x needs the classic `SparkSession`; the api-module one has no
`sharedState`. */
+private[iceberg] object IcebergRefreshCacheShim {
+ def refreshCache(
+ session: org.apache.spark.sql.SparkSession,
+ rel: DataSourceV2Relation): Unit = {
+ val classic = session.asInstanceOf[SparkSession]
+ classic.sharedState.cacheManager.recacheByPlan(classic, rel)
+ }
+}
diff --git
a/spark/src/main/spark-4.0/org/apache/spark/sql/comet/IcebergWriteSummaryShim.scala
b/spark/src/main/spark-4.0/org/apache/spark/sql/comet/IcebergWriteSummaryShim.scala
new file mode 100644
index 0000000000..d6a78ca2ea
--- /dev/null
+++
b/spark/src/main/spark-4.0/org/apache/spark/sql/comet/IcebergWriteSummaryShim.scala
@@ -0,0 +1,33 @@
+/*
+ * 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.comet
+
+import org.apache.spark.sql.connector.write.{BatchWrite, WriterCommitMessage}
+import org.apache.spark.sql.execution.SparkPlan
+
+/** The `BatchWrite.commit(messages, summary)` overload only exists on Spark
4.1+. */
+private[comet] object IcebergWriteSummaryShim {
+ def commit(
+ batchWrite: BatchWrite,
+ messages: Array[WriterCommitMessage],
+ query: SparkPlan): Unit = {
+ batchWrite.commit(messages)
+ }
+}
diff --git
a/spark/src/main/spark-4.1+/org/apache/comet/iceberg/IcebergRefreshCacheShim.scala
b/spark/src/main/spark-4.1+/org/apache/comet/iceberg/IcebergRefreshCacheShim.scala
new file mode 100644
index 0000000000..bd14131b8a
--- /dev/null
+++
b/spark/src/main/spark-4.1+/org/apache/comet/iceberg/IcebergRefreshCacheShim.scala
@@ -0,0 +1,43 @@
+/*
+ * 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.classic.SparkSession
+import org.apache.spark.sql.execution.datasources.v2.{DataSourceV2Relation,
ExtractV2CatalogAndIdentifier}
+
+/** Spark 4.x needs the classic `SparkSession`; the api-module one has no
`sharedState`. */
+private[iceberg] object IcebergRefreshCacheShim {
+
+ // Spark 4.1 recaches catalog tables by name rather than by plan so that
cache entries built
+ // before a schema change are still refreshed.
+ def refreshCache(
+ session: org.apache.spark.sql.SparkSession,
+ rel: DataSourceV2Relation): Unit = {
+ val classic = session.asInstanceOf[SparkSession]
+ val cacheManager = classic.sharedState.cacheManager
+ rel match {
+ case ExtractV2CatalogAndIdentifier(catalog, ident) =>
+ val nameParts = catalog.name +: ident.namespace.toSeq :+ ident.name
+ cacheManager.recacheTableOrView(classic, nameParts, includeTimeTravel
= false)
+ case _ =>
+ cacheManager.recacheByPlan(classic, rel)
+ }
+ }
+}
diff --git
a/spark/src/main/spark-4.1+/org/apache/spark/sql/comet/IcebergWriteSummaryShim.scala
b/spark/src/main/spark-4.1+/org/apache/spark/sql/comet/IcebergWriteSummaryShim.scala
new file mode 100644
index 0000000000..59aaad3f8c
--- /dev/null
+++
b/spark/src/main/spark-4.1+/org/apache/spark/sql/comet/IcebergWriteSummaryShim.scala
@@ -0,0 +1,55 @@
+/*
+ * 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.comet
+
+import org.apache.spark.sql.connector.write.{BatchWrite, MergeSummaryImpl,
WriterCommitMessage}
+import org.apache.spark.sql.execution.SparkPlan
+import org.apache.spark.sql.execution.adaptive.AdaptiveSparkPlanHelper
+import org.apache.spark.sql.execution.datasources.v2.MergeRowsExec
+
+/**
+ * Spark 4.1 collects MERGE metrics from the executed plan and hands them to
+ * `BatchWrite.commit(messages, summary)`; Iceberg 1.11+ records them in the
snapshot summary.
+ */
+private[comet] object IcebergWriteSummaryShim extends AdaptiveSparkPlanHelper {
+ def commit(
+ batchWrite: BatchWrite,
+ messages: Array[WriterCommitMessage],
+ query: SparkPlan): Unit = {
+ collectFirst(query) { case m: MergeRowsExec => m } match {
+ case Some(mergeRows) =>
+ val metrics = mergeRows.metrics
+ def metricValue(name: String): Long =
metrics.get(name).map(_.value).getOrElse(-1L)
+ batchWrite.commit(
+ messages,
+ MergeSummaryImpl(
+ metricValue("numTargetRowsCopied"),
+ metricValue("numTargetRowsDeleted"),
+ metricValue("numTargetRowsUpdated"),
+ metricValue("numTargetRowsInserted"),
+ metricValue("numTargetRowsMatchedUpdated"),
+ metricValue("numTargetRowsMatchedDeleted"),
+ metricValue("numTargetRowsNotMatchedBySourceUpdated"),
+ metricValue("numTargetRowsNotMatchedBySourceDeleted")))
+ case None =>
+ batchWrite.commit(messages)
+ }
+ }
+}
diff --git
a/spark/src/main/spark-4.x/org/apache/comet/iceberg/IcebergDriverMetricsShim.scala
b/spark/src/main/spark-4.x/org/apache/comet/iceberg/IcebergDriverMetricsShim.scala
new file mode 100644
index 0000000000..46f9abf35e
--- /dev/null
+++
b/spark/src/main/spark-4.x/org/apache/comet/iceberg/IcebergDriverMetricsShim.scala
@@ -0,0 +1,27 @@
+/*
+ * 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.connector.metric.CustomTaskMetric
+import org.apache.spark.sql.connector.write.Write
+
+object IcebergDriverMetricsShim {
+ def reportDriverMetrics(write: Write): Array[CustomTaskMetric] =
write.reportDriverMetrics()
+}
diff --git
a/spark/src/main/spark-4.x/org/apache/comet/iceberg/IcebergReplaceDataShim.scala
b/spark/src/main/spark-4.x/org/apache/comet/iceberg/IcebergReplaceDataShim.scala
new file mode 100644
index 0000000000..c860580005
--- /dev/null
+++
b/spark/src/main/spark-4.x/org/apache/comet/iceberg/IcebergReplaceDataShim.scala
@@ -0,0 +1,34 @@
+/*
+ * 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.plans.logical.ReplaceData
+
+/**
+ * Spark 4.x: `ReplaceData` carries `projections: ReplaceDataProjections` --
the rewritten row
+ * stream is prefixed with an operation column (5=WRITE,
6=WRITE_WITH_METADATA), and we have to
+ * apply `dataProj` / `metadataProj` before handing rows to the underlying
`DataWriter`.
+ */
+private[iceberg] object IcebergReplaceDataShim {
+ def extractProjections(rd: ReplaceData): Option[ReplaceDataDispatchInfo] = {
+ val p = rd.projections
+ Some(ReplaceDataDispatchInfo(p.rowProjection, p.metadataProjection))
+ }
+}
diff --git
a/spark/src/main/spark-4.x/org/apache/spark/sql/comet/IcebergDataWriterShim.scala
b/spark/src/main/spark-4.x/org/apache/spark/sql/comet/IcebergDataWriterShim.scala
new file mode 100644
index 0000000000..e30bdff4f4
--- /dev/null
+++
b/spark/src/main/spark-4.x/org/apache/spark/sql/comet/IcebergDataWriterShim.scala
@@ -0,0 +1,32 @@
+/*
+ * 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.comet
+
+import org.apache.spark.sql.catalyst.InternalRow
+import org.apache.spark.sql.connector.write.DataWriter
+
+private[comet] object IcebergDataWriterShim {
+
+ def writeWithMetadata(
+ writer: DataWriter[InternalRow],
+ metadata: InternalRow,
+ row: InternalRow): Unit =
+ writer.write(metadata, row)
+}
diff --git
a/spark/src/test/scala/org/apache/comet/CometIcebergRewriteActionSuite.scala
b/spark/src/test/scala/org/apache/comet/CometIcebergRewriteActionSuite.scala
index af2508aedc..19a4d6a3ce 100644
--- a/spark/src/test/scala/org/apache/comet/CometIcebergRewriteActionSuite.scala
+++ b/spark/src/test/scala/org/apache/comet/CometIcebergRewriteActionSuite.scala
@@ -204,7 +204,7 @@ class CometIcebergRewriteActionSuite extends CometTestBase
with CometIcebergTest
s"Expected >= 1 input file rewritten, got $rewrittenCount")
}
- val rewritePlans = plans.filter(_.hasNode("AppendData"))
+ val rewritePlans = plans.filter(isRewriteWrite)
assert(rewritePlans.nonEmpty, "Expected at least one rewrite plan")
assertReadsAreComet(rewritePlans)
@@ -266,11 +266,11 @@ class CometIcebergRewriteActionSuite extends
CometTestBase with CometIcebergTest
assertDataPreserved(rowsBefore, rowsAfterById, filesBefore,
filesAfter)
rc.verifyDataAfter(rowsAfterFileOrder)
- val rewritePlans = plans.filter(_.hasNode("AppendData"))
+ val rewritePlans = plans.filter(isRewriteWrite)
assert(
rewritePlans.nonEmpty,
- "Expected at least one captured plan with AppendData but got
none.\n" +
- dumpPlans(plans))
+ "Expected at least one captured plan with AppendData or
IcebergCommit but got " +
+ "none.\n" + dumpPlans(plans))
rc.verifyPlans(rewritePlans)
} catch {
case _: ClassNotFoundException =>
@@ -329,6 +329,11 @@ class CometIcebergRewriteActionSuite extends CometTestBase
with CometIcebergTest
// -- Plan assertions -------------------------------------------------------
+ // The per-group rewrite write plans as Spark's AppendData, or as Comet's
IcebergCommit when
+ // COMET_ICEBERG_WRITE_SPLIT_OPERATOR_ENABLED is on.
+ private def isRewriteWrite(plan: CapturedPlan): Boolean =
+ plan.hasNode("AppendData") || plan.hasNode("IcebergCommit")
+
private def assertReadsAreComet(plans: Seq[CapturedPlan]): Unit = {
assertOperator(plans, "CometIcebergNativeScan")
val withFallback =
diff --git
a/spark/src/test/scala/org/apache/comet/CometIcebergWriteActionSuite.scala
b/spark/src/test/scala/org/apache/comet/CometIcebergWriteActionSuite.scala
new file mode 100644
index 0000000000..41f596e5ea
--- /dev/null
+++ b/spark/src/test/scala/org/apache/comet/CometIcebergWriteActionSuite.scala
@@ -0,0 +1,979 @@
+/*
+ * 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
+
+import java.io.File
+import java.util.concurrent.{CountDownLatch, TimeUnit}
+
+import scala.collection.mutable
+import scala.concurrent.{Await, Future}
+import scala.concurrent.ExecutionContext.Implicits.global
+import scala.concurrent.duration.DurationInt
+
+import org.apache.spark.{CometListenerBusUtils, SparkConf}
+import org.apache.spark.sql.CometTestBase
+import org.apache.spark.sql.Row
+import org.apache.spark.sql.comet.{IcebergCommitExec, IcebergWriteExec}
+import org.apache.spark.sql.connector.catalog.InMemoryTableCatalog
+import org.apache.spark.sql.execution.{QueryExecution, SparkPlan}
+import org.apache.spark.sql.execution.adaptive.AdaptiveSparkPlanHelper
+import org.apache.spark.sql.types.{DoubleType, IntegerType, StringType,
StructField, StructType}
+import org.apache.spark.sql.util.QueryExecutionListener
+
+import org.apache.comet.CometSparkSessionExtensions.{isSpark35Plus,
isSpark41Plus}
+
+private case class WriteSnapshot(snapshotDelta: Long, plans: Seq[SparkPlan])
+
+class CometIcebergWriteActionSuite
+ extends CometTestBase
+ with AdaptiveSparkPlanHelper
+ with CometIcebergTestBase {
+
+ override protected def sparkConf: SparkConf = {
+ super.sparkConf
+ .set(CometConf.COMET_ICEBERG_WRITE_SPLIT_OPERATOR_ENABLED.key, "true")
+ .set(
+ "spark.sql.extensions",
+ "org.apache.iceberg.spark.extensions.IcebergSparkSessionExtensions")
+ }
+
+ test("AppendData unpartitioned INSERT INTO routes through two-op") {
+ assume(icebergAvailable, "Iceberg not available in classpath")
+ withIcebergCatalog { warehouseDir =>
+ createTable(warehouseDir, "append_unpart", partitionSpec = "")
+ val snapshot = captureWrite("append_unpart") {
+ spark.sql(
+ "INSERT INTO cat.db.append_unpart VALUES " +
+ "(1, 'us-east', 10.5), (2, 'us-west', 20.3), (3, 'eu', 30.7)")
+ }
+ assertExactlyOneCommit(snapshot)
+ assertRows("append_unpart", expectedIds = Seq(1, 2, 3))
+ }
+ }
+
+ test("AppendData partitioned INSERT INTO routes through two-op") {
+ assume(icebergAvailable, "Iceberg not available in classpath")
+ withIcebergCatalog { warehouseDir =>
+ createTable(warehouseDir, "append_part", partitionSpec = "PARTITIONED BY
(region)")
+ val snapshot = captureWrite("append_part") {
+ spark.sql(
+ "INSERT INTO cat.db.append_part VALUES " +
+ "(1, 'us-east', 10.5), (2, 'us-east', 20.3), (3, 'eu', 30.7)")
+ }
+ assertExactlyOneCommit(snapshot)
+ assertRows("append_part", expectedIds = Seq(1, 2, 3))
+ }
+ }
+
+ test("AppendData INSERT FROM SELECT survives the intervening exchange/sort")
{
+ assume(icebergAvailable, "Iceberg not available in classpath")
+ withIcebergCatalog { warehouseDir =>
+ createTable(warehouseDir, "src", partitionSpec = "")
+ createTable(warehouseDir, "append_from_select", partitionSpec =
"PARTITIONED BY (region)")
+ spark.sql(
+ "INSERT INTO cat.db.src VALUES " +
+ "(1, 'us-east', 10.5), (2, 'us-west', 20.3), (3, 'eu', 30.7)")
+
+ val snapshot = captureWrite("append_from_select") {
+ spark.sql(
+ "INSERT INTO cat.db.append_from_select " +
+ "SELECT id, region, amount FROM cat.db.src ORDER BY id")
+ }
+ assertExactlyOneCommit(snapshot)
+ assertRows("append_from_select", expectedIds = Seq(1, 2, 3))
+ }
+ }
+
+ test("AppendData on an empty source still emits a single commit") {
+ assume(icebergAvailable, "Iceberg not available in classpath")
+ withIcebergCatalog { warehouseDir =>
+ createTable(warehouseDir, "empty_target", partitionSpec = "")
+ val snapshot = captureWrite("empty_target") {
+ spark.sql(
+ "INSERT INTO cat.db.empty_target SELECT id, region, amount " +
+ "FROM (SELECT 1 AS id, 'r' AS region, 1.0 AS amount) WHERE id < 0")
+ }
+ assertExactlyOneCommit(snapshot)
+ assertRows("empty_target", expectedIds = Seq.empty)
+ }
+ }
+
+ test("AppendData from a zero-partition RDD still runs one write task and
commits") {
+ assume(icebergAvailable, "Iceberg not available in classpath")
+ withIcebergCatalog { warehouseDir =>
+ createTable(warehouseDir, "zero_part", partitionSpec = "")
+ val schema = StructType(
+ Seq(
+ StructField("id", IntegerType),
+ StructField("region", StringType),
+ StructField("amount", DoubleType)))
+ val emptyDf = spark.createDataFrame(spark.sparkContext.emptyRDD[Row],
schema)
+ assert(emptyDf.rdd.getNumPartitions == 0, "test requires a genuinely
zero-partition input")
+
+ val snapshot = captureWrite("zero_part") {
+ emptyDf.writeTo(s"$catalog.$ns.zero_part").append()
+ }
+ assertExactlyOneCommit(snapshot)
+ val commitNodes = snapshot.plans.flatMap { plan =>
+ collectWithSubqueries(plan) { case c: IcebergCommitExec => c }
+ }
+ assert(
+ commitNodes.exists(_.metrics("numCommittedMessages").value == 1),
+ "expected the dummy single-partition write task to produce one commit
message")
+ assertRows("zero_part", expectedIds = Seq.empty)
+ }
+ }
+
+ test("AQE re-plan of the writer subtree writes and commits exactly once") {
+ assume(icebergAvailable, "Iceberg not available in classpath")
+ withIcebergCatalog { warehouseDir =>
+ createTable(warehouseDir, "aqe_replan", partitionSpec = "")
+ val session = spark
+ import session.implicits._
+ (1 to 100)
+ .map(i => (i, s"r${i % 4}", i.toDouble))
+ .toDF("id", "region", "amount")
+ .createOrReplaceTempView("aqe_replan_left")
+ (1 to 100)
+ .map(i => (i, i * 10.0))
+ .toDF("id", "bonus")
+ .createOrReplaceTempView("aqe_replan_right")
+
+ // Broadcast is disabled at static planning time, so the initial plan
under the writer
+ // joins with a shuffle. AQE's runtime stats then re-plan it to a
broadcast join, which
+ // re-emits the writer subtree via IcebergWriteLogical mid-execution.
+ val snapshot = captureWrite("aqe_replan") {
+ withSQLConf(
+ "spark.sql.adaptive.enabled" -> "true",
+ "spark.sql.autoBroadcastJoinThreshold" -> "-1",
+ "spark.sql.adaptive.autoBroadcastJoinThreshold" -> "10m") {
+ spark.sql(
+ "INSERT INTO cat.db.aqe_replan " +
+ "SELECT l.id, l.region, l.amount + r.bonus " +
+ "FROM aqe_replan_left l JOIN aqe_replan_right r ON l.id = r.id")
+ }
+ }
+ assertExactlyOneCommit(snapshot)
+ val broadcastJoins = snapshot.plans.flatMap { plan =>
+ collectWithSubqueries(plan) {
+ case j if j.nodeName.contains("BroadcastHashJoin") => j
+ }
+ }
+ assert(
+ broadcastJoins.nonEmpty,
+ "expected AQE to re-plan the static shuffle join to a broadcast join.
Plans:\n" +
+ snapshot.plans.mkString("\n--\n"))
+ assertRows("aqe_replan", expectedIds = 1 to 100)
+ }
+ }
+
+ test("partitioned write under AQE keeps the clustered writer working across
the shuffle") {
+ assume(icebergAvailable, "Iceberg not available in classpath")
+ withIcebergCatalog { warehouseDir =>
+ createTable(warehouseDir, "aqe_part", partitionSpec = "PARTITIONED BY
(region)")
+ val session = spark
+ import session.implicits._
+ (1 to 500)
+ .map(i => (i, s"r${i % 8}", i.toDouble))
+ .toDF("id", "region", "amount")
+ .createOrReplaceTempView("aqe_part_src")
+
+ val snapshot = captureWrite("aqe_part") {
+ withSQLConf(
+ "spark.sql.adaptive.enabled" -> "true",
+ "spark.sql.shuffle.partitions" -> "8") {
+ spark.sql(s"INSERT INTO $catalog.$ns.aqe_part SELECT * FROM
aqe_part_src")
+ }
+ }
+ assertExactlyOneCommit(snapshot)
+ val exchanges = snapshot.plans.flatMap { plan =>
+ collectWithSubqueries(plan) { case e if
e.nodeName.contains("Exchange") => e }
+ }
+ assert(
+ exchanges.nonEmpty,
+ "expected the clustered partitioned write to shuffle its input.
Plans:\n" +
+ snapshot.plans.mkString("\n--\n"))
+ assertRows("aqe_part", expectedIds = 1 to 500)
+ }
+ }
+
+ test("multi-partition partitioned write collects one commit message per
task") {
+ assume(icebergAvailable, "Iceberg not available in classpath")
+ withIcebergCatalog { warehouseDir =>
+ createTable(warehouseDir, "multi_task", partitionSpec = "PARTITIONED BY
(region)")
+ val session = spark
+ import session.implicits._
+ (1 to 100)
+ .map(i => (i, s"r${i % 4}", i.toDouble))
+ .toDF("id", "region", "amount")
+ .createOrReplaceTempView("multi_task_src")
+
+ val snapshot = captureWrite("multi_task") {
+ withSQLConf(
+ "spark.sql.adaptive.coalescePartitions.enabled" -> "false",
+ "spark.sql.shuffle.partitions" -> "4") {
+ spark.sql(s"INSERT INTO $catalog.$ns.multi_task SELECT * FROM
multi_task_src")
+ }
+ }
+ assertExactlyOneCommit(snapshot)
+ val commitNodes = snapshot.plans.flatMap { plan =>
+ collectWithSubqueries(plan) { case c: IcebergCommitExec => c }
+ }
+ assert(
+ commitNodes.exists(_.metrics("numCommittedMessages").value >= 2),
+ "expected multiple task commit messages, got " +
+ s"${commitNodes.map(_.metrics("numCommittedMessages").value)}")
+ assertRows("multi_task", expectedIds = 1 to 100)
+ }
+ }
+
+ // Spark 4.1 recaches v2 catalog tables by name, so a cached table must pick
up an append
+ // that follows a schema-changing DDL; plan-based recaching on older Sparks
cannot match the
+ // pre-DDL cache entry, matching stock behaviour there.
+ test("cached table reflects an append that follows a schema change on Spark
4.1+") {
+ assume(icebergAvailable, "Iceberg not available in classpath")
+ assume(isSpark41Plus, "name-based cache refresh needs Spark 4.1+")
+ withIcebergCatalog { warehouseDir =>
+ createTable(warehouseDir, "cache_refresh", partitionSpec = "")
+ spark.sql(s"INSERT INTO $catalog.$ns.cache_refresh VALUES (1, 'us-east',
10.0)")
+ spark.sql(s"CACHE TABLE $catalog.$ns.cache_refresh")
+ assert(spark.sql(s"SELECT * FROM $catalog.$ns.cache_refresh").count() ==
1)
+
+ try {
+ spark.sql(s"ALTER TABLE $catalog.$ns.cache_refresh ADD COLUMN extra
INT")
+ spark.sql(s"INSERT INTO $catalog.$ns.cache_refresh VALUES (2, 'eu',
20.0, -1)")
+ // Sorted client-side: ORDER BY over an InMemoryTableScan trips an
unrelated,
+ // pre-existing Comet shuffle-planning bug regardless of this config.
+ val ids = spark
+ .sql(s"SELECT id FROM $catalog.$ns.cache_refresh")
+ .collect()
+ .map(_.getInt(0))
+ .sorted
+ .toSeq
+ assert(ids == Seq(1, 2), s"cached table missed the post-DDL append,
got $ids")
+ } finally {
+ spark.sql(s"UNCACHE TABLE IF EXISTS $catalog.$ns.cache_refresh")
+ }
+ }
+ }
+
+ test("OverwriteByExpression replaces existing rows via two-op") {
+ assume(icebergAvailable, "Iceberg not available in classpath")
+ withIcebergCatalog { warehouseDir =>
+ createTable(warehouseDir, "overwrite_static", partitionSpec = "")
+ spark.sql(
+ "INSERT INTO cat.db.overwrite_static VALUES " +
+ "(1, 'old', 1.0), (2, 'old', 2.0), (3, 'old', 3.0)")
+
+ val snapshot = captureWrite("overwrite_static") {
+ withSQLConf("spark.sql.sources.partitionOverwriteMode" -> "STATIC") {
+ spark.sql(
+ "INSERT OVERWRITE cat.db.overwrite_static VALUES " +
+ "(10, 'new', 100.0), (11, 'new', 110.0)")
+ }
+ }
+ assertExactlyOneCommit(snapshot)
+ assertRows("overwrite_static", expectedIds = Seq(10, 11))
+ }
+ }
+
+ test("OverwritePartitionsDynamic replaces only touched partitions") {
+ assume(icebergAvailable, "Iceberg not available in classpath")
+ withIcebergCatalog { warehouseDir =>
+ createTable(warehouseDir, "overwrite_dynamic", partitionSpec =
"PARTITIONED BY (region)")
+ spark.sql(
+ "INSERT INTO cat.db.overwrite_dynamic VALUES " +
+ "(1, 'us-east', 1.0), (2, 'us-west', 2.0), (3, 'eu', 3.0)")
+
+ val snapshot = captureWrite("overwrite_dynamic") {
+ withSQLConf("spark.sql.sources.partitionOverwriteMode" -> "DYNAMIC") {
+ spark.sql("INSERT OVERWRITE cat.db.overwrite_dynamic VALUES (10,
'us-east', 100.0)")
+ }
+ }
+ assertExactlyOneCommit(snapshot)
+ val ids = spark
+ .sql("SELECT id FROM cat.db.overwrite_dynamic ORDER BY id")
+ .collect()
+ .map(_.getInt(0))
+ .toSeq
+ assert(ids == Seq(2, 3, 10), s"expected (2,3,10), got $ids")
+ }
+ }
+
+ test("ReplaceData (CoW DELETE) on a row predicate goes through two-op") {
+ assume(icebergAvailable, "Iceberg not available in classpath")
+ withIcebergCatalog { warehouseDir =>
+ createTable(
+ warehouseDir,
+ "cow_delete",
+ partitionSpec = "",
+ properties = Some("'write.delete.mode'='copy-on-write'"))
+ withSQLConf(CometConf.COMET_ICEBERG_WRITE_SPLIT_OPERATOR_ENABLED.key ->
"false") {
+ coalesceInsert(
+ "cow_delete",
+ Seq((1, "us-east", 10.0), (2, "us-west", 20.0), (3, "eu", 30.0), (4,
"us-east", 40.0)))
+ }
+
+ val snapshot = captureWrite("cow_delete") {
+ spark.sql("DELETE FROM cat.db.cow_delete WHERE id = 2")
+ }
+ assertExactlyOneCommit(snapshot)
+ assertRows("cow_delete", expectedIds = Seq(1, 3, 4))
+ }
+ }
+
+ test("ReplaceData (CoW UPDATE) routes through two-op") {
+ assume(icebergAvailable, "Iceberg not available in classpath")
+ withIcebergCatalog { warehouseDir =>
+ createTable(
+ warehouseDir,
+ "cow_update",
+ partitionSpec = "",
+ properties = Some("'write.update.mode'='copy-on-write'"))
+ coalesceInsert(
+ "cow_update",
+ Seq((1, "us-east", 10.0), (2, "us-west", 20.0), (3, "eu", 30.0)))
+
+ val snapshot = captureWrite("cow_update") {
+ spark.sql("UPDATE cat.db.cow_update SET amount = amount * 2 WHERE id =
2")
+ }
+ assertExactlyOneCommit(snapshot)
+ val r = spark
+ .sql("SELECT id, amount FROM cat.db.cow_update WHERE id = 2")
+ .collect()
+ assert(r.length == 1 && r(0).getDouble(1) == 40.0, s"got ${r.toSeq}")
+ }
+ }
+
+ test("ReplaceData (CoW MERGE) with matched and unmatched legs routes through
two-op") {
+ assume(icebergAvailable, "Iceberg not available in classpath")
+ withIcebergCatalog { warehouseDir =>
+ createTable(
+ warehouseDir,
+ "cow_merge",
+ partitionSpec = "",
+ properties = Some("'write.merge.mode'='copy-on-write'"))
+ coalesceInsert("cow_merge", Seq((1, "us-east", 10.0), (2, "us-west",
20.0)))
+
+ val snapshot = captureWrite("cow_merge") {
+ spark.sql("""
+ |MERGE INTO cat.db.cow_merge t
+ |USING (SELECT 2 AS id, 'us-west' AS region, 200.0 AS amount UNION
ALL
+ | SELECT 3 AS id, 'eu' AS region, 30.0 AS amount) s
+ |ON t.id = s.id
+ |WHEN MATCHED THEN UPDATE SET t.amount = s.amount
+ |WHEN NOT MATCHED THEN INSERT (id, region, amount) VALUES (s.id,
s.region, s.amount)
+ |""".stripMargin)
+ }
+ assertExactlyOneCommit(snapshot)
+ assertRows("cow_merge", expectedIds = Seq(1, 2, 3))
+ }
+ }
+
+ // Spark 4.1 hands MERGE metrics to `BatchWrite.commit(messages, summary)`
and Iceberg 1.11+
+ // records them in the snapshot summary.
+ test("CoW MERGE records merge metrics in the snapshot summary on Spark
4.1+") {
+ assume(icebergAvailable, "Iceberg not available in classpath")
+ assume(isSpark41Plus, "MERGE write summaries need Spark 4.1+")
+ withIcebergCatalog { warehouseDir =>
+ createTable(
+ warehouseDir,
+ "merge_summary",
+ partitionSpec = "",
+ properties = Some("'write.merge.mode'='copy-on-write'"))
+ coalesceInsert("merge_summary", Seq((1, "us-east", 10.0), (2, "us-west",
20.0)))
+
+ spark.sql(s"""
+ |MERGE INTO $catalog.$ns.merge_summary t
+ |USING (SELECT 2 AS id, 'us-west' AS region, 200.0 AS amount UNION ALL
+ | SELECT 3 AS id, 'eu' AS region, 30.0 AS amount) s
+ |ON t.id = s.id
+ |WHEN MATCHED THEN UPDATE SET t.amount = s.amount
+ |WHEN NOT MATCHED THEN INSERT (id, region, amount) VALUES (s.id,
s.region, s.amount)
+ |""".stripMargin)
+
+ val summary = spark
+ .sql(s"SELECT summary FROM $catalog.$ns.merge_summary.snapshots " +
+ "ORDER BY committed_at DESC LIMIT 1")
+ .collect()(0)
+ .getMap[String, String](0)
+ val expected = Map(
+ "spark.merge-into.num-target-rows-copied" -> "1",
+ "spark.merge-into.num-target-rows-updated" -> "1",
+ "spark.merge-into.num-target-rows-inserted" -> "1")
+ expected.foreach { case (key, value) =>
+ assert(
+ summary.get(key).contains(value),
+ s"expected $key=$value in snapshot summary, got $summary")
+ }
+ }
+ }
+
+ test("failed write job aborts and leaves the table unchanged") {
+ assume(icebergAvailable, "Iceberg not available in classpath")
+ withIcebergCatalog { warehouseDir =>
+ createTable(warehouseDir, "task_fail", partitionSpec = "")
+ coalesceInsert("task_fail", Seq((1, "us-east", 10.0)))
+ val session = spark
+ import session.implicits._
+ (1 to 10)
+ .map(i => (i, s"r$i", i.toDouble))
+ .toDF("id", "region", "amount")
+ .createOrReplaceTempView("task_fail_src")
+ spark.udf.register(
+ "boom_on_seven",
+ (id: Int) => {
+ if (id == 7) throw new RuntimeException("boom")
+ id
+ })
+
+ val before = countSnapshots("task_fail")
+ val e = intercept[Exception] {
+ spark.sql(
+ s"INSERT INTO $catalog.$ns.task_fail " +
+ "SELECT boom_on_seven(id), region, amount FROM task_fail_src")
+ }
+ assert(
+ exceptionChain(e).exists(_.getMessage != null) &&
+ exceptionChain(e).exists(t =>
Option(t.getMessage).exists(_.contains("boom"))),
+ s"expected the injected task failure to surface, got $e")
+ assert(countSnapshots("task_fail") == before, "failed write must not
create a snapshot")
+ assertRows("task_fail", expectedIds = Seq(1))
+ }
+ }
+
+ test("commit-time validation still sees a conflicting concurrent append") {
+ assume(icebergAvailable, "Iceberg not available in classpath")
+ withIcebergCatalog { warehouseDir =>
+ createTable(
+ warehouseDir,
+ "conflict",
+ partitionSpec = "",
+ properties = Some(
+
"'write.delete.mode'='copy-on-write','write.delete.isolation-level'='serializable'"))
+ coalesceInsert("conflict", Seq((1, "us-east", 10.0), (2, "us-west",
20.0), (3, "eu", 30.0)))
+
+ ConflictGate.reset()
+ spark.udf.register(
+ "conflict_gate",
+ (id: Int) => {
+ ConflictGate.enter()
+ id
+ })
+
+ val before = countSnapshots("conflict")
+ // The gate UDF blocks the DELETE's tasks after its scan snapshot is
pinned, so the
+ // append below is guaranteed to land between the scan and the commit.
Runtime group
+ // filtering is disabled so the conflict is detected by Iceberg's
commit-time
+ // validation rather than aborted earlier by the runtime file filter.
+
withSQLConf("spark.sql.optimizer.runtime.rowLevelOperationGroupFilter.enabled"
-> "false") {
+ val delete = Future {
+ spark.sql(s"DELETE FROM $catalog.$ns.conflict WHERE
conflict_gate(id) = 2")
+ }
+ assert(
+ ConflictGate.awaitScanStarted(2, TimeUnit.MINUTES),
+ "DELETE never started scanning; gate UDF was not invoked")
+ spark.sql(s"INSERT INTO $catalog.$ns.conflict VALUES (2, 'us-west',
99.0)")
+ ConflictGate.releaseWrite()
+
+ val e = intercept[Exception] {
+ Await.result(delete, 2.minutes)
+ }
+ assert(
+ exceptionChain(e).exists(t =>
+ Option(t.getMessage).exists(_.toLowerCase.contains("conflict"))),
+ s"expected Iceberg commit-time validation to fail the DELETE, got
$e")
+ }
+ assert(
+ countSnapshots("conflict") == before + 1,
+ "only the concurrent append may commit; the DELETE must not")
+ val ids = spark
+ .sql(s"SELECT id FROM $catalog.$ns.conflict ORDER BY id")
+ .collect()
+ .map(_.getInt(0))
+ .toSeq
+ assert(ids == Seq(1, 2, 2, 3), s"expected (1,2,2,3), got $ids")
+ }
+ }
+
+ test("non-Iceberg V2 write plans through Spark unchanged with the config
on") {
+ withSQLConf(
+ "spark.sql.catalog.testcat" -> classOf[InMemoryTableCatalog].getName,
+ CometConf.COMET_ENABLED.key -> "true",
+ CometConf.COMET_EXEC_ENABLED.key -> "true") {
+ spark.sql("CREATE TABLE testcat.tbl (id INT, region STRING, amount
DOUBLE)")
+ try {
+ val plans = capturePlans {
+ spark.sql("INSERT INTO testcat.tbl VALUES (1, 'us-east', 10.5)")
+ }
+ val (commits, writes) = collectIcebergWriteOps(plans)
+ assert(commits.isEmpty, s"unexpected IcebergCommitExec on a
non-Iceberg write: $commits")
+ assert(writes.isEmpty, s"unexpected IcebergWriteExec on a non-Iceberg
write: $writes")
+ val ids = spark.sql("SELECT id FROM
testcat.tbl").collect().map(_.getInt(0)).toSeq
+ assert(ids == Seq(1), s"expected (1), got $ids")
+ } finally {
+ spark.sql("DROP TABLE IF EXISTS testcat.tbl")
+ }
+ }
+ }
+
+ test("sanity check: Spark's default DELETE path works against a Hadoop
catalog") {
+ assume(icebergAvailable, "Iceberg not available in classpath")
+ withIcebergCatalog { warehouseDir =>
+ withSQLConf(CometConf.COMET_ICEBERG_WRITE_SPLIT_OPERATOR_ENABLED.key ->
"false") {
+ createTable(
+ warehouseDir,
+ "spark_cow_delete",
+ partitionSpec = "",
+ properties = Some("'write.delete.mode'='copy-on-write'"))
+ coalesceInsert(
+ "spark_cow_delete",
+ Seq((1, "us-east", 10.0), (2, "us-west", 20.0), (3, "eu", 30.0), (4,
"us-east", 40.0)))
+ spark.sql("DELETE FROM cat.db.spark_cow_delete WHERE id = 2")
+ assertRows("spark_cow_delete", expectedIds = Seq(1, 3, 4))
+ }
+ }
+ }
+
+ test("disabled config falls through to Spark's V2ExistingTableWriteExec") {
+ assume(icebergAvailable, "Iceberg not available in classpath")
+ withIcebergCatalog { warehouseDir =>
+ createTable(warehouseDir, "disabled_conf", partitionSpec = "")
+
+ val snapshot = captureWrite("disabled_conf") {
+ withSQLConf(CometConf.COMET_ICEBERG_WRITE_SPLIT_OPERATOR_ENABLED.key
-> "false") {
+ spark.sql("INSERT INTO cat.db.disabled_conf VALUES (1, 'us-east',
10.5)")
+ }
+ }
+ val (commits, writes) = collectIcebergWriteOps(snapshot.plans)
+ assert(commits.isEmpty, s"unexpected IcebergCommitExec: $commits")
+ assert(writes.isEmpty, s"unexpected IcebergWriteExec: $writes")
+ assertRows("disabled_conf", expectedIds = Seq(1))
+ }
+ }
+
+ test("Comet-written rows round-trip through Spark's reader unchanged") {
+ assume(icebergAvailable, "Iceberg not available in classpath")
+ withIcebergCatalog { warehouseDir =>
+ createTable(warehouseDir, "parity_comet", partitionSpec = "PARTITIONED
BY (region)")
+ createTable(warehouseDir, "parity_spark", partitionSpec = "PARTITIONED
BY (region)")
+
+ spark.sql(
+ "INSERT INTO cat.db.parity_comet VALUES " +
+ "(1, 'us', 1.5), (2, 'eu', 2.5), (3, 'ap', 3.5), (4, 'us', 4.5)")
+
+ withSQLConf(CometConf.COMET_ICEBERG_WRITE_SPLIT_OPERATOR_ENABLED.key ->
"false") {
+ spark.sql(
+ "INSERT INTO cat.db.parity_spark VALUES " +
+ "(1, 'us', 1.5), (2, 'eu', 2.5), (3, 'ap', 3.5), (4, 'us', 4.5)")
+ }
+
+ val cometRows: Array[Row] = spark
+ .sql("SELECT id, region, amount FROM cat.db.parity_comet ORDER BY id")
+ .collect()
+ val sparkRows: Array[Row] = spark
+ .sql("SELECT id, region, amount FROM cat.db.parity_spark ORDER BY id")
+ .collect()
+ assert(cometRows.toSeq == sparkRows.toSeq, s"$cometRows vs $sparkRows")
+ }
+ }
+
+ test("write custom metrics are registered on the committer only") {
+ assume(icebergAvailable, "Iceberg not available in classpath")
+ withIcebergCatalog { warehouseDir =>
+ createTable(warehouseDir, "metrics_once", partitionSpec = "")
+ val snapshot = captureWrite("metrics_once") {
+ spark.sql(s"INSERT INTO $catalog.$ns.metrics_once VALUES (1,
'us-east', 10.5)")
+ }
+ assertExactlyOneCommit(snapshot)
+ val (commits, writes) = collectIcebergWriteOps(snapshot.plans)
+ writes.foreach { w =>
+ assert(
+ w.metrics.keySet == Set("numOutputRows"),
+ s"IcebergWriteExec must not re-register the write's custom metrics,
got ${w.metrics.keySet}")
+ }
+ // Iceberg publishes write custom metrics from 1.9 on; older versions
expose none.
+ commits.head.metrics.get("addedDataFiles").foreach { m =>
+ assert(m.value == 1L, s"expected addedDataFiles=1 on the committer,
got ${m.value}")
+ }
+ }
+ }
+
+ test("AppendData round-trips all primitive types unchanged") {
+ assume(icebergAvailable, "Iceberg not available in classpath")
+ withIcebergCatalog { _ =>
+ spark.sql(s"""
+ CREATE TABLE $catalog.$ns.prim_types (
+ b BOOLEAN, ti TINYINT, si SMALLINT, i INT, bi BIGINT,
+ f FLOAT, d DOUBLE, dec DECIMAL(18,4),
+ s STRING, bin BINARY, dt DATE, ts TIMESTAMP
+ ) USING iceberg""")
+ val literals =
+ "(true, CAST(1 AS TINYINT), CAST(2 AS SMALLINT), 3, CAST(4 AS BIGINT),
" +
+ "CAST(1.5 AS FLOAT), 2.5D, CAST(12.3456 AS DECIMAL(18,4)), " +
+ "'hello', X'0102', DATE'2024-01-15', TIMESTAMP'2024-01-15
10:00:00'), " +
+ "(false, CAST(NULL AS TINYINT), CAST(NULL AS SMALLINT), CAST(NULL AS
INT), " +
+ "CAST(NULL AS BIGINT), CAST(NULL AS FLOAT), CAST(NULL AS DOUBLE), " +
+ "CAST(NULL AS DECIMAL(18,4)), CAST(NULL AS STRING), CAST(NULL AS
BINARY), " +
+ "CAST(NULL AS DATE), CAST(NULL AS TIMESTAMP))"
+ val snapshot = captureWrite("prim_types") {
+ spark.sql(s"INSERT INTO $catalog.$ns.prim_types VALUES $literals")
+ }
+ assertExactlyOneCommit(snapshot)
+ checkAnswer(
+ spark.sql(s"SELECT * FROM $catalog.$ns.prim_types"),
+ spark.sql(s"SELECT * FROM VALUES $literals"))
+ }
+ }
+
+ test("AppendData writes NULLs in every column") {
+ assume(icebergAvailable, "Iceberg not available in classpath")
+ withIcebergCatalog { warehouseDir =>
+ createTable(warehouseDir, "nulls_all", partitionSpec = "")
+ val snapshot = captureWrite("nulls_all") {
+ spark.sql(s"""INSERT INTO $catalog.$ns.nulls_all VALUES
+ (CAST(NULL AS INT), CAST(NULL AS STRING), CAST(NULL AS DOUBLE)),
+ (1, NULL, 1.0),
+ (2, 'r', NULL)""")
+ }
+ assertExactlyOneCommit(snapshot)
+ checkAnswer(
+ spark.sql(s"SELECT id, region, amount FROM $catalog.$ns.nulls_all"),
+ Seq(Row(null, null, null), Row(1, null, 1.0), Row(2, "r", null)))
+ }
+ }
+
+ test("AppendData round-trips STRUCT, ARRAY, and MAP columns") {
+ assume(icebergAvailable, "Iceberg not available in classpath")
+ withIcebergCatalog { _ =>
+ spark.sql(s"""
+ CREATE TABLE $catalog.$ns.nested_types (
+ id INT,
+ s STRUCT<a: INT, b: STRING>,
+ arr ARRAY<INT>,
+ m MAP<STRING, INT>
+ ) USING iceberg""")
+ val snapshot = captureWrite("nested_types") {
+ spark.sql(s"""INSERT INTO $catalog.$ns.nested_types VALUES
+ (1, named_struct('a', 10, 'b', 'x'), array(1, 2, 3), map('k1', 1,
'k2', 2)),
+ (2, named_struct('a', CAST(NULL AS INT), 'b', CAST(NULL AS STRING)),
+ CAST(NULL AS ARRAY<INT>), CAST(NULL AS MAP<STRING, INT>))""")
+ }
+ assertExactlyOneCommit(snapshot)
+ checkAnswer(
+ spark.sql(s"SELECT id, s.a, s.b, arr, m FROM
$catalog.$ns.nested_types"),
+ Seq(
+ Row(1, 10, "x", Seq(1, 2, 3), Map("k1" -> 1, "k2" -> 2)),
+ Row(2, null, null, null, null)))
+ }
+ }
+
+ test("CoW MERGE with only a matched DELETE leg routes through two-op") {
+ assume(icebergAvailable, "Iceberg not available in classpath")
+ withIcebergCatalog { warehouseDir =>
+ createTable(
+ warehouseDir,
+ "merge_del",
+ partitionSpec = "",
+ properties = Some("'write.merge.mode'='copy-on-write'"))
+ coalesceInsert("merge_del", Seq((1, "a", 1.0), (2, "b", 2.0), (3, "c",
3.0)))
+
+ val snapshot = captureWrite("merge_del") {
+ spark.sql(s"""
+ |MERGE INTO $catalog.$ns.merge_del t
+ |USING (SELECT 2 AS id UNION ALL SELECT 3 AS id) s
+ |ON t.id = s.id
+ |WHEN MATCHED THEN DELETE
+ |""".stripMargin)
+ }
+ assertExactlyOneCommit(snapshot)
+ assertRows("merge_del", expectedIds = Seq(1))
+ }
+ }
+
+ test("CoW MERGE with a NOT MATCHED BY SOURCE leg routes through two-op") {
+ assume(icebergAvailable, "Iceberg not available in classpath")
+ // Iceberg 1.5.x (the Spark 3.4 pairing) rejects the clause in its
extensions parser.
+ assume(isSpark35Plus, "NOT MATCHED BY SOURCE needs Iceberg 1.8+")
+ withIcebergCatalog { warehouseDir =>
+ createTable(
+ warehouseDir,
+ "merge_nmbs",
+ partitionSpec = "",
+ properties = Some("'write.merge.mode'='copy-on-write'"))
+ coalesceInsert("merge_nmbs", Seq((1, "a", 1.0), (2, "b", 2.0), (3, "c",
3.0)))
+
+ val snapshot = captureWrite("merge_nmbs") {
+ spark.sql(s"""
+ |MERGE INTO $catalog.$ns.merge_nmbs t
+ |USING (SELECT 2 AS id) s
+ |ON t.id = s.id
+ |WHEN MATCHED THEN UPDATE SET t.amount = t.amount + 100
+ |WHEN NOT MATCHED BY SOURCE THEN DELETE
+ |""".stripMargin)
+ }
+ assertExactlyOneCommit(snapshot)
+ assertRows("merge_nmbs", expectedIds = Seq(2))
+ checkAnswer(spark.sql(s"SELECT amount FROM $catalog.$ns.merge_nmbs"),
Seq(Row(102.0)))
+ }
+ }
+
+ test("AppendData to a bucket- and truncate-partitioned table routes through
two-op") {
+ assume(icebergAvailable, "Iceberg not available in classpath")
+ withIcebergCatalog { _ =>
+ spark.sql(s"""
+ CREATE TABLE $catalog.$ns.bucketed (
+ id INT, region STRING, amount DOUBLE
+ ) USING iceberg PARTITIONED BY (bucket(4, id), truncate(2, region))""")
+ val snapshot = captureWrite("bucketed") {
+ spark.sql(
+ s"INSERT INTO $catalog.$ns.bucketed VALUES " +
+ "(1, 'aa', 1.0), (2, 'ab', 2.0), (3, 'bb', 3.0), (4, 'bc', 4.0),
(5, 'cc', 5.0)")
+ }
+ assertExactlyOneCommit(snapshot)
+ assertRows("bucketed", expectedIds = Seq(1, 2, 3, 4, 5))
+ }
+ }
+
+ test("AppendData to a days()-partitioned table routes each day to its own
partition") {
+ assume(icebergAvailable, "Iceberg not available in classpath")
+ withIcebergCatalog { _ =>
+ spark.sql(s"""
+ CREATE TABLE $catalog.$ns.by_day (
+ id INT, ts TIMESTAMP, amount DOUBLE
+ ) USING iceberg PARTITIONED BY (days(ts))""")
+ val snapshot = captureWrite("by_day") {
+ spark.sql(
+ s"INSERT INTO $catalog.$ns.by_day VALUES " +
+ "(1, TIMESTAMP'2024-01-15 10:00:00', 1.0), " +
+ "(2, TIMESTAMP'2024-01-15 10:00:00', 2.0), " +
+ "(3, TIMESTAMP'2024-02-15 10:00:00', 3.0)")
+ }
+ assertExactlyOneCommit(snapshot)
+ val partitions = spark
+ .sql(s"SELECT count(*) FROM $catalog.$ns.by_day.partitions")
+ .collect()
+ .head
+ .getLong(0)
+ assert(partitions == 2, s"expected 2 day partitions, got $partitions")
+ val ids = spark
+ .sql(s"SELECT id FROM $catalog.$ns.by_day ORDER BY id")
+ .collect()
+ .map(_.getInt(0))
+ .toSeq
+ assert(ids == Seq(1, 2, 3), s"expected Seq(1, 2, 3), got $ids")
+ }
+ }
+
+ // On Spark 3.5+ the staged CTAS/RTAS operators run their inner append as
its own
+ // `AppendData` QueryExecution, which IcebergWriteStrategy intercepts like
any other append.
+ // Spark 3.4 writes inline inside the exec (no re-planning), so nothing is
intercepted there.
+ test("CTAS and RTAS write through the split operators on Spark 3.5+") {
+ assume(icebergAvailable, "Iceberg not available in classpath")
+ withIcebergCatalog { _ =>
+ val session = spark
+ import session.implicits._
+ (1 to 5)
+ .map(i => (i, s"r${i % 2}", i.toDouble))
+ .toDF("id", "region", "amount")
+ .createOrReplaceTempView("ctas_src")
+
+ def assertSplitUsage(plans: Seq[SparkPlan], statement: String): Unit = {
+ val (commits, writes) = collectIcebergWriteOps(plans)
+ if (isSpark35Plus) {
+ assert(
+ commits.nonEmpty && writes.nonEmpty,
+ s"expected the $statement inner append to plan through the split
operators: $plans")
+ } else {
+ assert(
+ commits.isEmpty && writes.isEmpty,
+ s"expected the $statement write to stay inside Spark's staged
exec: $plans")
+ }
+ }
+
+ val ctasPlans = capturePlans {
+ spark.sql(s"CREATE TABLE $catalog.$ns.ctas_tgt USING iceberg AS SELECT
* FROM ctas_src")
+ }
+ assertSplitUsage(ctasPlans, "CTAS")
+ assert(countSnapshots("ctas_tgt") == 1L, "CTAS must land exactly one
snapshot")
+ assertRows("ctas_tgt", expectedIds = Seq(1, 2, 3, 4, 5))
+
+ val rtasPlans = capturePlans {
+ (1 to 2)
+ .map(i => (i, s"r$i", i.toDouble))
+ .toDF("id", "region", "amount")
+ .writeTo(s"$catalog.$ns.ctas_tgt")
+ .using("iceberg")
+ .createOrReplace()
+ }
+ assertSplitUsage(rtasPlans, "RTAS")
+ assertRows("ctas_tgt", expectedIds = Seq(1, 2))
+ }
+ }
+
+ private val catalog = "cat"
+ private val ns = "db"
+
+ private def withIcebergCatalog(f: File => Unit): Unit = withTempIcebergDir {
warehouseDir =>
+ withSQLConf(
+ s"spark.sql.catalog.$catalog" -> "org.apache.iceberg.spark.SparkCatalog",
+ s"spark.sql.catalog.$catalog.type" -> "hadoop",
+ s"spark.sql.catalog.$catalog.warehouse" -> warehouseDir.getAbsolutePath,
+ CometConf.COMET_ENABLED.key -> "true",
+ CometConf.COMET_EXEC_ENABLED.key -> "true") {
+ f(warehouseDir)
+ }
+ }
+
+ private def createTable(
+ warehouseDir: File,
+ tableName: String,
+ partitionSpec: String,
+ properties: Option[String] = None): Unit = {
+ val props = properties.map(s => s" TBLPROPERTIES ($s)").getOrElse("")
+ spark.sql(s"""
+ CREATE TABLE $catalog.$ns.$tableName (
+ id INT,
+ region STRING,
+ amount DOUBLE
+ ) USING iceberg
+ $partitionSpec
+ $props
+ """)
+ }
+
+ private def coalesceInsert(tableName: String, rows: Seq[(Int, String,
Double)]): Unit = {
+ val session = spark
+ import session.implicits._
+ rows
+ .toDF("id", "region", "amount")
+ .coalesce(1)
+ .writeTo(s"$catalog.$ns.$tableName")
+ .append()
+ }
+
+ private def capturePlans(action: => Unit): Seq[SparkPlan] = {
+ val captured = mutable.Buffer.empty[SparkPlan]
+ val listener = new QueryExecutionListener {
+ override def onSuccess(funcName: String, qe: QueryExecution, durationNs:
Long): Unit = {
+ captured += qe.executedPlan
+ }
+ override def onFailure(funcName: String, qe: QueryExecution, exception:
Exception): Unit =
+ ()
+ }
+ spark.listenerManager.register(listener)
+ try {
+ action
+ CometListenerBusUtils.waitUntilEmpty(spark.sparkContext)
+ } finally {
+ spark.listenerManager.unregister(listener)
+ }
+ captured.toSeq
+ }
+
+ private def captureWrite(tableName: String)(action: => Unit): WriteSnapshot
= {
+ val before = countSnapshots(tableName)
+ val plans = capturePlans(action)
+ WriteSnapshot(countSnapshots(tableName) - before, plans)
+ }
+
+ private def countSnapshots(tableName: String): Long =
+ try {
+ spark
+ .sql(s"SELECT count(*) FROM $catalog.$ns.$tableName.snapshots")
+ .collect()
+ .head
+ .getLong(0)
+ } catch {
+ case _: Throwable => 0L
+ }
+
+ private def collectIcebergWriteOps(
+ plans: Seq[SparkPlan]): (Seq[IcebergCommitExec], Seq[IcebergWriteExec])
= {
+ val commits = plans.flatMap { plan =>
+ collectWithSubqueries(plan) { case c: IcebergCommitExec => c }
+ }
+ val writes = plans.flatMap { plan =>
+ collectWithSubqueries(plan) { case w: IcebergWriteExec => w }
+ }
+ (commits, writes)
+ }
+
+ private def assertExactlyOneCommit(snapshot: WriteSnapshot): Unit = {
+ assert(
+ snapshot.snapshotDelta == 1L,
+ s"expected exactly 1 new Iceberg snapshot, got
${snapshot.snapshotDelta}. Plans:\n" +
+ snapshot.plans.mkString("\n--\n"))
+ val (commits, writes) = collectIcebergWriteOps(snapshot.plans)
+ assert(
+ commits.nonEmpty,
+ s"expected >= 1 IcebergCommitExec in captured plans, got
${commits.size}. Plans:\n" +
+ snapshot.plans.mkString("\n--\n"))
+ assert(
+ writes.nonEmpty,
+ s"expected >= 1 IcebergWriteExec in captured plans, got ${writes.size}.
Plans:\n" +
+ snapshot.plans.mkString("\n--\n"))
+ }
+
+ private def assertRows(tableName: String, expectedIds: Seq[Int]): Unit = {
+ val ids = spark
+ .sql(s"SELECT id FROM $catalog.$ns.$tableName ORDER BY id")
+ .collect()
+ .map(_.getInt(0))
+ .toSeq
+ assert(ids == expectedIds, s"expected $expectedIds, got $ids")
+ }
+
+ private def exceptionChain(t: Throwable): Seq[Throwable] = {
+ val chain = mutable.Buffer.empty[Throwable]
+ var current = t
+ while (current != null && !chain.contains(current)) {
+ chain += current
+ current = current.getCause
+ }
+ chain.toSeq
+ }
+
+}
+
+/**
+ * Blocks the DELETE's write job between its scan-snapshot pin and its commit
so the test can
+ * inject a conflicting commit. Top-level so the UDF closure doesn't capture
the suite.
+ */
+private object ConflictGate {
+ @volatile private var scanStarted = new CountDownLatch(1)
+ @volatile private var writeReleased = new CountDownLatch(1)
+
+ def reset(): Unit = {
+ scanStarted = new CountDownLatch(1)
+ writeReleased = new CountDownLatch(1)
+ }
+
+ def enter(): Unit = {
+ scanStarted.countDown()
+ if (!writeReleased.await(2, TimeUnit.MINUTES)) {
+ throw new IllegalStateException("ConflictGate was never released")
+ }
+ }
+
+ def awaitScanStarted(timeout: Long, unit: TimeUnit): Boolean =
scanStarted.await(timeout, unit)
+
+ def releaseWrite(): Unit = writeReleased.countDown()
+}
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
