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/docs/source/user-guide/latest/iceberg-writes.md b/docs/source/user-guide/latest/iceberg-writes.md new file mode 100644 index 0000000000..593cdcb69b --- /dev/null +++ b/docs/source/user-guide/latest/iceberg-writes.md @@ -0,0 +1,92 @@ + + +# Iceberg Writes: Comet's Split-Operator Plan (Experimental) + +**This feature is experimental and enabled by default.** Set +`spark.comet.write.iceberg.splitOperator.enabled=false` to restore Spark's stock combined +write operator. + +## Overview + +Spark writes an Iceberg table through a single physical operator that combines data-file +writing with metadata writing, committing, and catalog validation. Because that operator sits +outside Spark's Adaptive Query Execution (AQE), the sub-query feeding the write — the scans, +projects, sorts, and exchanges producing the rows — cannot be re-planned at runtime. + +When `spark.comet.write.iceberg.splitOperator.enabled=true`, Comet rewrites eligible Iceberg +writes into two operators: + +1. **`IcebergWrite`** — writes the data files on the executors, exactly as iceberg-java does + today, and returns each task's serialized commit message. This operator and the sub-query + feeding it run inside AQE. +2. **`IcebergCommit`** — collects the commit messages on the driver and performs the normal + Iceberg commit (including commit-time validation), outside AQE, exactly once. + +Data files are still written by iceberg-java; only the plan shape changes. The split makes the +write's input visible to AQE and to Comet's columnar rules, and it is the groundwork for a +planned follow-up in which Comet writes the data files natively via +[iceberg-rust](https://github.com/apache/iceberg-rust). + +## Configuration + +Standard Comet + Iceberg setup (see [`iceberg.md`](iceberg.md)) is all that is required; the +split-operator plan is applied automatically. To turn it off: + +``` +# Standard Comet / Iceberg wiring +spark.plugins=org.apache.spark.CometPlugin +spark.sql.extensions=org.apache.iceberg.spark.extensions.IcebergSparkSessionExtensions +spark.sql.catalog.=org.apache.iceberg.spark.SparkCatalog +spark.sql.catalog..type=hadoop # or hive / glue / rest / ... +spark.sql.catalog..warehouse=... + +# Split-operator plan (experimental, on by default); set to false to opt out +spark.comet.write.iceberg.splitOperator.enabled=false +``` + +## Supported operations + +The split-operator plan supports the following operations on every Spark version Comet +supports: + +- `INSERT INTO` / DataFrame `append` (`AppendData`) +- `INSERT OVERWRITE`, static and dynamic (`OverwriteByExpression`, + `OverwritePartitionsDynamic`) +- Copy-on-write `DELETE` / `UPDATE` / `MERGE` (`ReplaceData`) + +The mechanism behind row-level DML differs by Spark version: on Spark 4.0+ the analyzer emits +operation-coded rows that Comet's writer dispatches through `ReplaceData`'s projections, while +on Spark 3.4/3.5 the rewritten rows are written as a plain row stream. The supported set of +operations is the same either way. + +## When Comet falls back to Spark's write operator + +The rewrite is skipped — and the write runs through Spark's stock combined operator — when: + +- `spark.comet.write.iceberg.splitOperator.enabled` is `false`; +- the write is not an Iceberg `SparkWrite` (any other V2 data source); +- the table uses merge-on-read: delta writes (Iceberg `WriteDelta`) are not intercepted; +- the write requires Spark's commit coordinator, which Comet's per-task commit protocol does + not use; +- Comet cannot reflect the Iceberg internals needed to build the two-operator plan (for + example an unrecognised write class or a `ReplaceData` projection it cannot map). + +In every fallback case the write is planned as if Comet were absent; there is no correctness +trade-off, only no plan change. 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 e9cf8ca010..95ff23eee2 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(true) + 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 ca4e1e4e79..f74f653fd2 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..23a619c010 --- /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.recacheByPlan(session, rel) + Some( + IcebergCommitExec( + batchWrite, + write, + refresh, + // `replaceDataDispatch` may project the data into the format the writer expects. + planLater(IcebergWriteLogical(query, batchWrite, write, replaceDataDispatch)))) + } +} 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..557c0be1e8 --- /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) + batchWrite.commit(messages) + 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..3bd03ffe5d --- /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.catalyst.plans.logical.LogicalPlan + +private[iceberg] object IcebergRefreshCacheShim { + def recacheByPlan(session: SparkSession, plan: LogicalPlan): Unit = { + session.sharedState.cacheManager.recacheByPlan(session, plan) + } +} 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-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/IcebergRefreshCacheShim.scala b/spark/src/main/spark-4.x/org/apache/comet/iceberg/IcebergRefreshCacheShim.scala new file mode 100644 index 0000000000..fb05ff1885 --- /dev/null +++ b/spark/src/main/spark-4.x/org/apache/comet/iceberg/IcebergRefreshCacheShim.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.plans.logical.LogicalPlan +import org.apache.spark.sql.classic.SparkSession + +/** Spark 4.x needs the classic `SparkSession`; the api-module one has no `sharedState`. */ +private[iceberg] object IcebergRefreshCacheShim { + def recacheByPlan(session: org.apache.spark.sql.SparkSession, plan: LogicalPlan): Unit = { + val classic = session.asInstanceOf[SparkSession] + classic.sharedState.cacheManager.recacheByPlan(classic, plan) + } +} 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/CometIcebergWriteActionSuite.scala b/spark/src/test/scala/org/apache/comet/CometIcebergWriteActionSuite.scala new file mode 100644 index 0000000000..57bd59b38a --- /dev/null +++ b/spark/src/test/scala/org/apache/comet/CometIcebergWriteActionSuite.scala @@ -0,0 +1,675 @@ +/* + * 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 + +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) + } + } + + 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)) + } + } + + 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() +}