comphead commented on code in PR #6564:
URL: https://github.com/apache/datafusion-comet/pull/6564#discussion_r4186038706
##########
spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometShuffleExchangeExec.scala:
##########
@@ -622,6 +635,19 @@ object CometShuffleExchangeExec
reasons += s"unsupported hash partitioning data type for native
shuffle: $dt"
}
}
+ // A typed Dataset conversion moves the shuffle above it from Comet's
columnar shuffle,
+ // which partitions with Spark's hash, to native shuffle. Native
shuffle hashes a decimal
+ // wider than 18 digits differently from Spark, so a join with an
input that is still on
+ // the columnar shuffle would put matching keys in different
partitions. Leave such a
+ // shuffle where it was. A single partition hashes nothing.
+ // TODO: remove once native hashing matches Spark for wide decimals
(#5994).
Review Comment:
#6005 rejects native hashing of wide decimal keys for every native shuffle,
and #6607 carries its own copy of this guard (`hashesDifferentlyFromSpark`).
Once #6005 lands, this block, `readsTypedDatasetConversion`, the
`native_shuffle.md` bullet and the join test are no longer needed. This TODO
names only #5994, so the guard would stay behind after #6005. Could it name
#6005 as well, so whoever lands that removes this guard in the same change?
#6005 also edits the same section of `native_shuffle.md`, so one of the two
will need a rebase.
##########
spark/src/main/scala/org/apache/comet/rules/CometExecRule.scala:
##########
@@ -1274,6 +1306,83 @@ case class CometExecRule(session: SparkSession)
private def hasEnabledHandler(op: SparkPlan): Boolean =
allExecs.get(op.getClass).exists(_.enabledConfig.forall(_.get(op.conf)))
+ /**
+ * Tags each `SerializeFromObjectExec` whose output an operator above it can
stop reading early
+ * with that operator, so `transform` does not convert it. Spark computes
the rows of a typed
+ * Dataset operation one at a time, as the operator above reads them, while
the conversion fills
+ * a whole Arrow batch first. Below a limit, a `mapPartitions` function such
as `_.take(1)`, or
+ * code reading `Dataset.rdd`, the conversion would run the user function on
rows Spark never
+ * reaches, and a function that throws on one of them would fail a query
that succeeds in Spark.
+ * An operator that reads all of its input before it returns a row ends the
search, since Spark
+ * computes every row below it anyway: an exchange, a sort, a hash
aggregate, or a top-k over
+ * input that is not already sorted.
+ *
+ * Conversion is bottom-up, so this runs first. TreeNode tags survive the
child copies made
+ * during transformUp, while an identity set would not.
+ */
+ private def tagPartiallyReadTypedDatasetOutputs(plan: SparkPlan): Unit = {
+ def visit(op: SparkPlan, partialReader: Option[String]): Unit = {
+ val childReader = op match {
+ case serialize: SerializeFromObjectExec =>
+ partialReader.foreach(
+ serialize.setTagValue(CometExecRule.TYPED_DATASET_PARTIAL_READER,
_))
+ partialReader
+ case _: CollectLimitExec | _: LocalLimitExec | _: GlobalLimitExec =>
Some("a limit")
+ // A top-k reads only its first rows when its input is already sorted.
+ case topK: TakeOrderedAndProjectExec
+ if SortOrder.orderingSatisfies(topK.child.outputOrdering,
topK.sortOrder) =>
+ Some("a limit")
Review Comment:
Can this arm ever tag a `SerializeFromObjectExec`? Typed operators report no
`outputOrdering` (`SerializeFromObjectExec` does not override it), so an
ordering that satisfies the top-k comes from a `SortExec`, and the `SortExec`
arm below resets the search before it reaches the serializer. I could not build
a plan where this arm changes a tag, and no test reaches it. If there is none,
dropping the arm, its comment and the `SortOrder` import leaves
`TakeOrderedAndProjectExec` as a plain barrier. If there is one, a test would
pin it.
##########
spark/src/main/scala/org/apache/comet/rules/CometExecRule.scala:
##########
@@ -1274,6 +1306,83 @@ case class CometExecRule(session: SparkSession)
private def hasEnabledHandler(op: SparkPlan): Boolean =
allExecs.get(op.getClass).exists(_.enabledConfig.forall(_.get(op.conf)))
+ /**
+ * Tags each `SerializeFromObjectExec` whose output an operator above it can
stop reading early
+ * with that operator, so `transform` does not convert it. Spark computes
the rows of a typed
+ * Dataset operation one at a time, as the operator above reads them, while
the conversion fills
+ * a whole Arrow batch first. Below a limit, a `mapPartitions` function such
as `_.take(1)`, or
+ * code reading `Dataset.rdd`, the conversion would run the user function on
rows Spark never
+ * reaches, and a function that throws on one of them would fail a query
that succeeds in Spark.
+ * An operator that reads all of its input before it returns a row ends the
search, since Spark
+ * computes every row below it anyway: an exchange, a sort, a hash
aggregate, or a top-k over
+ * input that is not already sorted.
+ *
+ * Conversion is bottom-up, so this runs first. TreeNode tags survive the
child copies made
+ * during transformUp, while an identity set would not.
+ */
+ private def tagPartiallyReadTypedDatasetOutputs(plan: SparkPlan): Unit = {
+ def visit(op: SparkPlan, partialReader: Option[String]): Unit = {
+ val childReader = op match {
+ case serialize: SerializeFromObjectExec =>
+ partialReader.foreach(
+ serialize.setTagValue(CometExecRule.TYPED_DATASET_PARTIAL_READER,
_))
+ partialReader
+ case _: CollectLimitExec | _: LocalLimitExec | _: GlobalLimitExec =>
Some("a limit")
+ // A top-k reads only its first rows when its input is already sorted.
+ case topK: TakeOrderedAndProjectExec
+ if SortOrder.orderingSatisfies(topK.child.outputOrdering,
topK.sortOrder) =>
+ Some("a limit")
+ case _: MapPartitionsExec => Some("a mapPartitions function")
+ case _: Exchange | _: SortExec | _: HashAggregateExec | _:
ObjectHashAggregateExec |
+ _: TakeOrderedAndProjectExec =>
+ None
+ case _ => partialReader
+ }
+ op.children.foreach(visit(_, childReader))
+ }
+ // `Dataset.rdd` reads the objects the plan produces, and the RDD's own
code decides how many
+ // of them to read, as `take(1)` does. The plan's root is the
`DeserializeToObjectExec` that
+ // `Dataset.rdd` adds. When the Dataset ends in a typed operation such as
`map`, Spark's
+ // `EliminateSerialization` drops that deserializer together with the
operation's serializer,
+ // so the root is the operation itself, which produces objects too, or a
typed filter or a
+ // project over it. A Dataset's own plan ends in rows, so no other plan
has such a root.
+ def producesObjects(op: SparkPlan): Boolean = op match {
+ case _: ObjectProducerExec => true
+ case _: FilterExec | _: ProjectExec => producesObjects(op.children.head)
+ case _ => false
+ }
+ visit(plan, if (producesObjects(plan)) Some("code reading Dataset.rdd")
else None)
+ }
+
+ /**
+ * 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:
A native shuffle placed directly over this conversion reads it through
`executeColumnar()` and `ColumnarBatchArrowReader`, which closes each batch.
`RowArrowReader` reuses its vectors, and closing a struct vector drops its
children, so a struct column fails from the second batch. #6607 hits the same
thing and fixes it by reading a `CometNativeArrowSource` child as an Arrow
stream.
I replayed this with the `CometExecRule`, `CometShuffleExchangeExec` and
`CometConf` from this head over a local build of `main` (Spark 4.1.3,
`spark.comet.convert.typedDataset.enabled=true`, `spark.comet.batchSize=16`):
```scala
case class Rec(a: Int, b: String)
case class Out(id: Int, inner: Rec)
spark.range(0, 200, 1, 2).map(i => Out(i.toInt, Rec(i.toInt,
s"s$i"))).repartition(col("id")).collect()
```
This fails with `CometNativeException: ... no more field nodes for field a`,
with AQE on and off. `orderBy("id")`, the left side of a shuffled join, and a
union that feeds a shuffle fail the same way. The plan is `CometExchange ...
CometNativeShuffle` over `CometSparkRowToColumnar` over `SerializeFromObject`.
The same `repartition` returns all 200 rows with one batch, with flat columns,
with an `array<string>` column, or with the conversion off.
No test here puts a native shuffle directly over the conversion. The struct
test reads only `id` downstream, so Spark's `ObjectSerializerPruning` removes
`inner` and `tags` from the serializer before the conversion sees them (a plain
`explain` of that query shows `SerializeFromObject [... AS id]`). Could you add
a struct case with a shuffle directly above the conversion and a small
`spark.comet.batchSize`? It will fail until #6607 lands, so please either land
that first or decline struct columns in `convertTypedDatasetOutput` until then.
##########
docs/source/user-guide/latest/datasources.md:
##########
@@ -80,6 +80,10 @@ string collations remain unsupported at this conversion
boundary. Source default
This includes row-backed `ExistingRDD` inputs when
`spark.comet.convert.rdd.enabled=true`. Spark
still produces the RDD rows; conversion lets eligible downstream operators
execute in Comet.
+The same types apply to the output of typed `Dataset` operations, such as
`map`, which Comet
+converts when `spark.comet.convert.typedDataset.enabled=true`. A column of any
other type keeps
+the operators above the typed operation on Spark.
Review Comment:
The "Other Spark inputs" list above names each `spark.comet.convert.*`
conversion of a Spark operator's output, and this one is missing from it. Could
it get a bullet there, for example `spark.comet.convert.typedDataset.enabled`:
the output of typed `Dataset` operations such as `map` and `mapPartitions`? The
sentence about column types can stay here. Someone looking for what Comet can
convert will not find the new config otherwise.
##########
spark/src/test/scala/org/apache/spark/sql/benchmark/CometTypedDatasetBenchmark.scala:
##########
@@ -0,0 +1,187 @@
+/*
+ * 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, 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 = 4L * 1024 * 1024
+ private val loKeys = 100
+ private val hiKeys = 1024 * 1024
+
+ import spark.implicits._
+
+ private case class Arm(name: String, confs: Seq[(String, String)], check:
SparkPlan => Unit)
Review Comment:
`Arm`, `runArm` and the check-then-time loop in `runCometBenchmark` repeat
`CometRangeBenchmark`, where the `Arm` signature is identical. #6607 adds a
third copy in `CometShuffleInputConversionBenchmark`. Could `Arm` and one
`runArms(name, numRows, arms, query)` helper move into `CometBenchmarkBase` in
whichever of the two PRs lands first, so this file keeps only its cases and
plan checks?
--
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]
