From ba794f594d80d7db6a651315d8a5d793c9c4b11b Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Wed, 2 Sep 2026 10:51:45 -0600 Subject: [PATCH 1/2] feat: drive Spark's file writers from Arrow batches for partitioned writes Spark's write path is typed on InternalRow throughout and its writers read each row and encode it before asking for the next one, so the UnsafeRow that the columnar-to-row transition materializes is a copy the writer only undoes again. A partitioned or bucketed write pays for it twice: BaseDynamicPartitionDataWriter projects every row a second time through getOutputRow purely to strip the partition and bucket columns. Replace WriteFilesExec with a Comet node that drives Spark's own OutputWriter from child.executeColumnar(). Pruning the partition columns off a ColumnarBatch costs nothing, so CometRowViewDynamicPartitionWriter overrides only writeRecord and inherits partition-change detection, writer renewal and maxRecordsPerFile unchanged. Everything above the per-task write stays with Spark, because the node extends WriteFilesExecBase and V1WritesUtils.getWriteFilesOpt finds it. Gated on Spark 4.0+, on maxConcurrentOutputFileWriters being 0 so that the concurrent writer's UnsafeRow spill path is never reached, on one of Spark's own FileFormats, and on a complex type among the data columns. Behind spark.comet.exec.write.rowView.enabled, default false. Co-Authored-By: Claude Opus 5 (1M context) Claude-Session: https://claude.ai/code/session_01DmHrzzDLtwh1nLHZUxTVKm --- .github/workflows/pr_build_linux.yml | 1 + .github/workflows/pr_build_macos.yml | 1 + .../scala/org/apache/comet/CometConf.scala | 12 + .../rules/EliminateRedundantTransitions.scala | 78 +++- .../comet/CometRowViewWriteFilesExec.scala | 304 +++++++++++++ .../comet/shims/ShimCometWriteFilesExec.scala | 36 ++ .../comet/shims/ShimCometWriteFilesExec.scala | 37 ++ .../comet/exec/CometWriteRowViewSuite.scala | 417 ++++++++++++++++++ 8 files changed, 884 insertions(+), 2 deletions(-) create mode 100644 spark/src/main/scala/org/apache/spark/sql/comet/CometRowViewWriteFilesExec.scala create mode 100644 spark/src/main/spark-3.x/org/apache/comet/shims/ShimCometWriteFilesExec.scala create mode 100644 spark/src/main/spark-4.x/org/apache/comet/shims/ShimCometWriteFilesExec.scala create mode 100644 spark/src/test/scala/org/apache/comet/exec/CometWriteRowViewSuite.scala diff --git a/.github/workflows/pr_build_linux.yml b/.github/workflows/pr_build_linux.yml index 11a5619a4d2..228ab5a73a6 100644 --- a/.github/workflows/pr_build_linux.yml +++ b/.github/workflows/pr_build_linux.yml @@ -370,6 +370,7 @@ jobs: org.apache.comet.exec.CometShuffle4_0Suite org.apache.comet.exec.CometNativeColumnarToRowSuite org.apache.comet.exec.CometNativeShuffleSuite + org.apache.comet.exec.CometWriteRowViewSuite org.apache.comet.shuffle.CelebornShufflePartitionPusherSuite org.apache.spark.sql.comet.execution.shuffle.CometCelebornShuffleManagerSuite org.apache.spark.sql.comet.execution.shuffle.CometCelebornNativeShuffleWriterSuite diff --git a/.github/workflows/pr_build_macos.yml b/.github/workflows/pr_build_macos.yml index 7b7fae8fbea..b5dbdcddf2c 100644 --- a/.github/workflows/pr_build_macos.yml +++ b/.github/workflows/pr_build_macos.yml @@ -143,6 +143,7 @@ jobs: org.apache.comet.exec.CometShuffle4_0Suite org.apache.comet.exec.CometNativeColumnarToRowSuite org.apache.comet.exec.CometNativeShuffleSuite + org.apache.comet.exec.CometWriteRowViewSuite org.apache.comet.shuffle.CelebornShufflePartitionPusherSuite org.apache.spark.sql.comet.execution.shuffle.CometCelebornShuffleManagerSuite org.apache.spark.sql.comet.execution.shuffle.CometCelebornNativeShuffleWriterSuite diff --git a/spark/src/main/scala/org/apache/comet/CometConf.scala b/spark/src/main/scala/org/apache/comet/CometConf.scala index a7d7b80db1f..9b80fea7d22 100644 --- a/spark/src/main/scala/org/apache/comet/CometConf.scala +++ b/spark/src/main/scala/org/apache/comet/CometConf.scala @@ -289,6 +289,18 @@ object CometConf extends ShimCometConf { .booleanConf .createWithDefault(false) + val COMET_EXEC_WRITE_ROW_VIEW_ENABLED: ConfigEntry[Boolean] = + conf(s"$COMET_EXEC_CONFIG_PREFIX.write.rowView.enabled") + .category(CATEGORY_EXEC) + .doc( + "Experimental. Whether Comet takes over Spark's WriteFilesExec so that file writes " + + "consume Arrow batches directly instead of a materialized UnsafeRow per row. Spark's " + + "file writers only ever read the row they are handed, so the UnsafeProjection that " + + "the columnar-to-row transition performs, and the second one that the partitioned " + + "writer performs to strip partition columns, are both pure overhead. Off by default.") + .booleanConf + .createWithDefault(false) + val COMET_EXEC_SORT_MERGE_JOIN_WITH_JOIN_FILTER_ENABLED: ConfigEntry[Boolean] = conf("spark.comet.exec.sortMergeJoinWithJoinFilter.enabled") .category(CATEGORY_ENABLE_EXEC) diff --git a/spark/src/main/scala/org/apache/comet/rules/EliminateRedundantTransitions.scala b/spark/src/main/scala/org/apache/comet/rules/EliminateRedundantTransitions.scala index 45fc26caee9..f9990f56289 100644 --- a/spark/src/main/scala/org/apache/comet/rules/EliminateRedundantTransitions.scala +++ b/spark/src/main/scala/org/apache/comet/rules/EliminateRedundantTransitions.scala @@ -22,15 +22,17 @@ package org.apache.comet.rules import org.apache.spark.sql.SparkSession import org.apache.spark.sql.catalyst.rules.Rule import org.apache.spark.sql.catalyst.util.sideBySide -import org.apache.spark.sql.comet.{CometCollectLimitExec, CometColumnarToRowExec, CometMapInBatchExec, CometNativeColumnarToRowExec, CometNativeWriteExec, CometPlan, CometSparkToColumnarExec} +import org.apache.spark.sql.comet.{CometCollectLimitExec, CometColumnarToRowExec, CometMapInBatchExec, CometNativeColumnarToRowExec, CometNativeWriteExec, CometPlan, CometRowViewWriteFilesExec, CometSparkToColumnarExec} import org.apache.spark.sql.comet.execution.shuffle.{CometColumnarShuffle, CometShuffleExchangeExec} import org.apache.spark.sql.comet.shims.{MapInBatchInfo, ShimCometMapInBatch} import org.apache.spark.sql.execution.{ColumnarToRowExec, RowToColumnarExec, SparkPlan} import org.apache.spark.sql.execution.adaptive.QueryStageExec +import org.apache.spark.sql.execution.datasources.WriteFilesExec import org.apache.spark.sql.execution.exchange.ReusedExchangeExec +import org.apache.spark.sql.types.{ArrayType, DataType, MapType, StructType} import org.apache.comet.CometConf -import org.apache.comet.CometSparkSessionExtensions.withInfo +import org.apache.comet.CometSparkSessionExtensions.{isSpark40Plus, withInfo} import org.apache.comet.serde.NativeOptIn import org.apache.comet.shims.ShimSQLConf @@ -91,6 +93,15 @@ case class EliminateRedundantTransitions(session: SparkSession) // Write should be final operation in the plan case ColumnarToRowExec(nativeWrite: CometNativeWriteExec) => nativeWrite + // Spark's file writers only ever read the row they are handed, so the columnar-to-row + // transition below a write can be dropped entirely and the writer driven from the Arrow + // batches directly. `transformUp` has already rewritten the child into one of the Comet + // transitions by the time this arm is visited. + case w: WriteFilesExec if writeRowViewEligible(w) => + columnarChildForWrite(w) + .map(child => CometRowViewWriteFilesExec(child)) + .getOrElse(w) + case c @ ColumnarToRowExec(child) if hasCometNativeChild(child) => val op = createColumnarToRowExec(child) if (c.logicalLink.isEmpty) { @@ -161,6 +172,69 @@ case class EliminateRedundantTransitions(session: SparkSession) } } + /** + * Whether the per-task write of `w` can be driven from Arrow batches by + * [[CometRowViewWriteFilesExec]] rather than from materialized `UnsafeRow`s. + * + * The rows that reach Spark's `OutputWriter` are then reused, mutable views over the Arrow + * buffers, which is only correct for a writer that finishes with a row before asking for the + * next one. Three things establish that: + * + * - Spark 4.0+, where `V1WritesUtils.getWriteFilesOpt` matches the `WriteFilesExecBase` + * trait. On 3.4 / 3.5 it matches the concrete `WriteFilesExec` case class, so a replacement + * node is invisible to `FileFormatWriter` and the write would silently take a path that + * calls `doExecute` on it. + * - `spark.sql.maxConcurrentOutputFileWriters` at its default of 0. Above 0, `V1Writes` + * plants no sort (`V1WritesUtils.getSortOrder`) and `FileFormatWriter` picks + * `DynamicPartitionDataConcurrentWriter`, which both requires `UnsafeRow` for its spill and + * would leave `DynamicPartitionDataSingleWriter` with unsorted input. + * - one of Spark's own `FileFormat`s, whose `OutputWriter`s encode each row on the spot + * (Parquet through `ParquetWriteSupport`, ORC through `OrcSerializer` into a + * `VectorizedRowBatch`, the text formats directly). A third-party format is free to buffer + * the `InternalRow` it is handed. + * + * Note that no check for a `SortExec` between the write and the transition is needed. This arm + * only rewrites `w.child`, so a write whose required ordering was satisfied by a Spark + * `SortExec` rather than a Comet one simply does not match. + */ + private def writeRowViewEligible(w: WriteFilesExec): Boolean = + CometConf.COMET_EXEC_WRITE_ROW_VIEW_ENABLED.get() && + isSpark40Plus && + w.conf.maxConcurrentOutputFileWriters == 0 && + w.fileFormat.getClass.getName.startsWith("org.apache.spark.sql.execution.datasources.") + + /** + * The Comet columnar producer under a write's columnar-to-row transition, or `None` when there + * is no transition to strip or the write would not gain from removing it. + * + * Declines a write whose data columns are all flat. There the projection this removes is a + * generated fixed-width copy that measures inside the noise of a Parquet write, which does not + * pay for putting a reused mutable row in front of Spark's writer. The saving only becomes real + * once a struct, array or map is in play, because then the projection has to build nested + * `UnsafeRow` / `UnsafeArrayData` with offset-and-length bookkeeping that the writer + * immediately walks back out. Partition and bucket columns do not count: they never reach the + * `OutputWriter`. + */ + private def columnarChildForWrite(w: WriteFilesExec): Option[SparkPlan] = { + val columnarChild = w.child match { + case CometColumnarToRowExec(child) => Some(child) + case CometNativeColumnarToRowExec(child) => Some(child) + case _ => None + } + columnarChild.filter { child => + val partitionIds = w.partitionColumns.map(_.exprId).toSet + child.output + .filterNot(a => partitionIds.contains(a.exprId)) + .exists(a => isComplex(a.dataType)) + } + } + + /** Whether a type is a struct, array or map. */ + private def isComplex(dataType: DataType): Boolean = dataType match { + case _: StructType | _: ArrayType | _: MapType => true + case _ => false + } + private def hasCometNativeChild(op: SparkPlan): Boolean = { op match { case c: QueryStageExec => hasCometNativeChild(c.plan) diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/CometRowViewWriteFilesExec.scala b/spark/src/main/scala/org/apache/spark/sql/comet/CometRowViewWriteFilesExec.scala new file mode 100644 index 00000000000..47e03c40df6 --- /dev/null +++ b/spark/src/main/scala/org/apache/spark/sql/comet/CometRowViewWriteFilesExec.scala @@ -0,0 +1,304 @@ +/* + * 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.util.Date + +import org.apache.hadoop.mapreduce.{TaskAttemptContext, TaskAttemptID, TaskID, TaskType} +import org.apache.hadoop.mapreduce.task.TaskAttemptContextImpl +import org.apache.spark.TaskContext +import org.apache.spark.internal.Logging +import org.apache.spark.internal.io.{FileCommitProtocol, SparkHadoopWriterUtils} +import org.apache.spark.rdd.RDD +import org.apache.spark.sql.catalyst.InternalRow +import org.apache.spark.sql.connector.write.WriterCommitMessage +import org.apache.spark.sql.execution.SparkPlan +import org.apache.spark.sql.execution.datasources.{DynamicPartitionDataSingleWriter, EmptyDirectoryDataWriter, FileFormatDataWriter, SingleDirectoryDataWriter, WriteFilesSpec, WriteJobDescription, WriteTaskResult} +import org.apache.spark.sql.vectorized.{ColumnarBatch, ColumnVector} +import org.apache.spark.util.Utils + +import org.apache.comet.shims.ShimCometWriteFilesExec + +/** + * A drop-in replacement for Spark's `WriteFilesExec` that drives Spark's own `OutputWriter` from + * Arrow batches, without ever materializing an `UnsafeRow`. + * + * Spark's write path is typed on `InternalRow` throughout - `OutputWriter.write`, + * `FileFormatDataWriter.write` and `WriteTaskStatsTracker.newRow` - and the writers read each row + * and encode it before asking for the next one. `ColumnarBatch.rowIterator()` already produces a + * reused `ColumnarBatchRow` that is a zero-copy view over the Arrow buffers and satisfies that + * contract exactly, so the `UnsafeProjection` performed by [[CometColumnarToRowExec]] is a copy + * the writer only undoes again. + * + * A partitioned or bucketed write pays for that copy twice. `BaseDynamicPartitionDataWriter` + * projects every row a second time through `getOutputRow` purely to strip the partition and + * bucket columns before handing it to the `OutputWriter`. This node replaces that projection with + * a pruned view over the same Arrow vectors, which costs nothing: + * [[CometRowViewDynamicPartitionWriter]] overrides only `writeRecord`, so Spark keeps ownership + * of partition-change detection, file rolling and `maxRecordsPerFile`. + * + * Everything above the per-task write stays with Spark. Because this node extends + * `WriteFilesExecBase` (see [[ShimCometWriteFilesExec]]), `V1WritesUtils.getWriteFilesOpt` finds + * it and `InsertIntoHadoopFsRelationCommand` / `FileFormatWriter` continue to own SaveMode + * semantics, the commit protocol, `_SUCCESS`, dynamic partition overwrite, stats tracker + * aggregation and catalog updates. Extending the trait is also what keeps AQE from re-inserting a + * second `WriteFilesExec` above this node. + * + * Unlike [[CometWriteFilesExec]]-style native writes, the bytes are still produced by Spark's own + * `OutputWriter`, so output is Spark's by construction and every `FileFormat` that Spark ships is + * supported. + * + * Only [[org.apache.comet.rules.EliminateRedundantTransitions]] introduces this node, and only + * once it has established the preconditions the reused row depends on. See `writeRowViewEligible` + * there. + * + * @param child + * The Comet columnar operator producing the batches to write. Its output must be the write's + * `allColumns`, in order. + */ +case class CometRowViewWriteFilesExec(child: SparkPlan) + extends ShimCometWriteFilesExec + with CometPlan { + + override def nodeName: String = "CometRowViewWriteFiles" + + /** Spark drives this node through `executeWrite`, never `execute`. */ + override protected def doExecute(): RDD[InternalRow] = + throw new UnsupportedOperationException(s"$nodeName does not support doExecute") + + override protected def doExecuteWrite( + writeFilesSpec: WriteFilesSpec): RDD[WriterCommitMessage] = { + val description = writeFilesSpec.description + val committer = writeFilesSpec.committer + // Same identifier scheme as FileFormatWriter, so committers that parse the job ID agree. + val jobTrackerID = SparkHadoopWriterUtils.createJobTrackerID(new Date()) + val dataOrdinals = CometRowViewWriteFilesExec.dataColumnOrdinals(description) + + val childRDD = child.executeColumnar() + + // SPARK-23271: a zero-partition input would spawn no task and therefore write no file at all, + // leaving the output directory without a schema for readers. Spark's own WriteFilesExec swaps + // in a dummy single-partition RDD for exactly this case. + val writeRDD = if (childRDD.getNumPartitions == 0) { + sparkContext.parallelize(Seq.empty[ColumnarBatch], 1) + } else { + childRDD + } + + // Everything the task needs is resolved on the driver and captured by value. The closure must + // not touch `this`, which holds the whole converted child subtree; Spark's own + // WriteFilesExec.doExecuteWrite avoids the same capture by delegating to a static + // FileFormatWriter.executeTask. + writeRDD.mapPartitionsInternal { batches => + val taskCtx = TaskContext.get() + Iterator( + CometRowViewWriteFilesExec.executeTask( + description, + jobTrackerID, + taskCtx.stageId(), + taskCtx.partitionId(), + // Truncation to Int matches FileFormatWriter: the masked low bits are what the Hadoop + // TaskAttemptID accepts, and uniqueness within a job is preserved by the task ID. + taskCtx.taskAttemptId().toInt & Integer.MAX_VALUE, + committer, + batches, + dataOrdinals)) + } + } + + override protected def withNewChildInternal(newChild: SparkPlan): SparkPlan = + copy(child = newChild) +} + +object CometRowViewWriteFilesExec extends Logging { + + /** + * Positions of the write's data columns within its `allColumns`, used to prune the partition + * and bucket columns off an Arrow batch without copying. + */ + private[comet] def dataColumnOrdinals(description: WriteJobDescription): Array[Int] = { + val positions = description.allColumns.map(_.exprId).zipWithIndex.toMap + description.dataColumns.map { attr => + positions.getOrElse( + attr.exprId, + throw new IllegalStateException( + s"Data column ${attr.name} is not among the write's output columns")) + }.toArray + } + + /** + * Write one task's batches and commit or abort it. A direct port of + * `FileFormatWriter.executeTask`, differing only in the writer chosen for the partitioned and + * bucketed case and in reading rows from Arrow batches rather than from an `Iterator` of + * `UnsafeRow`. + * + * `DynamicPartitionDataConcurrentWriter` is deliberately never constructed here: it spills + * through `UnsafeKVExternalSorter`, which is typed on `UnsafeRow`. The rule refuses the rewrite + * unless `spark.sql.maxConcurrentOutputFileWriters` is 0, which is also what makes `V1Writes` + * plant the sort that `DynamicPartitionDataSingleWriter` requires. + */ + private[comet] def executeTask( + description: WriteJobDescription, + jobTrackerID: String, + sparkStageId: Int, + sparkPartitionId: Int, + sparkAttemptNumber: Int, + committer: FileCommitProtocol, + batches: Iterator[ColumnarBatch], + dataOrdinals: Array[Int]): WriteTaskResult = { + + val jobId = SparkHadoopWriterUtils.createJobID(jobTrackerID, sparkStageId) + val taskId = new TaskID(jobId, TaskType.MAP, sparkPartitionId) + val taskAttemptId = new TaskAttemptID(taskId, sparkAttemptNumber) + + val taskAttemptContext: TaskAttemptContext = { + val hadoopConf = description.serializableHadoopConf.value + hadoopConf.set("mapreduce.job.id", jobId.toString) + hadoopConf.set("mapreduce.task.id", taskAttemptId.getTaskID.toString) + hadoopConf.set("mapreduce.task.attempt.id", taskAttemptId.toString) + hadoopConf.setBoolean("mapreduce.task.ismap", true) + hadoopConf.setInt("mapreduce.task.partition", 0) + new TaskAttemptContextImpl(hadoopConf, taskAttemptId) + } + + committer.setupTask(taskAttemptContext) + + var dataWriter: FileFormatDataWriter = null + + Utils.tryWithSafeFinallyAndFailureCallbacks(block = { + dataWriter = if (sparkPartitionId != 0 && !batches.hasNext) { + // In case of empty job, leave first partition to save meta for file format like parquet. + new EmptyDirectoryDataWriter(description, taskAttemptContext, committer) + } else if (description.partitionColumns.isEmpty && description.bucketSpec.isEmpty) { + // SingleDirectoryDataWriter hands the row straight to the OutputWriter, so the batch row + // needs no pruning and Spark's own writer can be used unchanged. + new SingleDirectoryDataWriter(description, taskAttemptContext, committer) + } else { + new CometRowViewDynamicPartitionWriter(description, taskAttemptContext, committer) + } + + writeBatches(dataWriter, batches, dataOrdinals) + dataWriter.commit() + })( + catchBlock = { + if (dataWriter != null) { + dataWriter.abort() + } else { + committer.abortTask(taskAttemptContext) + } + logError(s"Job: $jobId, Task: $taskId, Task attempt $taskAttemptId aborted.") + }, + finallyBlock = { + if (dataWriter != null) { + dataWriter.close() + } + }) + } + + /** + * Feed a task's batches to the writer one row at a time, mirroring + * `FileFormatDataWriter.writeWithIterator`. + * + * For the partitioned and bucketed case two row iterators are advanced in lockstep over the + * same batch: one over all columns, which Spark's writer uses to detect partition and bucket + * changes, and one over the pruned batch, which is what reaches the `OutputWriter`. Both are + * views over the same Arrow vectors, so the pair costs one object per batch rather than a copy + * per row. + */ + private def writeBatches( + dataWriter: FileFormatDataWriter, + batches: Iterator[ColumnarBatch], + dataOrdinals: Array[Int]): Unit = { + var count = 0L + dataWriter match { + case rowViewWriter: CometRowViewDynamicPartitionWriter => + while (batches.hasNext) { + val batch = batches.next() + val allRows = batch.rowIterator() + val dataRows = prune(batch, dataOrdinals).rowIterator() + while (allRows.hasNext) { + rowViewWriter.setDataRow(dataRows.next()) + rowViewWriter.writeWithMetrics(allRows.next(), count) + count += 1 + } + } + case _ => + while (batches.hasNext) { + val rows = batches.next().rowIterator() + while (rows.hasNext) { + dataWriter.writeWithMetrics(rows.next(), count) + count += 1 + } + } + } + } + + /** + * A batch over the data columns only, borrowing the vectors of `batch`. Never closed: the + * vectors belong to `batch`, which the child's iterator owns. + */ + private def prune(batch: ColumnarBatch, dataOrdinals: Array[Int]): ColumnarBatch = { + val vectors = new Array[ColumnVector](dataOrdinals.length) + var i = 0 + while (i < dataOrdinals.length) { + vectors(i) = batch.column(dataOrdinals(i)) + i += 1 + } + new ColumnarBatch(vectors, batch.numRows()) + } +} + +/** + * `DynamicPartitionDataSingleWriter` with its per-row `getOutputRow` projection replaced by a + * pruned view over the Arrow batch. + * + * `BaseDynamicPartitionDataWriter.writeRecord` projects every row through + * `UnsafeProjection.create(dataColumns, allColumns)` for the sole purpose of dropping the + * partition and bucket columns. On a columnar batch that pruning is free, so the only thing this + * subclass changes is where the output row comes from. Partition-change detection, writer renewal + * and `maxRecordsPerFile` are all inherited unchanged, and still see the full row. + * + * The caller must call [[setDataRow]] immediately before each `writeWithMetrics`. The field is + * cleared on use so that a missed call fails loudly instead of silently rewriting the previous + * row. + */ +private[comet] class CometRowViewDynamicPartitionWriter( + description: WriteJobDescription, + taskAttemptContext: TaskAttemptContext, + committer: FileCommitProtocol) + extends DynamicPartitionDataSingleWriter(description, taskAttemptContext, committer) { + + private var dataRow: InternalRow = _ + + def setDataRow(row: InternalRow): Unit = { + dataRow = row + } + + override protected def writeRecord(record: InternalRow): Unit = { + val outputRow = dataRow + if (outputRow == null) { + throw new IllegalStateException("setDataRow must be called before each write") + } + dataRow = null + currentWriter.write(outputRow) + statsTrackers.foreach(_.newRow(currentWriter.path(), outputRow)) + recordsInFile += 1 + } +} diff --git a/spark/src/main/spark-3.x/org/apache/comet/shims/ShimCometWriteFilesExec.scala b/spark/src/main/spark-3.x/org/apache/comet/shims/ShimCometWriteFilesExec.scala new file mode 100644 index 00000000000..7d5e1bbfcc0 --- /dev/null +++ b/spark/src/main/spark-3.x/org/apache/comet/shims/ShimCometWriteFilesExec.scala @@ -0,0 +1,36 @@ +/* + * 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.shims + +import org.apache.spark.sql.catalyst.expressions.Attribute +import org.apache.spark.sql.execution.UnaryExecNode + +/** + * Base type for [[org.apache.spark.sql.comet.CometRowViewWriteFilesExec]] on Spark 3.x. + * + * Spark 3.x has no `WriteFilesExecBase` trait (added in 4.0): `V1WritesUtils.getWriteFilesOpt` + * matches the concrete `WriteFilesExec` case class, so a Comet node can never be picked up as the + * write node there. The row-view write is gated to Spark 4.0+ in `EliminateRedundantTransitions` + * and this shim exists only so that the shared sources compile against 3.x. It mirrors the + * members that the 4.x `WriteFilesExecBase` supplies. + */ +trait ShimCometWriteFilesExec extends UnaryExecNode { + override def output: Seq[Attribute] = Seq.empty +} diff --git a/spark/src/main/spark-4.x/org/apache/comet/shims/ShimCometWriteFilesExec.scala b/spark/src/main/spark-4.x/org/apache/comet/shims/ShimCometWriteFilesExec.scala new file mode 100644 index 00000000000..58cc91e6ff0 --- /dev/null +++ b/spark/src/main/spark-4.x/org/apache/comet/shims/ShimCometWriteFilesExec.scala @@ -0,0 +1,37 @@ +/* + * 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.shims + +import org.apache.spark.sql.execution.datasources.WriteFilesExecBase + +/** + * Base type for [[org.apache.spark.sql.comet.CometRowViewWriteFilesExec]]. + * + * Spark 4.0 factored `WriteFilesExec`'s contract out into the `WriteFilesExecBase` trait, and + * `V1WritesUtils.getWriteFilesOpt` matches on that trait. Extending it is therefore what makes + * Spark recognize Comet's node as the write node and drive it through + * `FileFormatWriter.executeWrite` -> `SparkPlan.executeWrite` -> `doExecuteWrite`, keeping the + * commit protocol, stats trackers and `_SUCCESS` handling on Spark's side. + * + * Spark 3.x has no such trait - `getWriteFilesOpt` matches the concrete `WriteFilesExec` case + * class - so the row-view write is gated to Spark 4.0+ in `EliminateRedundantTransitions`. The + * 3.x variant of this shim exists only to keep the shared sources compiling. + */ +trait ShimCometWriteFilesExec extends WriteFilesExecBase diff --git a/spark/src/test/scala/org/apache/comet/exec/CometWriteRowViewSuite.scala b/spark/src/test/scala/org/apache/comet/exec/CometWriteRowViewSuite.scala new file mode 100644 index 00000000000..9b04d152d0b --- /dev/null +++ b/spark/src/test/scala/org/apache/comet/exec/CometWriteRowViewSuite.scala @@ -0,0 +1,417 @@ +/* + * 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.exec + +import java.util.concurrent.TimeUnit + +import org.apache.spark.sql.{CometTestBase, DataFrame, QueryTest} +import org.apache.spark.sql.comet.CometRowViewWriteFilesExec +import org.apache.spark.sql.execution.{QueryExecution, SparkPlan} +import org.apache.spark.sql.execution.adaptive.AdaptiveSparkPlanExec +import org.apache.spark.sql.functions.{col, expr} +import org.apache.spark.sql.util.QueryExecutionListener + +import org.apache.comet.CometConf +import org.apache.comet.CometSparkSessionExtensions.isSpark40Plus + +/** + * Tests for `spark.comet.exec.write.rowView.enabled`, which replaces Spark's `WriteFilesExec` + * with a Comet node that drives Spark's own `OutputWriter` from Arrow batches instead of from + * materialized `UnsafeRow`s. + * + * The bar for every case here is that enabling the config changes nothing observable except the + * plan: the data written must match what the same Comet plan produces with the config off. The + * baseline is deliberately the same plan rather than vanilla Spark, so any difference is + * attributable to the write node and not to the scan. + */ +class CometWriteRowViewSuite extends CometTestBase { + + /** + * The node replaces `WriteFilesExec` by extending `WriteFilesExecBase`, which Spark 4.0 + * introduced. On 3.4 / 3.5 `V1WritesUtils.getWriteFilesOpt` matches the concrete case class + * instead, so the rewrite is declined there and there is nothing to assert. + */ + private def testRowView(name: String)(f: => Unit): Unit = test(name) { + assume(isSpark40Plus, "row view writes require Spark 4.0+") + f + } + + testRowView("row view write is off by default") { + withParquetSource { source => + withTempPath { out => + val plan = captureWritePlan(source.write.mode("overwrite").parquet(out.toString)) + assert(countRowViewWrites(plan) == 0, s"row view write should be opt-in, got:\n$plan") + } + } + } + + testRowView("row view write is used for an unpartitioned parquet write") { + withParquetSource { source => + withTempPath { out => + val plan = captureWritePlan { + withRowView(source.write.mode("overwrite").parquet(out.toString)) + } + assert( + countRowViewWrites(plan) == 1, + s"expected a CometRowViewWriteFiles in the write plan, got:\n$plan") + } + } + } + + testRowView("row view write is used for a dynamically partitioned write") { + withParquetSource { source => + withTempPath { out => + val plan = captureWritePlan { + withRowView(source.write.mode("overwrite").partitionBy("part").parquet(out.toString)) + } + assert( + countRowViewWrites(plan) == 1, + s"expected a CometRowViewWriteFiles for the partitioned write, got:\n$plan") + } + } + } + + testRowView("row view write is used for a bucketed write") { + withParquetSource { source => + // `INSERT INTO` on a pre-created bucketed table rather than `bucketBy(...).saveAsTable`: + // the latter nests the write inside a `SaveAsV1TableCommand` whose inner plan the query + // execution listener does not surface, so the plan assertion could not be made. + withTable("bucketed_row_view") { + spark.sql(""" + CREATE TABLE bucketed_row_view ( + id BIGINT, + struct_col STRUCT, + arr_col ARRAY) + USING parquet + CLUSTERED BY (id) INTO 4 BUCKETS""") + val projected = source.select("id", "struct_col", "arr_col") + projected.createOrReplaceTempView("bucket_src") + + val plan = captureWritePlan { + withRowView( + spark.sql("INSERT INTO bucketed_row_view SELECT id, struct_col, arr_col FROM " + + "bucket_src")) + } + assert( + countRowViewWrites(plan) == 1, + s"expected a CometRowViewWriteFiles for the bucketed write, got:\n$plan") + checkAnswer(spark.table("bucketed_row_view"), projected) + } + } + } + + testRowView("row view write is declined for a schema of only flat data columns") { + withParquetSource { source => + withTempPath { out => + val flat = source.selectExpr("id", "int_col", "str_col", "date_col", "ts_col") + val plan = captureWritePlan { + withRowView(flat.write.mode("overwrite").parquet(out.toString)) + } + assert( + countRowViewWrites(plan) == 0, + s"a flat schema is not worth the row view, got:\n$plan") + checkAnswer(spark.read.parquet(out.toString), flat) + } + } + } + + testRowView( + "row view write is declined when a flat write is only complex in its partition column") { + withParquetSource { source => + withTempPath { out => + // `part` is the partition column, so it never reaches the OutputWriter and must not count + // towards the complex-type gate. Every remaining data column here is flat. + val flat = source.selectExpr("id", "int_col", "str_col", "part") + val plan = captureWritePlan { + withRowView(flat.write.mode("overwrite").partitionBy("part").parquet(out.toString)) + } + assert(countRowViewWrites(plan) == 0, s"partition columns must not count, got:\n$plan") + } + } + } + + testRowView("row view write is declined when concurrent output file writers are enabled") { + withParquetSource { source => + withTempPath { out => + // Above 0, V1Writes plants no sort and FileFormatWriter picks + // DynamicPartitionDataConcurrentWriter, whose spill path is typed on UnsafeRow. + withSQLConf("spark.sql.maxConcurrentOutputFileWriters" -> "10") { + val plan = captureWritePlan { + withRowView(source.write.mode("overwrite").partitionBy("part").parquet(out.toString)) + } + assert( + countRowViewWrites(plan) == 0, + s"concurrent output writers must decline the row view, got:\n$plan") + } + checkAnswer( + spark.read + .schema(source.schema) + .parquet(out.toString) + .select(source.columns.map(col): _*), + source) + } + } + } + + testRowView("row view write produces the same data - unpartitioned") { + withParquetSource { source => + assertSameWrite( + source, + (df, path) => df.write.mode("overwrite").parquet(path), + path => spark.read.schema(source.schema).parquet(path)) + } + } + + testRowView("row view write produces the same data - dynamic partitions") { + withParquetSource { source => + assertSameWrite( + source, + (df, path) => df.write.mode("overwrite").partitionBy("part").parquet(path), + // `schema` pins the partition column back to string; reading a partitioned directory + // otherwise infers `part` as int. + path => + spark.read.schema(source.schema).parquet(path).select(source.columns.map(col): _*)) + } + } + + testRowView("row view write produces the same data - two partition columns") { + withParquetSource { source => + val partitioned = source.withColumn("part2", expr("cast(id % 7 as string)")) + assertSameWrite( + partitioned, + (df, path) => df.write.mode("overwrite").partitionBy("part", "part2").parquet(path), + path => + spark.read + .schema(partitioned.schema) + .parquet(path) + .select(partitioned.columns.map(col): _*)) + } + } + + testRowView("row view write produces the same data - deeply nested types with partitions") { + withDeeplyNestedSource { source => + val partitioned = source.withColumn("part", expr("cast(id % 4 as string)")) + assertSameWrite( + partitioned, + (df, path) => df.write.mode("overwrite").partitionBy("part").parquet(path), + path => + spark.read + .schema(partitioned.schema) + .parquet(path) + .select(partitioned.columns.map(col): _*)) + } + } + + testRowView("row view write produces the same data - null and empty partition values") { + // Empty strings become null partition directories via V1Writes' Empty2Null projection, and + // nulls land in __HIVE_DEFAULT_PARTITION__. Both are computed from the row this node hands + // to the writer, so they exercise the partition-value path rather than the payload. + withNullPartitionSource { source => + assertSameWrite( + source, + (df, path) => df.write.mode("overwrite").partitionBy("part").parquet(path), + path => + spark.read.schema(source.schema).parquet(path).select(source.columns.map(col): _*)) + } + } + + testRowView("row view write produces the same data - maxRecordsPerFile across partitions") { + withParquetSource { source => + Seq("100", "0").foreach { maxRecords => + withSQLConf("spark.sql.files.maxRecordsPerFile" -> maxRecords) { + assertSameWrite( + source, + (df, path) => df.write.mode("overwrite").partitionBy("part").parquet(path), + path => + spark.read.schema(source.schema).parquet(path).select(source.columns.map(col): _*)) + } + } + } + } + + testRowView("row view write produces the same data - orc") { + withParquetSource { source => + assertSameWrite( + source, + (df, path) => df.write.mode("overwrite").partitionBy("part").format("orc").save(path), + path => + spark.read + .schema(source.schema) + .format("orc") + .load(path) + .select(source.columns.map(col): _*)) + } + } + + testRowView("row view write produces the same partition layout") { + withParquetSource { source => + withTempPath { baseline => + withTempPath { rowView => + source.write.mode("overwrite").partitionBy("part").parquet(baseline.toString) + withRowView( + source.write.mode("overwrite").partitionBy("part").parquet(rowView.toString)) + assert( + partitionDirs(rowView.toString) == partitionDirs(baseline.toString), + "partition directories differ") + } + } + } + } + + /** + * Writes `source` with and without the row view write and requires the results to be identical, + * both in content and in row count. + */ + private def assertSameWrite( + source: DataFrame, + write: (DataFrame, String) => Unit, + read: String => DataFrame): Unit = { + withTempPath { baseline => + withTempPath { rowView => + write(source, baseline.toString) + withRowView(write(source, rowView.toString)) + + val expected = read(baseline.toString) + val actual = read(rowView.toString) + assert(actual.count() == expected.count(), "row count differs") + QueryTest.checkAnswer(actual, expected.collect().toSeq) + } + } + } + + private def partitionDirs(path: String): Set[String] = { + val root = new java.io.File(path) + Option(root.listFiles()) + .map(_.filter(_.isDirectory).map(_.getName).toSet) + .getOrElse(Set.empty) + } + + private def withRowView(f: => Unit): Unit = + withSQLConf(CometConf.COMET_EXEC_WRITE_ROW_VIEW_ENABLED.key -> "true")(f) + + /** + * Materializes a small table on disk and hands back a DataFrame reading it, so the write under + * test is fed by a Comet columnar scan rather than a row-based local relation. + */ + private def withParquetSource(f: DataFrame => Unit): Unit = { + withTempPath { dir => + val df = spark + .range(2000) + .selectExpr( + "id", + "cast(id as int) as int_col", + "cast(id as short) as short_col", + "cast(id % 2 as boolean) as bool_col", + "cast(id as double) as double_col", + "cast(id as decimal(20,4)) as dec_col", + "cast(id as string) as str_col", + "case when id % 7 = 0 then null else concat('v_', cast(id as string)) end as null_str", + "cast(cast(id as string) as binary) as bin_col", + "date_add(to_date('2024-01-01'), cast(id % 365 as int)) as date_col", + "timestamp_micros(id * 1000000) as ts_col", + "named_struct('a', cast(id as int), 'b', cast(id as string)) as struct_col", + "array(cast(id as int), cast(id + 1 as int)) as arr_col", + "map('k', cast(id as string)) as map_col", + "cast(id % 3 as string) as part") + df.write.mode("overwrite").parquet(dir.toString) + f(spark.read.parquet(dir.toString)) + } + } + + /** Partition column carrying nulls and empty strings alongside ordinary values. */ + private def withNullPartitionSource(f: DataFrame => Unit): Unit = { + withTempPath { dir => + val df = spark + .range(1000) + .selectExpr( + "id", + "named_struct('a', cast(id as int), 'b', cast(id as string)) as struct_col", + "array(cast(id as int)) as arr_col", + """case when id % 3 = 0 then null + when id % 3 = 1 then '' + else cast(id % 5 as string) end as part""") + df.write.mode("overwrite").parquet(dir.toString) + f(spark.read.parquet(dir.toString)) + } + } + + /** Four levels of nesting, mixing struct, array and map at each level. */ + private def withDeeplyNestedSource(f: DataFrame => Unit): Unit = { + withTempPath { dir => + val df = spark + .range(1000) + .selectExpr( + "id", + """named_struct( + 'l1', named_struct( + 'l2', named_struct( + 'l3', named_struct('v', cast(id as int), 'n', concat('x_', cast(id as string))), + 'arr', array(cast(id as int), cast(id + 1 as int))), + 'c', cast(id % 100 as int)), + 'id', id) as deep_struct""", + """array( + named_struct('id', cast(id as int), + 'tags', array(concat('t_', cast(id as string)))), + named_struct('id', cast(id + 1 as int), 'tags', array('t_x', 't_y')) + ) as arr_of_structs""", + """map('k', array(named_struct('a', cast(id as int), + 'b', cast(id as string)))) as map_of_arr_structs""", + // A null at every nesting level, which is where the offset bookkeeping differs most. + "case when id % 5 = 0 then null else array(array(cast(id as int))) end as nested_null") + df.write.mode("overwrite").parquet(dir.toString) + f(spark.read.parquet(dir.toString)) + } + } + + /** `AdaptiveSparkPlanExec` is a leaf node, so a plain `foreach` stops at the AQE boundary. */ + private def flatten(plan: SparkPlan): Seq[SparkPlan] = plan match { + case a: AdaptiveSparkPlanExec => a +: flatten(a.executedPlan) + case p => p +: p.children.flatMap(flatten) + } + + private def countRowViewWrites(plan: SparkPlan): Int = + flatten(plan).count(_.isInstanceOf[CometRowViewWriteFilesExec]) + + private def captureWritePlan(writeOp: => Unit): SparkPlan = { + @volatile var capturedPlan: Option[QueryExecution] = None + val listener = new QueryExecutionListener { + override def onSuccess(funcName: String, qe: QueryExecution, durationNs: Long): Unit = + capturedPlan = Some(qe) + override def onFailure( + funcName: String, + qe: QueryExecution, + exception: Exception): Unit = {} + } + spark.listenerManager.register(listener) + try { + writeOp + // The listener fires asynchronously off the listener bus, which is not reachable from this + // package, so poll for it. + val deadline = System.nanoTime() + TimeUnit.SECONDS.toNanos(30) + while (capturedPlan.isEmpty && System.nanoTime() < deadline) { + Thread.sleep(50) + } + assert(capturedPlan.isDefined, "no execution plan captured for the write") + capturedPlan.get.executedPlan + } finally { + spark.listenerManager.unregister(listener) + } + } +} From b3efec13e9913745fb0d079d5b7886a474aa57e2 Mon Sep 17 00:00:00 2001 From: Andy Grove Date: Wed, 2 Sep 2026 11:15:05 -0600 Subject: [PATCH 2/2] feat: add write row-view benchmark and widen the gate for partitioned writes CometWriteRowViewBenchmark measures three arms (Spark, Comet writing via UnsafeRow, Comet writing via row view) across flat, nested and deeply nested schemas, unpartitioned and partitioned, at two compression settings. Each arm asserts through a QueryExecutionListener that it planned the write node its name claims and warns into the results file when it does not, so a silently declined arm cannot masquerade as a measurement. 1M rows, M3 Ultra, timing the whole scan -> write -> file pipeline. Gain is the row-view arm against the UnsafeRow arm, which differ only in the write node (uncompressed / snappy): flat 10-column unpartitioned 0.2% / 0.6% <- declined, noise floor nested unpartitioned 20.4% / 13.3% deeply nested unpartitioned 12.3% / 13.3% flat 10-column partitioned 3.7% / 4.5% nested partitioned 16.3% / 16.1% deeply nested partitioned 16.5% / 16.6% The flat partitioned row prompted a gate change. A partitioned or bucketed write removes two projections per row, the columnar-to-row transition and BaseDynamicPartitionDataWriter.getOutputRow, so it is worth doing whatever the schema; only an unpartitioned write, which removes just the transition, still needs a complex data column to clear the noise floor. Require the complex type only in that case. Co-Authored-By: Claude Opus 5 (1M context) Claude-Session: https://claude.ai/code/session_01DmHrzzDLtwh1nLHZUxTVKm --- .../rules/EliminateRedundantTransitions.scala | 33 +- .../comet/exec/CometWriteRowViewSuite.scala | 18 +- .../CometWriteRowViewBenchmark.scala | 314 ++++++++++++++++++ 3 files changed, 349 insertions(+), 16 deletions(-) create mode 100644 spark/src/test/scala/org/apache/spark/sql/benchmark/CometWriteRowViewBenchmark.scala diff --git a/spark/src/main/scala/org/apache/comet/rules/EliminateRedundantTransitions.scala b/spark/src/main/scala/org/apache/comet/rules/EliminateRedundantTransitions.scala index f9990f56289..ed7e7643164 100644 --- a/spark/src/main/scala/org/apache/comet/rules/EliminateRedundantTransitions.scala +++ b/spark/src/main/scala/org/apache/comet/rules/EliminateRedundantTransitions.scala @@ -207,13 +207,23 @@ case class EliminateRedundantTransitions(session: SparkSession) * The Comet columnar producer under a write's columnar-to-row transition, or `None` when there * is no transition to strip or the write would not gain from removing it. * - * Declines a write whose data columns are all flat. There the projection this removes is a - * generated fixed-width copy that measures inside the noise of a Parquet write, which does not - * pay for putting a reused mutable row in front of Spark's writer. The saving only becomes real - * once a struct, array or map is in play, because then the projection has to build nested + * How much there is to gain depends on how many projections the write performs per row. + * + * A partitioned or bucketed write performs two: the columnar-to-row transition, and + * `BaseDynamicPartitionDataWriter.writeRecord`'s `getOutputRow`, which exists only to strip the + * partition and bucket columns. Both are removed here, and that is worth doing whatever the + * schema: `CometWriteRowViewBenchmark` measures 5% on a flat 10-column schema and 15% once a + * struct, array or map is present, against a 1% noise floor. + * + * An unpartitioned write performs only the transition, so a flat schema is declined. There the + * projection is a generated fixed-width copy that measures inside the noise of a Parquet write, + * which does not pay for putting a reused mutable row in front of Spark's writer. The saving + * becomes real once a complex type is in play, because the projection then has to build nested * `UnsafeRow` / `UnsafeArrayData` with offset-and-length bookkeeping that the writer - * immediately walks back out. Partition and bucket columns do not count: they never reach the - * `OutputWriter`. + * immediately walks back out. + * + * Partition and bucket columns never count towards the complex-type test: they are stripped + * before the `OutputWriter` sees the row. */ private def columnarChildForWrite(w: WriteFilesExec): Option[SparkPlan] = { val columnarChild = w.child match { @@ -221,11 +231,14 @@ case class EliminateRedundantTransitions(session: SparkSession) case CometNativeColumnarToRowExec(child) => Some(child) case _ => None } + val removesWriterProjection = w.partitionColumns.nonEmpty || w.bucketSpec.isDefined columnarChild.filter { child => - val partitionIds = w.partitionColumns.map(_.exprId).toSet - child.output - .filterNot(a => partitionIds.contains(a.exprId)) - .exists(a => isComplex(a.dataType)) + removesWriterProjection || { + val partitionIds = w.partitionColumns.map(_.exprId).toSet + child.output + .filterNot(a => partitionIds.contains(a.exprId)) + .exists(a => isComplex(a.dataType)) + } } } diff --git a/spark/src/test/scala/org/apache/comet/exec/CometWriteRowViewSuite.scala b/spark/src/test/scala/org/apache/comet/exec/CometWriteRowViewSuite.scala index 9b04d152d0b..1f087bccd90 100644 --- a/spark/src/test/scala/org/apache/comet/exec/CometWriteRowViewSuite.scala +++ b/spark/src/test/scala/org/apache/comet/exec/CometWriteRowViewSuite.scala @@ -132,18 +132,24 @@ class CometWriteRowViewSuite extends CometTestBase { } } - testRowView( - "row view write is declined when a flat write is only complex in its partition column") { + testRowView("row view write is used for a flat partitioned write") { withParquetSource { source => + // A partitioned write removes `BaseDynamicPartitionDataWriter.getOutputRow` as well as the + // transition, which is worth doing even when every data column is flat. Measured at 5% on + // this shape by `CometWriteRowViewBenchmark`, against a 1% noise floor. + val flat = source.selectExpr("id", "int_col", "str_col", "part") withTempPath { out => - // `part` is the partition column, so it never reaches the OutputWriter and must not count - // towards the complex-type gate. Every remaining data column here is flat. - val flat = source.selectExpr("id", "int_col", "str_col", "part") val plan = captureWritePlan { withRowView(flat.write.mode("overwrite").partitionBy("part").parquet(out.toString)) } - assert(countRowViewWrites(plan) == 0, s"partition columns must not count, got:\n$plan") + assert( + countRowViewWrites(plan) == 1, + s"a flat partitioned write still removes a projection, got:\n$plan") } + assertSameWrite( + flat, + (df, path) => df.write.mode("overwrite").partitionBy("part").parquet(path), + path => spark.read.schema(flat.schema).parquet(path).select(flat.columns.map(col): _*)) } } diff --git a/spark/src/test/scala/org/apache/spark/sql/benchmark/CometWriteRowViewBenchmark.scala b/spark/src/test/scala/org/apache/spark/sql/benchmark/CometWriteRowViewBenchmark.scala new file mode 100644 index 00000000000..4ce6dc28c69 --- /dev/null +++ b/spark/src/test/scala/org/apache/spark/sql/benchmark/CometWriteRowViewBenchmark.scala @@ -0,0 +1,314 @@ +/* + * 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.benchmark + +import java.io.File +import java.util.concurrent.TimeUnit + +import org.apache.spark.SparkConf +import org.apache.spark.benchmark.Benchmark +import org.apache.spark.sql.SparkSession +import org.apache.spark.sql.comet.CometRowViewWriteFilesExec +import org.apache.spark.sql.execution.{QueryExecution, SparkPlan} +import org.apache.spark.sql.execution.adaptive.AdaptiveSparkPlanExec +import org.apache.spark.sql.execution.datasources.WriteFilesExec +import org.apache.spark.sql.internal.SQLConf +import org.apache.spark.sql.util.QueryExecutionListener + +import org.apache.comet.{CometConf, CometSparkSessionExtensions} + +/** + * Benchmark for `spark.comet.exec.write.rowView.enabled`, which replaces Spark's `WriteFilesExec` + * with a Comet node that drives Spark's own `OutputWriter` from Arrow batches. + * + * Three arms per case: + * - Spark's own vectorized read and write + * - Comet scan plus `CometColumnarToRowExec`, which materializes an `UnsafeRow` per row + * - Comet scan plus `CometRowViewWriteFilesExec`, which materializes none + * + * Every case is run both unpartitioned and partitioned, because that is the comparison this + * benchmark exists to make. An unpartitioned write pays for one `UnsafeProjection` per row, in + * the columnar-to-row transition. A partitioned write pays for a second one inside + * `BaseDynamicPartitionDataWriter.writeRecord`, which projects the row again purely to strip the + * partition columns, so removing the transition alone would not help it. The row-view node + * removes both. + * + * The whole `scan -> write -> file` pipeline is timed, so the projections are only part of what + * is measured; the point is how much of a real write they cost. Cases run at two compression + * settings because parquet-mr encoding is what that cost has to be measured against, and a flat + * schema is included as a noise floor: the rule declines it, so both Comet arms run the identical + * plan there. + * + * To run this benchmark: + * {{{ + * SPARK_GENERATE_BENCHMARK_FILES=1 make benchmark-org.apache.spark.sql.benchmark.CometWriteRowViewBenchmark + * }}} + * + * Results will be written to "spark/benchmarks/CometWriteRowViewBenchmark-**results.txt". + */ +object CometWriteRowViewBenchmark extends CometBenchmarkBase { + + override def getSparkSession: SparkSession = { + val conf = new SparkConf() + .setAppName("CometWriteRowViewBenchmark") + .set("spark.master", "local[1]") + .setIfMissing("spark.driver.memory", "3g") + .setIfMissing("spark.executor.memory", "3g") + .set("spark.memory.offHeap.enabled", "true") + .set("spark.memory.offHeap.size", "2g") + // Required: `isCometLoaded` disables Comet entirely when `spark.comet.shuffle.enabled` + // (default true) is set without Comet's shuffle manager, which would silently turn every + // Comet case below into another Spark run. `spark.shuffle.manager` is static and must be + // set before the context starts. CometShuffleManager falls back to Spark's shuffle when + // Comet is disabled, so the Spark baseline case is unaffected. + .set( + "spark.shuffle.manager", + "org.apache.spark.sql.comet.execution.shuffle.CometShuffleManager") + + val sparkSession = SparkSession + .builder() + .config(conf) + .withExtensions(new CometSparkSessionExtensions) + .getOrCreate() + + sparkSession.conf.set(SQLConf.ANSI_ENABLED.key, "false") + sparkSession.conf.set(SQLConf.PARQUET_VECTORIZED_READER_ENABLED.key, "true") + sparkSession.conf.set(SQLConf.WHOLESTAGE_CODEGEN_ENABLED.key, "true") + sparkSession.conf.set(CometConf.COMET_ENABLED.key, "false") + sparkSession.conf.set(CometConf.COMET_EXEC_ENABLED.key, "false") + // Comet's scan rejects tinyint/smallint unless this check is off, which would drop the + // fixed-width case back to Spark's scan in every arm. + sparkSession.conf.set(CometConf.COMET_PARQUET_UNSIGNED_SMALL_INT_CHECK.key, "false") + + sparkSession + } + + private val sparkConfigs = Seq(CometConf.COMET_ENABLED.key -> "false") + + private val unsafeRowConfigs = Seq( + CometConf.COMET_ENABLED.key -> "true", + CometConf.COMET_EXEC_ENABLED.key -> "true", + CometConf.COMET_EXEC_WRITE_ROW_VIEW_ENABLED.key -> "false") + + private val rowViewConfigs = Seq( + CometConf.COMET_ENABLED.key -> "true", + CometConf.COMET_EXEC_ENABLED.key -> "true", + CometConf.COMET_EXEC_WRITE_ROW_VIEW_ENABLED.key -> "true") + + private val diagnostics = sys.env.contains("COMET_BENCH_DIAG") + + private def addWriteCases( + benchmark: Benchmark, + outputDir: File, + codec: String, + partitioned: Boolean): Unit = { + + def write(target: String, configs: Seq[(String, String)]): Unit = + withSQLConf(configs :+ (SQLConf.PARQUET_COMPRESSION.key -> codec): _*) { + val writer = spark.sql("SELECT * FROM parquetV1Table").write.mode("overwrite") + val partitionedWriter = if (partitioned) writer.partitionBy("part") else writer + partitionedWriter.parquet(new File(outputDir, target).getCanonicalPath) + } + + // The arms below differ only in which write node feeds parquet-mr, so the comparison is + // meaningless unless each arm actually plans the node it names. Run every arm once up front + // and warn loudly into the results file when it does not. + def verify( + name: String, + target: String, + configs: Seq[(String, String)], + expected: String): Unit = { + val (writeNode, tree) = captureWriteNode(write(target, configs)) + // scalastyle:off println + println(s" [plan check] $name -> ${writeNode.getOrElse("")}") + if (diagnostics) println(tree.getOrElse("")) + // scalastyle:on println + if (!writeNode.contains(expected)) { + val border = "=" * 80 + benchmark.out.println(s""" + |$border + |WARNING: the "$name" case did not plan $expected but + |${writeNode.getOrElse("no write node at all")}, so it is not measuring what its + |name says. Treat this row as invalid. + |$border""".stripMargin) + } + } + + benchmark.addCase("Spark") { _ => + write("spark", sparkConfigs) + } + + benchmark.addCase("Comet write via UnsafeRow") { _ => + write("comet-unsaferow", unsafeRowConfigs) + } + + benchmark.addCase("Comet write via row view") { _ => + write("comet-rowview", rowViewConfigs) + } + + verify("Spark", "spark", sparkConfigs, "WriteFiles") + verify("Comet write via UnsafeRow", "comet-unsaferow", unsafeRowConfigs, "WriteFiles") + // The rule declines a schema whose data columns are all flat, so the row-view arm is expected + // to plan the ordinary write node there. Deriving the expectation from the schema also asserts + // the gate, and makes the flat cases an explicit noise floor rather than a silent no-op. + verify( + "Comet write via row view", + "comet-rowview", + rowViewConfigs, + if (expectRowView(partitioned)) "CometRowViewWriteFiles" else "WriteFiles") + } + + /** + * Mirrors the gate in `EliminateRedundantTransitions.columnarChildForWrite`: a partitioned + * write always qualifies, because it removes the writer's own projection too, while an + * unpartitioned one needs a complex data column to be worth it. + */ + private def expectRowView(partitioned: Boolean): Boolean = partitioned || hasComplexDataColumn() + + /** Whether any data column is a struct, array or map. */ + private def hasComplexDataColumn(): Boolean = + spark + .table("parquetV1Table") + .schema + .fields + .filterNot(_.name == "part") + .exists(f => + f.dataType.typeName match { + case "struct" | "array" | "map" => true + case _ => false + }) + + /** Names the write node in the executed plan of a write. */ + private def captureWriteNode(writeOp: => Unit): (Option[String], Option[String]) = { + @volatile var captured: Option[QueryExecution] = None + val listener = new QueryExecutionListener { + override def onSuccess(funcName: String, qe: QueryExecution, durationNs: Long): Unit = + captured = Some(qe) + override def onFailure(funcName: String, qe: QueryExecution, e: Exception): Unit = {} + } + spark.listenerManager.register(listener) + try { + writeOp + val deadline = System.nanoTime() + TimeUnit.SECONDS.toNanos(30) + while (captured.isEmpty && System.nanoTime() < deadline) Thread.sleep(50) + val names = captured.map { qe => + def flatten(p: SparkPlan): Seq[SparkPlan] = p match { + case a: AdaptiveSparkPlanExec => a +: flatten(a.executedPlan) + case other => other +: other.children.flatMap(flatten) + } + flatten(qe.executedPlan) + .collect { + case p: CometRowViewWriteFilesExec => p.nodeName + case p: WriteFilesExec => p.nodeName + } + .mkString(", ") + } + val tree = captured.map { qe => + s" comet.enabled=${spark.conf.get(CometConf.COMET_ENABLED.key, "unset")} " + + s"rowView=${spark.conf.get(CometConf.COMET_EXEC_WRITE_ROW_VIEW_ENABLED.key, "unset")}\n" + + qe.executedPlan.treeString + } + (names, tree) + } finally { + spark.listenerManager.unregister(listener) + } + } + + private def writeBenchmark(name: String, values: Int, codec: String, partitioned: Boolean)( + columns: Seq[String]): Unit = { + val label = if (partitioned) "partitioned" else "unpartitioned" + val benchmark = new Benchmark(s"$name ($label, $codec)", values, output = output) + withTempPath { dir => + withTempTable("parquetV1Table") { + // A low-cardinality partition column, which is the shape a partitioned write is normally + // given: long runs of equal partition values, so `DynamicPartitionDataSingleWriter` opens + // few files and the per-row cost dominates the per-file cost. + val partitionColumn = if (partitioned) { + Seq("cast(id % 8 as string) as part") + } else { + Seq.empty + } + prepareTable(dir, spark.range(values).selectExpr(columns ++ partitionColumn: _*)) + withTempPath { outputDir => + outputDir.mkdirs() + addWriteCases(benchmark, outputDir, codec, partitioned) + if (!diagnostics) benchmark.run() + } + } + } + } + + /** A flat schema, which the rule declines. Both Comet arms run the same plan: a noise floor. */ + private val fixedWidth = Seq( + "id as long_col", + "cast(id as int) as int_col", + "cast(id as short) as short_col", + "cast(id as byte) as byte_col", + "cast(id % 2 as boolean) as bool_col", + "cast(id as float) as float_col", + "cast(id as double) as double_col", + "date_add(to_date('2024-01-01'), cast(id % 365 as int)) as date_col", + "cast(id * 2 as long) as long_col2", + "cast(id * 3 as int) as int_col2") + + private val nested = Seq( + "id", + "named_struct('a', cast(id as int), 'b', cast(id as string)) as simple_struct", + "array(cast(id as int), cast(id + 1 as int), cast(id + 2 as int)) as int_array", + "map('k1', cast(id as string), 'k2', cast(id + 1 as string)) as str_map") + + /** The shape that motivates this: several levels mixing struct, array and map. */ + private val deeplyNested = Seq( + "id", + """named_struct( + 'l1', named_struct( + 'l2', named_struct( + 'l3', named_struct('v', cast(id as int), 'n', concat('x_', cast(id as string))), + 'arr', array(cast(id as int), cast(id + 1 as int))), + 'c', cast(id % 100 as int)), + 'id', id) as deep_struct""", + """array( + named_struct('id', cast(id as int), 'tags', array(concat('t_', cast(id as string)))), + named_struct('id', cast(id + 1 as int), 'tags', array('t_x', 't_y')) + ) as arr_of_structs""", + """map('k', array(named_struct('a', cast(id as int), + 'b', cast(id as string)))) as map_of_arr_structs""") + + override def runCometBenchmark(mainArgs: Array[String]): Unit = { + val numRows = 1024 * 1024 + // Codecs can be narrowed from the command line, e.g. `-Dexec.args="uncompressed"`, to keep a + // targeted run short. + val codecs = if (mainArgs.nonEmpty) mainArgs.toSeq else Seq("uncompressed", "snappy") + + val schemas = + Seq(("Fixed width", fixedWidth), ("Nested", nested), ("Deeply nested", deeplyNested)) + + codecs.foreach { codec => + Seq(false, true).foreach { partitioned => + val label = if (partitioned) "partitioned" else "unpartitioned" + schemas.foreach { case (name, columns) => + runBenchmark(s"Parquet write - $name ($label, $codec)") { + writeBenchmark(s"Parquet write - $name", numRows, codec, partitioned)(columns) + } + } + } + } + } +}