andygrove commented on code in PR #6564:
URL: https://github.com/apache/datafusion-comet/pull/6564#discussion_r4173768291
##########
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:
Thanks for the repro. It reproduced exactly here: 9 rows instead of 100,
with `CometNativeShuffle` on the left and `CometColumnarShuffle` on the right.
Fixed in 1ec0e66265. A hash shuffle whose stage starts at a typed Dataset
conversion now stays on Comet's columnar shuffle when a key is or contains a
decimal wider than 18 digits and the shuffle has more than one partition, so
the conversion no longer moves it off the path the other side of the join uses.
That's the rule #6005 applies to every native shuffle, limited to the shuffles
this feature moves, so it becomes redundant once #6005 lands. Your join is now
a test (`a join on wide decimal keys with an input that is not converted`),
which fails without the guard. The rule is also listed in `native_shuffle.md`.
##########
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:
Fixed in 1ec0e66265. `numRows` is now `4L * 1024 * 1024`, and `./mvnw -B
test-compile -Pspark-3.5 -Pstrict-warnings -DskipTests` passes locally.
--
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]
