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: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
