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..ed7e7643164 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,82 @@ 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. + * + * 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 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 { + case CometColumnarToRowExec(child) => Some(child) + case CometNativeColumnarToRowExec(child) => Some(child) + case _ => None + } + val removesWriterProjection = w.partitionColumns.nonEmpty || w.bucketSpec.isDefined + columnarChild.filter { child => + removesWriterProjection || { + 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..1f087bccd90 --- /dev/null +++ b/spark/src/test/scala/org/apache/comet/exec/CometWriteRowViewSuite.scala @@ -0,0 +1,423 @@ +/* + * 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 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 => + val plan = captureWritePlan { + withRowView(flat.write.mode("overwrite").partitionBy("part").parquet(out.toString)) + } + 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): _*)) + } + } + + 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) + } + } +} 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) + } + } + } + } + } +}