diff --git a/.github/workflows/pr_build_linux.yml b/.github/workflows/pr_build_linux.yml index 0a71607ead..5e055ed152 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 2335588b70..b91079f469 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 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 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 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 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 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 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 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..36a27bbbbc --- /dev/null +++ b/docs/source/user-guide/latest/iceberg-writes.md @@ -0,0 +1,95 @@ + + +# 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.=org.apache.iceberg.spark.SparkCatalog +spark.sql.catalog..type=hadoop # or hive / glue / rest / ... +spark.sql.catalog..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 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 Writes S3 Credential Providers Kubernetes Guide 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 f867d4bc9b..aa90e511a1 100644 --- a/spark/src/main/scala/org/apache/comet/iceberg/IcebergReflection.scala +++ b/spark/src/main/scala/org/apache/comet/iceberg/IcebergReflection.scala @@ -54,6 +54,10 @@ object IcebergReflection extends Logging { val SPARK_SCHEMA_UTIL = "org.apache.iceberg.spark.SparkSchemaUtil" val TABLE = "org.apache.iceberg.Table" val PARTITIONING = "org.apache.iceberg.Partitioning" + val SPARK_WRITE = "org.apache.iceberg.spark.source.SparkWrite" + + // Iceberg 1.5.2 uses its own `ReplaceIcebergData` due to lack of `ReplaceData` in Spark 3.4. + val REPLACE_ICEBERG_DATA = "org.apache.spark.sql.catalyst.plans.logical.ReplaceIcebergData" } /** @@ -128,6 +132,45 @@ object IcebergReflection extends Logging { val UNKNOWN = "unknown" } + /** Loads a class, returning `None` when it's absent (e.g. Iceberg not on the classpath). */ + private def tryLoadClass(name: String): Option[Class[_]] = + try Some(loadClass(name)) + catch { case _: ClassNotFoundException => None } + + private lazy val sparkWriteClassOpt: Option[Class[_]] = tryLoadClass(ClassNames.SPARK_WRITE) + + /** Whether `write` is an Iceberg `SparkWrite` (false if Iceberg isn't on the classpath). */ + def isIcebergSparkWrite(write: Any): Boolean = + sparkWriteClassOpt.exists(_.isInstance(write)) + + def isReplaceIcebergData(plan: Any): Boolean = + plan != null && plan.getClass.getName == ClassNames.REPLACE_ICEBERG_DATA + + private def reflectField(plan: Any, fieldName: String): Option[AnyRef] = + try { + val field = plan.getClass.getDeclaredField(fieldName) + field.setAccessible(true) + Option(field.get(plan)) + } catch { + case e: Exception => + logError( + s"Iceberg reflection failure: $fieldName on ${plan.getClass.getName}: ${e.getMessage}") + None + } + + def extractReplaceIcebergDataFields(plan: Any): Option[(AnyRef, AnyRef, AnyRef, AnyRef)] = { + if (!isReplaceIcebergData(plan)) return None + for { + table <- reflectField(plan, "table") + query <- reflectField(plan, "query") + originalTable <- reflectField(plan, "originalTable") + write <- reflectField( + plan, + "write" + ) // Option[Write]; field can be Some(null) so kept AnyRef + } yield (table, query, originalTable, write) + } + /** * Loads a class using the thread context classloader first, then falls back to the system * classloader. 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..1a56e51002 --- /dev/null +++ b/spark/src/main/scala/org/apache/comet/iceberg/IcebergWriteLogical.scala @@ -0,0 +1,44 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +package org.apache.comet.iceberg + +import org.apache.spark.sql.catalyst.expressions.{Attribute, AttributeReference} +import org.apache.spark.sql.catalyst.plans.logical.{LogicalPlan, UnaryNode} +import org.apache.spark.sql.comet.IcebergWriteExec +import org.apache.spark.sql.connector.write.{BatchWrite, Write} +import org.apache.spark.sql.types.BinaryType + +/** Logical anchor for the writer. See `IcebergWriteStrategy` for the rationale. */ +case class IcebergWriteLogical( + child: LogicalPlan, + // Driver-side only: AQE re-planning is driver-local and write commands aren't cached. + @transient batchWrite: BatchWrite, + @transient write: Write, + replaceDataDispatch: Option[ReplaceDataDispatchInfo] = None) + extends UnaryNode { + + // Owns the commit-message attribute so the physical writer keeps the same exprId across + // AQE re-plans. + override val output: Seq[Attribute] = Seq( + AttributeReference(IcebergWriteExec.CommitMessageColumn, BinaryType, nullable = false)()) + + override protected def withNewChildInternal(newChild: LogicalPlan): IcebergWriteLogical = + copy(child = newChild) +} 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..a7f409f988 --- /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, write, replaceDataDispatch) => + Seq(IcebergWriteExec(batchWrite, write, l.output, planLater(child), replaceDataDispatch)) + case _ => Nil + } + } + + private def matchedSparkWrite( + table: org.apache.spark.sql.catalyst.analysis.NamedRelation, + write: Option[Write], + query: LogicalPlan, + replaceDataDispatch: Option[ReplaceDataDispatchInfo]): Option[SparkPlan] = { + table match { + case rel: DataSourceV2Relation => + write.flatMap { w => + if (IcebergReflection.isIcebergSparkWrite(w)) { + buildTwoOp(w, rel, query, replaceDataDispatch) + } else { + None + } + } + case _ => None + } + } + + /** + * Builds the two-op tree. The committer and writer share one `BatchWrite` (also reused across + * AQE re-plans): `toBatch()` returns a fresh instance per call, but the committer's commit-time + * validation must see the same instance the writer wrote through, hence we store it. The + * writer's child is wrapped in [[IcebergWriteLogical]] so AQE re-emits only the data-writing + * operator on each re-plan as opposed to multiple new commit operators. + * + * Iceberg's `SparkWrite` never asks for Spark's commit coordinator, so the + * `useCommitCoordinator` fallback below is defensive coverage in case a future Iceberg version + * changes that; the split writer's per-task commit protocol does not use it. + */ + private def buildTwoOp( + write: Write, + rel: DataSourceV2Relation, + query: LogicalPlan, + replaceDataDispatch: Option[ReplaceDataDispatchInfo]): Option[SparkPlan] = { + val batchWrite = write.toBatch + if (batchWrite.useCommitCoordinator()) { + return None + } + // To mirror Spark ReplaceData semantics we invalidate our cache of the state of + // `originalTable`. + val refresh: () => Unit = () => IcebergRefreshCacheShim.refreshCache(session, rel) + Some( + IcebergCommitExec( + batchWrite, + write, + refresh, + // `replaceDataDispatch` may project the data into the format the writer expects. + planLater(IcebergWriteLogical(query, batchWrite, write, 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..791a6623a8 --- /dev/null +++ b/spark/src/main/scala/org/apache/spark/sql/comet/IcebergWriteExec.scala @@ -0,0 +1,202 @@ +/* + * 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, Write, WriterCommitMessage} +import org.apache.spark.sql.execution.{SparkPlan, UnaryExecNode} +import org.apache.spark.sql.execution.metric.{CustomMetrics, 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, + // Driver-side only; used to declare the write's custom task metrics. + @transient write: Write, + 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) + + override lazy val metrics: Map[String, SQLMetric] = + Map("numOutputRows" -> SQLMetrics.createMetric(sparkContext, "number of output rows")) ++ + write + .supportedCustomMetrics() + .map(m => m.name -> SQLMetrics.createV2CustomMetric(sparkContext, m)) + + 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 customMetrics = metrics.filter { case (name, _) => name != "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, + customMetrics, + 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, + customMetrics: Map[String, SQLMetric], + projection: UnsafeProjection, + replaceDataDispatch: Option[ReplaceDataDispatchInfo]): Iterator[InternalRow] = { + val iterWithMetrics = new IteratorWithMetrics(iter, writer, customMetrics, rowsMetric) + val message = Utils.tryWithSafeFinallyAndFailureCallbacks(block = { + if (replaceDataDispatch.isDefined) { + runReplaceDataWriter(writer, iterWithMetrics, replaceDataDispatch.get) + } else { + while (iterWithMetrics.hasNext) { + writer.write(iterWithMetrics.next()) + } + } + CustomMetrics.updateMetrics(writer.currentMetricsValues.toSeq, customMetrics) + writer.commit() + })( + catchBlock = { + writer.abort() + }, + finallyBlock = { + writer.close() + }) + + Iterator.single(projection(InternalRow(serializeMessage(message))).copy()) + } + + private class IteratorWithMetrics( + iter: Iterator[InternalRow], + dataWriter: DataWriter[InternalRow], + customMetrics: Map[String, SQLMetric], + rowsMetric: SQLMetric) + extends Iterator[InternalRow] { + private var count = 0L + + override def hasNext: Boolean = iter.hasNext + + override def next(): InternalRow = { + if (count % CustomMetrics.NUM_ROWS_PER_UPDATE == 0) { + CustomMetrics.updateMetrics(dataWriter.currentMetricsValues.toSeq, customMetrics) + } + count += 1 + 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 + + // Spark has different `DataWriter#write` methods across versions. + @transient private lazy val dataWriterWriteWithMetadataMethod + : Option[java.lang.reflect.Method] = + try Some(classOf[DataWriter[_]].getMethod("write", classOf[Object], classOf[Object])) + catch { case _: NoSuchMethodException => None } + + 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) + val writeWithMetadata = dataWriterWriteWithMetadataMethod.getOrElse( + throw new UnsupportedOperationException( + "DataWriter.write(metadata, row) is not available in this Spark version but the " + + s"analyzer emitted operation code $WRITE_WITH_METADATA_OPERATION")) + writeWithMetadata.invoke(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/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/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..6ffb82fd0d --- /dev/null +++ b/spark/src/test/scala/org/apache/comet/CometIcebergWriteActionSuite.scala @@ -0,0 +1,746 @@ +/* + * 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.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") + } + } + + 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() +}