0lai0 commented on code in PR #6038: URL: https://github.com/apache/datafusion-comet/pull/6038#discussion_r4053298963
########## spark/src/test/scala/org/apache/spark/sql/benchmark/CometIcebergWriteBenchmark.scala: ########## @@ -0,0 +1,486 @@ +/* + * 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 scala.collection.mutable +import scala.concurrent.duration._ + +import org.apache.spark.CometListenerBusUtils +import org.apache.spark.SparkConf +import org.apache.spark.benchmark.Benchmark +import org.apache.spark.sql.SparkSession +import org.apache.spark.sql.comet.CometIcebergWriteExec +import org.apache.spark.sql.execution.{QueryExecution, SparkPlan} +import org.apache.spark.sql.internal.SQLConf +import org.apache.spark.sql.util.QueryExecutionListener + +import org.apache.comet.CometConf +import org.apache.comet.CometSparkSessionExtensions +import org.apache.comet.iceberg.IcebergReflection + +/** + * Benchmark of writes into an Iceberg table with Comet's native (iceberg-rust) writer on and off. + * + * Four cases are measured, covering the writer implementations an Iceberg write can reach: + * + * 1. an unpartitioned `INSERT INTO ... SELECT`, which writes one file per task with no + * exchange; + * 1. the same insert into a partitioned table, where Iceberg's default hash distribution + * clusters each partition onto one task and the clustered writer keeps one file open; + * 1. the same insert into a partitioned table configured for the fanout writer, which holds a + * file open per partition instead of requiring the exchange; + * 1. a copy-on-write `DELETE`, where the write is a rewrite of every file that holds a matching + * row, and so is a read and a write of the whole table rather than of new rows. + * + * Each case is measured under three configurations, so that the writer's own contribution can be + * read off the table the way the ad hoc measurements in + * [[https://github.com/apache/datafusion-comet/pull/5361 #5361]] reported it: + * + * 1. `Spark` - stock Spark, the baseline the `Relative` column is computed against. + * 1. `Comet scan` - Comet reads the Parquet source, iceberg-java still writes the data files. + * 1. `Comet scan + native write` - the per-task write is delegated to iceberg-rust. + * + * The first-to-second step is therefore the scan speedup and the second-to-third step is the + * writer speedup. The `Relative` column compares each case against stock Spark, so the writer's + * own factor has to be divided out of the two Comet rows rather than read directly. + * + * What the numbers do and do not cover: + * - The timed statement includes the driver-side Iceberg commit, which all three configurations + * pay equally. It dilutes the writer's factor rather than inflating it. + * - The warehouse is a local temporary directory, so no object-store latency is included. + * - The two writers choose different file roll points (see the accepted divergences in + * `docs/source/user-guide/latest/iceberg-writes.md`), so the resulting file layouts are not + * expected to match. Only wall clock and row counts are compared. + * - `Rate` and `Per Row` are always computed against the corpus size, so for the `DELETE` case + * they describe the rows the statement passed over, not the far smaller number it removed. + * - The corpus below is this benchmark's own, not the one #5361 measured. The factors are the + * same comparison repeated on different data, not a continuation of that PR's numbers. + * + * To run this benchmark: + * {{{ + * SPARK_GENERATE_BENCHMARK_FILES=1 make benchmark-org.apache.spark.sql.benchmark.CometIcebergWriteBenchmark + * }}} + * Results will be written to "spark/benchmarks/CometIcebergWriteBenchmark-**results.txt". + */ +object CometIcebergWriteBenchmark extends CometBenchmarkBase { + + private val catalog = defaultIcebergCatalog + private val namespace = "db" + private val targetTable = s"$catalog.$namespace.write_target" + + /** One null in eight in every column, matching the corpus of the other Iceberg benchmarks. */ + private val nullStride = 8 + + /** + * Source file count, and therefore the number of write tasks an unpartitioned insert runs. + * Pinned rather than left to `spark.sql.files.maxPartitionBytes` so that a change in corpus + * size or in the machine's core count does not silently change the parallelism and make two + * recorded runs incomparable. + */ + private val sourceFiles = 8 + + /** + * The partition column of the two partitioned cases. Its eight values and its nulls give nine + * partitions: few enough that the clustered writer holds one file open at a time, and more than + * the source file count so that the fanout writer has something to fan out over. + */ + private val partitionColumn = "c_str_dict" + + /** + * The rows the copy-on-write case deletes. One row in a hundred, spread evenly, so that every + * data file holds at least one of them and the rewrite covers the whole table - which is the + * shape of copy-on-write that the writer's speed actually decides. A predicate matching a + * contiguous range would instead measure how well the scan planner pruned files. + */ + private val deletePredicate = "PMOD(c_long, 100) = 0" + + /** + * How a case is configured and what its plan must contain. `expectComet` and + * `expectNativeWrite` are what [[verifyArm]] checks before anything is timed: a configuration + * that silently fell back measures a different engine than its name claims. + */ + private case class Arm( + name: String, + confs: Seq[(String, String)], + expectComet: Boolean, + expectNativeWrite: Boolean) + + private val arms: Seq[Arm] = Seq( + Arm( + "Spark", + Seq(CometConf.COMET_ENABLED.key -> "false"), + expectComet = false, + expectNativeWrite = false), + Arm( + "Comet scan", + Seq(CometConf.COMET_ENABLED.key -> "true", CometConf.COMET_EXEC_ENABLED.key -> "true"), + expectComet = true, + expectNativeWrite = false), + Arm( + "Comet scan + native write", + Seq( + CometConf.COMET_ENABLED.key -> "true", + CometConf.COMET_EXEC_ENABLED.key -> "true", + // The native writer requires the split-operator plan; enabling it alone is a no-op. + CometConf.COMET_ICEBERG_WRITE_SPLIT_OPERATOR_ENABLED.key -> "true", + CometConf.COMET_ICEBERG_NATIVE_WRITE_ENABLED.key -> "true"), + expectComet = true, + expectNativeWrite = true)) + + /** + * One timed statement and the table it needs. + * + * @param prePopulate + * whether the table has to hold the corpus before the statement runs, as the copy-on-write + * case does and the inserts do not. + * @param expectShuffle + * whether the statement's plan must contain an exchange, where that is what separates this + * case from another one. `None` where the case's identity does not rest on it. The clustered + * and fanout cases differ only in how Iceberg distributes rows to the writers, so if that + * ever stopped holding, the two would quietly become the same measurement printed twice. + */ + private case class Workload( + title: String, + partitionSpec: String, + properties: Seq[String], + prePopulate: Boolean, + statement: String, + expectedRowsAfter: Long, + expectShuffle: Option[Boolean]) + + private def workloads(values: Int): Seq[Workload] = { + val insert = s"INSERT INTO $targetTable SELECT * FROM parquetV1Table" + val partitioned = s"PARTITIONED BY ($partitionColumn)" + val deleted = + spark.sql(s"SELECT count(*) FROM parquetV1Table WHERE $deletePredicate").head().getLong(0) + + Seq( + Workload( + "unpartitioned INSERT INTO ... SELECT", + partitionSpec = "", + properties = Nil, + prePopulate = false, + statement = insert, + expectedRowsAfter = values.toLong, + expectShuffle = None), + Workload( + "partitioned INSERT INTO ... SELECT, clustered writer", + partitionSpec = partitioned, + // No `write.distribution-mode`: Iceberg defaults a partitioned table to hash, which is the + // clustered writer this case is named for. Spelling it out would hide a change of default. + properties = Nil, Review Comment: Thanks @sunchao for review. Fixed, Two things were needed: - Pinned `'write.spark.fanout.enabled'='false'`, which selects the clustered writer. - Added `ALTER TABLE ... WRITE DISTRIBUTED BY PARTITION LOCALLY ORDERED BY <partition column>` via a new `Workload.orderBy`. The local sort is what the clustered writer needs so it doesn't error on interleaved partitions, and it makes `hasOrdering` true so the writer choice no longer rests on the fanout default. `verifyArm` now checks for a `Sort` node through `Workload.expectSort` — the plan-visible, version-independent mark of that required ordering, which the exchange alone couldn't show. The clustered row measured fanout with hash distribution before this. Re-running and will update the description. -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected] --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
