diff --git a/.github/workflows/pr_build_linux.yml b/.github/workflows/pr_build_linux.yml index 31cc3c34da9..768e7809903 100644 --- a/.github/workflows/pr_build_linux.yml +++ b/.github/workflows/pr_build_linux.yml @@ -303,6 +303,7 @@ jobs: - name: "scans" value: | org.apache.comet.parquet.CometParquetWriterSuite + org.apache.comet.exec.CometWriteRowViewSuite org.apache.comet.parquet.ParquetReadV1Suite org.apache.comet.parquet.ParquetReadV2Suite org.apache.comet.parquet.ParquetReadFromFakeHadoopFsSuite diff --git a/.github/workflows/pr_build_macos.yml b/.github/workflows/pr_build_macos.yml index 7b7fae8fbea..cca9ebb9ee8 100644 --- a/.github/workflows/pr_build_macos.yml +++ b/.github/workflows/pr_build_macos.yml @@ -119,6 +119,7 @@ jobs: - name: "scans" value: | org.apache.comet.parquet.CometParquetWriterSuite + org.apache.comet.exec.CometWriteRowViewSuite org.apache.comet.parquet.ParquetReadV1Suite org.apache.comet.parquet.ParquetReadV2Suite org.apache.comet.parquet.ParquetReadFromFakeHadoopFsSuite diff --git a/spark/src/main/scala/org/apache/comet/CometConf.scala b/spark/src/main/scala/org/apache/comet/CometConf.scala index a7d7b80db1f..814f5afc058 100644 --- a/spark/src/main/scala/org/apache/comet/CometConf.scala +++ b/spark/src/main/scala/org/apache/comet/CometConf.scala @@ -289,6 +289,19 @@ object CometConf extends ShimCometConf { .booleanConf .createWithDefault(false) + val COMET_WRITE_ROW_VIEW_ENABLED: ConfigEntry[Boolean] = + conf(s"$COMET_EXEC_CONFIG_PREFIX.write.rowView.enabled") + .category(CATEGORY_EXEC) + .doc( + "Whether to feed Spark's file write path zero-copy row views over Arrow batches " + + "instead of materializing an UnsafeRow per row. Nothing on the write path requires " + + "an UnsafeRow, so the copy is undone immediately by the writer. Only applies to " + + "unpartitioned, unbucketed writes of a schema containing a struct, array or map, " + + "where the writer is guaranteed not to retain rows and the saving is large enough " + + "to be worth measuring. This feature is experimental.") + .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..c021ae6d84f 100644 --- a/spark/src/main/scala/org/apache/comet/rules/EliminateRedundantTransitions.scala +++ b/spark/src/main/scala/org/apache/comet/rules/EliminateRedundantTransitions.scala @@ -22,12 +22,15 @@ 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, CometColumnarToRowViewExec, CometMapInBatchExec, CometNativeColumnarToRowExec, CometNativeWriteExec, CometPlan, 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.command.DataWritingCommandExec +import org.apache.spark.sql.execution.datasources.{InsertIntoHadoopFsRelationCommand, 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 @@ -91,6 +94,29 @@ case class EliminateRedundantTransitions(session: SparkSession) // Write should be final operation in the plan case ColumnarToRowExec(nativeWrite: CometNativeWriteExec) => nativeWrite + + // Spark's file writers consume `InternalRow` and never need an `UnsafeRow`, so the + // materializing transition below a write can be swapped for a zero-copy row view over the + // Arrow batch. `transformUp` has already rewritten the child into one of the Comet + // transitions by the time these arms are visited. + // + // `plannedWrite` (the default since Spark 3.4) puts the transition under `WriteFilesExec`; + // with it disabled the write command executes its child directly. Both are handled, and both + // are gated on the write being unpartitioned and unbucketed (`rowViewSafeForWrite`) and on + // the schema containing a complex type (`rowView`). + case w: WriteFilesExec + if rowViewSafeForWrite( + w.partitionColumns.nonEmpty, + w.bucketSpec.isDefined, + w.fileFormat.getClass.getName) => + rowView(w.child).map(c => w.withNewChildren(Seq(c))).getOrElse(w) + case d @ DataWritingCommandExec(cmd: InsertIntoHadoopFsRelationCommand, _) + if rowViewSafeForWrite( + cmd.partitionColumns.nonEmpty || cmd.staticPartitions.nonEmpty, + cmd.bucketSpec.isDefined, + cmd.fileFormat.getClass.getName) => + rowView(d.child).map(c => d.withNewChildren(Seq(c))).getOrElse(d) + case c @ ColumnarToRowExec(child) if hasCometNativeChild(child) => val op = createColumnarToRowExec(child) if (c.logicalLink.isEmpty) { @@ -169,6 +195,64 @@ case class EliminateRedundantTransitions(session: SparkSession) } } + /** + * Whether a write can consume the reused, mutable rows produced by + * [[CometColumnarToRowViewExec]] rather than materialized `UnsafeRow`s. + * + * The row view is only correct for a consumer that finishes with a row before pulling the next + * one. That holds for `SingleDirectoryDataWriter`, which is what `FileFormatWriter` picks when + * there are no partition and no bucket columns; it writes the row straight through to the + * `OutputWriter`, and `BasicWriteTaskStatsTracker.newRow` ignores the row entirely. The + * partitioned and bucketed writers do not qualify: + * + * - `FileFormatWriter` requires an ordering on the partition/bucket columns, so a `SortExec` + * sits between this transition and the writer, and `UnsafeExternalSorter` needs + * `UnsafeRow`. + * - `DynamicPartitionDataConcurrentWriter` spills through `UnsafeKVExternalSorter.insertKV`, + * which is typed on `UnsafeRow`. + * + * The format check keeps this to Spark's own `FileFormat` implementations. Their + * `OutputWriter`s encode each row on the spot (Parquet through `ParquetWriteSupport`, ORC + * through `OrcSerializer` into a `VectorizedRowBatch`, the text formats directly), whereas a + * third-party format is free to buffer the `InternalRow` it is handed. + */ + private def rowViewSafeForWrite( + partitioned: Boolean, + bucketed: Boolean, + fileFormat: String): Boolean = + CometConf.COMET_WRITE_ROW_VIEW_ENABLED.get() && + !partitioned && !bucketed && + fileFormat.startsWith("org.apache.spark.sql.execution.datasources.") + + /** + * Rewrites a Comet columnar-to-row transition into the zero-copy row view. Returns `None` for + * anything else, which leaves the plan untouched - notably for the `WriteFilesExec` under a + * `DataWritingCommandExec`, and for a write whose input was never columnar to begin with. + * + * Also declines a schema of nothing but flat types. There the `UnsafeProjection` this replaces + * is a generated fixed-width copy that measures at 0-2% of a Parquet write, which does not pay + * for the reused-mutable-row hazard. 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 `ParquetWriteSupport` immediately walks back out. See + * `CometParquetWriteBenchmark` for the measurements behind this cut-off. + */ + private def rowView(plan: SparkPlan): Option[SparkPlan] = plan match { + case CometColumnarToRowExec(child) if hasComplexType(child.schema) => + Some(CometColumnarToRowViewExec(child)) + case CometNativeColumnarToRowExec(child) if hasComplexType(child.schema) => + Some(CometColumnarToRowViewExec(child)) + case _ => None + } + + /** Whether the schema has a struct, array or map anywhere in it. */ + private def hasComplexType(schema: StructType): Boolean = { + def isComplex(dataType: DataType): Boolean = dataType match { + case _: StructType | _: ArrayType | _: MapType => true + case _ => false + } + schema.fields.exists(f => isComplex(f.dataType)) + } + /** * If the given plan is a Comet ColumnarToRow transition, returns the columnar child the Python * UDF operator can consume directly. By the time this rule runs the earlier diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/CometColumnarToRowViewExec.scala b/spark/src/main/scala/org/apache/spark/sql/comet/CometColumnarToRowViewExec.scala new file mode 100644 index 00000000000..6966b60e68c --- /dev/null +++ b/spark/src/main/scala/org/apache/spark/sql/comet/CometColumnarToRowViewExec.scala @@ -0,0 +1,92 @@ +/* + * 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 scala.jdk.CollectionConverters._ + +import org.apache.spark.rdd.RDD +import org.apache.spark.sql.catalyst.InternalRow +import org.apache.spark.sql.catalyst.expressions.{Attribute, SortOrder} +import org.apache.spark.sql.catalyst.plans.physical.Partitioning +import org.apache.spark.sql.execution.{ColumnarToRowTransition, SparkPlan} +import org.apache.spark.sql.execution.metric.{SQLMetric, SQLMetrics} +import org.apache.spark.util.Utils + +/** + * A columnar-to-row transition that hands the consumer `ColumnarBatch.rowIterator()` directly + * instead of materializing an `UnsafeRow` per row. + * + * `ColumnarBatch.rowIterator()` returns a single mutable `ColumnarBatchRow` that is advanced over + * the batch, so each row is a zero-copy view over the underlying Arrow buffers. That makes this + * operator roughly free, but it is only correct for a consumer that fully consumes a row before + * requesting the next one and never retains a reference to it. + * + * Spark's file write path is such a consumer: `OutputWriter.write`, `FileFormatDataWriter.write` + * and `WriteTaskStatsTracker.newRow` all take a plain `InternalRow`, and `ParquetWriteSupport` + * reads fields through `SpecializedGetters` and encodes them immediately. Nothing there requires + * an `UnsafeRow`, so the `UnsafeProjection` performed by [[CometColumnarToRowExec]] is a copy + * that the writer only undoes again. + * + * This is deliberately NOT a `CodegenSupport` node: whole-stage codegen would generate an + * `UnsafeRowWriter` loop and reintroduce exactly the copy this operator exists to avoid. + * + * Only [[org.apache.comet.rules.EliminateRedundantTransitions]] introduces this operator, and + * only where it has proven the consumer is a non-retaining one. Do not use it as a general + * replacement for [[CometColumnarToRowExec]]. + * + * @param child + * The child plan that produces columnar batches + */ +case class CometColumnarToRowViewExec(child: SparkPlan) + extends ColumnarToRowTransition + with CometPlan { + + // supportsColumnar requires to be only called on driver side, see also SPARK-37779. + assert(Utils.isInRunningSparkTask || child.supportsColumnar) + + override def output: Seq[Attribute] = child.output + + override def outputPartitioning: Partitioning = child.outputPartitioning + + override def outputOrdering: Seq[SortOrder] = child.outputOrdering + + override def nodeName: String = "CometColumnarToRowView" + + override lazy val metrics: Map[String, SQLMetric] = Map( + "numOutputRows" -> SQLMetrics.createMetric(sparkContext, "number of output rows"), + "numInputBatches" -> SQLMetrics.createMetric(sparkContext, "number of input batches")) + + override def doExecute(): RDD[InternalRow] = { + val numOutputRows = longMetric("numOutputRows") + val numInputBatches = longMetric("numInputBatches") + child.executeColumnar().mapPartitionsInternal { batches => + batches.flatMap { batch => + numInputBatches += 1 + numOutputRows += batch.numRows() + // `flatMap` does not advance to the next batch until this row iterator is exhausted, so + // the Arrow buffers the returned rows point at stay live for as long as the rows are used. + batch.rowIterator().asScala + } + } + } + + override def withNewChildInternal(newChild: SparkPlan): SparkPlan = + copy(child = newChild) +} 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..c599b3b24f4 --- /dev/null +++ b/spark/src/test/scala/org/apache/comet/exec/CometWriteRowViewSuite.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.comet.exec + +import java.util.concurrent.TimeUnit + +import scala.util.Random + +import org.apache.spark.sql.{CometTestBase, DataFrame, QueryTest} +import org.apache.spark.sql.comet.CometColumnarToRowViewExec +import org.apache.spark.sql.execution.{QueryExecution, SparkPlan} +import org.apache.spark.sql.execution.adaptive.AdaptiveSparkPlanExec +import org.apache.spark.sql.util.QueryExecutionListener + +import org.apache.comet.CometConf +import org.apache.comet.testing.{DataGenOptions, FuzzDataGenerator, SchemaGenOptions} + +/** + * Tests for `spark.comet.exec.write.rowView.enabled`, which replaces the UnsafeRow-materializing + * columnar-to-row transition below a write with a zero-copy row view over the Arrow batch. + * + * The bar for every case here is that enabling the config changes nothing observable except the + * plan: the bytes written must match what the same write produces without it. + */ +class CometWriteRowViewSuite extends CometTestBase { + + test("row view is used for an unpartitioned parquet write") { + withParquetSource { source => + withTempPath { out => + val plan = captureWritePlan { + withRowView(source.write.mode("overwrite").parquet(out.toString)) + } + assert( + countRowViews(plan) == 1, + s"expected a CometColumnarToRowView in the write plan, got:\n$plan") + } + } + } + + test("row view is off by default") { + withParquetSource { source => + withTempPath { out => + val plan = captureWritePlan(source.write.mode("overwrite").parquet(out.toString)) + assert(countRowViews(plan) == 0, s"row view should be opt-in, got:\n$plan") + } + } + } + + test("row view is not used for a schema of only flat types") { + 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(countRowViews(plan) == 0, s"a flat schema is not worth the row view, got:\n$plan") + checkAnswer(spark.read.parquet(out.toString), flat) + } + } + } + + test("row view is used when a complex type is nested below a flat top level") { + withParquetSource { source => + withTempPath { out => + // The gate looks at top-level fields, so this asserts the common shape where the only + // complex column sits alongside flat ones. + val mixed = source.selectExpr("id", "str_col", "struct_col") + val plan = captureWritePlan { + withRowView(mixed.write.mode("overwrite").parquet(out.toString)) + } + assert(countRowViews(plan) == 1, s"expected the row view, got:\n$plan") + checkAnswer(spark.read.parquet(out.toString), mixed) + } + } + } + + test("row view writes the same data - deeply nested types") { + withDeeplyNestedSource { source => + withTempPath { out => + val plan = captureWritePlan { + withRowView(source.write.mode("overwrite").parquet(out.toString)) + } + assert(countRowViews(plan) == 1, s"expected the row view, got:\n$plan") + } + assertSameWrite(source, "parquet") + } + } + + test("row view writes the same data - primitives, strings and nulls") { + withParquetSource { source => + assertSameWrite(source, "parquet") + } + } + + test("row view writes the same data - fuzz generated flat schema") { + withFuzzSource( + SchemaGenOptions(generateArray = false, generateStruct = false, generateMap = false)) { + source => assertSameWrite(source, "parquet") + } + } + + test("row view writes the same data - fuzz generated nested schema") { + withFuzzSource( + SchemaGenOptions(generateArray = true, generateStruct = true, generateMap = true)) { + source => assertSameWrite(source, "parquet") + } + } + + test("row view writes the same data - orc and json") { + withParquetSource { source => + assertSameWrite(source, "orc") + assertSameWrite(source, "json") + } + } + + test("row view respects maxRecordsPerFile") { + withParquetSource { source => + Seq("100", "0").foreach { maxRecords => + withTempPath { out => + withSQLConf("spark.sql.files.maxRecordsPerFile" -> maxRecords) { + withRowView(source.write.mode("overwrite").parquet(out.toString)) + } + checkAnswer(spark.read.parquet(out.toString), source) + } + } + } + } + + test("row view is not used for partitioned or bucketed writes") { + withParquetSource { source => + withTempPath { out => + val plan = captureWritePlan { + withRowView(source.write.mode("overwrite").partitionBy("part").parquet(out.toString)) + } + assert(countRowViews(plan) == 0, s"partitioned write must not use the row view:\n$plan") + // `schema` pins the partition column back to string; reading a partitioned directory + // otherwise infers `part` as int. + checkAnswer(spark.read.schema(source.schema).parquet(out.toString), source) + } + } + + withParquetSource { source => + withTable("bucketed") { + val plan = captureWritePlan { + withRowView( + source.write + .mode("overwrite") + .bucketBy(4, "id") + .format("parquet") + .saveAsTable("bucketed")) + } + assert(countRowViews(plan) == 0, s"bucketed write must not use the row view:\n$plan") + checkAnswer(spark.table("bucketed"), source) + } + } + } + + /** + * Writes `source` with and without the row view and requires the results to be identical, both + * in content and in row count. The baseline is written by the same Comet plan with the config + * off, so any difference is attributable to the transition and not to the scan. + */ + private def assertSameWrite(source: DataFrame, format: String): Unit = { + withTempPath { baseline => + withTempPath { rowView => + source.write.mode("overwrite").format(format).save(baseline.toString) + withRowView(source.write.mode("overwrite").format(format).save(rowView.toString)) + + val expected = spark.read.schema(source.schema).format(format).load(baseline.toString) + val actual = spark.read.schema(source.schema).format(format).load(rowView.toString) + assert(actual.count() == expected.count(), s"row count differs for $format") + QueryTest.checkAnswer(actual, expected.collect().toSeq) + } + } + } + + private def withRowView(f: => Unit): Unit = + withSQLConf(CometConf.COMET_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)) + } + } + + /** 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)) + } + } + + private def withFuzzSource(schemaOptions: SchemaGenOptions)(f: DataFrame => Unit): Unit = { + withTempPath { dir => + val schema = FuzzDataGenerator.generateSchema(schemaOptions) + val df = FuzzDataGenerator.generateDataFrame( + new Random(42), + spark, + schema, + 1000, + DataGenOptions(generateNegativeZero = false)) + withSQLConf(CometConf.COMET_EXEC_ENABLED.key -> "false") { + 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 countRowViews(plan: SparkPlan): Int = + flatten(plan).count(_.isInstanceOf[CometColumnarToRowViewExec]) + + 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/CometParquetWriteBenchmark.scala b/spark/src/test/scala/org/apache/spark/sql/benchmark/CometParquetWriteBenchmark.scala new file mode 100644 index 00000000000..fa8c58d5bbd --- /dev/null +++ b/spark/src/test/scala/org/apache/spark/sql/benchmark/CometParquetWriteBenchmark.scala @@ -0,0 +1,335 @@ +/* + * 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.{CometColumnarToRowExec, CometColumnarToRowViewExec, CometNativeColumnarToRowExec} +import org.apache.spark.sql.execution.{ColumnarToRowExec, QueryExecution, SparkPlan} +import org.apache.spark.sql.execution.adaptive.AdaptiveSparkPlanExec +import org.apache.spark.sql.internal.SQLConf +import org.apache.spark.sql.types.{ArrayType, MapType, StructType} +import org.apache.spark.sql.util.QueryExecutionListener + +import org.apache.comet.{CometConf, CometSparkSessionExtensions} + +/** + * Benchmark for the columnar-to-row transition that sits below a Parquet write: + * - Spark's own vectorized read plus `ColumnarToRowExec` + * - Comet scan plus the JVM `CometColumnarToRowExec` (materializes an UnsafeRow per row) + * - Comet scan plus the native `CometNativeColumnarToRowExec` + * - Comet scan plus `CometColumnarToRowViewExec` (zero-copy row views over the Arrow batch) + * + * Each case measures `scan -> transition -> ParquetWriteSupport -> file`, so the transition is + * only part of what is timed; the point of the comparison is how much of a real write the + * UnsafeRow materialization costs. Cases are run at two compression settings because the encoding + * cost is what the transition cost has to be measured against. + * + * To run this benchmark: + * {{{ + * SPARK_GENERATE_BENCHMARK_FILES=1 make benchmark-org.apache.spark.sql.benchmark.CometParquetWriteBenchmark + * }}} + * + * Results will be written to "spark/benchmarks/CometParquetWriteBenchmark-**results.txt". + */ +object CometParquetWriteBenchmark extends CometBenchmarkBase { + + override def getSparkSession: SparkSession = { + val conf = new SparkConf() + .setAppName("CometParquetWriteBenchmark") + .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 def addWriteCases(benchmark: Benchmark, outputDir: File, codec: String): Unit = { + val query = "SELECT * FROM parquetV1Table" + + def write(target: String, configs: Seq[(String, String)]): Unit = + withSQLConf(configs :+ (SQLConf.PARQUET_COMPRESSION.key -> codec): _*) { + spark + .sql(query) + .write + .mode("overwrite") + .parquet(new File(outputDir, target).getCanonicalPath) + } + + // The four cases below differ only in which columnar-to-row transition feeds the writer, so + // the comparison is meaningless unless each arm actually plans the transition it names. Run + // every arm once up front and fail loudly when it does not. + def verify( + name: String, + target: String, + configs: Seq[(String, String)], + expected: String): Unit = { + val (transition, tree) = captureTransition(write(target, configs)) + // scalastyle:off println + println(s" [plan check] $name -> ${transition.getOrElse("")}") + if (diagnostics) println(tree.getOrElse("")) + // scalastyle:on println + if (!transition.contains(expected)) { + val border = "=" * 80 + benchmark.out.println(s"""\n$border + |WARNING: the "$name" case did not plan $expected but + |${transition.getOrElse("no transition 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 JVM C2R (UnsafeRow)") { _ => + write("comet-jvm", jvmC2RConfigs) + } + + benchmark.addCase("Comet native C2R (UnsafeRow)") { _ => + write("comet-native", nativeC2RConfigs) + } + + benchmark.addCase("Comet row view (zero-copy)") { _ => + write("comet-rowview", rowViewConfigs) + } + + verify("Spark", "spark", sparkConfigs, "ColumnarToRow") + verify("Comet JVM C2R (UnsafeRow)", "comet-jvm", jvmC2RConfigs, "CometColumnarToRow") + verify( + "Comet native C2R (UnsafeRow)", + "comet-native", + nativeC2RConfigs, + "CometNativeColumnarToRow") + // The rule declines a schema of only flat types, so the row-view arm is expected to plan the + // ordinary transition there. Deriving the expectation from the schema also asserts the gate. + verify( + "Comet row view (zero-copy)", + "comet-rowview", + rowViewConfigs, + if (hasComplexType) "CometColumnarToRowView" else "CometColumnarToRow") + } + + /** Mirrors the complex-type gate in `EliminateRedundantTransitions.rowView`. */ + private def hasComplexType: Boolean = + spark.table("parquetV1Table").schema.fields.exists { f => + f.dataType.isInstanceOf[StructType] || f.dataType.isInstanceOf[ArrayType] || + f.dataType.isInstanceOf[MapType] + } + + private val sparkConfigs = Seq(CometConf.COMET_ENABLED.key -> "false") + + private val jvmC2RConfigs = Seq( + CometConf.COMET_ENABLED.key -> "true", + CometConf.COMET_EXEC_ENABLED.key -> "true", + CometConf.COMET_NATIVE_COLUMNAR_TO_ROW_ENABLED.key -> "false", + CometConf.COMET_WRITE_ROW_VIEW_ENABLED.key -> "false") + + private val nativeC2RConfigs = Seq( + CometConf.COMET_ENABLED.key -> "true", + CometConf.COMET_EXEC_ENABLED.key -> "true", + CometConf.COMET_NATIVE_COLUMNAR_TO_ROW_ENABLED.key -> "true", + CometConf.COMET_WRITE_ROW_VIEW_ENABLED.key -> "false") + + private val rowViewConfigs = Seq( + CometConf.COMET_ENABLED.key -> "true", + CometConf.COMET_EXEC_ENABLED.key -> "true", + CometConf.COMET_NATIVE_COLUMNAR_TO_ROW_ENABLED.key -> "false", + CometConf.COMET_WRITE_ROW_VIEW_ENABLED.key -> "true") + + /** Names the columnar-to-row transition in the executed plan of a write. */ + private def captureTransition(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: ColumnarToRowExec => p.nodeName + case p: CometColumnarToRowExec => p.nodeName + case p: CometNativeColumnarToRowExec => p.nodeName + case p: CometColumnarToRowViewExec => p.nodeName + } + .mkString(", ") + } + val tree = captured.map { qe => + s" comet.enabled=${spark.conf.get(CometConf.COMET_ENABLED.key, "unset")} " + + s"exec.enabled=${spark.conf.get(CometConf.COMET_EXEC_ENABLED.key, "unset")} " + + s"rowView=${spark.conf.get(CometConf.COMET_WRITE_ROW_VIEW_ENABLED.key, "unset")}\n" + + qe.executedPlan.treeString + } + (names, tree) + } finally { + spark.listenerManager.unregister(listener) + } + } + + private val diagnostics = sys.env.contains("COMET_BENCH_DIAG") + + private def writeBenchmark(name: String, values: Int, codec: String)( + columns: Seq[String]): Unit = { + val benchmark = new Benchmark(s"$name ($codec)", values, output = output) + withTempPath { dir => + withTempTable("parquetV1Table") { + prepareTable(dir, spark.range(values).selectExpr(columns: _*)) + withTempPath { outputDir => + outputDir.mkdirs() + addWriteCases(benchmark, outputDir, codec) + if (!diagnostics) benchmark.run() + } + } + } + } + + 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 strings = Seq( + "id", + "concat('short_', cast(id % 100 as string)) as short_str", + "concat('medium_string_value_', cast(id as string), '_with_more_content') as medium_str", + "repeat(concat('long_', cast(id as string)), 10) as long_str") + + private val wide = (0 until 50).map { i => + i % 5 match { + case 0 => s"cast(id + $i as int) as int_col_$i" + case 1 => s"cast(id + $i as long) as long_col_$i" + case 2 => s"cast(id + $i as double) as double_col_$i" + case 3 => s"concat('str_${i}_', cast(id as string)) as str_col_$i" + case 4 => s"cast((id + $i) % 2 as boolean) as bool_col_$i" + } + } + + 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") + + /** One struct, nested `depth` levels deep, to show how the saving scales with depth. */ + private def structOfDepth(depth: Int): Seq[String] = { + def build(level: Int): String = + if (level == depth) { + "named_struct('v', cast(id as int), 'n', concat('x_', cast(id as string)))" + } else { + s"named_struct('c', cast(id % 100 as int), 'inner', ${build(level + 1)})" + } + Seq("id", s"${build(1)} as deep_struct") + } + + /** 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", "zstd") + + codecs.foreach { codec => + runBenchmark(s"Parquet write - Fixed width ($codec)") { + writeBenchmark("Parquet write - Fixed width", numRows, codec)(fixedWidth) + } + runBenchmark(s"Parquet write - Strings ($codec)") { + writeBenchmark("Parquet write - Strings", numRows, codec)(strings) + } + runBenchmark(s"Parquet write - Wide 50 columns ($codec)") { + writeBenchmark("Parquet write - Wide 50 columns", numRows, codec)(wide) + } + runBenchmark(s"Parquet write - Nested ($codec)") { + writeBenchmark("Parquet write - Nested", numRows, codec)(nested) + } + runBenchmark(s"Parquet write - Deeply nested ($codec)") { + writeBenchmark("Parquet write - Deeply nested", numRows, codec)(deeplyNested) + } + Seq(1, 2, 4, 8).foreach { depth => + runBenchmark(s"Parquet write - Struct depth $depth ($codec)") { + writeBenchmark(s"Parquet write - Struct depth $depth", numRows, codec)( + structOfDepth(depth)) + } + } + } + } +}