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/04 16:59:14 UTC
[PR] feat: use native shuffle for row input by converting it to Arrow [datafusion-comet]
andygrove opened a new pull request, #6607:
URL: https://github.com/apache/datafusion-comet/pull/6607
## Which issue does this PR close?
Closes #6598.
## Rationale for this change
When a shuffle's child is a Spark row-based operator, Comet uses its JVM columnar shuffle. #6598 found that converting the child's rows with `CometSparkToColumnarExec` and using native shuffle instead is faster. This adds that behind `spark.comet.convert.shuffleInput.enabled`, off by default, next to the other `spark.comet.convert.*` settings (#6593). It reuses the `ArrowWriter` that #6566 sped up.
## What changes are included in this PR?
- `CometShuffleExchangeExec.convertsInputForNativeShuffle` decides whether a shuffle takes the conversion. The config has to be on, and the shuffle's child must be a Spark row-based plan that would otherwise go to the JVM columnar shuffle. `CometSparkToColumnarExec` has to support the child's columns, and native shuffle has to support the partitioning and the columns. The checks are `shuffleSupported`'s own, so the rule converts only a shuffle that will become native.
- A shuffle that hashes a decimal wider than 18 digits into more than one partition stays on the JVM columnar shuffle. Native shuffle doesn't hash those as Spark does (#5994), and the other input of a join may still be partitioned by Spark. This is the same guard as in #6564.
- The conversion leaves Celeborn alone for now, since it has no JVM columnar shuffle to replace.
- `CometExecRule` puts `CometSparkToColumnarExec` over the shuffle's child. As in #6564, it first applies Spark's own transition rule to the child's subtree, because Spark inserts no transitions below a `RowToColumnarTransition`. The generic fallback case no longer reports those transitions as unsupported operators when AQE runs the rule over the plan again.
- `revertRedundantColumnarShuffle` (#4004) also reverts a native shuffle over converted rows when both its parent and the converted operator are Spark aggregates. That's the same row → Arrow → shuffle → Arrow → row pattern.
- The shuffle tuning guide describes the setting.
The new `CometShuffleInputConversionBenchmark` hash-partitions 4M rows into 200 partitions and aggregates them. It checks every arm's plan and result before timing. Best times on an M3 Max, `local[1]`, release build:
| Case | Spark | Comet | Comet, converted | Comet (repeat) |
| ----------------------------------------------------- | ------- | ------- | ---------------- | -------------- |
| RDD rows: int, long, double | 2145 ms | 1973 ms | 1546 ms | 2014 ms |
| RDD rows: int, long, string, double, decimal(18,2) | 3079 ms | 3505 ms | 2627 ms | 3167 ms |
| `map` over a Parquet scan | 1936 ms | 1846 ms | 1145 ms | 1771 ms |
"Comet" is the JVM columnar shuffle. The repeat row shows the noise: other jobs were running on the machine.
This overlaps with #6564, which converts a typed Dataset operation's output the same way. The two share the transition handling and the fallback case for transitions, so whichever lands second has a small conflict in `CometExecRule`.
## How are these changes tested?
`CometShuffleInputConversionSuite` covers:
- a shuffle of RDD rows, with AQE on and off
- a shuffle over a typed `map` of a Comet scan, which checks that the Spark operators under the conversion read Comet operators through Comet's transitions and that no later pass reports those transitions as fallbacks
- aggregates whose partial aggregate runs on Spark
- the revert between two Spark aggregates, with the revert on and off
- a join with one input on native shuffle and one on the JVM columnar shuffle
- a join on wide decimal keys
- columns the conversion does not support
- range partitioning
- `spark.comet.shuffle.mode=jvm`
- the default
The suite turns off the conversion of leaf operators, which `CometTestBase` turns on and production leaves off.
I checked that the suite fails when each piece is removed: the conversion, the wide decimal guard, the transitions, the fallback case for transitions, and the revert of the new shape.
On Spark 4.1, the new suite passes along with `CometExecRuleSuite`, `CometShuffleSuite` and `DisableAQECometShuffleSuite` (192 tests). On Spark 3.5, the new suite and `CometExecRuleSuite` pass. The strict warnings build and the semantic scalafix check pass on Spark 3.5.
--
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: use native shuffle for row input by converting it to Arrow [datafusion-comet]
Posted by "andygrove (via GitHub)" <gi...@apache.org>.
andygrove merged PR #6607:
URL: https://github.com/apache/datafusion-comet/pull/6607
--
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: use native shuffle for row input by converting it to Arrow [datafusion-comet]
Posted by "comphead (via GitHub)" <gi...@apache.org>.
comphead commented on code in PR #6607:
URL: https://github.com/apache/datafusion-comet/pull/6607#discussion_r4197642414
##########
spark/src/test/scala/org/apache/spark/sql/comet/execution/shuffle/CometCelebornShufflePlanningSuite.scala:
##########
@@ -265,6 +265,21 @@ class CometCelebornShufflePlanningSuite extends CometTestBase {
}
}
+ test("the shuffle input conversion leaves a Spark child to Spark's shuffle") {
+ // The conversion only takes over from the JVM columnar shuffle, which Celeborn does not
+ // have, and Celeborn's native shuffle needs a native child.
+ withSQLConf(
+ CometConf.COMET_SHUFFLE_MODE.key -> "native",
Review Comment:
This test sets `spark.comet.shuffle.mode` to `native`. In that mode `columnarShuffleFailureReasons` already returns "Comet columnar shuffle not enabled", so `shuffleSupported` is `None` and `convertsInputForNativeShuffle` is false with or without Celeborn. Reading the code, the test would still pass if the Celeborn check in `shuffleSupported` were removed. The suite default is `auto`, and there the Celeborn check is the only thing that keeps a Spark child on Spark's shuffle.
Could this test run in `auto` so that it covers the exclusion? I expect it to pass there but have not run it.
--
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: use native shuffle for row input by converting it to Arrow [datafusion-comet]
Posted by "andygrove (via GitHub)" <gi...@apache.org>.
andygrove commented on PR #6607:
URL: https://github.com/apache/datafusion-comet/pull/6607#issuecomment-6021061724
I've updated the description. The stream read bullet now covers only the range sampling change, the overlap paragraph is gone, and the `CometExecRule` bullet credits the fallback case for transitions to #6564. It also mentions the revert's new INFO text. While I was there I changed "hash aggregates" to "aggregates" in the revert bullet to match the docs, dropped the stream read and the fallback case from the list of changes the tests catch, and updated the test counts at c90dcfca5.
The `Lint Java (Spark 4.0, JDK 21)` failure wasn't from the merge. In its log, the download of `scala-compiler-2.13.16.jar` starts and never finishes, and `scala-compile-first` then fails on `comet-common` with `ClassNotFoundException: scala.tools.nsc.Main`, before any source is compiled. The JDK 17 job downloaded the same jar and passed, and on c90dcfca5 the JDK 21 job passed too.
I've applied `run-all-spark-profiles`. The Comet suites on 3.4, 3.5, 4.0 and 4.2 are running now.
--
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: use native shuffle for row input by converting it to Arrow [datafusion-comet]
Posted by "andygrove (via GitHub)" <gi...@apache.org>.
andygrove commented on code in PR #6607:
URL: https://github.com/apache/datafusion-comet/pull/6607#discussion_r4186685038
##########
spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometShuffleExchangeExec.scala:
##########
@@ -144,9 +146,16 @@ case class CometShuffleExchangeExec(
ctx.perPartitionByKey,
positionalRoundRobin.isDefined)
case None =>
- // Non-native child (e.g. CometSparkToColumnarExec): no subtree to inline. The dep gets
- // built via the convenience overload below; we just need a real RDD of batches.
- child.executeColumnar()
+ child match {
+ // Native reads the source's Arrow stream, as it does for a native operator's input. The
+ // batches from the source's `executeColumnar` share vectors that it reuses for the next
+ // batch, and the convenience overload below closes each batch once native has it.
+ // Closing a struct vector drops its children, so the next batch would lose them.
+ case source: CometNativeArrowSource => source.doExecuteAsArrowStream()
Review Comment:
Agreed, it's in #6689 now. The leaf conversion half of the struct test moved to `CometNativeShuffleSuite` there, and it also covers a local table scan, which fails the same way (#6685). This PR keeps the same change until #6689 lands, and then I'll merge main. For the backport call: the code path dates from #4572, so 1.0.0 has the bug as well, and every conversion that reaches it is off by default, so it isn't a 1.1.0 regression. I've left it without a backport label for now.
--
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: use native shuffle for row input by converting it to Arrow [datafusion-comet]
Posted by "andygrove (via GitHub)" <gi...@apache.org>.
andygrove commented on code in PR #6607:
URL: https://github.com/apache/datafusion-comet/pull/6607#discussion_r4197877768
##########
spark/src/test/scala/org/apache/spark/sql/comet/execution/shuffle/CometCelebornShufflePlanningSuite.scala:
##########
@@ -265,6 +265,21 @@ class CometCelebornShufflePlanningSuite extends CometTestBase {
}
}
+ test("the shuffle input conversion leaves a Spark child to Spark's shuffle") {
+ // The conversion only takes over from the JVM columnar shuffle, which Celeborn does not
+ // have, and Celeborn's native shuffle needs a native child.
+ withSQLConf(
+ CometConf.COMET_SHUFFLE_MODE.key -> "native",
Review Comment:
Confirmed: with the Celeborn block in `shuffleSupported` removed, the test still passed in `native`. Running it in `auto` alone wouldn't catch that either, though. Under Celeborn, `shuffleSupported` first calls `cometCelebornShuffleFallbackReason` through `isCometShuffleEnabledReason`, and that already turns down `auto` and `jvm` with "Celeborn-backed Comet shuffle requires spark.comet.shuffle.mode=native". The `auto` version passed with the block removed, and failed only once both checks were gone.
c90dcfca5 runs the test in `auto`, `jvm` and `native`, and checks that the fallback reason names Celeborn, which in `native` only the block gives; without it the reason is "Comet columnar shuffle not enabled". The test now fails if the block is removed, if both Celeborn checks are removed, since `auto` then converts, or if the gate stops requiring `CometColumnarShuffle`. Removing only the mode checks in `cometCelebornShuffleFallbackReason` fails eight other tests in the suite.
--
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: use native shuffle for row input by converting it to Arrow [datafusion-comet]
Posted by "andygrove (via GitHub)" <gi...@apache.org>.
andygrove commented on code in PR #6607:
URL: https://github.com/apache/datafusion-comet/pull/6607#discussion_r4184364143
##########
spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometShuffleExchangeExec.scala:
##########
@@ -524,6 +525,45 @@ object CometShuffleExchangeExec
None
}
+ /**
+ * Whether a shuffle whose child is a Spark row-based plan can convert the child's rows to Arrow
+ * with `CometSparkToColumnarExec` and use native shuffle, instead of the JVM columnar shuffle.
+ * See [[CometConf.COMET_CONVERT_FROM_SHUFFLE_INPUT_ENABLED]]. The checks are those that
+ * [[shuffleSupported]] would make for the converted plan, so `CometExecRule` only converts a
+ * shuffle that will become native. Pure: does not tag the node.
+ */
+ def convertsInputForNativeShuffle(s: ShuffleExchangeExec): Boolean = {
+ val conf = s.conf
+ CometConf.COMET_CONVERT_FROM_SHUFFLE_INPUT_ENABLED.get(conf) &&
+ // A decision an earlier pass recorded stands, as in shuffleSupported.
+ !hasFallbackReason(s) &&
+ isCometShuffleEnabledReason(s).isEmpty &&
+ CometConf.COMET_SHUFFLE_CONVERT_FROM_SPARK_PLAN_ENABLED.get(conf) &&
+ // Not exercised with Celeborn, which has no JVM columnar shuffle to replace.
+ !isCometCelebornShuffleManagerEnabled(conf) &&
+ !isCometPlan(s.child) &&
+ !s.child.supportsColumnar &&
+ !isShuffleOperator(s.child) &&
+ !stageContainsDPPScan(s) &&
+ CometSparkToColumnarExec.isSchemaSupported(s.child.schema, ListBuffer.empty) &&
+ nativeShuffleFailureReasons(s).isEmpty &&
+ !hashesWideDecimal(s)
Review Comment:
Confirmed with your follow-up case: with the conversion on, the join on `hash(lk) === hash(rk)` returned 0 rows to Spark's 800. a20a4a5e9c has the guard look for a string anywhere in a hash key's expression rather than only at the key's type, so a key computed from a string stays on the JVM columnar shuffle too. A wide decimal still counts only as the key's own type, since Comet's `hash` of a wide decimal already falls back to Spark.
The string join test now also joins on the hashes of the strings, and fails without the change. It uses your two keys instead of four. Both shuffles write the strings with invalid UTF-8 replaced (#4764), and with four keys two of the hashes share a partition, where the join then matches them. That happens with the conversion off as well.
--
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: use native shuffle for row input by converting it to Arrow [datafusion-comet]
Posted by "andygrove (via GitHub)" <gi...@apache.org>.
andygrove closed pull request #6607: feat: use native shuffle for row input by converting it to Arrow
URL: https://github.com/apache/datafusion-comet/pull/6607
--
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: use native shuffle for row input by converting it to Arrow [datafusion-comet]
Posted by "andygrove (via GitHub)" <gi...@apache.org>.
andygrove commented on code in PR #6607:
URL: https://github.com/apache/datafusion-comet/pull/6607#discussion_r4186480163
##########
spark/src/main/scala/org/apache/comet/rules/CometExecRule.scala:
##########
@@ -1274,6 +1291,30 @@ 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 Spark operator feeds to a shuffle to Arrow, so that the shuffle runs as
+ * native shuffle instead of the JVM columnar shuffle. See
+ * [[CometConf.COMET_CONVERT_FROM_SHUFFLE_INPUT_ENABLED]], and
+ * [[CometShuffleExchangeExec.convertsInputForNativeShuffle]] for when it applies.
+ *
+ * Spark inserts the columnar transitions after this rule, but it does not look below a
+ * `RowToColumnarTransition` such as `CometSparkToColumnarExec`. Without a transition, a Spark
+ * operator in the child's subtree that reads a Comet operator would do so 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 convertShuffleInput(s: ShuffleExchangeExec): SparkPlan = {
+ val child =
+ ApplyColumnarRulesAndInsertTransitions(Seq.empty, outputsColumnar = false).apply(s.child)
Review Comment:
Confirmed. With `transitionRevert.enabled=true` and `maxTransitions=0` over the typed `map` case, the task failed casting Spark's `OffHeapColumnVector` to `CometVector` in `reconcileStreamSchema`. 73c727c3b reverts the exchange with its stage: when `RevertNativeForTransitionHeavyStages` reverts the stage below a native shuffle whose child is the conversion, the shuffle goes back to the JVM columnar shuffle that the conversion replaced, and `EliminateRedundantTransitions` drops the `RowToColumnarExec` under it as it does for any JVM columnar shuffle. The gate now only converts a shuffle that `shuffleSupported` would give the JVM columnar shuffle, so that shuffle is always available there. Your test is in the suite, with AQE on and off, and fails without the change.
##########
spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometShuffleExchangeExec.scala:
##########
@@ -524,6 +548,58 @@ object CometShuffleExchangeExec
None
}
+ /**
+ * Whether a shuffle whose child is a Spark row-based plan can convert the child's rows to Arrow
+ * with `CometSparkToColumnarExec` and use native shuffle, instead of the JVM columnar shuffle.
+ * See [[CometConf.COMET_CONVERT_FROM_SHUFFLE_INPUT_ENABLED]]. The checks are those that
+ * [[shuffleSupported]] would make for the converted plan, so `CometExecRule` only converts a
+ * shuffle that will become native. Pure: does not tag the node.
+ */
+ def convertsInputForNativeShuffle(s: ShuffleExchangeExec): Boolean = {
Review Comment:
Updated all three in ef1d8ec60. The first two now give `spark.comet.convert.shuffleInput.enabled` as the exception, and the plan construction step says that a `CometNativeArrowSource` child exports its own stream, which `prepareArrowStreamShuffleDependency` hands to native, while any other input goes through the `prepareShuffleDependency(rdd, ...)` overload. The tuning guide still quoted the revert's old log message, which this PR changed, so that's fixed too.
##########
spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometShuffleExchangeExec.scala:
##########
@@ -524,6 +548,58 @@ object CometShuffleExchangeExec
None
}
+ /**
+ * Whether a shuffle whose child is a Spark row-based plan can convert the child's rows to Arrow
+ * with `CometSparkToColumnarExec` and use native shuffle, instead of the JVM columnar shuffle.
+ * See [[CometConf.COMET_CONVERT_FROM_SHUFFLE_INPUT_ENABLED]]. The checks are those that
+ * [[shuffleSupported]] would make for the converted plan, so `CometExecRule` only converts a
+ * shuffle that will become native. Pure: does not tag the node.
+ */
+ def convertsInputForNativeShuffle(s: ShuffleExchangeExec): Boolean = {
+ val conf = s.conf
+ CometConf.COMET_CONVERT_FROM_SHUFFLE_INPUT_ENABLED.get(conf) &&
+ // A decision an earlier pass recorded stands, as in shuffleSupported.
+ !hasFallbackReason(s) &&
+ isCometShuffleEnabledReason(s).isEmpty &&
+ CometConf.COMET_SHUFFLE_CONVERT_FROM_SPARK_PLAN_ENABLED.get(conf) &&
+ // Not exercised with Celeborn, which has no JVM columnar shuffle to replace.
+ !isCometCelebornShuffleManagerEnabled(conf) &&
+ !isCometPlan(s.child) &&
+ !s.child.supportsColumnar &&
+ !isShuffleOperator(s.child) &&
+ !stageContainsDPPScan(s) &&
+ CometSparkToColumnarExec.isSchemaSupported(s.child.schema, ListBuffer.empty) &&
+ // Arrow holds the time part of an interval in nanoseconds, so the conversion overflows on a
+ // calendar interval with more microseconds than that can hold. The JVM columnar shuffle
+ // leaves calendar intervals to Spark's shuffle as well.
+ !s.child.schema.existsRecursively(_.isInstanceOf[CalendarIntervalType]) &&
+ nativeShuffleFailureReasons(s).isEmpty &&
+ !hashesDifferentlyFromSpark(s)
Review Comment:
Done in 21651aa1d, as you laid it out: the config, `!isCometPlan(s.child)`, `shuffleSupported(s).contains(CometColumnarShuffle)`, the schema check, `nativeShuffleFailureReasons(s).isEmpty` and `!hashesDifferentlyFromSpark(s)`. The `native` mode case wasn't intended. The setting is only meant to take over from the JVM columnar shuffle, so `native` mode now keeps Spark's shuffle for a Spark child, and the shuffle mode test checks `native` as well as `jvm`. The predicate now tags the node when `shuffleSupported` finds no Comet shuffle, but the next case in `CometExecRule` would record the same reasons anyway.
##########
spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometShuffleExchangeExec.scala:
##########
@@ -227,12 +236,27 @@ case class CometShuffleExchangeExec(
ctx,
positionalRoundRobin))
case None =>
- CometShuffleExchangeExec.prepareShuffleDependency(
- inputRDD.asInstanceOf[RDD[ColumnarBatch]],
- child.output,
- outputPartitioning,
- serializer,
- metrics)
+ child match {
+ case _: CometNativeArrowSource =>
+ CometShuffleExchangeExec.prepareArrowStreamShuffleDependency(
+ inputRDD.asInstanceOf[RDD[ArrowArrayStream]],
+ // The range partitioner samples rows, so it needs them as batches.
+ outputPartitioning match {
+ case _: RangePartitioning => Some(child.executeColumnar())
Review Comment:
Done in 54defa139. When the conversion's child is row-based, the sampling job now reads those rows through `conversion.child.execute()` and projects only the sort keys, as `prepareJVMShuffleDependency` does. `prepareNativeShuffleDependency` takes its sampling input as rows, and the other callers turn their batches into rows as it did before. The old sampling job also ran through the conversion's metrics, so a range-partitioned shuffle reported twice the rows it wrote. The range partitioning test now checks the count (2d5e9cf45).
The benchmark has a range case now, over all five columns. Best times in ms for Spark, Comet, Comet converted and the Comet repeat were 2955, 3296, 3133 and 3079, so range partitioning comes out about even with the JVM columnar shuffle, where the hash cases gain 15 to 30%. The sampling change didn't make a difference I could measure above the noise, which was around 5% on a busy machine. The tuning guide now says that range partitioning comes out about even (80253c8fe).
##########
spark/src/test/scala/org/apache/comet/exec/CometShuffleInputConversionSuite.scala:
##########
@@ -0,0 +1,310 @@
+/*
+ * 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.comet.exec
+
+import org.apache.spark.sql.{CometTestBase, DataFrame, Row}
+import org.apache.spark.sql.comet.CometSparkToColumnarExec
+import org.apache.spark.sql.comet.execution.shuffle.{CometColumnarShuffle, CometNativeShuffle, CometShuffleExchangeExec}
+import org.apache.spark.sql.execution.{ColumnarToRowExec, ColumnarToRowTransition, SparkPlan}
+import org.apache.spark.sql.functions.{array, avg, col, count, hash, length, max, min, size, sum}
+import org.apache.spark.sql.internal.SQLConf
+import org.apache.spark.sql.types.{ArrayType, BinaryType, CalendarIntervalType, DataTypes, DecimalType, IntegerType, LongType, StringType, StructType}
+import org.apache.spark.unsafe.types.CalendarInterval
+
+import org.apache.comet.{CometConf, ExtendedExplainInfo}
+
+// Top-level, so the encoder needs no outer pointer.
+case class ShuffleInputRec(a: Int, b: String)
+
+/** Tests for [[CometConf.COMET_CONVERT_FROM_SHUFFLE_INPUT_ENABLED]]. */
+class CometShuffleInputConversionSuite extends CometTestBase {
+
+ import testImplicits._
+
+ /**
+ * `CometTestBase` turns on the conversion of leaf operators such as RDD scans, which would
+ * convert the scans below the shuffles here before this conversion sees them. It is off by
+ * default.
+ */
+ private def withoutLeafConversion(f: => Unit): Unit =
+ withSQLConf(sparkToArrowConversionConfs(enabled = false): _*)(f)
+
+ /** Defines a test that runs with the shuffle input conversion enabled. */
+ private def convertTest(name: String)(f: => Unit): Unit =
+ test(name) {
+ withoutLeafConversion {
+ withSQLConf(CometConf.COMET_CONVERT_FROM_SHUFFLE_INPUT_ENABLED.key -> "true")(f)
+ }
+ }
+
+ private val rowSchema = new StructType()
+ .add("k", IntegerType)
+ .add("l", LongType)
+ .add("s", DataTypes.StringType)
+ .add("d", DataTypes.DoubleType)
+ .add("m", DecimalType(18, 2))
+
+ /**
+ * `rows` rows from an RDD, so the shuffle above them reads a Spark operator: the conversion of
+ * RDD scans, `spark.comet.convert.rdd.enabled`, is off by default.
+ */
+ private def rowsDf(rows: Int = 1000): DataFrame = {
+ val data = (0 until rows).map { i =>
+ Row(
+ i % 23,
+ i.toLong,
+ if (i % 11 == 0) null else s"value-$i",
+ i + 0.5d,
+ java.math.BigDecimal.valueOf(i * 7919L % 1000000L, 2))
+ }
+ spark.createDataFrame(spark.sparkContext.parallelize(data, 4), rowSchema)
+ }
+
+ /** Rows with an `array<int>` column, which `CometSparkToColumnarExec` does not convert. */
+ private def arraysDf(rows: Int = 100): DataFrame = {
+ val schema = new StructType().add("k", IntegerType).add("xs", ArrayType(IntegerType))
+ val data = (0 until rows).map(i => Row(i % 23, Seq(i, i + 1)))
+ spark.createDataFrame(spark.sparkContext.parallelize(data, 4), schema)
+ }
+
+ private def conversions(plan: SparkPlan): Seq[CometSparkToColumnarExec] =
+ collectWithSubqueries(plan) { case c: CometSparkToColumnarExec => c }
+
+ private def cometShuffles(plan: SparkPlan): Seq[CometShuffleExchangeExec] =
+ collectWithSubqueries(plan) { case s: CometShuffleExchangeExec => s }
+
+ /** The native shuffles that read rows `CometSparkToColumnarExec` converted. */
+ private def convertedShuffles(plan: SparkPlan): Seq[CometShuffleExchangeExec] =
+ cometShuffles(plan).filter { s =>
+ s.shuffleType == CometNativeShuffle && s.child.isInstanceOf[CometSparkToColumnarExec]
+ }
+
+ /**
+ * Spark inserts no columnar transitions below a `CometSparkToColumnarExec`, so the rule adds
+ * them for the operators under the conversion. Without one, an operator reads its Comet child
+ * through `CometExec.doExecute`, Spark's interpreted columnar-to-row path, which gives the
+ * right answer slowly, so only the plan shows it.
+ */
+ private def assertRowOperatorsReadThroughTransitions(plan: SparkPlan): Unit = {
+ val converted = conversions(plan)
+ assert(converted.nonEmpty, plan)
+ converted.foreach { conversion =>
+ // `InputAdapter` and `WholeStageCodegenExec` report their child's `supportsColumnar`.
+ val bare = conversion.child.collect {
+ case p
+ if !p.supportsColumnar && !p.isInstanceOf[ColumnarToRowTransition] &&
+ p.children.exists(_.supportsColumnar) =>
+ p
+ }
+ assert(bare.isEmpty, s"row operators read a columnar child without a transition:\n$plan")
+ assert(
+ conversion.child.collect { case c: ColumnarToRowExec => c }.isEmpty,
+ s"expected Comet's columnar-to-row transitions, not Spark's:\n$plan")
+ }
+ // The transitions are added on the first of the rule's passes under AQE. A later pass must
+ // not report them as operators Comet failed to convert.
+ val reasons = new ExtendedExplainInfo().getFallbackReasons(plan)
+ assert(!reasons.exists(_.contains("ColumnarToRow")), reasons)
+ }
+
+ convertTest("the shuffle of a Spark operator's rows runs as native shuffle") {
+ Seq("true", "false").foreach { aqe =>
+ withSQLConf(SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> aqe) {
+ val (_, plan) = checkSparkAnswer(
+ rowsDf()
+ .repartition(10, col("k"))
+ .groupBy("k")
+ .agg(sum("l"), sum(length(col("s"))), sum("d"), sum("m")))
+ assert(convertedShuffles(plan).length == 1, s"AQE $aqe:\n$plan")
+ assert(
+ !cometShuffles(plan).exists(_.shuffleType == CometColumnarShuffle),
+ s"AQE $aqe:\n$plan")
+ }
+ }
+ }
+
+ convertTest("Spark operators over Comet operators read them through transitions") {
+ withParquetTable((0 until 200).map(i => (i, (i % 13).toString)), "tbl") {
+ Seq("true", "false").foreach { aqe =>
+ withSQLConf(SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> aqe) {
+ val ds = spark.sql("SELECT _1 AS a, _2 AS b FROM tbl").as[ShuffleInputRec]
+ // The shuffle reads the typed operation's SerializeFromObject.
+ val (_, repartitioned) = checkSparkAnswer(
+ ds.map(r => ShuffleInputRec(r.a % 13, r.b))
+ .repartition(7, col("a"))
+ .groupBy("a")
+ .agg(count("b")))
+ assert(convertedShuffles(repartitioned).length == 1, s"AQE $aqe:\n$repartitioned")
+ assertRowOperatorsReadThroughTransitions(repartitioned)
+ // The shuffle reads a Spark partial aggregate over the typed operation. The final
+ // aggregate stays on Spark too, so the shuffle would go back to Spark's own.
+ withSQLConf(CometConf.COMET_SHUFFLE_REVERT_REDUNDANT_COLUMNAR_ENABLED.key -> "false") {
+ val (_, aggregated) =
+ checkSparkAnswer(ds.map(r => ShuffleInputRec(r.a % 7, r.b)).groupBy("a").count())
+ assert(convertedShuffles(aggregated).length == 1, s"AQE $aqe:\n$aggregated")
+ assertRowOperatorsReadThroughTransitions(aggregated)
+ }
+ }
+ }
+ }
+ }
+
+ convertTest("aggregates whose partial aggregate runs on Spark") {
Review Comment:
Added in ddb428bb3. The shuffle did convert here, as it turns out: `min(s)` on a string makes Spark plan sort aggregates, and the revert in the next test only matches hash aggregates, so the shuffle stayed native between two Spark `SortAggregateExec`s. The test now asserts each plan for two shapes: hash aggregates whose final aggregate runs natively over the Spark partial, and the original aggregates on Spark sort aggregates.
##########
spark/src/test/scala/org/apache/comet/exec/CometShuffleInputConversionSuite.scala:
##########
@@ -0,0 +1,310 @@
+/*
+ * 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.comet.exec
+
+import org.apache.spark.sql.{CometTestBase, DataFrame, Row}
+import org.apache.spark.sql.comet.CometSparkToColumnarExec
+import org.apache.spark.sql.comet.execution.shuffle.{CometColumnarShuffle, CometNativeShuffle, CometShuffleExchangeExec}
+import org.apache.spark.sql.execution.{ColumnarToRowExec, ColumnarToRowTransition, SparkPlan}
+import org.apache.spark.sql.functions.{array, avg, col, count, hash, length, max, min, size, sum}
+import org.apache.spark.sql.internal.SQLConf
+import org.apache.spark.sql.types.{ArrayType, BinaryType, CalendarIntervalType, DataTypes, DecimalType, IntegerType, LongType, StringType, StructType}
+import org.apache.spark.unsafe.types.CalendarInterval
+
+import org.apache.comet.{CometConf, ExtendedExplainInfo}
+
+// Top-level, so the encoder needs no outer pointer.
+case class ShuffleInputRec(a: Int, b: String)
+
+/** Tests for [[CometConf.COMET_CONVERT_FROM_SHUFFLE_INPUT_ENABLED]]. */
+class CometShuffleInputConversionSuite extends CometTestBase {
+
+ import testImplicits._
+
+ /**
+ * `CometTestBase` turns on the conversion of leaf operators such as RDD scans, which would
+ * convert the scans below the shuffles here before this conversion sees them. It is off by
+ * default.
+ */
+ private def withoutLeafConversion(f: => Unit): Unit =
+ withSQLConf(sparkToArrowConversionConfs(enabled = false): _*)(f)
+
+ /** Defines a test that runs with the shuffle input conversion enabled. */
+ private def convertTest(name: String)(f: => Unit): Unit =
+ test(name) {
+ withoutLeafConversion {
+ withSQLConf(CometConf.COMET_CONVERT_FROM_SHUFFLE_INPUT_ENABLED.key -> "true")(f)
+ }
+ }
+
+ private val rowSchema = new StructType()
+ .add("k", IntegerType)
+ .add("l", LongType)
+ .add("s", DataTypes.StringType)
+ .add("d", DataTypes.DoubleType)
+ .add("m", DecimalType(18, 2))
+
+ /**
+ * `rows` rows from an RDD, so the shuffle above them reads a Spark operator: the conversion of
+ * RDD scans, `spark.comet.convert.rdd.enabled`, is off by default.
+ */
+ private def rowsDf(rows: Int = 1000): DataFrame = {
+ val data = (0 until rows).map { i =>
+ Row(
+ i % 23,
+ i.toLong,
+ if (i % 11 == 0) null else s"value-$i",
+ i + 0.5d,
+ java.math.BigDecimal.valueOf(i * 7919L % 1000000L, 2))
+ }
+ spark.createDataFrame(spark.sparkContext.parallelize(data, 4), rowSchema)
+ }
+
+ /** Rows with an `array<int>` column, which `CometSparkToColumnarExec` does not convert. */
+ private def arraysDf(rows: Int = 100): DataFrame = {
+ val schema = new StructType().add("k", IntegerType).add("xs", ArrayType(IntegerType))
+ val data = (0 until rows).map(i => Row(i % 23, Seq(i, i + 1)))
+ spark.createDataFrame(spark.sparkContext.parallelize(data, 4), schema)
+ }
+
+ private def conversions(plan: SparkPlan): Seq[CometSparkToColumnarExec] =
+ collectWithSubqueries(plan) { case c: CometSparkToColumnarExec => c }
+
+ private def cometShuffles(plan: SparkPlan): Seq[CometShuffleExchangeExec] =
+ collectWithSubqueries(plan) { case s: CometShuffleExchangeExec => s }
+
+ /** The native shuffles that read rows `CometSparkToColumnarExec` converted. */
+ private def convertedShuffles(plan: SparkPlan): Seq[CometShuffleExchangeExec] =
+ cometShuffles(plan).filter { s =>
+ s.shuffleType == CometNativeShuffle && s.child.isInstanceOf[CometSparkToColumnarExec]
+ }
+
+ /**
+ * Spark inserts no columnar transitions below a `CometSparkToColumnarExec`, so the rule adds
+ * them for the operators under the conversion. Without one, an operator reads its Comet child
+ * through `CometExec.doExecute`, Spark's interpreted columnar-to-row path, which gives the
+ * right answer slowly, so only the plan shows it.
+ */
+ private def assertRowOperatorsReadThroughTransitions(plan: SparkPlan): Unit = {
+ val converted = conversions(plan)
+ assert(converted.nonEmpty, plan)
+ converted.foreach { conversion =>
+ // `InputAdapter` and `WholeStageCodegenExec` report their child's `supportsColumnar`.
+ val bare = conversion.child.collect {
+ case p
+ if !p.supportsColumnar && !p.isInstanceOf[ColumnarToRowTransition] &&
+ p.children.exists(_.supportsColumnar) =>
+ p
+ }
+ assert(bare.isEmpty, s"row operators read a columnar child without a transition:\n$plan")
+ assert(
+ conversion.child.collect { case c: ColumnarToRowExec => c }.isEmpty,
+ s"expected Comet's columnar-to-row transitions, not Spark's:\n$plan")
+ }
+ // The transitions are added on the first of the rule's passes under AQE. A later pass must
+ // not report them as operators Comet failed to convert.
+ val reasons = new ExtendedExplainInfo().getFallbackReasons(plan)
+ assert(!reasons.exists(_.contains("ColumnarToRow")), reasons)
+ }
+
+ convertTest("the shuffle of a Spark operator's rows runs as native shuffle") {
+ Seq("true", "false").foreach { aqe =>
+ withSQLConf(SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> aqe) {
+ val (_, plan) = checkSparkAnswer(
+ rowsDf()
+ .repartition(10, col("k"))
+ .groupBy("k")
+ .agg(sum("l"), sum(length(col("s"))), sum("d"), sum("m")))
+ assert(convertedShuffles(plan).length == 1, s"AQE $aqe:\n$plan")
+ assert(
+ !cometShuffles(plan).exists(_.shuffleType == CometColumnarShuffle),
+ s"AQE $aqe:\n$plan")
+ }
+ }
+ }
+
+ convertTest("Spark operators over Comet operators read them through transitions") {
+ withParquetTable((0 until 200).map(i => (i, (i % 13).toString)), "tbl") {
+ Seq("true", "false").foreach { aqe =>
+ withSQLConf(SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> aqe) {
+ val ds = spark.sql("SELECT _1 AS a, _2 AS b FROM tbl").as[ShuffleInputRec]
+ // The shuffle reads the typed operation's SerializeFromObject.
+ val (_, repartitioned) = checkSparkAnswer(
+ ds.map(r => ShuffleInputRec(r.a % 13, r.b))
+ .repartition(7, col("a"))
+ .groupBy("a")
+ .agg(count("b")))
+ assert(convertedShuffles(repartitioned).length == 1, s"AQE $aqe:\n$repartitioned")
+ assertRowOperatorsReadThroughTransitions(repartitioned)
+ // The shuffle reads a Spark partial aggregate over the typed operation. The final
+ // aggregate stays on Spark too, so the shuffle would go back to Spark's own.
+ withSQLConf(CometConf.COMET_SHUFFLE_REVERT_REDUNDANT_COLUMNAR_ENABLED.key -> "false") {
+ val (_, aggregated) =
+ checkSparkAnswer(ds.map(r => ShuffleInputRec(r.a % 7, r.b)).groupBy("a").count())
+ assert(convertedShuffles(aggregated).length == 1, s"AQE $aqe:\n$aggregated")
+ assertRowOperatorsReadThroughTransitions(aggregated)
+ }
+ }
+ }
+ }
+ }
+
+ convertTest("aggregates whose partial aggregate runs on Spark") {
+ checkSparkAnswer(
+ rowsDf()
+ .groupBy((col("k") % 7).as("g"))
+ .agg(sum("l"), count("s"), avg("d"), sum("m"), max("l"), min("s")))
+ }
+
+ convertTest("a shuffle between two Spark aggregates goes back to Spark's shuffle") {
+ withSQLConf(CometConf.COMET_EXEC_AGGREGATE_ENABLED.key -> "false") {
+ val (_, plan) = checkSparkAnswer(rowsDf().groupBy("k").agg(sum("l")))
+ assert(cometShuffles(plan).isEmpty, plan)
+ withSQLConf(CometConf.COMET_SHUFFLE_REVERT_REDUNDANT_COLUMNAR_ENABLED.key -> "false") {
+ val (_, kept) = checkSparkAnswer(rowsDf().groupBy("k").agg(sum("l")))
+ assert(convertedShuffles(kept).length == 1, kept)
+ }
+ }
+ }
+
+ convertTest("a join with an input that stays on the JVM columnar shuffle") {
Review Comment:
Done in ddb428bb3. There's now a test for each key type the gate admits: boolean, the four integer types, float and double with NaN and both zeros, decimal(18, 2), date, timestamp, timestamp_ntz and binary, each with boundary values and NULL keys. Each test first compares the `spark_partition_id()` of every row after a `repartition` with Spark's, which covers the NULL keys that the join filters out, and then runs the join with one input on each shuffle. They share a `withShuffledJoins` helper with the wide decimal and string tests. To check that the placement comparison would notice a difference, I took out the wide decimal guard and added a `decimal(38, 0)` key, and it failed on the partition ids.
The Celeborn case is in `CometCelebornShufflePlanningSuite` (21651aa1d), and fails if the gate stops excluding Celeborn. The test over several batches now has `array<string>` and `map<string,string>` columns with NULL elements, values and rows. Without the stream read it fails on the map too (`no more field nodes for field key`), not only on the struct.
--
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: use native shuffle for row input by converting it to Arrow [datafusion-comet]
Posted by "andygrove (via GitHub)" <gi...@apache.org>.
andygrove commented on code in PR #6607:
URL: https://github.com/apache/datafusion-comet/pull/6607#discussion_r4186482074
##########
spark/src/test/scala/org/apache/comet/exec/CometShuffleInputConversionSuite.scala:
##########
@@ -0,0 +1,310 @@
+/*
+ * 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.comet.exec
+
+import org.apache.spark.sql.{CometTestBase, DataFrame, Row}
+import org.apache.spark.sql.comet.CometSparkToColumnarExec
+import org.apache.spark.sql.comet.execution.shuffle.{CometColumnarShuffle, CometNativeShuffle, CometShuffleExchangeExec}
+import org.apache.spark.sql.execution.{ColumnarToRowExec, ColumnarToRowTransition, SparkPlan}
+import org.apache.spark.sql.functions.{array, avg, col, count, hash, length, max, min, size, sum}
+import org.apache.spark.sql.internal.SQLConf
+import org.apache.spark.sql.types.{ArrayType, BinaryType, CalendarIntervalType, DataTypes, DecimalType, IntegerType, LongType, StringType, StructType}
+import org.apache.spark.unsafe.types.CalendarInterval
+
+import org.apache.comet.{CometConf, ExtendedExplainInfo}
+
+// Top-level, so the encoder needs no outer pointer.
+case class ShuffleInputRec(a: Int, b: String)
+
+/** Tests for [[CometConf.COMET_CONVERT_FROM_SHUFFLE_INPUT_ENABLED]]. */
+class CometShuffleInputConversionSuite extends CometTestBase {
+
+ import testImplicits._
+
+ /**
+ * `CometTestBase` turns on the conversion of leaf operators such as RDD scans, which would
+ * convert the scans below the shuffles here before this conversion sees them. It is off by
+ * default.
+ */
+ private def withoutLeafConversion(f: => Unit): Unit =
+ withSQLConf(sparkToArrowConversionConfs(enabled = false): _*)(f)
+
+ /** Defines a test that runs with the shuffle input conversion enabled. */
+ private def convertTest(name: String)(f: => Unit): Unit =
+ test(name) {
+ withoutLeafConversion {
+ withSQLConf(CometConf.COMET_CONVERT_FROM_SHUFFLE_INPUT_ENABLED.key -> "true")(f)
+ }
+ }
+
+ private val rowSchema = new StructType()
+ .add("k", IntegerType)
+ .add("l", LongType)
+ .add("s", DataTypes.StringType)
+ .add("d", DataTypes.DoubleType)
+ .add("m", DecimalType(18, 2))
+
+ /**
+ * `rows` rows from an RDD, so the shuffle above them reads a Spark operator: the conversion of
+ * RDD scans, `spark.comet.convert.rdd.enabled`, is off by default.
+ */
+ private def rowsDf(rows: Int = 1000): DataFrame = {
+ val data = (0 until rows).map { i =>
+ Row(
+ i % 23,
+ i.toLong,
+ if (i % 11 == 0) null else s"value-$i",
+ i + 0.5d,
+ java.math.BigDecimal.valueOf(i * 7919L % 1000000L, 2))
+ }
+ spark.createDataFrame(spark.sparkContext.parallelize(data, 4), rowSchema)
+ }
+
+ /** Rows with an `array<int>` column, which `CometSparkToColumnarExec` does not convert. */
+ private def arraysDf(rows: Int = 100): DataFrame = {
+ val schema = new StructType().add("k", IntegerType).add("xs", ArrayType(IntegerType))
+ val data = (0 until rows).map(i => Row(i % 23, Seq(i, i + 1)))
+ spark.createDataFrame(spark.sparkContext.parallelize(data, 4), schema)
+ }
+
+ private def conversions(plan: SparkPlan): Seq[CometSparkToColumnarExec] =
+ collectWithSubqueries(plan) { case c: CometSparkToColumnarExec => c }
+
+ private def cometShuffles(plan: SparkPlan): Seq[CometShuffleExchangeExec] =
+ collectWithSubqueries(plan) { case s: CometShuffleExchangeExec => s }
+
+ /** The native shuffles that read rows `CometSparkToColumnarExec` converted. */
+ private def convertedShuffles(plan: SparkPlan): Seq[CometShuffleExchangeExec] =
+ cometShuffles(plan).filter { s =>
+ s.shuffleType == CometNativeShuffle && s.child.isInstanceOf[CometSparkToColumnarExec]
+ }
+
+ /**
+ * Spark inserts no columnar transitions below a `CometSparkToColumnarExec`, so the rule adds
+ * them for the operators under the conversion. Without one, an operator reads its Comet child
+ * through `CometExec.doExecute`, Spark's interpreted columnar-to-row path, which gives the
+ * right answer slowly, so only the plan shows it.
+ */
+ private def assertRowOperatorsReadThroughTransitions(plan: SparkPlan): Unit = {
+ val converted = conversions(plan)
+ assert(converted.nonEmpty, plan)
+ converted.foreach { conversion =>
+ // `InputAdapter` and `WholeStageCodegenExec` report their child's `supportsColumnar`.
+ val bare = conversion.child.collect {
+ case p
+ if !p.supportsColumnar && !p.isInstanceOf[ColumnarToRowTransition] &&
+ p.children.exists(_.supportsColumnar) =>
+ p
+ }
+ assert(bare.isEmpty, s"row operators read a columnar child without a transition:\n$plan")
+ assert(
+ conversion.child.collect { case c: ColumnarToRowExec => c }.isEmpty,
+ s"expected Comet's columnar-to-row transitions, not Spark's:\n$plan")
+ }
+ // The transitions are added on the first of the rule's passes under AQE. A later pass must
+ // not report them as operators Comet failed to convert.
+ val reasons = new ExtendedExplainInfo().getFallbackReasons(plan)
+ assert(!reasons.exists(_.contains("ColumnarToRow")), reasons)
+ }
+
+ convertTest("the shuffle of a Spark operator's rows runs as native shuffle") {
+ Seq("true", "false").foreach { aqe =>
+ withSQLConf(SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> aqe) {
+ val (_, plan) = checkSparkAnswer(
+ rowsDf()
+ .repartition(10, col("k"))
+ .groupBy("k")
+ .agg(sum("l"), sum(length(col("s"))), sum("d"), sum("m")))
+ assert(convertedShuffles(plan).length == 1, s"AQE $aqe:\n$plan")
+ assert(
+ !cometShuffles(plan).exists(_.shuffleType == CometColumnarShuffle),
+ s"AQE $aqe:\n$plan")
+ }
+ }
+ }
+
+ convertTest("Spark operators over Comet operators read them through transitions") {
+ withParquetTable((0 until 200).map(i => (i, (i % 13).toString)), "tbl") {
+ Seq("true", "false").foreach { aqe =>
+ withSQLConf(SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> aqe) {
+ val ds = spark.sql("SELECT _1 AS a, _2 AS b FROM tbl").as[ShuffleInputRec]
+ // The shuffle reads the typed operation's SerializeFromObject.
+ val (_, repartitioned) = checkSparkAnswer(
+ ds.map(r => ShuffleInputRec(r.a % 13, r.b))
+ .repartition(7, col("a"))
+ .groupBy("a")
+ .agg(count("b")))
+ assert(convertedShuffles(repartitioned).length == 1, s"AQE $aqe:\n$repartitioned")
+ assertRowOperatorsReadThroughTransitions(repartitioned)
+ // The shuffle reads a Spark partial aggregate over the typed operation. The final
+ // aggregate stays on Spark too, so the shuffle would go back to Spark's own.
+ withSQLConf(CometConf.COMET_SHUFFLE_REVERT_REDUNDANT_COLUMNAR_ENABLED.key -> "false") {
+ val (_, aggregated) =
+ checkSparkAnswer(ds.map(r => ShuffleInputRec(r.a % 7, r.b)).groupBy("a").count())
+ assert(convertedShuffles(aggregated).length == 1, s"AQE $aqe:\n$aggregated")
+ assertRowOperatorsReadThroughTransitions(aggregated)
+ }
+ }
+ }
+ }
+ }
+
+ convertTest("aggregates whose partial aggregate runs on Spark") {
+ checkSparkAnswer(
+ rowsDf()
+ .groupBy((col("k") % 7).as("g"))
+ .agg(sum("l"), count("s"), avg("d"), sum("m"), max("l"), min("s")))
+ }
+
+ convertTest("a shuffle between two Spark aggregates goes back to Spark's shuffle") {
+ withSQLConf(CometConf.COMET_EXEC_AGGREGATE_ENABLED.key -> "false") {
+ val (_, plan) = checkSparkAnswer(rowsDf().groupBy("k").agg(sum("l")))
+ assert(cometShuffles(plan).isEmpty, plan)
+ withSQLConf(CometConf.COMET_SHUFFLE_REVERT_REDUNDANT_COLUMNAR_ENABLED.key -> "false") {
+ val (_, kept) = checkSparkAnswer(rowsDf().groupBy("k").agg(sum("l")))
+ assert(convertedShuffles(kept).length == 1, kept)
+ }
+ }
+ }
+
+ convertTest("a join with an input that stays on the JVM columnar shuffle") {
Review Comment:
Done in ddb428bb3. There's now a test for each key type the gate admits: boolean, the four integer types, float and double with NaN and both zeros, decimal(18, 2), date, timestamp, timestamp_ntz and binary, each with boundary values and NULL keys. Each test first compares the `spark_partition_id()` of every row after a `repartition` with Spark's, which covers the NULL keys that the join filters out, and then runs the join with one input on each shuffle. They share a `withShuffledJoins` helper with the wide decimal and string tests. To check that the placement comparison would notice a difference, I took out the wide decimal guard and added a `decimal(38, 0)` key, and it failed on the partition ids.
The Celeborn case is in `CometCelebornShufflePlanningSuite` (21651aa1d), and fails if the gate stops excluding Celeborn. The test over several batches now has `array<string>` and `map<string,string>` columns with NULL elements, values and rows next to the struct. Without the stream read the struct is what breaks: a map or an array column on its own passes over several batches, NULLs included, as #6685 says.
--
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: use native shuffle for row input by converting it to Arrow [datafusion-comet]
Posted by "andygrove (via GitHub)" <gi...@apache.org>.
andygrove commented on code in PR #6607:
URL: https://github.com/apache/datafusion-comet/pull/6607#discussion_r4179215168
##########
spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometShuffleExchangeExec.scala:
##########
@@ -524,6 +525,45 @@ object CometShuffleExchangeExec
None
}
+ /**
+ * Whether a shuffle whose child is a Spark row-based plan can convert the child's rows to Arrow
+ * with `CometSparkToColumnarExec` and use native shuffle, instead of the JVM columnar shuffle.
+ * See [[CometConf.COMET_CONVERT_FROM_SHUFFLE_INPUT_ENABLED]]. The checks are those that
+ * [[shuffleSupported]] would make for the converted plan, so `CometExecRule` only converts a
+ * shuffle that will become native. Pure: does not tag the node.
+ */
+ def convertsInputForNativeShuffle(s: ShuffleExchangeExec): Boolean = {
+ val conf = s.conf
+ CometConf.COMET_CONVERT_FROM_SHUFFLE_INPUT_ENABLED.get(conf) &&
+ // A decision an earlier pass recorded stands, as in shuffleSupported.
+ !hasFallbackReason(s) &&
+ isCometShuffleEnabledReason(s).isEmpty &&
+ CometConf.COMET_SHUFFLE_CONVERT_FROM_SPARK_PLAN_ENABLED.get(conf) &&
+ // Not exercised with Celeborn, which has no JVM columnar shuffle to replace.
+ !isCometCelebornShuffleManagerEnabled(conf) &&
+ !isCometPlan(s.child) &&
+ !s.child.supportsColumnar &&
+ !isShuffleOperator(s.child) &&
+ !stageContainsDPPScan(s) &&
+ CometSparkToColumnarExec.isSchemaSupported(s.child.schema, ListBuffer.empty) &&
+ nativeShuffleFailureReasons(s).isEmpty &&
+ !hashesWideDecimal(s)
Review Comment:
Confirmed with your case: with the conversion on, the join returned 0 rows to Spark's 400. 9795b61f52 keeps a shuffle that hashes a string, at any nesting depth, on the JVM columnar shuffle, the way the wide decimal guard already worked. Preserving the raw-byte hash would mean native shuffle hashing the bytes before the import replaces invalid UTF-8, which is more than this PR should take on, so I went with the fallback. Your join is now a test in `CometShuffleInputConversionSuite`, and the config doc and the tuning guide say that a string key stays on columnar shuffle. The typed `map` test and the benchmark case that hashed a string now hash a number.
##########
spark/src/main/scala/org/apache/comet/rules/CometExecRule.scala:
##########
@@ -1256,6 +1273,30 @@ 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 Spark operator feeds to a shuffle to Arrow, so that the shuffle runs as
+ * native shuffle instead of the JVM columnar shuffle. See
+ * [[CometConf.COMET_CONVERT_FROM_SHUFFLE_INPUT_ENABLED]], and
+ * [[CometShuffleExchangeExec.convertsInputForNativeShuffle]] for when it applies.
+ *
+ * Spark inserts the columnar transitions after this rule, but it does not look below a
+ * `RowToColumnarTransition` such as `CometSparkToColumnarExec`. Without a transition, a Spark
+ * operator in the child's subtree that reads a Comet operator would do so 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 convertShuffleInput(s: ShuffleExchangeExec): SparkPlan = {
+ val child =
+ ApplyColumnarRulesAndInsertTransitions(Seq.empty, outputsColumnar = false).apply(s.child)
+ convertToComet(child, CometSparkToColumnarExec)
Review Comment:
Confirmed: your 24-row case failed with `no more field nodes for field v` at batch size 7. 16a1f73aa4 has native shuffle read a `CometNativeArrowSource` child through `doExecuteAsArrowStream`, as a native operator already does with such an input, so its batches no longer go through `ColumnarBatchArrowReader`. Range partitioning still samples from `executeColumnar`, but the sampling doesn't close the batches. The new test runs hash and range partitioning over a struct column that fills several batches, with leaf conversion on and off. With leaf conversion on, the same shuffle over a converted RDD scan failed the same way before this change, so `spark.comet.sparkToColumnar.enabled` gets the fix too.
##########
spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometShuffleExchangeExec.scala:
##########
@@ -524,6 +525,45 @@ object CometShuffleExchangeExec
None
}
+ /**
+ * Whether a shuffle whose child is a Spark row-based plan can convert the child's rows to Arrow
+ * with `CometSparkToColumnarExec` and use native shuffle, instead of the JVM columnar shuffle.
+ * See [[CometConf.COMET_CONVERT_FROM_SHUFFLE_INPUT_ENABLED]]. The checks are those that
+ * [[shuffleSupported]] would make for the converted plan, so `CometExecRule` only converts a
+ * shuffle that will become native. Pure: does not tag the node.
+ */
+ def convertsInputForNativeShuffle(s: ShuffleExchangeExec): Boolean = {
+ val conf = s.conf
+ CometConf.COMET_CONVERT_FROM_SHUFFLE_INPUT_ENABLED.get(conf) &&
+ // A decision an earlier pass recorded stands, as in shuffleSupported.
+ !hasFallbackReason(s) &&
+ isCometShuffleEnabledReason(s).isEmpty &&
+ CometConf.COMET_SHUFFLE_CONVERT_FROM_SPARK_PLAN_ENABLED.get(conf) &&
+ // Not exercised with Celeborn, which has no JVM columnar shuffle to replace.
+ !isCometCelebornShuffleManagerEnabled(conf) &&
+ !isCometPlan(s.child) &&
+ !s.child.supportsColumnar &&
+ !isShuffleOperator(s.child) &&
+ !stageContainsDPPScan(s) &&
+ CometSparkToColumnarExec.isSchemaSupported(s.child.schema, ListBuffer.empty) &&
Review Comment:
Confirmed, your row failed with `long overflow` in `IntervalMonthDayNanoWriter`. d8d7ce4b41 leaves an input with a calendar interval anywhere in its schema out of the conversion, so that shuffle stays on Spark's shuffle as it did before. There's a test with your value, and the config doc and the tuning guide list calendar intervals among the columns the conversion rules out.
--
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: use native shuffle for row input by converting it to Arrow [datafusion-comet]
Posted by "sunchao (via GitHub)" <gi...@apache.org>.
sunchao commented on code in PR #6607:
URL: https://github.com/apache/datafusion-comet/pull/6607#discussion_r4178941107
##########
spark/src/main/scala/org/apache/comet/rules/CometExecRule.scala:
##########
@@ -1256,6 +1273,30 @@ 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 Spark operator feeds to a shuffle to Arrow, so that the shuffle runs as
+ * native shuffle instead of the JVM columnar shuffle. See
+ * [[CometConf.COMET_CONVERT_FROM_SHUFFLE_INPUT_ENABLED]], and
+ * [[CometShuffleExchangeExec.convertsInputForNativeShuffle]] for when it applies.
+ *
+ * Spark inserts the columnar transitions after this rule, but it does not look below a
+ * `RowToColumnarTransition` such as `CometSparkToColumnarExec`. Without a transition, a Spark
+ * operator in the child's subtree that reads a Comet operator would do so 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 convertShuffleInput(s: ShuffleExchangeExec): SparkPlan = {
+ val child =
+ ApplyColumnarRulesAndInsertTransitions(Seq.empty, outputsColumnar = false).apply(s.child)
+ convertToComet(child, CometSparkToColumnarExec)
Review Comment:
[P2] Keep the converted input’s struct vectors valid across batches. This new plan shape sends `CometSparkToColumnarExec` through `executeColumnar()` and then `ColumnarBatchArrowReader`, which closes batches backed by `RowArrowReader`’s reusable vectors. A struct payload succeeds for one batch but fails when the next batch is loaded. This affects ordinary struct columns even when the partition key is an integer. Consume the conversion’s Arrow stream directly or repair the bridge’s ownership handling before admitting these inputs.
Evidence: On Spark 4.1.3, create 8,200 RDD rows in one input partition with schema `k: int, payload: struct<v: long>`, then run `repartition(3, col("k")).collect()` with leaf conversion disabled and shuffle-input conversion enabled. At batch size 8,192, the job fails with `C Data interface error: java.lang.IllegalArgumentException: no more field nodes for field v: Int(64, true) and vector []`. Spark, the supplied base, and conversion-disabled head return all 8,200 rows. A 24-row reproduction also passes with one batch and fails with batch size 7.
##########
spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometShuffleExchangeExec.scala:
##########
@@ -524,6 +525,45 @@ object CometShuffleExchangeExec
None
}
+ /**
+ * Whether a shuffle whose child is a Spark row-based plan can convert the child's rows to Arrow
+ * with `CometSparkToColumnarExec` and use native shuffle, instead of the JVM columnar shuffle.
+ * See [[CometConf.COMET_CONVERT_FROM_SHUFFLE_INPUT_ENABLED]]. The checks are those that
+ * [[shuffleSupported]] would make for the converted plan, so `CometExecRule` only converts a
+ * shuffle that will become native. Pure: does not tag the node.
+ */
+ def convertsInputForNativeShuffle(s: ShuffleExchangeExec): Boolean = {
+ val conf = s.conf
+ CometConf.COMET_CONVERT_FROM_SHUFFLE_INPUT_ENABLED.get(conf) &&
+ // A decision an earlier pass recorded stands, as in shuffleSupported.
+ !hasFallbackReason(s) &&
+ isCometShuffleEnabledReason(s).isEmpty &&
+ CometConf.COMET_SHUFFLE_CONVERT_FROM_SPARK_PLAN_ENABLED.get(conf) &&
+ // Not exercised with Celeborn, which has no JVM columnar shuffle to replace.
+ !isCometCelebornShuffleManagerEnabled(conf) &&
+ !isCometPlan(s.child) &&
+ !s.child.supportsColumnar &&
+ !isShuffleOperator(s.child) &&
+ !stageContainsDPPScan(s) &&
+ CometSparkToColumnarExec.isSchemaSupported(s.child.schema, ListBuffer.empty) &&
Review Comment:
[P2] Retain fallback for calendar intervals whose range the Arrow conversion cannot represent. `isSchemaSupported` accepts `CalendarIntervalType`, and the native shuffle checks admit it, so this predicate also converts inputs that previously stayed on Spark shuffle. Spark stores a signed 64-bit microsecond component, but `IntervalMonthDayNanoWriter` multiplies it by 1,000 using `Math.multiplyExact`. Valid Spark intervals therefore become query failures. Reject calendar intervals recursively for this conversion until the full Spark range is supported.
Evidence: On Spark 4.1.3, create one RDD row with schema `k: int, i: CalendarIntervalType` and value `Row(1, new CalendarInterval(0, 0, 10800000000000000L))`; execute `repartition(2, col("k")).collect()`. Spark, the supplied base, and conversion-disabled head return the row. Enabling `spark.comet.convert.shuffleInput.enabled` fails with `ArithmeticException: long overflow` at `IntervalMonthDayNanoWriter.setValueUnsafe`, `ArrowWriters.scala:1623`. The base plan retains Spark’s exchange.
##########
spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometShuffleExchangeExec.scala:
##########
@@ -524,6 +525,45 @@ object CometShuffleExchangeExec
None
}
+ /**
+ * Whether a shuffle whose child is a Spark row-based plan can convert the child's rows to Arrow
+ * with `CometSparkToColumnarExec` and use native shuffle, instead of the JVM columnar shuffle.
+ * See [[CometConf.COMET_CONVERT_FROM_SHUFFLE_INPUT_ENABLED]]. The checks are those that
+ * [[shuffleSupported]] would make for the converted plan, so `CometExecRule` only converts a
+ * shuffle that will become native. Pure: does not tag the node.
+ */
+ def convertsInputForNativeShuffle(s: ShuffleExchangeExec): Boolean = {
+ val conf = s.conf
+ CometConf.COMET_CONVERT_FROM_SHUFFLE_INPUT_ENABLED.get(conf) &&
+ // A decision an earlier pass recorded stands, as in shuffleSupported.
+ !hasFallbackReason(s) &&
+ isCometShuffleEnabledReason(s).isEmpty &&
+ CometConf.COMET_SHUFFLE_CONVERT_FROM_SPARK_PLAN_ENABLED.get(conf) &&
+ // Not exercised with Celeborn, which has no JVM columnar shuffle to replace.
+ !isCometCelebornShuffleManagerEnabled(conf) &&
+ !isCometPlan(s.child) &&
+ !s.child.supportsColumnar &&
+ !isShuffleOperator(s.child) &&
+ !stageContainsDPPScan(s) &&
+ CometSparkToColumnarExec.isSchemaSupported(s.child.schema, ListBuffer.empty) &&
+ nativeShuffleFailureReasons(s).isEmpty &&
+ !hashesWideDecimal(s)
Review Comment:
[P1] Preserve Spark’s raw-byte hash when converting string partition keys. With `spark.comet.convert.shuffleInput.enabled=true`, a string-key join can convert one input while an `array<int>` payload keeps the other on JVM shuffle. Spark hashes `UTF8String` bytes directly, whereas the converted input passes through `ScanExec.import_column`, which replaces invalid UTF-8 before native hashing. Equal Spark keys then reach different partitions and matching rows disappear. Preserve the original partition hash before sanitization, or retain JVM shuffle for affected inputs.
Evidence: Spark 4.1.3 reproduction: create 40 RDD rows `(Array((0x80 + i % 4).toByte), i.toLong)`, cast the binary column to string on both join inputs, and retain an `array<int>` payload only on the right. Disable leaf conversion, AQE, and broadcast joins; use 10 shuffle partitions. Spark, the supplied base, and head with conversion disabled return 400 rows. Head with conversion enabled returns 0. Reproduced by `FullShuffleReviewSuite` in `/tmp/pr6607-review-full-l2ehudzv/head-tests.log`; the executed plan contains one native and one JVM columnar exchange.
--
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: use native shuffle for row input by converting it to Arrow [datafusion-comet]
Posted by "comphead (via GitHub)" <gi...@apache.org>.
comphead commented on code in PR #6607:
URL: https://github.com/apache/datafusion-comet/pull/6607#discussion_r4185774420
##########
spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometShuffleExchangeExec.scala:
##########
@@ -524,6 +548,58 @@ object CometShuffleExchangeExec
None
}
+ /**
+ * Whether a shuffle whose child is a Spark row-based plan can convert the child's rows to Arrow
+ * with `CometSparkToColumnarExec` and use native shuffle, instead of the JVM columnar shuffle.
+ * See [[CometConf.COMET_CONVERT_FROM_SHUFFLE_INPUT_ENABLED]]. The checks are those that
+ * [[shuffleSupported]] would make for the converted plan, so `CometExecRule` only converts a
+ * shuffle that will become native. Pure: does not tag the node.
+ */
+ def convertsInputForNativeShuffle(s: ShuffleExchangeExec): Boolean = {
+ val conf = s.conf
+ CometConf.COMET_CONVERT_FROM_SHUFFLE_INPUT_ENABLED.get(conf) &&
+ // A decision an earlier pass recorded stands, as in shuffleSupported.
+ !hasFallbackReason(s) &&
+ isCometShuffleEnabledReason(s).isEmpty &&
+ CometConf.COMET_SHUFFLE_CONVERT_FROM_SPARK_PLAN_ENABLED.get(conf) &&
+ // Not exercised with Celeborn, which has no JVM columnar shuffle to replace.
+ !isCometCelebornShuffleManagerEnabled(conf) &&
+ !isCometPlan(s.child) &&
+ !s.child.supportsColumnar &&
+ !isShuffleOperator(s.child) &&
+ !stageContainsDPPScan(s) &&
+ CometSparkToColumnarExec.isSchemaSupported(s.child.schema, ListBuffer.empty) &&
+ // Arrow holds the time part of an interval in nanoseconds, so the conversion overflows on a
+ // calendar interval with more microseconds than that can hold. The JVM columnar shuffle
+ // leaves calendar intervals to Spark's shuffle as well.
+ !s.child.schema.existsRecursively(_.isInstanceOf[CalendarIntervalType]) &&
+ nativeShuffleFailureReasons(s).isEmpty &&
+ !hashesDifferentlyFromSpark(s)
Review Comment:
Most of these conjuncts restate `shuffleSupported`: an earlier fallback reason, shuffle enabled, Celeborn, `convertFromSparkPlan`, shuffle operator, columnar child and DPP scan. The config doc says the setting replaces "a shuffle that would use Comet's JVM columnar shuffle", which is `shuffleSupported(s).contains(CometColumnarShuffle)`. With that, only `!isCometPlan(s.child)` (first, so Comet children skip the call), the schema check, `nativeShuffleFailureReasons(s).isEmpty` and `hashesDifferentlyFromSpark(s)` remain. The calendar interval conjunct goes too, since `columnarShuffleFailureReasons` already passes `allowCalendarInterval = false`, and a gate added to `shuffleSupported` later applies here without a second edit.
One case changes. With `spark.comet.shuffle.mode=native`, a Spark child keeps Spark's shuffle, but this predicate converts it when the setting is on. Is that intended? The config doc does not say so.
##########
spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometShuffleExchangeExec.scala:
##########
@@ -524,6 +548,58 @@ object CometShuffleExchangeExec
None
}
+ /**
+ * Whether a shuffle whose child is a Spark row-based plan can convert the child's rows to Arrow
+ * with `CometSparkToColumnarExec` and use native shuffle, instead of the JVM columnar shuffle.
+ * See [[CometConf.COMET_CONVERT_FROM_SHUFFLE_INPUT_ENABLED]]. The checks are those that
+ * [[shuffleSupported]] would make for the converted plan, so `CometExecRule` only converts a
+ * shuffle that will become native. Pure: does not tag the node.
+ */
+ def convertsInputForNativeShuffle(s: ShuffleExchangeExec): Boolean = {
Review Comment:
This changes a selection rule that the contributor docs state as fact. `native_shuffle.md` says the child must be a `CometPlan` and that "Row-based Spark operators require JVM shuffle" (lines 48 and 49). `jvm_shuffle.md` says JVM shuffle is "the only option" for a Spark row-based child (lines 45 to 47). The plan construction step in `native_shuffle.md` (lines 157 to 161) still lists `CometSparkToColumnarExec` under the `prepareShuffleDependency(rdd, ...)` overload, but a `CometNativeArrowSource` child now reads its own Arrow stream through `prepareArrowStreamShuffleDependency`.
Could those three places be updated in this PR, with `spark.comet.convert.shuffleInput.enabled` as the exception to the first two?
##########
spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometShuffleExchangeExec.scala:
##########
@@ -227,12 +236,27 @@ case class CometShuffleExchangeExec(
ctx,
positionalRoundRobin))
case None =>
- CometShuffleExchangeExec.prepareShuffleDependency(
- inputRDD.asInstanceOf[RDD[ColumnarBatch]],
- child.output,
- outputPartitioning,
- serializer,
- metrics)
+ child match {
+ case _: CometNativeArrowSource =>
+ CometShuffleExchangeExec.prepareArrowStreamShuffleDependency(
+ inputRDD.asInstanceOf[RDD[ArrowArrayStream]],
+ // The range partitioner samples rows, so it needs them as batches.
+ outputPartitioning match {
+ case _: RangePartitioning => Some(child.executeColumnar())
Review Comment:
For range partitioning, the sampling job now converts every column of every row to Arrow through `CometSparkToColumnarExec`, only to read the sort keys back with `rowIterator()`. `prepareJVMShuffleDependency`, which this replaces, samples just the key projection straight from the rows. The benchmark only hash-partitions, and the tuning guide says the native route is usually faster, so I could not tell what range partitioning costs.
Could the sampling input be the child's rows (an `RDD[InternalRow]`), or could the benchmark get a range case? I have not measured this.
##########
spark/src/main/scala/org/apache/comet/rules/CometExecRule.scala:
##########
@@ -1274,6 +1291,30 @@ 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 Spark operator feeds to a shuffle to Arrow, so that the shuffle runs as
+ * native shuffle instead of the JVM columnar shuffle. See
+ * [[CometConf.COMET_CONVERT_FROM_SHUFFLE_INPUT_ENABLED]], and
+ * [[CometShuffleExchangeExec.convertsInputForNativeShuffle]] for when it applies.
+ *
+ * Spark inserts the columnar transitions after this rule, but it does not look below a
+ * `RowToColumnarTransition` such as `CometSparkToColumnarExec`. Without a transition, a Spark
+ * operator in the child's subtree that reads a Comet operator would do so 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 convertShuffleInput(s: ShuffleExchangeExec): SparkPlan = {
+ val child =
+ ApplyColumnarRulesAndInsertTransitions(Seq.empty, outputsColumnar = false).apply(s.child)
Review Comment:
This inserts `ColumnarToRowExec` nodes into the stage under the native exchange, and `RevertNativeForTransitionHeavyStages` counts them. Reading the rule, with `spark.comet.exec.transitionRevert.enabled=true` and more than `maxTransitions` of them, `revertToSpark` strips `CometSparkToColumnarExec` but the exchange keeps `shuffleType = CometNativeShuffle`. `revertStageIfNeeded` then wraps the stage in Spark's `RowToColumnarExec` because the exchange supports columnar output, and `EliminateRedundantTransitions` only removes that wrapper for `CometColumnarShuffle`. I would expect `executeColumnar()` to return Spark vectors and the cast to `CometVector` in `reconcileStreamSchema` to fail when the task runs.
Could you add a test with `transitionRevert.enabled=true` and `maxTransitions=0` over the typed `map` case? If it fails, reverting the exchange with its stage, or skipping the revert for this shape, would avoid it.
##########
spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometShuffleExchangeExec.scala:
##########
@@ -144,9 +146,16 @@ case class CometShuffleExchangeExec(
ctx.perPartitionByKey,
positionalRoundRobin.isDefined)
case None =>
- // Non-native child (e.g. CometSparkToColumnarExec): no subtree to inline. The dep gets
- // built via the convenience overload below; we just need a real RDD of batches.
- child.executeColumnar()
+ child match {
+ // Native reads the source's Arrow stream, as it does for a native operator's input. The
+ // batches from the source's `executeColumnar` share vectors that it reuses for the next
+ // batch, and the convenience overload below closes each batch once native has it.
+ // Closing a struct vector drops its children, so the next batch would lose them.
+ case source: CometNativeArrowSource => source.doExecuteAsArrowStream()
Review Comment:
This stream read is a fix on its own. As the PR description says, native shuffle over a converted leaf failed on a struct column after the first batch, and `branch-1.1` still has `child.executeColumnar()` at this spot. After the squash merge, that fix cannot be cherry-picked to the release branches without the opt-in feature.
Could it go in a separate PR first, with the `leafConversion = true` half of the struct test, so it can be backported as the backporting guide describes?
##########
spark/src/test/scala/org/apache/comet/exec/CometShuffleInputConversionSuite.scala:
##########
@@ -0,0 +1,310 @@
+/*
+ * 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.comet.exec
+
+import org.apache.spark.sql.{CometTestBase, DataFrame, Row}
+import org.apache.spark.sql.comet.CometSparkToColumnarExec
+import org.apache.spark.sql.comet.execution.shuffle.{CometColumnarShuffle, CometNativeShuffle, CometShuffleExchangeExec}
+import org.apache.spark.sql.execution.{ColumnarToRowExec, ColumnarToRowTransition, SparkPlan}
+import org.apache.spark.sql.functions.{array, avg, col, count, hash, length, max, min, size, sum}
+import org.apache.spark.sql.internal.SQLConf
+import org.apache.spark.sql.types.{ArrayType, BinaryType, CalendarIntervalType, DataTypes, DecimalType, IntegerType, LongType, StringType, StructType}
+import org.apache.spark.unsafe.types.CalendarInterval
+
+import org.apache.comet.{CometConf, ExtendedExplainInfo}
+
+// Top-level, so the encoder needs no outer pointer.
+case class ShuffleInputRec(a: Int, b: String)
+
+/** Tests for [[CometConf.COMET_CONVERT_FROM_SHUFFLE_INPUT_ENABLED]]. */
+class CometShuffleInputConversionSuite extends CometTestBase {
+
+ import testImplicits._
+
+ /**
+ * `CometTestBase` turns on the conversion of leaf operators such as RDD scans, which would
+ * convert the scans below the shuffles here before this conversion sees them. It is off by
+ * default.
+ */
+ private def withoutLeafConversion(f: => Unit): Unit =
+ withSQLConf(sparkToArrowConversionConfs(enabled = false): _*)(f)
+
+ /** Defines a test that runs with the shuffle input conversion enabled. */
+ private def convertTest(name: String)(f: => Unit): Unit =
+ test(name) {
+ withoutLeafConversion {
+ withSQLConf(CometConf.COMET_CONVERT_FROM_SHUFFLE_INPUT_ENABLED.key -> "true")(f)
+ }
+ }
+
+ private val rowSchema = new StructType()
+ .add("k", IntegerType)
+ .add("l", LongType)
+ .add("s", DataTypes.StringType)
+ .add("d", DataTypes.DoubleType)
+ .add("m", DecimalType(18, 2))
+
+ /**
+ * `rows` rows from an RDD, so the shuffle above them reads a Spark operator: the conversion of
+ * RDD scans, `spark.comet.convert.rdd.enabled`, is off by default.
+ */
+ private def rowsDf(rows: Int = 1000): DataFrame = {
+ val data = (0 until rows).map { i =>
+ Row(
+ i % 23,
+ i.toLong,
+ if (i % 11 == 0) null else s"value-$i",
+ i + 0.5d,
+ java.math.BigDecimal.valueOf(i * 7919L % 1000000L, 2))
+ }
+ spark.createDataFrame(spark.sparkContext.parallelize(data, 4), rowSchema)
+ }
+
+ /** Rows with an `array<int>` column, which `CometSparkToColumnarExec` does not convert. */
+ private def arraysDf(rows: Int = 100): DataFrame = {
+ val schema = new StructType().add("k", IntegerType).add("xs", ArrayType(IntegerType))
+ val data = (0 until rows).map(i => Row(i % 23, Seq(i, i + 1)))
+ spark.createDataFrame(spark.sparkContext.parallelize(data, 4), schema)
+ }
+
+ private def conversions(plan: SparkPlan): Seq[CometSparkToColumnarExec] =
+ collectWithSubqueries(plan) { case c: CometSparkToColumnarExec => c }
+
+ private def cometShuffles(plan: SparkPlan): Seq[CometShuffleExchangeExec] =
+ collectWithSubqueries(plan) { case s: CometShuffleExchangeExec => s }
+
+ /** The native shuffles that read rows `CometSparkToColumnarExec` converted. */
+ private def convertedShuffles(plan: SparkPlan): Seq[CometShuffleExchangeExec] =
+ cometShuffles(plan).filter { s =>
+ s.shuffleType == CometNativeShuffle && s.child.isInstanceOf[CometSparkToColumnarExec]
+ }
+
+ /**
+ * Spark inserts no columnar transitions below a `CometSparkToColumnarExec`, so the rule adds
+ * them for the operators under the conversion. Without one, an operator reads its Comet child
+ * through `CometExec.doExecute`, Spark's interpreted columnar-to-row path, which gives the
+ * right answer slowly, so only the plan shows it.
+ */
+ private def assertRowOperatorsReadThroughTransitions(plan: SparkPlan): Unit = {
+ val converted = conversions(plan)
+ assert(converted.nonEmpty, plan)
+ converted.foreach { conversion =>
+ // `InputAdapter` and `WholeStageCodegenExec` report their child's `supportsColumnar`.
+ val bare = conversion.child.collect {
+ case p
+ if !p.supportsColumnar && !p.isInstanceOf[ColumnarToRowTransition] &&
+ p.children.exists(_.supportsColumnar) =>
+ p
+ }
+ assert(bare.isEmpty, s"row operators read a columnar child without a transition:\n$plan")
+ assert(
+ conversion.child.collect { case c: ColumnarToRowExec => c }.isEmpty,
+ s"expected Comet's columnar-to-row transitions, not Spark's:\n$plan")
+ }
+ // The transitions are added on the first of the rule's passes under AQE. A later pass must
+ // not report them as operators Comet failed to convert.
+ val reasons = new ExtendedExplainInfo().getFallbackReasons(plan)
+ assert(!reasons.exists(_.contains("ColumnarToRow")), reasons)
+ }
+
+ convertTest("the shuffle of a Spark operator's rows runs as native shuffle") {
+ Seq("true", "false").foreach { aqe =>
+ withSQLConf(SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> aqe) {
+ val (_, plan) = checkSparkAnswer(
+ rowsDf()
+ .repartition(10, col("k"))
+ .groupBy("k")
+ .agg(sum("l"), sum(length(col("s"))), sum("d"), sum("m")))
+ assert(convertedShuffles(plan).length == 1, s"AQE $aqe:\n$plan")
+ assert(
+ !cometShuffles(plan).exists(_.shuffleType == CometColumnarShuffle),
+ s"AQE $aqe:\n$plan")
+ }
+ }
+ }
+
+ convertTest("Spark operators over Comet operators read them through transitions") {
+ withParquetTable((0 until 200).map(i => (i, (i % 13).toString)), "tbl") {
+ Seq("true", "false").foreach { aqe =>
+ withSQLConf(SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> aqe) {
+ val ds = spark.sql("SELECT _1 AS a, _2 AS b FROM tbl").as[ShuffleInputRec]
+ // The shuffle reads the typed operation's SerializeFromObject.
+ val (_, repartitioned) = checkSparkAnswer(
+ ds.map(r => ShuffleInputRec(r.a % 13, r.b))
+ .repartition(7, col("a"))
+ .groupBy("a")
+ .agg(count("b")))
+ assert(convertedShuffles(repartitioned).length == 1, s"AQE $aqe:\n$repartitioned")
+ assertRowOperatorsReadThroughTransitions(repartitioned)
+ // The shuffle reads a Spark partial aggregate over the typed operation. The final
+ // aggregate stays on Spark too, so the shuffle would go back to Spark's own.
+ withSQLConf(CometConf.COMET_SHUFFLE_REVERT_REDUNDANT_COLUMNAR_ENABLED.key -> "false") {
+ val (_, aggregated) =
+ checkSparkAnswer(ds.map(r => ShuffleInputRec(r.a % 7, r.b)).groupBy("a").count())
+ assert(convertedShuffles(aggregated).length == 1, s"AQE $aqe:\n$aggregated")
+ assertRowOperatorsReadThroughTransitions(aggregated)
+ }
+ }
+ }
+ }
+ }
+
+ convertTest("aggregates whose partial aggregate runs on Spark") {
+ checkSparkAnswer(
+ rowsDf()
+ .groupBy((col("k") % 7).as("g"))
+ .agg(sum("l"), count("s"), avg("d"), sum("m"), max("l"), min("s")))
+ }
+
+ convertTest("a shuffle between two Spark aggregates goes back to Spark's shuffle") {
+ withSQLConf(CometConf.COMET_EXEC_AGGREGATE_ENABLED.key -> "false") {
+ val (_, plan) = checkSparkAnswer(rowsDf().groupBy("k").agg(sum("l")))
+ assert(cometShuffles(plan).isEmpty, plan)
+ withSQLConf(CometConf.COMET_SHUFFLE_REVERT_REDUNDANT_COLUMNAR_ENABLED.key -> "false") {
+ val (_, kept) = checkSparkAnswer(rowsDf().groupBy("k").agg(sum("l")))
+ assert(convertedShuffles(kept).length == 1, kept)
+ }
+ }
+ }
+
+ convertTest("a join with an input that stays on the JVM columnar shuffle") {
Review Comment:
The join's safety argument is that native shuffle hashes an admitted key as Spark does, but only an `int` key is covered. Could this run over the other key types the gate admits (`long`, `float` and `double` with NaN, `date`, `timestamp`, `decimal(18, 2)`, `binary`, `boolean`), with NULL keys and boundary values? The three join tests repeat the same `withSQLConf` block, so one helper taking the key type would add this without more boilerplate.
Two more cases would help: the Celeborn exclusion (`CometCelebornShufflePlanningSuite` has the pattern for Spark children), and `array<string>` and `map<string,string>` columns with NULLs through the multi-batch stream path, which `CometSparkToColumnarExec` accepts.
##########
spark/src/test/scala/org/apache/comet/exec/CometShuffleInputConversionSuite.scala:
##########
@@ -0,0 +1,310 @@
+/*
+ * 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.comet.exec
+
+import org.apache.spark.sql.{CometTestBase, DataFrame, Row}
+import org.apache.spark.sql.comet.CometSparkToColumnarExec
+import org.apache.spark.sql.comet.execution.shuffle.{CometColumnarShuffle, CometNativeShuffle, CometShuffleExchangeExec}
+import org.apache.spark.sql.execution.{ColumnarToRowExec, ColumnarToRowTransition, SparkPlan}
+import org.apache.spark.sql.functions.{array, avg, col, count, hash, length, max, min, size, sum}
+import org.apache.spark.sql.internal.SQLConf
+import org.apache.spark.sql.types.{ArrayType, BinaryType, CalendarIntervalType, DataTypes, DecimalType, IntegerType, LongType, StringType, StructType}
+import org.apache.spark.unsafe.types.CalendarInterval
+
+import org.apache.comet.{CometConf, ExtendedExplainInfo}
+
+// Top-level, so the encoder needs no outer pointer.
+case class ShuffleInputRec(a: Int, b: String)
+
+/** Tests for [[CometConf.COMET_CONVERT_FROM_SHUFFLE_INPUT_ENABLED]]. */
+class CometShuffleInputConversionSuite extends CometTestBase {
+
+ import testImplicits._
+
+ /**
+ * `CometTestBase` turns on the conversion of leaf operators such as RDD scans, which would
+ * convert the scans below the shuffles here before this conversion sees them. It is off by
+ * default.
+ */
+ private def withoutLeafConversion(f: => Unit): Unit =
+ withSQLConf(sparkToArrowConversionConfs(enabled = false): _*)(f)
+
+ /** Defines a test that runs with the shuffle input conversion enabled. */
+ private def convertTest(name: String)(f: => Unit): Unit =
+ test(name) {
+ withoutLeafConversion {
+ withSQLConf(CometConf.COMET_CONVERT_FROM_SHUFFLE_INPUT_ENABLED.key -> "true")(f)
+ }
+ }
+
+ private val rowSchema = new StructType()
+ .add("k", IntegerType)
+ .add("l", LongType)
+ .add("s", DataTypes.StringType)
+ .add("d", DataTypes.DoubleType)
+ .add("m", DecimalType(18, 2))
+
+ /**
+ * `rows` rows from an RDD, so the shuffle above them reads a Spark operator: the conversion of
+ * RDD scans, `spark.comet.convert.rdd.enabled`, is off by default.
+ */
+ private def rowsDf(rows: Int = 1000): DataFrame = {
+ val data = (0 until rows).map { i =>
+ Row(
+ i % 23,
+ i.toLong,
+ if (i % 11 == 0) null else s"value-$i",
+ i + 0.5d,
+ java.math.BigDecimal.valueOf(i * 7919L % 1000000L, 2))
+ }
+ spark.createDataFrame(spark.sparkContext.parallelize(data, 4), rowSchema)
+ }
+
+ /** Rows with an `array<int>` column, which `CometSparkToColumnarExec` does not convert. */
+ private def arraysDf(rows: Int = 100): DataFrame = {
+ val schema = new StructType().add("k", IntegerType).add("xs", ArrayType(IntegerType))
+ val data = (0 until rows).map(i => Row(i % 23, Seq(i, i + 1)))
+ spark.createDataFrame(spark.sparkContext.parallelize(data, 4), schema)
+ }
+
+ private def conversions(plan: SparkPlan): Seq[CometSparkToColumnarExec] =
+ collectWithSubqueries(plan) { case c: CometSparkToColumnarExec => c }
+
+ private def cometShuffles(plan: SparkPlan): Seq[CometShuffleExchangeExec] =
+ collectWithSubqueries(plan) { case s: CometShuffleExchangeExec => s }
+
+ /** The native shuffles that read rows `CometSparkToColumnarExec` converted. */
+ private def convertedShuffles(plan: SparkPlan): Seq[CometShuffleExchangeExec] =
+ cometShuffles(plan).filter { s =>
+ s.shuffleType == CometNativeShuffle && s.child.isInstanceOf[CometSparkToColumnarExec]
+ }
+
+ /**
+ * Spark inserts no columnar transitions below a `CometSparkToColumnarExec`, so the rule adds
+ * them for the operators under the conversion. Without one, an operator reads its Comet child
+ * through `CometExec.doExecute`, Spark's interpreted columnar-to-row path, which gives the
+ * right answer slowly, so only the plan shows it.
+ */
+ private def assertRowOperatorsReadThroughTransitions(plan: SparkPlan): Unit = {
+ val converted = conversions(plan)
+ assert(converted.nonEmpty, plan)
+ converted.foreach { conversion =>
+ // `InputAdapter` and `WholeStageCodegenExec` report their child's `supportsColumnar`.
+ val bare = conversion.child.collect {
+ case p
+ if !p.supportsColumnar && !p.isInstanceOf[ColumnarToRowTransition] &&
+ p.children.exists(_.supportsColumnar) =>
+ p
+ }
+ assert(bare.isEmpty, s"row operators read a columnar child without a transition:\n$plan")
+ assert(
+ conversion.child.collect { case c: ColumnarToRowExec => c }.isEmpty,
+ s"expected Comet's columnar-to-row transitions, not Spark's:\n$plan")
+ }
+ // The transitions are added on the first of the rule's passes under AQE. A later pass must
+ // not report them as operators Comet failed to convert.
+ val reasons = new ExtendedExplainInfo().getFallbackReasons(plan)
+ assert(!reasons.exists(_.contains("ColumnarToRow")), reasons)
+ }
+
+ convertTest("the shuffle of a Spark operator's rows runs as native shuffle") {
+ Seq("true", "false").foreach { aqe =>
+ withSQLConf(SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> aqe) {
+ val (_, plan) = checkSparkAnswer(
+ rowsDf()
+ .repartition(10, col("k"))
+ .groupBy("k")
+ .agg(sum("l"), sum(length(col("s"))), sum("d"), sum("m")))
+ assert(convertedShuffles(plan).length == 1, s"AQE $aqe:\n$plan")
+ assert(
+ !cometShuffles(plan).exists(_.shuffleType == CometColumnarShuffle),
+ s"AQE $aqe:\n$plan")
+ }
+ }
+ }
+
+ convertTest("Spark operators over Comet operators read them through transitions") {
+ withParquetTable((0 until 200).map(i => (i, (i % 13).toString)), "tbl") {
+ Seq("true", "false").foreach { aqe =>
+ withSQLConf(SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> aqe) {
+ val ds = spark.sql("SELECT _1 AS a, _2 AS b FROM tbl").as[ShuffleInputRec]
+ // The shuffle reads the typed operation's SerializeFromObject.
+ val (_, repartitioned) = checkSparkAnswer(
+ ds.map(r => ShuffleInputRec(r.a % 13, r.b))
+ .repartition(7, col("a"))
+ .groupBy("a")
+ .agg(count("b")))
+ assert(convertedShuffles(repartitioned).length == 1, s"AQE $aqe:\n$repartitioned")
+ assertRowOperatorsReadThroughTransitions(repartitioned)
+ // The shuffle reads a Spark partial aggregate over the typed operation. The final
+ // aggregate stays on Spark too, so the shuffle would go back to Spark's own.
+ withSQLConf(CometConf.COMET_SHUFFLE_REVERT_REDUNDANT_COLUMNAR_ENABLED.key -> "false") {
+ val (_, aggregated) =
+ checkSparkAnswer(ds.map(r => ShuffleInputRec(r.a % 7, r.b)).groupBy("a").count())
+ assert(convertedShuffles(aggregated).length == 1, s"AQE $aqe:\n$aggregated")
+ assertRowOperatorsReadThroughTransitions(aggregated)
+ }
+ }
+ }
+ }
+ }
+
+ convertTest("aggregates whose partial aggregate runs on Spark") {
Review Comment:
This is the only test here without a plan assertion. If the final aggregate stays on Spark, the next test shows the shuffle goes back to Spark's, so this one could pass without the conversion in play. Could it assert `convertedShuffles(plan).nonEmpty`?
--
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