sunchao commented on code in PR #6564: URL: https://github.com/apache/datafusion-comet/pull/6564#discussion_r4173623238
########## spark/src/test/scala/org/apache/spark/sql/benchmark/CometTypedDatasetBenchmark.scala: ########## @@ -0,0 +1,192 @@ +/* + * 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 org.apache.spark.benchmark.Benchmark +import org.apache.spark.sql.{DataFrame, Dataset, Encoder, Encoders, Row} +import org.apache.spark.sql.catalyst.expressions.aggregate.Partial +import org.apache.spark.sql.comet.{CometHashAggregateExec, CometPlan, CometSparkToColumnarExec} +import org.apache.spark.sql.execution.SparkPlan +import org.apache.spark.sql.functions.{col, count, length, lit, sum} +import org.apache.spark.sql.internal.SQLConf + +import org.apache.comet.CometConf + +// Top-level, so the encoders need no outer pointer, which is the ordinary user shape. +case class TypedDatasetBenchRec(a: Long, b: String) + +case class TypedDatasetBenchWide(a: Long, b: String, c: Long, d: String) + +/** + * Compares three ways to run a query over the output of a typed Dataset operation: + * + * - Spark: Comet disabled. + * - Comet: the default. The typed operation runs in Spark, and Comet takes over again at the + * shuffle above it, so the operators in between stay on Spark. + * - Comet, converted: `spark.comet.convert.typedDataset.enabled`, which converts the output of + * the typed operation to Arrow so the operators above it run natively. + * + * The cases sweep how much work sits above the typed operation, from an aggregate over 100 groups + * that Spark's whole-stage codegen fuses with the operation to one over a million groups. Every + * arm's result and plan are checked before it is timed, and the Comet arm runs again at the end + * of each case to show the noise. To run this benchmark: + * {{{ + * SPARK_GENERATE_BENCHMARK_FILES=1 make benchmark-org.apache.spark.sql.benchmark.CometTypedDatasetBenchmark + * }}} + * Results will be written to "spark/benchmarks/CometTypedDatasetBenchmark-**results.txt". + */ +object CometTypedDatasetBenchmark extends CometBenchmarkBase { + + private val numRows = 4 * 1024 * 1024 Review Comment: [P2] Declare `numRows` as a `Long` so the required strict Spark 3.5 compilation succeeds. It is currently inferred as `Int`, but both `spark.range(numRows)` and `new Benchmark(name, numRows, ...)` require `Long`. The strict profile promotes those implicit numeric-widening warnings to errors, causing this PR’s CI build to fail. Using `4L * 1024 * 1024` fixes both call sites. Evidence: Exact-head CI job https://github.com/apache/datafusion-comet/actions/runs/37130634250/job/111225450414 ran `./mvnw -B test-compile -Pspark-3.5 -Pstrict-warnings -DskipTests` and failed with `CometTypedDatasetBenchmark.scala:161: implicit numeric widening` and the same error at line 178. The log ends with `two errors found` and a failed `scala-maven-plugin:testCompile` goal. ########## spark/src/main/scala/org/apache/comet/rules/CometExecRule.scala: ########## @@ -1166,6 +1178,35 @@ case class CometExecRule(session: SparkSession) private def hasEnabledHandler(op: SparkPlan): Boolean = allExecs.get(op.getClass).exists(_.enabledConfig.forall(_.get(op.conf))) + /** + * Converts the rows a typed Dataset operation produces to Arrow, so the operators above it can + * run natively. See [[CometConf.COMET_CONVERT_FROM_TYPED_DATASET_ENABLED]]. + * + * Spark inserts the columnar transitions after this rule, but it does not look below a + * `RowToColumnarTransition` such as `CometSparkToColumnarExec`. That is harmless above a leaf. + * Here the typed operation's own operators sit below the conversion, and without a transition + * they would read a Comet child through `CometExec.doExecute`, Spark's interpreted + * columnar-to-row path. So the subtree gets its transitions now, from Spark's own rule, and + * `EliminateRedundantTransitions` later replaces each one over a Comet child with Comet's own. + * Spark's rule leaves existing transitions alone, which matters because this rule runs over the + * same plan twice under AQE. + */ + private def convertTypedDatasetOutput(op: SerializeFromObjectExec): SparkPlan = { + val unsupported = op.output.filterNot(a => + CometSparkToColumnarExec.isTypeSupported(a.dataType, a.name, ListBuffer.empty)) + if (unsupported.nonEmpty) { + withFallbackReason( + op, + "Comet cannot convert the output of a typed Dataset operation to Arrow because it does " + + "not support the type of these columns: " + + unsupported.map(a => s"${a.name}: ${a.dataType.simpleString}").mkString(", ")) + } else { + val withTransitions = + ApplyColumnarRulesAndInsertTransitions(Seq.empty, outputsColumnar = false).apply(op) + convertToComet(withTransitions, CometSparkToColumnarExec).getOrElse(withTransitions) Review Comment: [P1] Prevent conversion from creating incompatible wide-decimal shuffles. With this feature enabled, a supported typed input switches to native shuffle while another typed input containing a retained `array<int>` column stays on JVM shuffle. Joining their `decimal(38,18)` keys with AQE disabled then silently loses matching rows: my reproduction returns 9 instead of 100. Conversion-disabled Comet returns all 100. This newly exposes the existing decimal hash difference as incorrect query results. Keep affected exchanges on JVM shuffle until native wide-decimal hashing matches Spark, and cover this mixed-path join. Evidence: Reproduced on Spark 4.1.3/JDK 17 in CometTestBase with `spark.sql.adaptive.enabled=false`, `spark.sql.autoBroadcastJoinThreshold=-1`, `spark.sql.shuffle.partitions=10`, and `spark.comet.shuffle.mode=auto`. Define `case class L(k: java.math.BigDecimal, v: Long)` and `case class R(k: java.math.BigDecimal, xs: Seq[Int])`. Build `l = spark.range(0,100,1,2).map(i => L(new java.math.BigDecimal(i), i)).alias("l")` and `r = spark.range(0,100,1,2).map(i => R(new java.math.BigDecimal(i), Seq(i.toInt))).alias("r")`. Collect `l.join(r, col("l.k") === col("r.k")).select(col("l.v"), col("r.xs"))`. Spark and conversion-disabled Comet return 100 rows. Conversion-enabled Comet returns 9. The executed plan contains left `CometNativeShuffle` and right `CometColumnarShuffle`. Setting shuffle mode to `jvm` restores 100 rows. Spark hashes wide decimals using `unscaledValue().toByteArray`; native `hash_array_decimal!` uses fixed-width `to_le_bytes()`. -- 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]
