mbutrovich commented on code in PR #4658: URL: https://github.com/apache/datafusion-comet/pull/4658#discussion_r3647994772
########## spark/src/main/scala/org/apache/spark/sql/comet/IcebergWriteExec.scala: ########## @@ -0,0 +1,177 @@ +/* + * 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 java.io.{ByteArrayInputStream, ByteArrayOutputStream, ObjectInputStream, ObjectOutputStream} + +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, AttributeReference} +import org.apache.spark.sql.catalyst.expressions.UnsafeProjection +import org.apache.spark.sql.catalyst.plans.physical.{Distribution, UnspecifiedDistribution} +import org.apache.spark.sql.connector.write.{BatchWrite, DataWriter, PhysicalWriteInfoImpl, WriterCommitMessage} +import org.apache.spark.sql.execution.{SparkPlan, UnaryExecNode} +import org.apache.spark.sql.execution.metric.{SQLMetric, SQLMetrics} +import org.apache.spark.sql.types.{BinaryType, StructField, StructType} +import org.apache.spark.util.Utils + +import org.apache.comet.iceberg.ReplaceDataDispatchInfo + +/** + * Executor-side file writer for Comet's split-operator Iceberg V2 write. + */ +case class IcebergWriteExec( + // `batchWrite` only stored driver side, only the writer factory is shipped to executors. + @transient batchWrite: BatchWrite, + child: SparkPlan, + replaceDataDispatch: Option[ReplaceDataDispatchInfo] = None) + extends UnaryExecNode { + + override def output: Seq[Attribute] = Seq( + AttributeReference(IcebergWriteExec.CommitMessageColumn, BinaryType, nullable = false)()) + + // 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")) + + override protected def doExecute(): RDD[InternalRow] = { + val rdd = child.execute() + val factory = batchWrite.createBatchWriterFactory(PhysicalWriteInfoImpl(rdd.getNumPartitions)) + require( Review Comment: Is this `require` unreachable? `buildTwoOp` already returns `None` when `useCommitCoordinator()` is true (`IcebergWriteStrategy.scala` L107), and the AQE re-plan case derives from that same guarded path. ########## spark/src/main/scala/org/apache/spark/sql/comet/IcebergWriteExec.scala: ########## @@ -0,0 +1,177 @@ +/* + * 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 java.io.{ByteArrayInputStream, ByteArrayOutputStream, ObjectInputStream, ObjectOutputStream} + +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, AttributeReference} +import org.apache.spark.sql.catalyst.expressions.UnsafeProjection +import org.apache.spark.sql.catalyst.plans.physical.{Distribution, UnspecifiedDistribution} +import org.apache.spark.sql.connector.write.{BatchWrite, DataWriter, PhysicalWriteInfoImpl, WriterCommitMessage} +import org.apache.spark.sql.execution.{SparkPlan, UnaryExecNode} +import org.apache.spark.sql.execution.metric.{SQLMetric, SQLMetrics} +import org.apache.spark.sql.types.{BinaryType, StructField, StructType} +import org.apache.spark.util.Utils + +import org.apache.comet.iceberg.ReplaceDataDispatchInfo + +/** + * Executor-side file writer for Comet's split-operator Iceberg V2 write. + */ +case class IcebergWriteExec( + // `batchWrite` only stored driver side, only the writer factory is shipped to executors. + @transient batchWrite: BatchWrite, + child: SparkPlan, + replaceDataDispatch: Option[ReplaceDataDispatchInfo] = None) + extends UnaryExecNode { + + override def output: Seq[Attribute] = Seq( + AttributeReference(IcebergWriteExec.CommitMessageColumn, BinaryType, nullable = false)()) + + // 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")) + + override protected def doExecute(): RDD[InternalRow] = { + val rdd = child.execute() + val factory = batchWrite.createBatchWriterFactory(PhysicalWriteInfoImpl(rdd.getNumPartitions)) + require( + !batchWrite.useCommitCoordinator(), + "Comet's Iceberg write path does not currently support BatchWrite implementations that " + + "require Spark's commit coordinator; received: " + batchWrite.getClass.getName) + + val rowsMetric = longMetric("numOutputRows") + val schemaTypes = output.map(_.dataType).toArray + val capturedReplaceDataDispatch = replaceDataDispatch + rdd.mapPartitionsInternal { iter => + val partId = TaskContext.getPartitionId() + val taskId = TaskContext.get().taskAttemptId() + val writer = factory.createWriter(partId, taskId) + val projection = UnsafeProjection.create(schemaTypes) + IcebergWriteExec.runWriter( + writer, + iter, + rowsMetric, + projection, + capturedReplaceDataDispatch) + } + } + + override protected def withNewChildInternal(newChild: SparkPlan): IcebergWriteExec = + copy(child = newChild) + + override def nodeName: String = "IcebergWrite" +} + +object IcebergWriteExec { + + val CommitMessageColumn: String = "iceberg_commit_message" + + val OutputSchema: StructType = StructType( + Seq(StructField(CommitMessageColumn, BinaryType, nullable = false))) + + /** Writes data files and returns the serialised Iceberg commit message. */ + def runWriter( + writer: DataWriter[InternalRow], + iter: Iterator[InternalRow], + rowsMetric: SQLMetric, + projection: UnsafeProjection, + replaceDataDispatch: Option[ReplaceDataDispatchInfo]): Iterator[InternalRow] = { + val message = Utils.tryWithSafeFinallyAndFailureCallbacks(block = { + if (replaceDataDispatch.isDefined) { + runReplaceDataWriter(writer, iter, replaceDataDispatch.get, rowsMetric) + } else { + while (iter.hasNext) { + writer.write(iter.next()) + rowsMetric.add(1L) + } + } + writer.commit() + })( + catchBlock = { + writer.abort() + }, + finallyBlock = { + writer.close() + }) + + Iterator.single(projection(InternalRow(serializeMessage(message))).copy()) + } + + // Mirrors Spark RowDeltaUtils, which is private and changes location across versions. + 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] = { + val bos = new ByteArrayOutputStream() + val oos = new ObjectOutputStream(bos) + try oos.writeObject(message) + finally oos.close() + bos.toByteArray + } + + def deserializeMessage(bytes: Array[Byte]): WriterCommitMessage = { + val bis = new ByteArrayInputStream(bytes) + val ois = new ObjectInputStream(bis) + try ois.readObject().asInstanceOf[WriterCommitMessage] + finally ois.close() + } Review Comment: Would Spark's serializer bound to Spark's classloader be safer here than the raw `ObjectOutputStream` / `ObjectInputStream` round-trip? My concern is that raw `ObjectInputStream` resolves classes with the latest user-defined loader and may fail to find Iceberg classes under `--packages`, REPL, or child-classloader isolation. Comet already uses `SparkEnv.get.serializer` for this elsewhere (`CometShuffleDependency.scala` L58, `CometBlockStoreShuffleReader.scala` L50). What do you think about `Utils.serialize` / `Utils.deserialize` with `Utils.getContextOrSparkClassLoader`? It might also let us drop one of two serialization passes, since the message is Java-serialized here and then serialized again by Spark's result serializer during `executeCollect`. ########## spark/src/main/scala/org/apache/spark/sql/comet/IcebergCommitExec.scala: ########## @@ -0,0 +1,77 @@ +/* + * 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, WriterCommitMessage} +import org.apache.spark.sql.execution.{SparkPlan, UnaryExecNode} +import org.apache.spark.sql.execution.datasources.v2.V2CommandExec +import org.apache.spark.sql.execution.metric.{SQLMetric, SQLMetrics} + +/** + * Driver-side committer for Comet's split-operator Iceberg V2 write. + */ +case class IcebergCommitExec( + // Neither of these fields are serialized, this is all run on the driver. + @transient batchWrite: BatchWrite, + @transient refreshCache: () => Unit, + 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")) + + override protected def run(): Seq[InternalRow] = { + val messages: Array[WriterCommitMessage] = child.executeCollect().map { row => + IcebergWriteExec.deserializeMessage(row.getBinary(0)) + } + 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 Review Comment: Should we call `write.reportDriverMetrics()` and post the results after a successful commit, the way `V2ExistingTableWriteExec.run` does in its finally block? On Iceberg v4.0 the `Write` wires an `InMemoryMetricsReporter` into the table for this (`SparkWrite.java` v4.0 L145-148, L287-289). Without it, my read is that Iceberg driver-side write metrics disappear when the split is on. ########## spark/src/main/scala/org/apache/comet/CometSparkSessionExtensions.scala: ########## @@ -99,6 +100,7 @@ class CometSparkSessionExtensions extensions.injectQueryStagePrepRule { session => CometExecRule(session) } injectQueryStageOptimizerRuleShim(extensions, CometPlanAdaptiveDynamicPruningFilters) injectQueryStageOptimizerRuleShim(extensions, CometReuseSubquery) + extensions.injectPlannerStrategy { session => IcebergWriteStrategy(session) } Review Comment: `injectPlannerStrategy` runs `IcebergWriteStrategy` ahead of `DataSourceV2Strategy` for every query in the session, not only Iceberg writes. Could we add a test that a non-Iceberg V2 write plans through Spark unchanged with the config on? The strategy returns `Nil` for those, and I don't think that fall-through is currently exercised. ########## spark/src/test/scala/org/apache/comet/CometIcebergWriteActionSuite.scala: ########## @@ -0,0 +1,437 @@ +/* + * 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 scala.collection.mutable + +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.execution.{QueryExecution, SparkPlan} +import org.apache.spark.sql.execution.adaptive.AdaptiveSparkPlanHelper +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) + } + } Review Comment: Could we add a multi-partition partitioned write so message collection and clustering see more than one task (several tests use `coalesce(1)`)? A concurrent-writer test that drives Iceberg's commit-time conflict validation (`validateNoConflictingData`) would also be valuable. My understanding is that the shared `BatchWrite` instance exists to make that validation see the writer's scan state, and I don't think anything currently tests it. ########## spark/src/test/scala/org/apache/comet/CometIcebergWriteActionSuite.scala: ########## @@ -0,0 +1,437 @@ +/* + * 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 scala.collection.mutable + +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.execution.{QueryExecution, SparkPlan} +import org.apache.spark.sql.execution.adaptive.AdaptiveSparkPlanHelper +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("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("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("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 captureWrite(tableName: String)(action: => Unit): WriteSnapshot = { + val before = countSnapshots(tableName) + 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 + try CometListenerBusUtils.waitUntilEmpty(spark.sparkContext) + catch { case _: java.util.concurrent.TimeoutException => () } + } finally { + spark.listenerManager.unregister(listener) + } + val after = countSnapshots(tableName) + WriteSnapshot(after - before, captured.toSeq) + } + + 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")) Review Comment: `assertExactlyOneCommit` proves the commit ran once, but does it prove files were written once? And I don't see a failure test. Could we add a case that injects a task failure and a case that injects a commit failure, asserting the table is unchanged and `abort` ran? The abort path feels like the entire risk surface of a split write, and right now it looks untested. A speculative-execution case (`spark.speculation=true`) asserting a single commit and documenting the orphan-file behavior would also be reassuring. ########## spark/src/main/scala/org/apache/spark/sql/comet/IcebergCommitExec.scala: ########## @@ -0,0 +1,77 @@ +/* + * 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, WriterCommitMessage} +import org.apache.spark.sql.execution.{SparkPlan, UnaryExecNode} +import org.apache.spark.sql.execution.datasources.v2.V2CommandExec +import org.apache.spark.sql.execution.metric.{SQLMetric, SQLMetrics} + +/** + * Driver-side committer for Comet's split-operator Iceberg V2 write. + */ +case class IcebergCommitExec( + // Neither of these fields are serialized, this is all run on the driver. + @transient batchWrite: BatchWrite, + @transient refreshCache: () => Unit, + 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")) + + override protected def run(): Seq[InternalRow] = { Review Comment: Could we add a comment stating why `run()` executes at most once? As I understand it the exactly-once guarantee rests on `V2CommandExec` memoizing `run()` via its `result` lazy val, plus the writer executing once inside the AQE boundary. That feels non-obvious and load-bearing, and I worry a future change to the base class or plan shape could reintroduce a double commit with no failing test to catch it. ########## spark/src/main/scala/org/apache/spark/sql/comet/IcebergWriteExec.scala: ########## @@ -0,0 +1,177 @@ +/* + * 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 java.io.{ByteArrayInputStream, ByteArrayOutputStream, ObjectInputStream, ObjectOutputStream} + +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, AttributeReference} +import org.apache.spark.sql.catalyst.expressions.UnsafeProjection +import org.apache.spark.sql.catalyst.plans.physical.{Distribution, UnspecifiedDistribution} +import org.apache.spark.sql.connector.write.{BatchWrite, DataWriter, PhysicalWriteInfoImpl, WriterCommitMessage} +import org.apache.spark.sql.execution.{SparkPlan, UnaryExecNode} +import org.apache.spark.sql.execution.metric.{SQLMetric, SQLMetrics} +import org.apache.spark.sql.types.{BinaryType, StructField, StructType} +import org.apache.spark.util.Utils + +import org.apache.comet.iceberg.ReplaceDataDispatchInfo + +/** + * Executor-side file writer for Comet's split-operator Iceberg V2 write. + */ +case class IcebergWriteExec( + // `batchWrite` only stored driver side, only the writer factory is shipped to executors. + @transient batchWrite: BatchWrite, + child: SparkPlan, + replaceDataDispatch: Option[ReplaceDataDispatchInfo] = None) + extends UnaryExecNode { + + override def output: Seq[Attribute] = Seq( + AttributeReference(IcebergWriteExec.CommitMessageColumn, BinaryType, nullable = false)()) + + // 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")) + + override protected def doExecute(): RDD[InternalRow] = { + val rdd = child.execute() + val factory = batchWrite.createBatchWriterFactory(PhysicalWriteInfoImpl(rdd.getNumPartitions)) + require( + !batchWrite.useCommitCoordinator(), + "Comet's Iceberg write path does not currently support BatchWrite implementations that " + + "require Spark's commit coordinator; received: " + batchWrite.getClass.getName) + + val rowsMetric = longMetric("numOutputRows") + val schemaTypes = output.map(_.dataType).toArray + val capturedReplaceDataDispatch = replaceDataDispatch + rdd.mapPartitionsInternal { iter => + val partId = TaskContext.getPartitionId() + val taskId = TaskContext.get().taskAttemptId() + val writer = factory.createWriter(partId, taskId) + val projection = UnsafeProjection.create(schemaTypes) + IcebergWriteExec.runWriter( + writer, + iter, + rowsMetric, + projection, + capturedReplaceDataDispatch) + } + } + + override protected def withNewChildInternal(newChild: SparkPlan): IcebergWriteExec = + copy(child = newChild) + + override def nodeName: String = "IcebergWrite" +} + +object IcebergWriteExec { + + val CommitMessageColumn: String = "iceberg_commit_message" + + val OutputSchema: StructType = StructType( + Seq(StructField(CommitMessageColumn, BinaryType, nullable = false))) + + /** Writes data files and returns the serialised Iceberg commit message. */ + def runWriter( + writer: DataWriter[InternalRow], + iter: Iterator[InternalRow], + rowsMetric: SQLMetric, + projection: UnsafeProjection, + replaceDataDispatch: Option[ReplaceDataDispatchInfo]): Iterator[InternalRow] = { + val message = Utils.tryWithSafeFinallyAndFailureCallbacks(block = { + if (replaceDataDispatch.isDefined) { + runReplaceDataWriter(writer, iter, replaceDataDispatch.get, rowsMetric) + } else { + while (iter.hasNext) { + writer.write(iter.next()) + rowsMetric.add(1L) + } + } + writer.commit() + })( + catchBlock = { + writer.abort() + }, + finallyBlock = { + writer.close() + }) + + Iterator.single(projection(InternalRow(serializeMessage(message))).copy()) + } + + // Mirrors Spark RowDeltaUtils, which is private and changes location across versions. + 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] = { + val bos = new ByteArrayOutputStream() + val oos = new ObjectOutputStream(bos) + try oos.writeObject(message) + finally oos.close() + bos.toByteArray + } + + def deserializeMessage(bytes: Array[Byte]): WriterCommitMessage = { + val bis = new ByteArrayInputStream(bytes) + val ois = new ObjectInputStream(bis) + try ois.readObject().asInstanceOf[WriterCommitMessage] + finally ois.close() + } + + private def runReplaceDataWriter( + writer: DataWriter[InternalRow], + iter: Iterator[InternalRow], + dispatch: ReplaceDataDispatchInfo, + rowsMetric: SQLMetric): Unit = { + val rowProjection = dispatch.rowProjection + val metadataProjection = dispatch.metadataProjection.orNull + while (iter.hasNext) { + val row = iter.next() + rowsMetric.add(1L) + 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: " + Review Comment: I believe codes 5 (WRITE) and 6 (WRITE_WITH_METADATA) are correct and complete for Spark 4.0+ copy-on-write ReplaceData (DELETE emits 6, UPDATE emits 6, MERGE emits 5 and 6), and that they do not exist before 4.0. This path looks safe only because `IcebergReplaceDataShim` (spark-3.x) returns `None`, routing 3.x through the plain `write(row)` loop. Could we state that invariant in a comment here? And would it be worth adding the suite to CI for all four Spark lines so the version split cannot regress unnoticed? ########## spark/src/main/scala/org/apache/comet/iceberg/IcebergWriteLogical.scala: ########## @@ -0,0 +1,38 @@ +/* + * 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 +import org.apache.spark.sql.catalyst.plans.logical.{LogicalPlan, UnaryNode} +import org.apache.spark.sql.connector.write.BatchWrite + +/** Logical anchor for the writer. See `IcebergWriteStrategy` for the rationale. */ +case class IcebergWriteLogical( + child: LogicalPlan, + // Driver-side only: AQE re-planning is driver-local and write commands aren't cached. + @transient batchWrite: BatchWrite, + replaceDataDispatch: Option[ReplaceDataDispatchInfo] = None) + extends UnaryNode { + + override def output: Seq[Attribute] = Nil Review Comment: `output` is `Nil` here while the physical `IcebergWriteExec.output` is a single binary column. Is that intentional? If so, could we add a comment explaining why `Nil` is safe? I'm wary that a logical anchor whose output disagrees with its physical node could break any rule that reasons over `logicalLink.output`, and this is the node AQE re-optimizes against. ########## spark/src/main/scala/org/apache/spark/sql/comet/IcebergWriteExec.scala: ########## @@ -0,0 +1,177 @@ +/* + * 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 java.io.{ByteArrayInputStream, ByteArrayOutputStream, ObjectInputStream, ObjectOutputStream} + +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, AttributeReference} +import org.apache.spark.sql.catalyst.expressions.UnsafeProjection +import org.apache.spark.sql.catalyst.plans.physical.{Distribution, UnspecifiedDistribution} +import org.apache.spark.sql.connector.write.{BatchWrite, DataWriter, PhysicalWriteInfoImpl, WriterCommitMessage} +import org.apache.spark.sql.execution.{SparkPlan, UnaryExecNode} +import org.apache.spark.sql.execution.metric.{SQLMetric, SQLMetrics} +import org.apache.spark.sql.types.{BinaryType, StructField, StructType} +import org.apache.spark.util.Utils + +import org.apache.comet.iceberg.ReplaceDataDispatchInfo + +/** + * Executor-side file writer for Comet's split-operator Iceberg V2 write. + */ +case class IcebergWriteExec( + // `batchWrite` only stored driver side, only the writer factory is shipped to executors. + @transient batchWrite: BatchWrite, + child: SparkPlan, + replaceDataDispatch: Option[ReplaceDataDispatchInfo] = None) + extends UnaryExecNode { + + override def output: Seq[Attribute] = Seq( + AttributeReference(IcebergWriteExec.CommitMessageColumn, BinaryType, nullable = false)()) + + // 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")) + + override protected def doExecute(): RDD[InternalRow] = { + val rdd = child.execute() + val factory = batchWrite.createBatchWriterFactory(PhysicalWriteInfoImpl(rdd.getNumPartitions)) + require( + !batchWrite.useCommitCoordinator(), + "Comet's Iceberg write path does not currently support BatchWrite implementations that " + + "require Spark's commit coordinator; received: " + batchWrite.getClass.getName) + + val rowsMetric = longMetric("numOutputRows") + val schemaTypes = output.map(_.dataType).toArray + val capturedReplaceDataDispatch = replaceDataDispatch + rdd.mapPartitionsInternal { iter => + val partId = TaskContext.getPartitionId() + val taskId = TaskContext.get().taskAttemptId() + val writer = factory.createWriter(partId, taskId) + val projection = UnsafeProjection.create(schemaTypes) + IcebergWriteExec.runWriter( + writer, + iter, + rowsMetric, + projection, + capturedReplaceDataDispatch) + } + } + + override protected def withNewChildInternal(newChild: SparkPlan): IcebergWriteExec = + copy(child = newChild) + + override def nodeName: String = "IcebergWrite" +} + +object IcebergWriteExec { + + val CommitMessageColumn: String = "iceberg_commit_message" + + val OutputSchema: StructType = StructType( + Seq(StructField(CommitMessageColumn, BinaryType, nullable = false))) + + /** Writes data files and returns the serialised Iceberg commit message. */ + def runWriter( + writer: DataWriter[InternalRow], + iter: Iterator[InternalRow], + rowsMetric: SQLMetric, + projection: UnsafeProjection, + replaceDataDispatch: Option[ReplaceDataDispatchInfo]): Iterator[InternalRow] = { + val message = Utils.tryWithSafeFinallyAndFailureCallbacks(block = { + if (replaceDataDispatch.isDefined) { + runReplaceDataWriter(writer, iter, replaceDataDispatch.get, rowsMetric) + } else { + while (iter.hasNext) { + writer.write(iter.next()) + rowsMetric.add(1L) + } + } + writer.commit() Review Comment: Should we be flushing the writer's `dataWriter.currentMetricsValues` during and after iteration? Spark does this in `WritingSparkTask` via `IteratorWithMetrics`, and it looks like that is how Iceberg surfaces per-task write metrics. Since `runWriter` only increments `numOutputRows`, my read is that the Iceberg task write metrics are lost when the split is enabled. Am I missing where they get picked up? ########## spark/src/main/scala/org/apache/spark/sql/comet/IcebergWriteExec.scala: ########## @@ -0,0 +1,177 @@ +/* + * 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 java.io.{ByteArrayInputStream, ByteArrayOutputStream, ObjectInputStream, ObjectOutputStream} + +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, AttributeReference} +import org.apache.spark.sql.catalyst.expressions.UnsafeProjection +import org.apache.spark.sql.catalyst.plans.physical.{Distribution, UnspecifiedDistribution} +import org.apache.spark.sql.connector.write.{BatchWrite, DataWriter, PhysicalWriteInfoImpl, WriterCommitMessage} +import org.apache.spark.sql.execution.{SparkPlan, UnaryExecNode} +import org.apache.spark.sql.execution.metric.{SQLMetric, SQLMetrics} +import org.apache.spark.sql.types.{BinaryType, StructField, StructType} +import org.apache.spark.util.Utils + +import org.apache.comet.iceberg.ReplaceDataDispatchInfo + +/** + * Executor-side file writer for Comet's split-operator Iceberg V2 write. + */ +case class IcebergWriteExec( + // `batchWrite` only stored driver side, only the writer factory is shipped to executors. + @transient batchWrite: BatchWrite, + child: SparkPlan, + replaceDataDispatch: Option[ReplaceDataDispatchInfo] = None) + extends UnaryExecNode { + + override def output: Seq[Attribute] = Seq( + AttributeReference(IcebergWriteExec.CommitMessageColumn, BinaryType, nullable = false)()) + + // 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")) + + override protected def doExecute(): RDD[InternalRow] = { + val rdd = child.execute() + val factory = batchWrite.createBatchWriterFactory(PhysicalWriteInfoImpl(rdd.getNumPartitions)) Review Comment: How does this behave when the child RDD has zero partitions? Spark's SPARK-23271 handling runs a single-partition job so the empty commit is still produced deterministically, as `writeWithV2` does. `mapPartitionsInternal` over a zero-partition RDD produces zero messages and relies on `commit([])` behaving identically, which I'm not sure is guaranteed across write types. Could we add a test with a genuinely zero-partition child? The current empty-source test (L96) uses `WHERE id < 0`, which I believe yields one empty partition, not zero. ########## spark/src/main/scala/org/apache/comet/iceberg/IcebergWriteStrategy.scala: ########## @@ -0,0 +1,120 @@ +/* + * 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 IcebergWriteLogical(child, batchWrite, replaceDataDispatch) => + Seq(IcebergWriteExec(batchWrite, 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. + * + * Returns None, falling back to Spark's combined write operator, when the `BatchWrite` requires + * Spark's commit coordinator, which the split writer's per-task commit protocol does not use. + */ + private def buildTwoOp( + write: Write, + rel: DataSourceV2Relation, + query: LogicalPlan, + replaceDataDispatch: Option[ReplaceDataDispatchInfo]): Option[SparkPlan] = { + val batchWrite = write.toBatch + if (batchWrite.useCommitCoordinator()) { + return None Review Comment: Could we reword this comment? It reads as if the `useCommitCoordinator()` fallback is what keeps the split correct, but since every Iceberg `BatchWrite` returns false, this never fires for Iceberg. As written I think it implies an Iceberg code path that does not exist. Would it be clearer to say plainly that this is coverage for non-Iceberg V2 sinks? ########## spark/src/main/scala/org/apache/comet/iceberg/IcebergWriteStrategy.scala: ########## @@ -0,0 +1,120 @@ +/* + * 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 IcebergWriteLogical(child, batchWrite, replaceDataDispatch) => + Seq(IcebergWriteExec(batchWrite, 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. + * + * Returns None, falling back to Spark's combined write operator, when the `BatchWrite` requires + * Spark's commit coordinator, which the split writer's per-task commit protocol does not use. + */ + 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(rel) Review Comment: What do you think about passing the captured `session` into `recacheByPlan` instead of resolving `SparkSession.active` inside the shim? Binding the refresh callback to a specific session matches how Spark builds it in `DataSourceV2Strategy` and would remove a dependency on thread-local active-session state. ########## spark/src/main/scala/org/apache/comet/iceberg/IcebergWriteStrategy.scala: ########## @@ -0,0 +1,120 @@ +/* + * 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 IcebergWriteLogical(child, batchWrite, replaceDataDispatch) => + Seq(IcebergWriteExec(batchWrite, 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. + * + * Returns None, falling back to Spark's combined write operator, when the `BatchWrite` requires + * Spark's commit coordinator, which the split writer's per-task commit protocol does not use. + */ + 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(rel) + Some( + IcebergCommitExec( + batchWrite, + refresh, + // `replaceDataDispatch` may project the data into the format the writer expects. + planLater(IcebergWriteLogical(query, batchWrite, replaceDataDispatch)))) Review Comment: The writer sets `requiredChildDistribution = UnspecifiedDistribution` and depends on `V2Writes` having injected the repartition and local sort into `.query`. Iceberg uses the clustered writer for partitioned writes and throws on unclustered rows (`ClusteredWriter.java` L66-105). The AQE test (`CometIcebergWriteActionSuite.scala` L110) uses an unpartitioned table, so as far as I can tell coalesce and skew-split against a clustered partitioned write is never exercised. Could we add a partitioned write under AQE with a shuffle? ########## docs/source/user-guide/latest/iceberg-writes.md: ########## @@ -0,0 +1,85 @@ +<!--- + Licensed to the Apache Software Foundation (ASF) under one + or more contributor license agreements. See the NOTICE file + distributed with this work for additional information + regarding copyright ownership. The ASF licenses this file + to you under the Apache License, Version 2.0 (the + "License"); you may not use this file except in compliance + with the License. You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + + Unless required by applicable law or agreed to in writing, + software distributed under the License is distributed on an + "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + KIND, either express or implied. See the License for the + specific language governing permissions and limitations + under the License. +--> + +# Iceberg Writes: Comet's Split-Operator Plan (Experimental) + +**This feature is experimental and disabled by default.** Enable it only after validating it +against your own workloads. + +## Overview + +Spark writes an Iceberg table through a single physical operator that combines data-file +writing with metadata writing, committing, and catalog validation. Because that operator sits +outside Spark's Adaptive Query Execution (AQE), the sub-query feeding the write — the scans, +projects, sorts, and exchanges producing the rows — cannot be re-planned at runtime. + +When `spark.comet.write.iceberg.splitOperator.enabled=true`, Comet rewrites eligible Iceberg +writes into two operators: + +1. **`IcebergWrite`** — writes the data files on the executors, exactly as iceberg-java does + today, and returns each task's serialized commit message. This operator and the sub-query + feeding it run inside AQE. +2. **`IcebergCommit`** — collects the commit messages on the driver and performs the normal + Iceberg commit (including commit-time validation), outside AQE, exactly once. + +Data files are still written by iceberg-java; only the plan shape changes. The split makes the +write's input visible to AQE and to Comet's columnar rules, and it is the groundwork for a +planned follow-up in which Comet writes the data files natively via +[iceberg-rust](https://github.com/apache/iceberg-rust). + +## Configuration + +Standard Comet + Iceberg setup (see [`iceberg.md`](iceberg.md)) plus the write-side toggle: + +``` +# Standard Comet / Iceberg wiring +spark.plugins=org.apache.spark.CometPlugin +spark.sql.extensions=org.apache.iceberg.spark.extensions.IcebergSparkSessionExtensions +spark.sql.catalog.<name>=org.apache.iceberg.spark.SparkCatalog +spark.sql.catalog.<name>.type=hadoop # or hive / glue / rest / ... +spark.sql.catalog.<name>.warehouse=... + +# Split-operator plan (experimental, off by default) +spark.comet.write.iceberg.splitOperator.enabled=true +``` + +## Supported operations + +The split-operator plan is supported on every Spark version Comet supports, with identical +coverage on each: + Review Comment: Could we revisit the "identical coverage on every Spark version" claim? The copy-on-write ReplaceData path looks like it differs by version: 4.0+ uses operation-coded rows with projections, while 3.4/3.5 use a plain row stream. Would it be more accurate to state that the row-level DML mechanism differs by Spark version so the docs match the shim behavior? ########## spark/src/test/scala/org/apache/comet/CometIcebergWriteActionSuite.scala: ########## @@ -0,0 +1,437 @@ +/* + * 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 scala.collection.mutable + +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.execution.{QueryExecution, SparkPlan} +import org.apache.spark.sql.execution.adaptive.AdaptiveSparkPlanHelper +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("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("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("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 captureWrite(tableName: String)(action: => Unit): WriteSnapshot = { + val before = countSnapshots(tableName) + 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 + try CometListenerBusUtils.waitUntilEmpty(spark.sparkContext) + catch { case _: java.util.concurrent.TimeoutException => () } Review Comment: Should we stop swallowing `TimeoutException` from `waitUntilEmpty`? My worry is that a dropped or late listener event turns the commit-count assertion into a silent pass. Would it be better to fail on timeout, or to read the snapshot delta directly from the Iceberg metadata rather than depending on listener delivery? -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected] --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
