sunchao commented on code in PR #6110: URL: https://github.com/apache/datafusion-comet/pull/6110#discussion_r4240812029
########## spark/src/test/scala/org/apache/spark/sql/benchmark/CometUnserializableShuffleKeyBenchmark.scala: ########## @@ -0,0 +1,155 @@ +/* + * 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.comet.execution.shuffle.CometShuffleExchangeExec +import org.apache.spark.sql.execution.SparkPlan +import org.apache.spark.sql.execution.adaptive.QueryStageExec +import org.apache.spark.sql.execution.exchange.ShuffleExchangeExec +import org.apache.spark.sql.internal.SQLConf + +import org.apache.comet.CometConf + +/** + * What a range-partitioned exchange costs when its sort key has no Comet serde, which is the + * shape #5971 moves from Spark's shuffle to Comet's JVM columnar shuffle. + * + * The operator above such an exchange serializes the same key and fails the same check, so it + * stays on Spark. The exchange therefore sits between two Spark operators, and the columnar path + * pulls its rows through `ColumnarToRowExec`, converts them to Arrow, and hands them back as + * rows. This measures whether that round trip costs anything. + * + * Turning the codegen dispatcher off is only a stable way to obtain a key with no serde; a Hive + * UDF behaves the same. + * + * The baseline case switches Comet shuffle off, which reproduces the plan this query had before + * #5971: a plain Spark `Exchange` with the rest of the plan unchanged. Comparing the two inside + * one build isolates the exchange, where comparing two branches would also carry every unrelated + * commit between them. + * + * To run this benchmark: + * {{{ + * SPARK_GENERATE_BENCHMARK_FILES=1 make benchmark-org.apache.spark.sql.benchmark.CometUnserializableShuffleKeyBenchmark + * }}} + * Results will be written to + * "spark/benchmarks/CometUnserializableShuffleKeyBenchmark-**results.txt". + */ +object CometUnserializableShuffleKeyBenchmark extends CometBenchmarkBase { + + private val Rows = 8L * 1024 * 1024 + + /** Range partitioning, as in the reported shape, and hash partitioning for comparison. */ + private val rangeQuery = "SELECT * FROM parquetV1Table ORDER BY bump(a), b" + private val hashQuery = + "SELECT max(b) FROM parquetV1Table GROUP BY bump(a) ORDER BY 1 LIMIT 10" Review Comment: [P2] Use a hash query whose exchange key actually fails the removed serde probe. Spark computes `bump(a)` below the partial aggregate here, so the exchange hashes the integer attribute `_groupingexpression`, which already passes `exprToProto` on the base revision. Consequently, disabling Comet shuffle does not reproduce this query's pre-PR plan, and the reported hash speedups do not measure this change. `checkPlans` misses this because it only checks exchange names. This invalidates two of the comparisons used to weigh the existing range slowdown. Use a direct `repartition(..., bump(col("a")))` or `DISTRIBUTE BY bump(a)`, assert that the exchange key fails native serde with the dispatcher disabled, and rerun the comparison. Evidence: A compiled Spark 4.1.3 probe over 100 Parquet rows executed the exact benchmark SQL with AQE=false and AQE=true. Both produced `hashpartitioning(_groupingexpression#..., 4)` with `UDF_IN_EXCHANGE=false`. The control `SELECT * FROM parquetV1Table DISTRIBUTE BY bump(a)` produced a hash key containing `bump` with `UDF_IN_EXCHANGE=true`. Spark's `AggUtils.planAggregateWithoutDistinct` uses grouping attributes for the final aggregate's distribution. Comet's unchanged `CometAttributeReference` serde accepts this integer attribute, so the old columnar probe does not reject it. Reproduction source and output: `/tmp/review6110-planprobe-6f4_q1ri/Review6110PlanProbe.scala` and `/tmp/review6110-planprobe-6f4_q1ri/probe.log`. -- 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]
