You are viewing a plain text version of this content. The canonical link for it is here.
Posted to github@datafusion.apache.org by "andygrove (via GitHub)" <gi...@apache.org> on 2026/10/03 14:43:28 UTC
[PR] feat: run the operators above a typed Dataset operation natively [datafusion-comet]
andygrove opened a new pull request, #6564:
URL: https://github.com/apache/datafusion-comet/pull/6564
## Which issue does this PR close?
Closes #5710. Part of #5572. Supersedes #5714.
## Rationale for this change
A typed `Dataset` operation such as `ds.map(f)` plans a `DeserializeToObject` / `MapElements` / `SerializeFromObject` island that has to run in Spark, because it passes JVM objects between its operators. Comet only takes over again at the next shuffle, so whatever sits between the island and that shuffle stays on Spark too. For `ds.map(f).groupBy(k).agg(...)` that is the partial aggregate, which is expensive when there are many groups. A broadcast join on the probe side is in the same position.
#5714 tried to fuse the island into one projection in the JVM codegen dispatcher. Since #6542 the dispatcher only calls into Spark's own classes, so that approach is blocked. When I measured the alternative in this PR against it, the conversion got 93-100% of the fused speedup in the cases where either one pays off. It is also much simpler, and it covers every typed operation rather than only `map`.
## What changes are included in this PR?
Every typed operation ends in `SerializeFromObjectExec`, whose output is ordinary rows. With the new `spark.comet.convert.typedDataset.enabled`, `CometExecRule` puts a `CometSparkToColumnarExec` above it, so the operators above the typed operation can run natively. The operation itself, including the user function, runs in Spark exactly as before. This covers `map`, `flatMap`, `mapPartitions`, `groupByKey(...).mapGroups`, `cogroup`, and a Dataset built from an RDD of objects. If nothing native consumes the output, `EliminateRedundantTransitions` removes the conversion again, so a typed operation at the top of a plan is unchanged.
Spark inserts no columnar transitions below a `RowToColumnarTransition`, and `CometSparkToColumnarExec` is one. Above a leaf that does not matter. Here the typed operation's own operators sit below the conversion, and without transitions they read their Comet child through `CometExec.doExecute`, which is Spark's interpreted columnar-to-row path. That gives the right answer slowly, and in the benchmark it made whole queries 1.5-2.4x slower. The rule therefore applies Spark's own `ApplyColumnarRulesAndInsertTransitions` to the subtree, as `CometRule.buildPreview` already does, and `EliminateRedundantTransitions` swaps in Comet's columnar-to-row as usual. Spark's rule leaves existing transitions alone, which matters because `CometExecRule` runs over the same plan twice under AQE.
That second pass used to tag the inserted `ColumnarToRowExec` with "ColumnarToRow is not supported", so `CometExecRule` no longer reports a `ColumnarToRowTransition` as an operator it failed to convert. A column type that Spark-to-Comet conversion does not support, such as `array<int>`, keeps the operators above the typed operation on Spark and records a fallback reason that names the column.
It is off by default because whether it pays depends on the work above the typed operation, which the planner cannot see. Here is `CometTypedDatasetBenchmark` on an M3 Max with 4Mi rows, `local[1]`, AQE off and one shuffle partition. Times are best of the iterations in ms, and the last column repeats the default Comet arm as a noise check.
| case | Spark | Comet (default) | Comet, converted | default repeat |
| ---------------------------------- | ----- | --------------- | ---------------- | -------------- |
| map -> group by 100 keys | 249 | 248 | 313 | 221 |
| map -> group by 1M keys | 1244 | 1079 | **484** | 940 |
| map -> filter -> group by 100 keys | 180 | 153 | 319 | 153 |
| map -> filter -> group by 1M keys | 896 | 633 | **451** | 636 |
| map -> group by long key, 1M keys | 811 | 699 | **421** | 686 |
| map, 4 columns -> group by 1M keys | 1359 | 1101 | **682** | 1103 |
| mapPartitions -> group by 1M keys | 1289 | 1148 | **738** | 1136 |
With an aggregate over many groups above the typed operation, it is 1.4-2.2x faster than today's default. With a cheap aggregate it is slower, 0.5x after a selective filter, because Spark compiles the typed operation, the filter and a small aggregate into one loop, while the conversion writes every row to Arrow first.
Two behaviors worth knowing when it is on:
- An exception thrown by the user function reaches the driver as `CometNativeException: C Data interface error: java.lang.IllegalArgumentException: ...` rather than as the original exception. Every Spark-to-Arrow input wraps errors this way today (#6234), and the open #6243 fixes it for all of them.
- The shuffle above the typed operation, which today is Comet's columnar shuffle, becomes a native shuffle. That is the choice Comet makes for any native child. The native partitioner's divergence on decimals wider than 18 digits (#5994) applies to it like any other native shuffle.
The user guide's operator page and its section on Spark-to-Comet conversion types describe the new config. The contributor guide's paragraph on typed operators suggested the dispatcher fuse, so it now describes the conversion and why the fuse was dropped.
## How are these changes tested?
New `CometTypedDatasetSuite` has 10 tests, registered in both PR workflows. Most of them check the answer against Spark, that every operator above the typed operation is native, and that the conversion sits directly on `SerializeFromObjectExec`. They cover:
- `map` followed by an aggregate with AQE on and off, and `flatMap`, `mapPartitions`, `mapGroups`, `cogroup`, and an RDD of objects.
- A broadcast join above the typed operation running as `CometBroadcastHashJoinExec`.
- Decimal, `Option`, nested struct and `array<string>` fields, and an `array<int>` column that declines with its fallback reason.
- A typed operation at the top of the plan, where nothing is converted.
- The user function running exactly once per row.
- An exception thrown past the first batch failing the query with the original message.
- The feature being off by default.
Every converted plan is also checked for row operators reading a columnar child without a transition, and for spurious `ColumnarToRow` fallback reasons. Dropping the transition insertion fails 5 tests, and dropping the `ColumnarToRowTransition` change fails 4.
The suite passes locally on Spark 3.4, 3.5, 4.0, 4.1 and 4.2. On 4.1 these also pass: `CometExecSuite`, `CometExecRuleSuite`, `EliminateRedundantTransitionsSuite`, `RevertNativeForTransitionHeavyStagesSuite`, `CometInMemoryCacheSuite`, `CometRangeExecSuite`, `CometJoinSuite`, and both TPC-DS plan stability suites with no approved plan changes. Semantic scalafix (3.4), scalastyle, spotless and prettier are clean.
I also ran Spark's own typed `Dataset` suites from the 4.1.3 tests jar, with Comet enabled through system properties: `DatasetSuite`, `DatasetPrimitiveSuite`, `DatasetAggregatorSuite`, `DatasetOptimizationSuite`, `DatasetCacheSuite` and `DatasetSerializerRegistratorSuite`, 281 tests. With the conversion off and on, the same 278 pass and the same 3 fail: the two TIME-type tests, which need Spark's test-only confs, and `groupBy.as`, which asserts on Spark's own plan nodes. A query listener showed that 22 of the suites' queries ran with the conversion when it was on.
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] feat: run the operators above a typed Dataset operation natively [datafusion-comet]
Posted by "andygrove (via GitHub)" <gi...@apache.org>.
andygrove commented on code in PR #6564:
URL: https://github.com/apache/datafusion-comet/pull/6564#discussion_r4184778300
##########
spark/src/main/scala/org/apache/comet/rules/CometExecRule.scala:
##########
@@ -1256,6 +1288,78 @@ 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` plans a `DeserializeToObjectExec` at the root, and the RDD's own code
+ // decides how much of it to read, as `take(1)` does.
+ val rootReader = plan match {
+ case _: DeserializeToObjectExec => Some("code reading Dataset.rdd")
Review Comment:
Confirmed: with the conversion on, `map(f).filter(...).map(g).rdd.take(1)` failed on row 30, with `MapElements` at the root of the plan. 1e238811fe takes any root that produces objects as code reading `Dataset.rdd`: an `ObjectProducerExec`, or a filter or project over one. That covers the `DeserializeToObjectExec` that `Dataset.rdd` adds, the typed operation left at the root once `EliminateSerialization` drops it, and a typed filter Spark pulls above that operation. A Dataset's own plan ends in rows, so no other plan has such a root.
The RDD test now runs your query, and the same query with a typed filter at the end, with AQE on and off. Your query fails with the old check, and the typed filter one fails without the walk through the filter.
##########
spark/src/main/scala/org/apache/comet/rules/CometExecRule.scala:
##########
@@ -1256,6 +1268,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:
Fixed in 2737816792, and generalized in c202f51588. The conversion leaves a typed operation's output unconverted when a limit above it can stop reading early, unless an operator that reads all of its input, such as an exchange, a sort or a hash aggregate, sits in between. Your query is the test `a limit does not evaluate typed Dataset rows beyond the result`, with AQE on and off.
##########
spark/src/main/scala/org/apache/comet/rules/CometExecRule.scala:
##########
@@ -473,6 +498,23 @@ case class CometExecRule(session: SparkSession)
case op if shouldApplySparkToColumnar(conf, op) =>
convertToComet(op, CometSparkToColumnarExec).getOrElse(op)
+ // Typed Dataset operations (`map`, `flatMap`, `mapPartitions`, `mapGroups`, ...) pass JVM
+ // objects between their operators, so those stay on Spark. Each of them ends in
+ // `SerializeFromObjectExec`, though, whose output is ordinary rows, and converting those to
+ // Arrow lets the operators above the typed operation run natively.
+ case op: SerializeFromObjectExec
+ if CometConf.COMET_CONVERT_FROM_TYPED_DATASET_ENABLED.get(conf) =>
+ if (op
+ .getTagValue(CometExecRule.SKIP_TYPED_DATASET_CONVERSION_UNDER_LIMIT)
+ .isDefined) {
+ withFallbackReason(
+ op,
+ "Comet does not convert the output of a typed Dataset operation below a limit " +
+ "because Arrow batching could evaluate rows beyond Spark's row-level limit")
+ } else {
+ convertTypedDatasetOutput(op)
Review Comment:
Fixed in c202f51588. A `MapPartitionsExec` above the typed operation now counts as a reader that can stop early, as a limit does. Your query is the test `a mapPartitions function does not evaluate typed Dataset rows it never reads`, with AQE on and off.
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] feat: run the operators above a typed Dataset operation natively [datafusion-comet]
Posted by "sunchao (via GitHub)" <gi...@apache.org>.
sunchao commented on code in PR #6564:
URL: https://github.com/apache/datafusion-comet/pull/6564#discussion_r4178950799
##########
spark/src/main/scala/org/apache/comet/rules/CometExecRule.scala:
##########
@@ -1256,6 +1288,78 @@ 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` plans a `DeserializeToObjectExec` at the root, and the RDD's own code
+ // decides how much of it to read, as `take(1)` does.
+ val rootReader = plan match {
+ case _: DeserializeToObjectExec => Some("code reading Dataset.rdd")
Review Comment:
[P2] Preserve RDD laziness after Spark eliminates serialization. `Dataset.rdd` does not always retain a `DeserializeToObjectExec` root: when the Dataset ends with another `map`, Spark’s `EliminateSerialization` removes that deserializer and leaves `MapElementsExec`. Consequently, `map(f).filter(...).map(g).rdd.take(1)` bypasses this guard and converts the earlier serializer. If `f` throws on row 30, conversion-enabled Comet fails although Spark and conversion-disabled Comet return the first result successfully. Could the safeguard recognize these optimized RDD roots and preserve row-level consumption? Add this case alongside the current RDD regression test.
Evidence: Reproduced at this head on Spark 4.1.3/JDK 17: `spark.range(0, 100, 1, 1).map { i => if (i == 30L) throw new IllegalArgumentException("unexpected evaluation of row 30"); i + 1L }.filter(col("value") > 0L).map(i => i + 1L).rdd.take(1)`. With AQE both false and true, Spark and Comet with `spark.comet.convert.typedDataset.enabled=false` return `Array(2L)`. Enabling conversion throws `SparkException` caused by `CometNativeException` wrapping that `IllegalArgumentException`, with `RowArrowReader.loadNextBatch` in the stack. The RDD execution plan is rooted at `MapElements`, above `DeserializeToObject -> CometColumnarToRow -> CometFilter -> CometSparkRowToColumnar -> SerializeFromObject`. Spark’s `EliminateSerialization` rule removes the outer deserialize/serialize pair across all five supported versions. The six-configuration probe failed only in the two conversion-enabled cases.
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] feat: run the operators above a typed Dataset operation natively [datafusion-comet]
Posted by "comphead (via GitHub)" <gi...@apache.org>.
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] feat: run the operators above a typed Dataset operation natively [datafusion-comet]
Posted by "andygrove (via GitHub)" <gi...@apache.org>.
andygrove commented on PR #6564:
URL: https://github.com/apache/datafusion-comet/pull/6564#issuecomment-5988083467
@copilot resolve the merge conflicts on this branch.
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] feat: run the operators above a typed Dataset operation natively [datafusion-comet]
Posted by "andygrove (via GitHub)" <gi...@apache.org>.
andygrove commented on PR #6564:
URL: https://github.com/apache/datafusion-comet/pull/6564#issuecomment-5998880198
I reproduced the struct failure, but it isn't new in this PR, so I filed #6685 for it. Main fails the same way today with `spark.comet.convert.rdd.enabled` or `spark.comet.exec.localTableScan.enabled`, a struct column and a repartition, and the same code is in every release since 1.0.0. Native shuffle wraps the batches of any child that isn't a native operator in `ColumnarBatchArrowReader`, which closes the vectors the conversion reuses. It also fails at the default batch size once a partition has more than one batch. #6607 reads a `CometNativeArrowSource` child as an Arrow stream instead. With just that change applied on top of this branch, your repro, the orderBy and join variants, and the two leaf conversions all pass. Since this is already queued and the config is off by default, I'd rather let it land and add your struct-under-shuffle test to #6607, which will close #6685, when I merge main into it. A struct decline here would only come out again in #6607.
On the profiles, I ran `CometTypedDatasetSuite` against main plus this branch on 3.4, 4.0 and 4.2, and all 15 tests pass. With the earlier 3.5 and 4.1 runs, it's green on every profile. You're right about the top-k arm: `SerializeFromObjectExec` doesn't report an output ordering in any supported version, so I'll remove it. I'll do that, the #6005 TODO, the shared benchmark helper and the `datasources.md` bullet in #6607, since it touches the same code and lands next.
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] feat: run the operators above a typed Dataset operation natively [datafusion-comet]
Posted by "andygrove (via GitHub)" <gi...@apache.org>.
andygrove merged PR #6564:
URL: https://github.com/apache/datafusion-comet/pull/6564
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] feat: run the operators above a typed Dataset operation natively [datafusion-comet]
Posted by "sunchao (via GitHub)" <gi...@apache.org>.
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] feat: run the operators above a typed Dataset operation natively [datafusion-comet]
Posted by "sunchao (via GitHub)" <gi...@apache.org>.
sunchao commented on code in PR #6564:
URL: https://github.com/apache/datafusion-comet/pull/6564#discussion_r4178564733
##########
spark/src/main/scala/org/apache/comet/rules/CometExecRule.scala:
##########
@@ -1256,6 +1268,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:
[P2] Preserve short-circuiting when a limit consumes typed output. With `spark.comet.convert.typedDataset.enabled=true`, `map(...).limit(1)` fills an Arrow batch before returning its first row. A user function that throws on row 30 therefore fails the query, although Spark and conversion-disabled Comet return the first row successfully. The resulting plan is `CometCollectLimit -> CometSparkRowToColumnar -> SerializeFromObject`. Please preserve row-level limiting before batching where valid, or decline conversion for this pipeline, and add a regression test.
Evidence: Reproduced at the reviewed head in CometTestBase on Spark 4.1.3/JDK 17, with AQE both false and true: `spark.range(0, 100, 1, 1).map { i => if (i == 30L) throw new IllegalArgumentException("unexpected evaluation of row 30"); i + 1L }.toDF().limit(1).collect()`. Spark and Comet with typed conversion disabled return `[Row(1)]`. Enabling conversion throws `SparkException` caused by that `IllegalArgumentException`, with `RowArrowReader.loadNextBatch` in the stack. Spark’s `CollectLimitExec.executeCollect` calls `child.executeTake(limit)`, whereas the inserted Arrow reader consumes a batch before the limit can stop it. The six-configuration probe failed only in the two conversion-enabled cases.
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] feat: run the operators above a typed Dataset operation natively [datafusion-comet]
Posted by "andygrove (via GitHub)" <gi...@apache.org>.
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] feat: run the operators above a typed Dataset operation natively [datafusion-comet]
Posted by "sunchao (via GitHub)" <gi...@apache.org>.
sunchao commented on code in PR #6564:
URL: https://github.com/apache/datafusion-comet/pull/6564#discussion_r4178740950
##########
spark/src/main/scala/org/apache/comet/rules/CometExecRule.scala:
##########
@@ -473,6 +498,23 @@ case class CometExecRule(session: SparkSession)
case op if shouldApplySparkToColumnar(conf, op) =>
convertToComet(op, CometSparkToColumnarExec).getOrElse(op)
+ // Typed Dataset operations (`map`, `flatMap`, `mapPartitions`, `mapGroups`, ...) pass JVM
+ // objects between their operators, so those stay on Spark. Each of them ends in
+ // `SerializeFromObjectExec`, though, whose output is ordinary rows, and converting those to
+ // Arrow lets the operators above the typed operation run natively.
+ case op: SerializeFromObjectExec
+ if CometConf.COMET_CONVERT_FROM_TYPED_DATASET_ENABLED.get(conf) =>
+ if (op
+ .getTagValue(CometExecRule.SKIP_TYPED_DATASET_CONVERSION_UNDER_LIMIT)
+ .isDefined) {
+ withFallbackReason(
+ op,
+ "Comet does not convert the output of a typed Dataset operation below a limit " +
+ "because Arrow batching could evaluate rows beyond Spark's row-level limit")
+ } else {
+ convertTypedDatasetOutput(op)
Review Comment:
[P2] Preserve lazy consumption before downstream `MapPartitionsExec` consumers. The physical-limit guard does not cover iterator operations such as `mapPartitions(_.take(1))`. With typed conversion enabled, an intervening native filter batches the earlier typed output and evaluates row 30 before returning the first result. A user function throwing on row 30 therefore fails a query that succeeds in Spark and conversion-disabled Comet. Please decline conversion across these lazy iterator boundaries unless full consumption is guaranteed, or otherwise preserve row-level consumption, and add this regression case.
Evidence: Reproduced at this head on Spark 4.1.3/JDK 17: `spark.range(0, 100, 1, 1).map { i => if (i == 30L) throw new IllegalArgumentException("unexpected evaluation of row 30"); i + 1L }.filter(col("value") > 0L).mapPartitions(_.take(1)).toDF().collect()`. With AQE both false and true, Spark and Comet with `spark.comet.convert.typedDataset.enabled=false` return `[Row(1)]`. Setting it to true throws `SparkException` caused by `CometNativeException` wrapping that `IllegalArgumentException`. The failing plan contains `MapPartitions -> DeserializeToObject -> CometColumnarToRow -> CometFilter -> CometSparkRowToColumnar -> SerializeFromObject`, with no physical limit node. Spark’s `MapPartitionsExec` passes a lazy mapped iterator to the user function, whereas `RowArrowReader.loadNextBatch` consumes upstream rows to fill a batch. The six-configuration probe failed only in the two conversion-enabled cases.
--
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: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org