You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@datafusion.apache.org by gi...@apache.org on 2026/10/07 11:25:33 UTC
(datafusion-comet) branch main updated: feat: use native shuffle for row input by converting it to Arrow (#6607)
This is an automated email from the ASF dual-hosted git repository.
github-merge-queue[bot] pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/datafusion-comet.git
The following commit(s) were added to refs/heads/main by this push:
new e038cf05fa feat: use native shuffle for row input by converting it to Arrow (#6607)
e038cf05fa is described below
commit e038cf05fa665dc61fdecbdde1c07b8b6a1ce3cb
Author: Andy Grove <ag...@apache.org>
AuthorDate: Wed Oct 7 10:17:45 2026 +0000
feat: use native shuffle for row input by converting it to Arrow (#6607)
* feat: use native shuffle for row input by converting it to Arrow
When a shuffle's child is a Spark row-based operator, Comet uses its JVM
columnar shuffle. With the new spark.comet.convert.shuffleInput.enabled
(off by default), CometExecRule instead puts CometSparkToColumnarExec over
the child and uses native shuffle, where native shuffle supports the
partitioning and the conversion supports the columns. A shuffle that hashes
a decimal wider than 18 digits stays on the JVM columnar shuffle, because
native shuffle does not hash those as Spark does.
As for the typed Dataset conversion, the child's subtree gets its columnar
transitions from Spark's own rule first, since Spark inserts none below a
RowToColumnarTransition. The revert of a shuffle between two Spark
aggregates covers the new shape too.
* test: run CometShuffleInputConversionSuite in the PR builds
* fix: keep a shuffle that hashes a string on the JVM columnar shuffle
Spark's partitioner hashes a string's bytes as they are. Once
CometSparkToColumnarExec has converted the rows, the import into native
replaces invalid UTF-8 before native shuffle hashes the string, so such a
key can go to a different partition than Spark's partitioner sends it to.
A join whose other input stays on the JVM columnar shuffle then loses the
matches for that key. The conversion now leaves a shuffle that hashes a
string on the JVM columnar shuffle, as it does one that hashes a wide
decimal.
* fix: have native shuffle read a converted child's Arrow stream
A native shuffle over a CometSparkToColumnarExec wrapped the child's
executeColumnar batches in a ColumnarBatchArrowReader, which closes each
batch once native has it. Those batches share the vectors that the
conversion reuses for the next batch, and closing a struct vector drops
its children, so the next batch failed to import. Native shuffle now reads
such a child's Arrow stream, as a native operator does, which also fixes
the same failure with spark.comet.sparkToColumnar.enabled.
* fix: leave calendar intervals out of the shuffle input conversion
Arrow holds the time part of an interval in nanoseconds, so
IntervalMonthDayNanoWriter overflows on a calendar interval with more
microseconds than that can hold, which Spark accepts. The JVM columnar
shuffle leaves calendar intervals to Spark's shuffle, and the conversion
now does too.
* test: turn off RDD scan conversion with its own config after the merge of main
Since #6602, CometTestBase turns on the conversion of RDD scans with
spark.comet.convert.rdd.enabled rather than spark.comet.sparkToColumnar.enabled,
so setting the old key to false no longer kept the RDD scans in
CometShuffleInputConversionSuite on Spark. The scans were converted before
the shuffle input conversion saw them, and six tests failed. The suite now
turns off the conversions CometTestBase turns on, and the struct test and
the benchmark use the RDD scan config.
* fix: keep a shuffle whose key is computed from a string on the JVM columnar shuffle
The string guard looked only at the type of each hash key, so a key that
is a number computed from a string, such as hash(s), passed it. Native
shuffle evaluates such a key after the import into native has replaced
invalid UTF-8 in the string, so it can still send a row to a different
partition than Spark's partitioner does, and a join whose other input
stays on the JVM columnar shuffle loses the matches. The guard now looks
for a string anywhere in the key expression. A wide decimal still matters
only as the type of the key itself, because Comet's hash of a wide
decimal already falls back to Spark.
The string join test also joins on the hash of the strings, and uses two
keys instead of four. Both shuffles write the strings with invalid UTF-8
replaced, so the keys stay apart only while each key expression sends
them to different partitions, and two of the four hashes share one.
* fix: put a converted shuffle back on the JVM columnar shuffle when its stage is reverted
RevertNativeForTransitionHeavyStages strips CometSparkToColumnarExec with the rest of a
reverted stage, which left a native shuffle reading the batches of Spark's RowToColumnarExec,
and the task failed casting them to CometVector. The shuffle now goes back to the JVM columnar
shuffle that the conversion replaced, which reads the stage's rows.
* refactor: convert only the input of a shuffle that would use the JVM columnar shuffle
convertsInputForNativeShuffle now starts from shuffleSupported, whose checks it restated, and
which also rules out Celeborn and calendar intervals. With spark.comet.shuffle.mode=native, a
Spark child keeps Spark's shuffle, as the config doc says.
* perf: sample a conversion's input rows for range partitioning
The range partitioner reads only the sort keys, so it samples the rows the conversion reads
rather than converting all of their columns to Arrow first. The benchmark gets a
range-partitioned case.
* test: check where native shuffle puts each key type the conversion admits
Checks the partition of every row and a join against the JVM columnar shuffle for each hash key
type the conversion admits, with NULL and boundary values. Adds array<string> and
map<string,string> columns with NULLs to the test over several batches, and plan assertions to
the test of aggregates whose partial aggregate runs on Spark.
* docs: describe the shuffle input conversion in the contributor guide
Also corrects the revert's log message in the tuning guide.
* test: count a range-partitioned conversion's rows once, and range-partition every column in the benchmark
Before the sampling change, the range partitioner's sampling job ran through the conversion and
added its rows to the conversion's metrics, so a range-partitioned shuffle reported twice the
rows it wrote.
* docs: say that the conversion gains with hash partitioning and comes out even with range partitioning
Also rewords the contributor guide's rule for when native shuffle takes a Spark child.
* test: correct why the shuffle between the two sort aggregates stays native
Since #4565, the revert of a shuffle between two Spark aggregates takes sort aggregates too, so
the reason the test gave, that the revert covers hash aggregates only, no longer held. The
shuffle stays native because the final sort aggregate reads it through a sort.
* docs: name sort aggregates in the tuning guide's shuffle revert section
Since #4565, the revert of a shuffle between two Spark aggregates takes SortAggregateExec too, but
only where the final aggregate reads the shuffle directly. One with grouping keys reads it through
a sort, so that shuffle is not reverted. The last sentence also no longer calls the shuffle that
disabling the revert keeps a columnar shuffle, since with the shuffle input conversion it can be a
native one.
* docs: say that the shuffle revert takes any Spark aggregate
Since #4565, the revert of a shuffle between two Spark aggregates matches any BaseAggregateExec,
not only hash aggregates, so the config doc and the scaladoc of revertRedundantColumnarShuffle
no longer name hash aggregates alone. The scaladoc also says why the shuffle below a sort
aggregate with grouping keys stays.
* test: check the Celeborn exclusion of the shuffle input conversion in every shuffle mode
In native mode the JVM columnar shuffle is off, so the conversion stays off
with or without the Celeborn checks. Run the test in auto and jvm too, where
Celeborn is what keeps the Spark child on Spark's shuffle, and check that the
fallback reason names Celeborn, which in native mode only the Celeborn check
in shuffleSupported gives.
* fix: keep native shuffle over a stacked bridge when reverting a transition-heavy stage
---
.github/workflows/pr_build_linux.yml | 1 +
.github/workflows/pr_build_macos.yml | 1 +
docs/source/contributor-guide/jvm_shuffle.md | 6 +-
docs/source/contributor-guide/native_shuffle.md | 21 +-
docs/source/user-guide/latest/tuning/shuffle.md | 25 +-
.../main/scala/org/apache/comet/CometConf.scala | 30 +-
.../org/apache/comet/rules/CometExecRule.scala | 82 +++-
.../RevertNativeForTransitionHeavyStages.scala | 49 ++-
.../shuffle/CometShuffleExchangeExec.scala | 105 +++--
.../exec/CometShuffleInputConversionSuite.scala | 428 +++++++++++++++++++++
.../CometShuffleInputConversionBenchmark.scala | 192 +++++++++
.../CometCelebornShufflePlanningSuite.scala | 20 +
12 files changed, 877 insertions(+), 83 deletions(-)
diff --git a/.github/workflows/pr_build_linux.yml b/.github/workflows/pr_build_linux.yml
index d6d487eb14..69f0fa8961 100644
--- a/.github/workflows/pr_build_linux.yml
+++ b/.github/workflows/pr_build_linux.yml
@@ -486,6 +486,7 @@ jobs:
org.apache.comet.exec.CometShuffle4_0Suite
org.apache.comet.exec.CometNativeColumnarToRowSuite
org.apache.comet.exec.CometNativeShuffleSuite
+ org.apache.comet.exec.CometShuffleInputConversionSuite
org.apache.comet.shuffle.CelebornShufflePartitionPusherSuite
org.apache.spark.sql.comet.execution.shuffle.CometCelebornShuffleManagerSuite
org.apache.spark.sql.comet.execution.shuffle.CometCelebornNativeShuffleWriterSuite
diff --git a/.github/workflows/pr_build_macos.yml b/.github/workflows/pr_build_macos.yml
index bf6a86b8f4..0b52541994 100644
--- a/.github/workflows/pr_build_macos.yml
+++ b/.github/workflows/pr_build_macos.yml
@@ -190,6 +190,7 @@ jobs:
org.apache.comet.exec.CometShuffle4_0Suite
org.apache.comet.exec.CometNativeColumnarToRowSuite
org.apache.comet.exec.CometNativeShuffleSuite
+ org.apache.comet.exec.CometShuffleInputConversionSuite
org.apache.comet.shuffle.CelebornShufflePartitionPusherSuite
org.apache.spark.sql.comet.execution.shuffle.CometCelebornShuffleManagerSuite
org.apache.spark.sql.comet.execution.shuffle.CometCelebornNativeShuffleWriterSuite
diff --git a/docs/source/contributor-guide/jvm_shuffle.md b/docs/source/contributor-guide/jvm_shuffle.md
index 11014f3490..dca7144a01 100644
--- a/docs/source/contributor-guide/jvm_shuffle.md
+++ b/docs/source/contributor-guide/jvm_shuffle.md
@@ -43,8 +43,10 @@ JVM shuffle (`CometColumnarExchange`) is used instead of native shuffle (`CometE
1. **Shuffle mode is explicitly set to "jvm"**: When `spark.comet.shuffle.mode` is set to `jvm`.
2. **Child plan is not a Comet native operator**: When the child plan is a Spark row-based operator
- (not a `CometPlan`), JVM shuffle is the only option since native shuffle requires columnar input
- from Comet operators.
+ (not a `CometPlan`), JVM shuffle is used, since native shuffle requires columnar input from
+ Comet operators. The exception is `spark.comet.convert.shuffleInput.enabled`, which converts
+ the child's rows to Arrow with `CometSparkToColumnarExec` so that native shuffle can take the
+ shuffle instead. See [When Native Shuffle is Used](native_shuffle.md#when-native-shuffle-is-used).
3. **Unsupported partition key types**: `RangePartitioning` keys must be primitive, so a complex
range key always falls back here. `HashPartitioning` keys must be primitive only by default:
diff --git a/docs/source/contributor-guide/native_shuffle.md b/docs/source/contributor-guide/native_shuffle.md
index e9cf777441..704fa2a0ea 100644
--- a/docs/source/contributor-guide/native_shuffle.md
+++ b/docs/source/contributor-guide/native_shuffle.md
@@ -52,7 +52,13 @@ Native shuffle (`CometExchange`) is selected when all of the following condition
1. **Shuffle mode allows native**: `spark.comet.shuffle.mode` is `native` or `auto`.
2. **Child plan is a Comet native operator**: The child must be a `CometPlan` that produces
- columnar output. Row-based Spark operators require JVM shuffle.
+ columnar output. Row-based Spark operators require JVM shuffle, unless
+ `spark.comet.convert.shuffleInput.enabled` is set. Then a shuffle that JVM shuffle would take
+ uses native shuffle instead, provided native shuffle supports the partitioning and the
+ columns, and a `CometSparkToColumnarExec` converts the child's rows to Arrow for it. A shuffle
+ that hashes a string, a value computed from a string, or a decimal wider than 18 digits stays
+ on JVM shuffle, because native shuffle would not put every row in the partition Spark's
+ partitioner does (`CometShuffleExchangeExec.convertsInputForNativeShuffle`).
3. **Supported partitioning type**: Native shuffle supports:
- `HashPartitioning`
@@ -166,11 +172,14 @@ The native shuffle implementation is its own workspace crate, `datafusion-comet-
- The child plan's `nativeOp` directly, when `CometShuffleExchangeExec`'s child is a
`CometNativeExec` subtree. The upstream operators run inside the same `CometExecIterator`
as the writer, with no JVM-to-native batch boundary between them.
- - A synthetic `Scan("ShuffleWriterInput")` placeholder, when the dep was built via the
- convenience `prepareShuffleDependency(rdd, ...)` overload (used by
- `CometCollectLimitExec` and `CometTakeOrderedAndProjectExec`, or when the
- exchange's child is a non-native `CometPlan` such as `CometSparkToColumnarExec`). Native
- code reads `ColumnarBatch`es from the JVM input iterator via Arrow C Stream Interface.
+ - A synthetic `Scan("ShuffleWriterInput")` placeholder, when the exchange's child is a
+ non-native `CometPlan`, and for the dependencies that `CometCollectLimitExec` and
+ `CometTakeOrderedAndProjectExec` build. Native code reads the input through the Arrow C
+ Stream Interface. A `CometNativeArrowSource` child, such as `CometSparkToColumnarExec`,
+ exports its own Arrow stream from `doExecuteAsArrowStream()`, which
+ `prepareArrowStreamShuffleDependency` hands to native. Any other input is an
+ `RDD[ColumnarBatch]`, which the convenience `prepareShuffleDependency(rdd, ...)` overload
+ wraps in a stream.
2. **Native execution**: A single `CometExecIterator` per partition runs the unified plan.
diff --git a/docs/source/user-guide/latest/tuning/shuffle.md b/docs/source/user-guide/latest/tuning/shuffle.md
index 71c9c431ab..b7ad3d2f09 100644
--- a/docs/source/user-guide/latest/tuning/shuffle.md
+++ b/docs/source/user-guide/latest/tuning/shuffle.md
@@ -66,23 +66,34 @@ write side. To restrict columnar shuffle to cases where the child is already a C
as native Spark shuffles, which avoids the row to columnar conversion but means the downstream stage will also start
on Spark.
+Setting `spark.comet.convert.shuffleInput.enabled=true` makes such a shuffle use native shuffle instead. Comet
+converts the child's rows to Arrow, as the `spark.comet.convert.*` settings do for Spark scans, and shuffles the
+batches natively. With hash partitioning that is usually faster than columnar shuffle, while with range partitioning
+the two come out about even. This applies only where native shuffle supports the partitioning and Comet can convert
+all of the child's columns, which rules out calendar intervals, and arrays and maps other than `array<string>` and
+`map<string,string>`. A shuffle that hashes a decimal wider than 18 digits, or a string or a value computed from a
+string, stays on columnar shuffle, so that it partitions rows as Spark does. The converted batches are allocated on
+the JVM outside Comet's memory pools. The setting is disabled by default.
+
### Automatic Revert to Spark Shuffle
-When a Comet columnar shuffle ends up between a partial and a final aggregate that Comet could not convert (both
-remain Spark `HashAggregateExec` or `ObjectHashAggregateExec` operators), Comet reverts it to Spark's built-in shuffle.
+When a Comet columnar shuffle, or a native shuffle over rows that `spark.comet.convert.shuffleInput.enabled` converted,
+ends up between a partial and a final aggregate that Comet could not convert (both remain Spark `HashAggregateExec`,
+`ObjectHashAggregateExec` or `SortAggregateExec` operators), Comet reverts it to Spark's built-in shuffle.
Keeping columnar shuffle between the two row-based aggregates would add `row -> Arrow -> shuffle -> Arrow -> row`
-conversions with no Comet consumer on either side to benefit from columnar output. Other shuffles between non-Comet
-operators are not reverted.
+conversions with no Comet consumer on either side to benefit from columnar output. The final aggregate has to read the
+shuffle directly, which a `SortAggregateExec` with grouping keys does not: it reads the shuffle through a sort, so that
+shuffle is not reverted. Other shuffles between non-Comet operators are not reverted either.
This shifts the affected shuffles from Comet's off-heap memory pool back to the JVM execution memory pool. Clusters
tuned for a small JVM heap may see `ExternalSorter` spills on queries where this revert fires. Shuffle I/O may also
grow marginally because Spark's row-based serializer generally compresses less well than Comet's Arrow IPC format.
-Each revert is logged at `INFO` level on the driver as `Reverting Comet columnar shuffle to Spark shuffle between
-<parent> and <child>`, which lets you correlate any unexpected behavior with this optimization.
+Each revert is logged at `INFO` level on the driver as `Reverting Comet shuffle to Spark shuffle between <parent>
+and <child>`, which lets you correlate any unexpected behavior with this optimization.
This optimization is enabled by default and can be disabled by setting
-`spark.comet.shuffle.revertRedundantColumnar.enabled=false`, in which case Comet will keep the columnar shuffle
+`spark.comet.shuffle.revertRedundantColumnar.enabled=false`, in which case Comet will keep its shuffle
even when both of those aggregates run on Spark.
## Shuffle Compression
diff --git a/spark/src/main/scala/org/apache/comet/CometConf.scala b/spark/src/main/scala/org/apache/comet/CometConf.scala
index 3d08c72dba..15889d5327 100644
--- a/spark/src/main/scala/org/apache/comet/CometConf.scala
+++ b/spark/src/main/scala/org/apache/comet/CometConf.scala
@@ -271,6 +271,20 @@ object CometConf extends ShimCometConf {
.booleanConf
.createWithDefault(false)
+ val COMET_CONVERT_FROM_SHUFFLE_INPUT_ENABLED: ConfigEntry[Boolean] =
+ conf("spark.comet.convert.shuffleInput.enabled")
+ .category(CATEGORY_EXEC)
+ .doc(
+ "When enabled, a shuffle that would use Comet's JVM columnar shuffle because its input " +
+ "comes from a Spark operator converts that input to Arrow format and uses Comet " +
+ "native shuffle instead. This applies only where native shuffle supports the " +
+ "partitioning and Comet can convert all of the input's columns, which rules out " +
+ "calendar intervals. A shuffle that hashes a decimal wider than 18 digits, or a " +
+ "string or a value computed from a string, stays on the JVM columnar shuffle, so " +
+ "that it partitions rows as Spark does.")
+ .booleanConf
+ .createWithDefault(false)
+
val COMET_EXEC_ENABLED: ConfigEntry[Boolean] = conf(s"$COMET_EXEC_CONFIG_PREFIX.enabled")
.category(CATEGORY_EXEC)
.doc(
@@ -726,14 +740,14 @@ object CometConf extends ShimCometConf {
conf("spark.comet.shuffle.revertRedundantColumnar.enabled")
.withAlternative(s"$COMET_EXEC_CONFIG_PREFIX.shuffle.revertRedundantColumnar.enabled")
.category(CATEGORY_SHUFFLE)
- .doc(
- "When enabled, Comet reverts a `CometShuffleExchangeExec` with `CometColumnarShuffle` " +
- "back to Spark's `ShuffleExchangeExec` when both its parent and child are non-Comet " +
- "hash aggregate operators. This avoids a redundant " +
- "row -> Arrow -> shuffle -> Arrow -> row conversion when no Comet operator on either " +
- "side can consume columnar output. Disable to keep Comet columnar shuffle even in " +
- "that case, which preserves Comet's off-heap shuffle memory accounting at the cost of " +
- "the extra conversion.")
+ .doc("When enabled, Comet reverts a `CometShuffleExchangeExec` with `CometColumnarShuffle` " +
+ "back to Spark's `ShuffleExchangeExec` when both its parent and child are non-Comet " +
+ "aggregate operators. The same applies to a native shuffle whose input " +
+ s"`${COMET_CONVERT_FROM_SHUFFLE_INPUT_ENABLED.key}` converted. This avoids a redundant " +
+ "row -> Arrow -> shuffle -> Arrow -> row conversion when no Comet operator on either " +
+ "side can consume columnar output. Disable to keep the Comet shuffle even in " +
+ "that case, which preserves Comet's off-heap shuffle memory accounting at the cost of " +
+ "the extra conversion.")
.booleanConf
.createWithDefault(true)
diff --git a/spark/src/main/scala/org/apache/comet/rules/CometExecRule.scala b/spark/src/main/scala/org/apache/comet/rules/CometExecRule.scala
index e2712b01e3..d339f27965 100644
--- a/spark/src/main/scala/org/apache/comet/rules/CometExecRule.scala
+++ b/spark/src/main/scala/org/apache/comet/rules/CometExecRule.scala
@@ -205,11 +205,14 @@ case class CometExecRule(session: SparkSession, queryStagePrep: Boolean = false)
/**
* Revert any `CometShuffleExchangeExec` with `CometColumnarShuffle` whose parent and child are
- * both non-Comet `HashAggregateExec` / `ObjectHashAggregateExec` operators back to the original
- * Spark `ShuffleExchangeExec`. This is the partial-final-aggregate pattern where Comet couldn't
+ * both non-Comet aggregate operators (any `BaseAggregateExec`) back to the original Spark
+ * `ShuffleExchangeExec`. This is the partial-final-aggregate pattern where Comet couldn't
* convert either aggregate; keeping a columnar shuffle between them only adds
* row->arrow->shuffle->arrow->row conversion overhead with no Comet consumer on either side.
- * See https://github.com/apache/datafusion-comet/issues/4004.
+ * See https://github.com/apache/datafusion-comet/issues/4004. A native shuffle over rows that
+ * `convertShuffleInput` converted is the same pattern, and is reverted the same way. A
+ * `SortAggregateExec` with grouping keys reads the shuffle through the `SortExec` that
+ * `EnsureRequirements` adds, so it is not the shuffle's parent and the shuffle stays.
*
* The match is intentionally narrow (both sides must be row-based aggregates that remained JVM
* after the main transform pass). Running the revert post-transform means we only fire when the
@@ -230,26 +233,35 @@ case class CometExecRule(session: SparkSession, queryStagePrep: Boolean = false)
private def revertRedundantColumnarShuffle(plan: SparkPlan): SparkPlan = {
def isAggregate(p: SparkPlan): Boolean = p.isInstanceOf[BaseAggregateExec]
- def isRedundantShuffle(child: SparkPlan): Boolean = child match {
- case s: CometShuffleExchangeExec =>
- s.shuffleType == CometColumnarShuffle && isAggregate(s.child)
- case _ => false
+ // The Spark operator whose rows a Comet shuffle writes: the child of a JVM columnar shuffle,
+ // or the input of the conversion that `convertShuffleInput` put under a native shuffle.
+ def rowInput(s: CometShuffleExchangeExec): Option[SparkPlan] =
+ (s.shuffleType, s.child) match {
+ case (CometColumnarShuffle, child) => Some(child)
+ case (CometNativeShuffle, conversion: CometSparkToColumnarExec) => Some(conversion.child)
+ case _ => None
+ }
+
+ def redundantShuffleInput(child: SparkPlan): Option[SparkPlan] = child match {
+ case s: CometShuffleExchangeExec => rowInput(s).filter(isAggregate)
+ case _ => None
}
plan.transform {
- case op if isAggregate(op) && op.children.exists(isRedundantShuffle) =>
- val newChildren = op.children.map {
- case s: CometShuffleExchangeExec
- if s.shuffleType == CometColumnarShuffle && isAggregate(s.child) =>
- val reverted =
- s.originalPlan.withNewChildren(Seq(s.child)).asInstanceOf[ShuffleExchangeExec]
- reverted.setTagValue(CometExecRule.SKIP_COMET_SHUFFLE_TAG, ())
- logInfo(
- "Reverting Comet columnar shuffle to Spark shuffle between " +
- s"${op.getClass.getSimpleName} and ${s.child.getClass.getSimpleName} " +
- "(no Comet operator on either side to consume columnar output)")
- reverted
- case other => other
+ case op if isAggregate(op) && op.children.exists(redundantShuffleInput(_).isDefined) =>
+ val newChildren = op.children.map { child =>
+ (child, redundantShuffleInput(child)) match {
+ case (s: CometShuffleExchangeExec, Some(input)) =>
+ val reverted =
+ s.originalPlan.withNewChildren(Seq(input)).asInstanceOf[ShuffleExchangeExec]
+ reverted.setTagValue(CometExecRule.SKIP_COMET_SHUFFLE_TAG, ())
+ logInfo(
+ "Reverting Comet shuffle to Spark shuffle between " +
+ s"${op.getClass.getSimpleName} and ${input.getClass.getSimpleName} " +
+ "(no Comet operator on either side to consume columnar output)")
+ reverted
+ case _ => child
+ }
}
op.withNewChildren(newChildren)
}
@@ -608,6 +620,9 @@ case class CometExecRule(session: SparkSession, queryStagePrep: Boolean = false)
case s: ShuffleExchangeExec if shouldSkipCometShuffle(s) =>
preserveSparkAggregateBuffers(s)
+ case s: ShuffleExchangeExec if CometShuffleExchangeExec.convertsInputForNativeShuffle(s) =>
+ convertShuffleInput(s)
+
case s: ShuffleExchangeExec =>
convertToComet(s, CometShuffleExchangeExec)
.getOrElse(preserveSparkAggregateBuffers(s))
@@ -635,7 +650,8 @@ case class CometExecRule(session: SparkSession, queryStagePrep: Boolean = false)
op
case _: ColumnarToRowTransition =>
// A transition does no work of its own. This rule only meets one that
- // `convertTypedDatasetOutput` inserted on an earlier pass over the same plan.
+ // `convertTypedDatasetOutput` or `convertShuffleInput` inserted on an earlier pass
+ // over the same plan.
op
case _: WriteFilesExec =>
// The write is converted at the enclosing DataWritingCommandExec above: on Spark 3.x
@@ -1397,6 +1413,30 @@ case class CometExecRule(session: SparkSession, queryStagePrep: Boolean = false)
}
}
+ /**
+ * 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)
+ .flatMap(converted =>
+ convertToComet(s.withNewChildren(Seq(converted)), CometShuffleExchangeExec))
+ .getOrElse(
+ convertToComet(s, CometShuffleExchangeExec).getOrElse(preserveSparkAggregateBuffers(s)))
+ }
+
private def shouldApplySparkToColumnar(conf: SQLConf, op: SparkPlan): Boolean = {
// Only consider converting leaf nodes to columnar currently, so that all the following
// operators can have a chance to be converted to columnar. Leaf operators that output
diff --git a/spark/src/main/scala/org/apache/comet/rules/RevertNativeForTransitionHeavyStages.scala b/spark/src/main/scala/org/apache/comet/rules/RevertNativeForTransitionHeavyStages.scala
index 3f3c57f064..9ad2d00384 100644
--- a/spark/src/main/scala/org/apache/comet/rules/RevertNativeForTransitionHeavyStages.scala
+++ b/spark/src/main/scala/org/apache/comet/rules/RevertNativeForTransitionHeavyStages.scala
@@ -24,7 +24,7 @@ import org.apache.spark.sql.SparkSession
import org.apache.spark.sql.catalyst.expressions.aggregate.{Final, Partial, PartialMerge}
import org.apache.spark.sql.catalyst.rules.Rule
import org.apache.spark.sql.comet.{CometBaseAggregateExec, CometColumnarToRowExec, CometExec, CometNativeColumnarToRowExec, CometSparkToColumnarExec}
-import org.apache.spark.sql.comet.execution.shuffle.{CometNativeShuffle, CometShuffleExchangeExec}
+import org.apache.spark.sql.comet.execution.shuffle.{CometColumnarShuffle, CometNativeShuffle, CometShuffleExchangeExec}
import org.apache.spark.sql.execution.{ColumnarToRowExec, ColumnarToRowTransition, RowToColumnarExec, SparkPlan}
import org.apache.spark.sql.execution.adaptive.QueryStageExec
import org.apache.spark.sql.execution.exchange.{BroadcastExchangeLike, ShuffleExchangeLike}
@@ -63,9 +63,7 @@ case class RevertNativeForTransitionHeavyStages(session: SparkSession, wholePlan
plan match {
case _: BroadcastExchangeLike => plan
case exchange: ShuffleExchangeLike =>
- revertShuffleStageIfNeeded(exchange)
- .map(reverted => exchange.withNewChildren(Seq(reverted)))
- .getOrElse(plan)
+ revertShuffleStageIfNeeded(exchange).getOrElse(plan)
case _ =>
// Result stage: its output is collected as rows.
revertStageIfNeeded(plan, outputColumnar = false).getOrElse(plan)
@@ -74,20 +72,47 @@ case class RevertNativeForTransitionHeavyStages(session: SparkSession, wholePlan
private def applyForNonAQE(plan: SparkPlan): SparkPlan = {
val withRevertedStages = plan.transformUp { case exchange: ShuffleExchangeLike =>
- revertShuffleStageIfNeeded(exchange)
- .map(reverted => exchange.withNewChildren(Seq(reverted)))
- .getOrElse(exchange)
+ revertShuffleStageIfNeeded(exchange).getOrElse(exchange)
}
revertStageIfNeeded(withRevertedStages, outputColumnar = false)
.getOrElse(withRevertedStages)
}
- private def revertShuffleStageIfNeeded(exchange: ShuffleExchangeLike): Option[SparkPlan] = {
- val outputArrow = exchange match {
- case comet: CometShuffleExchangeExec => comet.shuffleType == CometNativeShuffle
- case _ => false
+ /**
+ * Reverts the stage below a shuffle if needed and returns the shuffle over the reverted stage.
+ * A native shuffle over rows that `spark.comet.convert.shuffleInput.enabled` converted loses
+ * the conversion with the rest of the stage, so it goes back to the JVM columnar shuffle that
+ * the conversion replaced, which reads the stage's rows. Any other native shuffle consumes
+ * Arrow-backed Comet vectors, so the reverted stage is bridged back to them.
+ */
+ private def revertShuffleStageIfNeeded(exchange: ShuffleExchangeLike): Option[SparkPlan] =
+ exchange match {
+ case s: CometShuffleExchangeExec
+ if s.shuffleType == CometNativeShuffle && convertsSparkRows(s.child) =>
+ revertStageIfNeeded(s.child, s.supportsColumnar).map { reverted =>
+ val columnar = s.copy(child = reverted, shuffleType = CometColumnarShuffle)
+ columnar.copyTagsFrom(s)
+ columnar
+ }
+ case _ =>
+ val outputArrow = exchange match {
+ case comet: CometShuffleExchangeExec => comet.shuffleType == CometNativeShuffle
+ case _ => false
+ }
+ revertStageIfNeeded(exchange.child, exchange.supportsColumnar, outputArrow)
+ .map(reverted => exchange.withNewChildren(Seq(reverted)))
}
- revertStageIfNeeded(exchange.child, exchange.supportsColumnar, outputArrow)
+
+ /**
+ * Whether `plan` is the conversion that `spark.comet.convert.shuffleInput.enabled` puts over a
+ * Spark operator. A Comet transition directly under a `CometSparkToColumnarExec` is a stacked
+ * bridge that the revert strips and then restores for a native shuffle, not that conversion.
+ */
+ private def convertsSparkRows(plan: SparkPlan): Boolean = plan match {
+ case conversion: CometSparkToColumnarExec =>
+ !conversion.child.isInstanceOf[CometNativeColumnarToRowExec] &&
+ !conversion.child.isInstanceOf[CometColumnarToRowExec]
+ case _ => false
}
/**
diff --git a/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometShuffleExchangeExec.scala b/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometShuffleExchangeExec.scala
index 3216b6a430..c5e5c7641e 100644
--- a/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometShuffleExchangeExec.scala
+++ b/spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometShuffleExchangeExec.scala
@@ -21,6 +21,7 @@ package org.apache.spark.sql.comet.execution.shuffle
import java.util.function.Supplier
+import scala.collection.mutable.ListBuffer
import scala.concurrent.Future
import scala.jdk.CollectionConverters._
@@ -214,17 +215,18 @@ case class CometShuffleExchangeExec(
val nativeChild = child.asInstanceOf[CometNativeExec]
// RangePartitioner needs real rows for sampling. Reuse the precomputed context so we
// don't re-walk the SparkPlan tree or re-broadcast the encryption Hadoop conf.
- val samplingRDD: Option[RDD[ColumnarBatch]] = outputPartitioning match {
+ val samplingRows: Option[RDD[InternalRow]] = outputPartitioning match {
case _: RangePartitioning =>
Some(
- nativeChild.executeColumnarWithContext(
- ctx,
- nativeChildMetricNode.withoutAggregateMetrics(nativeChild)))
+ CometShuffleExchangeExec.rowsOf(
+ nativeChild.executeColumnarWithContext(
+ ctx,
+ nativeChildMetricNode.withoutAggregateMetrics(nativeChild))))
case _ => None
}
CometShuffleExchangeExec.prepareNativeShuffleDependency(
inputRDD.asInstanceOf[CometNativeShuffleInputRDD],
- samplingRDD,
+ samplingRows,
child.output,
outputPartitioning,
serializer,
@@ -236,12 +238,19 @@ case class CometShuffleExchangeExec(
positionalRoundRobin))
case None =>
child match {
- case _: CometNativeArrowSource =>
+ case source: 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())
+ // The range partitioner samples only the sort keys, so a conversion of rows
+ // gives it the rows it reads, rather than all of their columns in Arrow.
+ case _: RangePartitioning =>
+ Some(source match {
+ case conversion: CometSparkToColumnarExec
+ if !conversion.child.supportsColumnar =>
+ conversion.child.execute()
+ case _ => CometShuffleExchangeExec.rowsOf(child.executeColumnar())
+ })
case _ => None
},
child.output,
@@ -547,6 +556,48 @@ object CometShuffleExchangeExec
None
}
+ /**
+ * Whether a shuffle that would use the JVM columnar shuffle, because its child is a Spark
+ * row-based plan, can convert the child's rows to Arrow with `CometSparkToColumnarExec` and use
+ * native shuffle instead. See [[CometConf.COMET_CONVERT_FROM_SHUFFLE_INPUT_ENABLED]]. The
+ * native checks are those that [[shuffleSupported]] would make for the converted plan, so
+ * `CometExecRule` only converts a shuffle that will become native. Like [[shuffleSupported]],
+ * tags the node when no Comet shuffle can take it.
+ */
+ def convertsInputForNativeShuffle(s: ShuffleExchangeExec): Boolean =
+ CometConf.COMET_CONVERT_FROM_SHUFFLE_INPUT_ENABLED.get(s.conf) &&
+ !isCometPlan(s.child) &&
+ // This also rules out Celeborn, which has no JVM columnar shuffle, and calendar intervals.
+ // Arrow holds the time part of an interval in nanoseconds, so the conversion would
+ // overflow on one with more microseconds than that can hold.
+ shuffleSupported(s).contains(CometColumnarShuffle) &&
+ CometSparkToColumnarExec.isSchemaSupported(s.child.schema, ListBuffer.empty) &&
+ nativeShuffleFailureReasons(s).isEmpty &&
+ !hashesDifferentlyFromSpark(s)
+
+ /**
+ * Whether the shuffle hashes into more than one partition a key that native shuffle, reading
+ * converted rows, may put in a different partition from Spark's partitioner. Native shuffle
+ * must not take over such a shuffle from the JVM columnar shuffle: a join with an input that is
+ * still partitioned by Spark would put matching keys in different partitions. The keys are:
+ *
+ * - A decimal wider than 18 digits, which native shuffle hashes differently (#5994).
+ * - A string, or any value computed from one, such as `hash(s)`. Spark reads the string's
+ * bytes as they are, but the import of the converted batch into native replaces invalid
+ * UTF-8 before native shuffle evaluates the key.
+ *
+ * TODO: allow wide decimals once native hashing matches Spark for them.
+ */
+ private def hashesDifferentlyFromSpark(s: ShuffleExchangeExec): Boolean =
+ s.outputPartitioning match {
+ case HashPartitioning(expressions, numPartitions) =>
+ numPartitions > 1 && expressions.exists { key =>
+ key.dataType.existsRecursively(DecimalType.isByteArrayDecimalType) ||
+ key.exists(_.dataType.existsRecursively(_.isInstanceOf[StringType]))
+ }
+ case _ => false
+ }
+
/**
* Whether the stage feeding a shuffle starts at a typed Dataset conversion (see
* [[CometConf.COMET_CONVERT_FROM_TYPED_DATASET_ENABLED]]). `CometExecRule` decides the shuffle
@@ -862,6 +913,10 @@ object CometShuffleExchangeExec
}
}
+ /** The rows of `batches`, for the range partitioner to sample. */
+ private def rowsOf(batches: RDD[ColumnarBatch]): RDD[InternalRow] =
+ batches.mapPartitionsInternal(_.flatMap(_.rowIterator().asScala))
+
/**
* Build a Comet native shuffle dependency around an existing `RDD[ColumnarBatch]` of real
* batches. Used by [[org.apache.spark.sql.comet.CometCollectLimitExec]] and
@@ -891,7 +946,7 @@ object CometShuffleExchangeExec
"ShuffleWriterInput")
prepareArrowStreamShuffleDependency(
streamRDD,
- Some(rdd),
+ Some(rowsOf(rdd)),
outputAttributes,
outputPartitioning,
serializer,
@@ -900,12 +955,12 @@ object CometShuffleExchangeExec
/**
* [[prepareShuffleDependency]] for input that native reads as one Arrow stream per partition,
- * such as a [[CometNativeArrowSource]]'s. `samplingRDD` gives the same rows as batches, and is
- * only required for [[RangePartitioning]].
+ * such as a [[CometNativeArrowSource]]'s. `samplingRows` gives the same rows, and is only
+ * required for [[RangePartitioning]].
*/
private def prepareArrowStreamShuffleDependency(
streamRDD: RDD[ArrowArrayStream],
- samplingRDD: Option[RDD[ColumnarBatch]],
+ samplingRows: Option[RDD[InternalRow]],
outputAttributes: Seq[Attribute],
outputPartitioning: Partitioning,
serializer: Serializer,
@@ -946,7 +1001,7 @@ object CometShuffleExchangeExec
// is `shuffleWriterMetrics` at the root with one empty leaf for the Scan child.
prepareNativeShuffleDependency(
thinRDD,
- samplingRDD,
+ samplingRows,
outputAttributes,
outputPartitioning,
serializer,
@@ -965,13 +1020,14 @@ object CometShuffleExchangeExec
* @param thinRDD
* scheduling-anchor RDD whose `compute` returns a [[CometNativeShuffleInputIterator]];
* produces no batches itself.
- * @param samplingRDD
- * regular columnar execution of the child, only required for [[RangePartitioning]] (sampling
- * needs real rows). `None` for hash / single / round-robin.
+ * @param samplingRows
+ * the child's rows, such as those of a regular columnar execution of the child, only required
+ * for [[RangePartitioning]] (sampling needs real rows). `None` for hash / single /
+ * round-robin.
*/
def prepareNativeShuffleDependency(
thinRDD: CometNativeShuffleInputRDD,
- samplingRDD: Option[RDD[ColumnarBatch]],
+ samplingRows: Option[RDD[InternalRow]],
outputAttributes: Seq[Attribute],
outputPartitioning: Partitioning,
serializer: Serializer,
@@ -1011,11 +1067,11 @@ object CometShuffleExchangeExec
// ShuffleExchangeExec::prepareShuffleDependency
val (partitioner, rangePartitionBounds) = outputPartitioning match {
case rangePartitioning: RangePartitioning =>
- // Sampling needs real rows; use the dedicated samplingRDD (a regular columnar execution
- // of the child). The thin RDD itself yields nothing.
- val samplingInput = samplingRDD.getOrElse(
+ // Sampling needs real rows; use the dedicated samplingRows (the child's rows). The thin
+ // RDD itself yields nothing.
+ val samplingInput = samplingRows.getOrElse(
throw new IllegalStateException(
- "RangePartitioning requires a samplingRDD on the native-shuffle path"))
+ "RangePartitioning requires samplingRows on the native-shuffle path"))
// Extract only fields used for sorting to avoid collecting large fields that does not
// affect sorting result when deciding partition bounds in RangePartitioner
val rddForSampling = samplingInput.mapPartitionsInternal { iter =>
@@ -1025,12 +1081,7 @@ object CometShuffleExchangeExec
// Internally, RangePartitioner runs a job on the RDD that samples keys to compute
// partition bounds. To get accurate samples, we need to copy the mutable keys.
- iter.flatMap { batch =>
- val rowIter = batch.rowIterator().asScala
- rowIter.map { row =>
- mutablePair.update(projection(row).copy(), null)
- }
- }
+ iter.map(row => mutablePair.update(projection(row).copy(), null))
}
// Construct ordering on extracted sort key.
diff --git a/spark/src/test/scala/org/apache/comet/exec/CometShuffleInputConversionSuite.scala b/spark/src/test/scala/org/apache/comet/exec/CometShuffleInputConversionSuite.scala
new file mode 100644
index 0000000000..0318113600
--- /dev/null
+++ b/spark/src/test/scala/org/apache/comet/exec/CometShuffleInputConversionSuite.scala
@@ -0,0 +1,428 @@
+/*
+ * 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 java.sql.{Date, Timestamp}
+import java.time.LocalDateTime
+
+import org.apache.spark.sql.{CometTestBase, DataFrame, Row}
+import org.apache.spark.sql.comet.{CometHashAggregateExec, CometPlan, 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.execution.aggregate.{HashAggregateExec, SortAggregateExec}
+import org.apache.spark.sql.functions.{array, avg, col, count, hash, length, max, min, size, spark_partition_id, sum}
+import org.apache.spark.sql.internal.SQLConf
+import org.apache.spark.sql.types.{ArrayType, BinaryType, BooleanType, ByteType, CalendarIntervalType, DataType, DataTypes, DateType, DecimalType, DoubleType, FloatType, IntegerType, LongType, MapType, ShortType, StringType, StructType, TimestampNTZType, TimestampType}
+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)
+ }
+
+ /** Runs `f` with each join planned as a sort-merge join of two shuffles into 10 partitions. */
+ private def withShuffledJoins(f: => Unit): Unit =
+ withSQLConf(
+ SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false",
+ SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1",
+ SQLConf.SHUFFLE_PARTITIONS.key -> "10")(f)
+
+ /**
+ * Values of each type the conversion admits as a hash key, with boundary values. The floating
+ * point values include NaN and both zeros, which a join normalizes.
+ */
+ private val hashKeys: Seq[(DataType, Seq[Any])] = Seq(
+ BooleanType -> Seq(true, false),
+ ByteType -> Seq(0.toByte, -1.toByte, Byte.MinValue, Byte.MaxValue),
+ ShortType -> Seq(0.toShort, -1.toShort, Short.MinValue, Short.MaxValue),
+ IntegerType -> Seq(0, -1, Int.MinValue, Int.MaxValue),
+ LongType -> Seq(0L, -1L, Long.MinValue, Long.MaxValue),
+ FloatType -> Seq(
+ 0.0f,
+ -0.0f,
+ 1.5f,
+ Float.NaN,
+ Float.MinPositiveValue,
+ Float.MaxValue,
+ Float.NegativeInfinity),
+ DoubleType -> Seq(
+ 0.0d,
+ -0.0d,
+ 1.5d,
+ Double.NaN,
+ Double.MinPositiveValue,
+ Double.MaxValue,
+ Double.NegativeInfinity),
+ DecimalType(18, 2) -> Seq("0.00", "-0.01", "9999999999999999.99", "-9999999999999999.99")
+ .map(new java.math.BigDecimal(_)),
+ DateType -> Seq("1970-01-01", "1969-12-31", "0001-01-01", "9999-12-31").map(Date.valueOf),
+ TimestampType -> Seq(
+ new Timestamp(0L),
+ new Timestamp(-1L),
+ Timestamp.valueOf("0001-01-01 00:00:00"),
+ Timestamp.valueOf("9999-12-31 23:59:59.999999")),
+ TimestampNTZType -> Seq(
+ LocalDateTime.of(1970, 1, 1, 0, 0),
+ LocalDateTime.of(1969, 12, 31, 23, 59, 59, 999999000),
+ LocalDateTime.of(1, 1, 1, 0, 0),
+ LocalDateTime.of(9999, 12, 31, 23, 59, 59, 999999000)),
+ BinaryType -> Seq(Array.emptyByteArray, Array[Byte](0), Array[Byte](-1, -128, 127)))
+
+ /** Four rows of each value and of a NULL, as column `k`, over four input partitions. */
+ private def keysDf(keyType: DataType, values: Seq[Any]): DataFrame = {
+ val keys = values :+ null
+ val data = (0 until keys.length * 4).map(i => Row(keys(i % keys.length), i.toLong))
+ val schema = new StructType().add("k", keyType).add("v", LongType)
+ spark.createDataFrame(spark.sparkContext.parallelize(data, 4), schema)
+ }
+
+ 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") {
+ // The partial aggregate reads Spark rows, so it runs on Spark, and native shuffle carries its
+ // buffers to the final aggregate. With these aggregates the final one runs natively.
+ val (_, split) = checkSparkAnswer(
+ rowsDf()
+ .groupBy((col("k") % 7).as("g"))
+ .agg(sum("l"), avg("d"), max("l"), min("l")))
+ assert(convertedShuffles(split).length == 1, split)
+ assert(collect(split) { case a: HashAggregateExec => a }.length == 1, split)
+ assert(collect(split) { case a: CometHashAggregateExec => a }.length == 1, split)
+ // The min of a string makes Spark plan sort aggregates, which stay on Spark, so the final
+ // aggregate reads the buffers back as rows. The shuffle stays native: the final aggregate
+ // reads it through a sort, and the revert in the next test only takes a shuffle that an
+ // aggregate reads directly.
+ val (_, sorted) = checkSparkAnswer(
+ rowsDf()
+ .groupBy((col("k") % 7).as("g"))
+ .agg(sum("l"), count("s"), avg("d"), sum("m"), max("l"), min("s")))
+ assert(convertedShuffles(sorted).length == 1, sorted)
+ assert(collect(sorted) { case a: SortAggregateExec => a }.length == 2, sorted)
+ }
+
+ 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 stage that the transition revert puts back on Spark keeps the JVM shuffle") {
+ // The typed operation reads the Comet scan through a transition. With no transitions allowed,
+ // the revert puts the stage below the shuffle back on Spark, the conversion included, and the
+ // shuffle goes back to the JVM columnar shuffle, which reads the stage's rows.
+ withParquetTable((0 until 200).map(i => (i, (i % 13).toString)), "tbl") {
+ Seq("true", "false").foreach { aqe =>
+ withSQLConf(
+ SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> aqe,
+ CometConf.COMET_EXEC_TRANSITION_REVERT_ENABLED.key -> "true",
+ CometConf.COMET_EXEC_TRANSITION_REVERT_MAX_TRANSITIONS.key -> "0") {
+ val ds = spark.sql("SELECT _1 AS a, _2 AS b FROM tbl").as[ShuffleInputRec]
+ val (_, plan) = checkSparkAnswer(
+ ds.map(r => ShuffleInputRec(r.a % 13, r.b))
+ .repartition(7, col("a"))
+ .groupBy("a")
+ .agg(max("b")))
+ assert(conversions(plan).isEmpty, s"AQE $aqe:\n$plan")
+ val reverted = cometShuffles(plan).filter(_.shuffleType == CometColumnarShuffle)
+ assert(reverted.length == 1, s"AQE $aqe:\n$plan")
+ assert(!reverted.head.child.exists(_.isInstanceOf[CometPlan]), s"AQE $aqe:\n$plan")
+ }
+ }
+ }
+ }
+
+ hashKeys.foreach { case (keyType, values) =>
+ convertTest(s"native shuffle puts ${keyType.simpleString} keys where Spark does") {
+ withShuffledJoins {
+ val keys = keysDf(keyType, values)
+ // The partition of each row, against Spark's partitioner, NULL keys included.
+ val (_, placed) = checkSparkAnswer(
+ keys.repartition(10, col("k")).select(col("k"), col("v"), spark_partition_id()))
+ assert(convertedShuffles(placed).length == 1, placed)
+ // The left input converts and the right one, with its array<bigint> column, does not, so
+ // matching keys meet only if both shuffles put them in the same partition.
+ val left = keys.select(col("k").as("lk"), col("v"))
+ val right = keys.select(col("k").as("rk"), array(col("v")).as("xs"))
+ val (_, joined) = checkSparkAnswer(left.join(right, col("lk") === col("rk")))
+ assert(convertedShuffles(joined).length == 1, joined)
+ assert(cometShuffles(joined).count(_.shuffleType == CometColumnarShuffle) == 1, joined)
+ }
+ }
+ }
+
+ convertTest("a shuffle that hashes a wide decimal stays on the JVM columnar shuffle") {
+ // Native shuffle hashes a decimal wider than 18 digits differently from Spark's partitioner
+ // (#5994). The right input, with its array<int> column, stays on the JVM columnar shuffle,
+ // so the left one has to as well, or matching keys land in different partitions.
+ withShuffledJoins {
+ val left = rowsDf(100).select(col("k").cast(DecimalType(38, 0)).as("lk"), col("l"))
+ val right = arraysDf().select(col("k").cast(DecimalType(38, 0)).as("rk"), col("xs"))
+ val df = left.join(right, col("lk") === col("rk"))
+ val (_, plan) = checkSparkAnswer(df)
+ assert(conversions(plan).isEmpty, plan)
+ checkCometExchange(df, 2, native = false)
+ }
+ }
+
+ convertTest("a shuffle that hashes a string stays on the JVM columnar shuffle") {
+ // Spark's partitioner hashes a string's bytes as they are, but native shuffle hashes them
+ // after the import into native has replaced invalid UTF-8. The right input, with its
+ // array<int> column, stays on the JVM columnar shuffle, so the left one has to as well, or
+ // matching keys with invalid UTF-8 land in different partitions. A key computed from a
+ // string, such as its hash, is no different: native shuffle would compute it from the
+ // replaced string. Both shuffles write the rows with invalid UTF-8 replaced, and the join
+ // compares what they wrote, so the two keys stay apart only because each key expression puts
+ // them in different partitions.
+ withShuffledJoins {
+ val schema = new StructType().add("b", BinaryType).add("v", LongType)
+ val data = (0 until 40).map { i =>
+ Row(Array((if (i % 2 == 0) 0x80 else 0x83).toByte), i.toLong)
+ }
+ val keys = spark.createDataFrame(spark.sparkContext.parallelize(data, 4), schema)
+ val left = keys.select(col("b").cast(StringType).as("lk"), col("v"))
+ val right = keys.select(
+ col("b").cast(StringType).as("rk"),
+ array(col("v").cast(IntegerType)).as("xs"))
+ Seq(col("lk") === col("rk"), hash(col("lk")) === hash(col("rk"))).foreach { condition =>
+ val df = left.join(right, condition)
+ val (_, plan) = checkSparkAnswer(df)
+ assert(conversions(plan).isEmpty, s"$condition:\n$plan")
+ checkCometExchange(df, 2, native = false)
+ }
+ }
+ }
+
+ convertTest("nested columns over several batches") {
+ // The rows of each input partition fill several batches, and the conversion reuses its
+ // vectors from one batch to the next. With the conversion of RDD scans on, native shuffle
+ // reads the converted scan the same way.
+ val schema = new StructType()
+ .add("k", IntegerType)
+ .add("payload", new StructType().add("v", LongType).add("s", StringType))
+ .add("xs", ArrayType(StringType))
+ .add("m", MapType(StringType, StringType))
+ val data = (0 until 24).map { i =>
+ Row(
+ i,
+ if (i % 5 == 4) null else Row(i.toLong, if (i % 3 == 0) null else s"s$i"),
+ if (i % 6 == 5) null else Seq.tabulate(i % 4)(j => if (j == 1) null else s"x$i-$j"),
+ if (i % 7 == 6) null
+ else Seq.tabulate(i % 3)(j => s"k$j" -> (if (j == 1) null else s"v$i-$j")).toMap)
+ }
+ Seq("false", "true").foreach { leafConversion =>
+ withSQLConf(
+ CometConf.COMET_BATCH_SIZE.key -> "7",
+ CometConf.COMET_CONVERT_FROM_RDD_ENABLED.key -> leafConversion) {
+ val df = spark.createDataFrame(spark.sparkContext.parallelize(data, 1), schema)
+ Seq(df.repartition(3, col("k")), df.repartitionByRange(3, col("k"))).foreach { shuffled =>
+ val (_, plan) = checkSparkAnswer(shuffled)
+ assert(convertedShuffles(plan).length == 1, s"leaf conversion $leafConversion:\n$plan")
+ }
+ }
+ }
+ }
+
+ convertTest("calendar intervals keep Spark's shuffle") {
+ // Arrow holds an interval's time part in nanoseconds, which can't hold every number of
+ // microseconds Spark can. The JVM columnar shuffle doesn't take calendar intervals either.
+ val schema = new StructType().add("k", IntegerType).add("i", CalendarIntervalType)
+ val data = Seq(Row(1, new CalendarInterval(0, 0, 10800000000000000L)))
+ val df = spark
+ .createDataFrame(spark.sparkContext.parallelize(data, 1), schema)
+ .repartition(2, col("k"))
+ val (_, plan) = checkSparkAnswer(df)
+ assert(cometShuffles(plan).isEmpty, plan)
+ }
+
+ convertTest("columns the conversion does not support keep the JVM columnar shuffle") {
+ val df = arraysDf().repartition(5, col("k")).groupBy("k").agg(sum(size(col("xs"))))
+ val (_, plan) = checkSparkAnswer(df)
+ assert(conversions(plan).isEmpty, plan)
+ assert(cometShuffles(plan).forall(_.shuffleType == CometColumnarShuffle), plan)
+ }
+
+ convertTest("range partitioning") {
+ val (_, plan) = checkSparkAnswer(rowsDf().orderBy(col("l").desc))
+ assert(convertedShuffles(plan).length == 1, plan)
+ // The range partitioner samples the rows that the conversion reads, without converting them,
+ // so the conversion only counts the rows that the shuffle writes.
+ assert(conversions(plan).head.metrics("numInputRows").value == 1000, plan)
+ }
+
+ convertTest("the conversion only takes over from the JVM columnar shuffle") {
+ // In `jvm` mode the shuffle keeps the JVM columnar shuffle. In `native` mode, where Comet has
+ // no shuffle for a Spark operator's rows, it keeps Spark's shuffle.
+ Seq("jvm" -> Seq(CometColumnarShuffle), "native" -> Seq.empty).foreach {
+ case (mode, shuffleTypes) =>
+ withSQLConf(CometConf.COMET_SHUFFLE_MODE.key -> mode) {
+ val (_, plan) =
+ checkSparkAnswer(rowsDf().repartition(10, col("k")).groupBy("k").count())
+ assert(conversions(plan).isEmpty, s"$mode:\n$plan")
+ assert(
+ cometShuffles(plan).map(_.shuffleType).distinct == shuffleTypes,
+ s"$mode:\n$plan")
+ }
+ }
+ }
+
+ test("off by default") {
+ withoutLeafConversion {
+ val (_, plan) = checkSparkAnswer(rowsDf().repartition(10, col("k")).groupBy("k").count())
+ assert(conversions(plan).isEmpty, plan)
+ assert(cometShuffles(plan).exists(_.shuffleType == CometColumnarShuffle), plan)
+ }
+ }
+}
diff --git a/spark/src/test/scala/org/apache/spark/sql/benchmark/CometShuffleInputConversionBenchmark.scala b/spark/src/test/scala/org/apache/spark/sql/benchmark/CometShuffleInputConversionBenchmark.scala
new file mode 100644
index 0000000000..3942f1859e
--- /dev/null
+++ b/spark/src/test/scala/org/apache/spark/sql/benchmark/CometShuffleInputConversionBenchmark.scala
@@ -0,0 +1,192 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.spark.sql.benchmark
+
+import org.apache.spark.benchmark.Benchmark
+import org.apache.spark.sql.{DataFrame, Row}
+import org.apache.spark.sql.comet.{CometPlan, CometSparkToColumnarExec}
+import org.apache.spark.sql.comet.execution.shuffle.{CometColumnarShuffle, CometNativeShuffle, CometShuffleExchangeExec}
+import org.apache.spark.sql.execution.SparkPlan
+import org.apache.spark.sql.functions.{col, count, length, lit, sum}
+import org.apache.spark.sql.internal.SQLConf
+import org.apache.spark.sql.types.{DataTypes, DecimalType, IntegerType, LongType, StructType}
+
+import org.apache.comet.CometConf
+
+// Top-level, so the encoder needs no outer pointer.
+case class ShuffleInputBenchRec(a: Long, b: String)
+
+/**
+ * Compares three ways to run a shuffle whose input comes from a Spark operator:
+ *
+ * - Spark: Comet disabled.
+ * - Comet: the default. Comet's JVM columnar shuffle buffers the rows and converts them to
+ * Arrow natively, one partition at a time.
+ * - Comet, converted: `spark.comet.convert.shuffleInput.enabled`, which converts the rows with
+ * `CometSparkToColumnarExec` and uses native shuffle.
+ *
+ * Each case partitions its rows into 200 partitions and aggregates them. Most hash-partition on a
+ * number, because the conversion leaves a shuffle that hashes a string on the JVM columnar
+ * shuffle, and one range-partitions. Every arm's result and plan are checked before it is timed,
+ * and the Comet arm runs again at the end of each case to show the noise. To run this benchmark:
+ * {{{
+ * SPARK_GENERATE_BENCHMARK_FILES=1 make benchmark-org.apache.spark.sql.benchmark.CometShuffleInputConversionBenchmark
+ * }}}
+ * Results will be written to
+ * "spark/benchmarks/CometShuffleInputConversionBenchmark-**results.txt".
+ */
+object CometShuffleInputConversionBenchmark extends CometBenchmarkBase {
+
+ private val numRows = 4L * 1024 * 1024
+ private val numPartitions = 200
+
+ import spark.implicits._
+
+ private case class Arm(name: String, confs: Seq[(String, String)], check: SparkPlan => Unit)
+
+ // No AQE, so every arm plans the same shape on every iteration.
+ private val planConfs = Seq(SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false")
+
+ private def cometConfs(convert: Boolean) = planConfs ++ Seq(
+ CometConf.COMET_ENABLED.key -> "true",
+ CometConf.COMET_EXEC_ENABLED.key -> "true",
+ // Spark scans the RDDs, as it does by default, so the shuffle reads a Spark operator.
+ CometConf.COMET_CONVERT_FROM_RDD_ENABLED.key -> "false",
+ CometConf.COMET_CONVERT_FROM_SHUFFLE_INPUT_ENABLED.key -> convert.toString)
+
+ private def shuffles(plan: SparkPlan): Seq[CometShuffleExchangeExec] =
+ collect(plan) { case s: CometShuffleExchangeExec => s }
+
+ private val arms = Seq(
+ Arm(
+ "Spark",
+ planConfs :+ (CometConf.COMET_ENABLED.key -> "false"),
+ plan => assert(collect(plan) { case c: CometPlan => c }.isEmpty, plan)),
+ Arm(
+ "Comet",
+ cometConfs(convert = false),
+ plan => assert(shuffles(plan).exists(_.shuffleType == CometColumnarShuffle), plan)),
+ Arm(
+ "Comet, converted",
+ cometConfs(convert = true),
+ plan =>
+ assert(
+ shuffles(plan).exists(s =>
+ s.shuffleType == CometNativeShuffle && s.child
+ .isInstanceOf[CometSparkToColumnarExec]),
+ plan)))
+
+ private val rowSchema = new StructType()
+ .add("k", IntegerType)
+ .add("l", LongType)
+ .add("s", DataTypes.StringType)
+ .add("d", DataTypes.DoubleType)
+ .add("m", DecimalType(18, 2))
+
+ /** `numRows` rows from an RDD, generated as the query runs. */
+ private def rddRows(): DataFrame = {
+ val rows = spark.sparkContext.range(0, numRows, 1, 1).map { id =>
+ val i = id.toInt
+ Row(
+ i % 1000,
+ id,
+ s"value-${i * 31 % 100003}",
+ i + 0.5d,
+ java.math.BigDecimal.valueOf(i * 7919L % 1000000000L, 2))
+ }
+ spark.createDataFrame(rows, rowSchema)
+ }
+
+ /** Totals the per-group results in `s`, so each case returns one row. */
+ private def total(grouped: DataFrame): DataFrame = grouped.agg(sum("s"), count(lit(1)))
+
+ private val cases: Seq[(String, () => DataFrame)] = Seq(
+ "RDD rows: int, long, double" -> (() =>
+ total(
+ rddRows()
+ .select("k", "l", "d")
+ .repartition(numPartitions, col("k"))
+ .groupBy("k")
+ .agg((sum("l") + sum("d")).as("s")))),
+ "RDD rows: int, long, string, double, decimal(18,2)" -> (() =>
+ total(
+ rddRows()
+ .repartition(numPartitions, col("k"))
+ .groupBy("k")
+ .agg((sum("l") + sum(length(col("s"))) + sum("d") + sum("m")).as("s")))),
+ "RDD rows: int, long, string, double, decimal(18,2), range partitioned" -> (() =>
+ total(
+ rddRows()
+ .repartitionByRange(numPartitions, col("l"))
+ .select((col("k") + col("l") + length(col("s")) + col("d").cast("long") + col("m"))
+ .as("s")))),
+ "map over a Parquet scan" -> (() =>
+ total(
+ spark
+ .table("parquetV1Table")
+ .select(col("id").as("a"), col("key").as("b"))
+ .as[ShuffleInputBenchRec]
+ .map(r => ShuffleInputBenchRec(r.a % 1000, r.b))
+ .repartition(numPartitions, col("a"))
+ .groupBy("a")
+ .agg(sum(length(col("b"))).as("s")))))
+
+ private def runArm(arm: Arm, query: () => DataFrame): (Seq[Row], SparkPlan) = {
+ var result: (Seq[Row], SparkPlan) = null
+ withSQLConf(arm.confs: _*) {
+ val df = query()
+ result = (df.collect().toSeq, df.queryExecution.executedPlan)
+ }
+ result
+ }
+
+ override def runCometBenchmark(mainArgs: Array[String]): Unit = {
+ withTempPath { dir =>
+ withTempTable("parquetV1Table") {
+ prepareTable(
+ dir,
+ spark.range(numRows).selectExpr("id", "CAST(id % 1000 AS STRING) AS key"))
+
+ cases.foreach { case (name, query) =>
+ runBenchmark(name) {
+ // Check every arm's plan and result before timing anything.
+ val results = arms.map(runArm(_, query))
+ arms.zip(results).foreach { case (arm, (rows, plan)) =>
+ arm.check(plan)
+ assert(
+ rows == results.head._1,
+ s"${arm.name} returned $rows, expected ${results.head._1}")
+ }
+
+ val benchmark = new Benchmark(name, numRows, output = output)
+ (arms :+ arms(1).copy(name = "Comet (repeat)")).foreach { arm =>
+ benchmark.addCase(arm.name) { _ =>
+ withSQLConf(arm.confs: _*) {
+ query().collect()
+ }
+ }
+ }
+ benchmark.run()
+ }
+ }
+ }
+ }
+ }
+}
diff --git a/spark/src/test/scala/org/apache/spark/sql/comet/execution/shuffle/CometCelebornShufflePlanningSuite.scala b/spark/src/test/scala/org/apache/spark/sql/comet/execution/shuffle/CometCelebornShufflePlanningSuite.scala
index a63560f337..1d217bf6a2 100644
--- a/spark/src/test/scala/org/apache/spark/sql/comet/execution/shuffle/CometCelebornShufflePlanningSuite.scala
+++ b/spark/src/test/scala/org/apache/spark/sql/comet/execution/shuffle/CometCelebornShufflePlanningSuite.scala
@@ -265,6 +265,26 @@ 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. Without Celeborn, auto and jvm
+ // would convert this child. Native mode has no JVM columnar shuffle either, so the reason
+ // is what shows that Celeborn kept Spark's shuffle there.
+ Seq("auto", "jvm", "native").foreach { mode =>
+ withSQLConf(
+ CometConf.COMET_SHUFFLE_MODE.key -> mode,
+ CometConf.COMET_CONVERT_FROM_SHUFFLE_INPUT_ENABLED.key -> "true") {
+ val leaf = sparkLeaf()
+ val exchange = ShuffleExchangeExec(HashPartitioning(leaf.output, 2), leaf)
+ assert(!CometShuffleExchangeExec.convertsInputForNativeShuffle(exchange), mode)
+ assert(reasons(exchange).exists(_.contains("Celeborn")), s"$mode: ${reasons(exchange)}")
+ val planned = CometExecRule(spark)(exchange)
+ assert(planned.isInstanceOf[ShuffleExchangeExec], planned)
+ assert(collect(planned) { case c: CometSparkToColumnarExec => c }.isEmpty, planned)
+ }
+ }
+ }
+
test("fallback is sticky after AQE reshapes the child or native mode becomes available") {
val child = nativeChild()
val exchange = ShuffleExchangeExec(SinglePartition, child)
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@datafusion.apache.org
For additional commands, e-mail: commits-help@datafusion.apache.org