andygrove commented on code in PR #6607:
URL: https://github.com/apache/datafusion-comet/pull/6607#discussion_r4186480163
##########
spark/src/main/scala/org/apache/comet/rules/CometExecRule.scala:
##########
@@ -1274,6 +1291,30 @@ case class CometExecRule(session: SparkSession)
private def hasEnabledHandler(op: SparkPlan): Boolean =
allExecs.get(op.getClass).exists(_.enabledConfig.forall(_.get(op.conf)))
+ /**
+ * Converts the rows a Spark operator feeds to a shuffle to Arrow, so that
the shuffle runs as
+ * native shuffle instead of the JVM columnar shuffle. See
+ * [[CometConf.COMET_CONVERT_FROM_SHUFFLE_INPUT_ENABLED]], and
+ * [[CometShuffleExchangeExec.convertsInputForNativeShuffle]] for when it
applies.
+ *
+ * Spark inserts the columnar transitions after this rule, but it does not
look below a
+ * `RowToColumnarTransition` such as `CometSparkToColumnarExec`. Without a
transition, a Spark
+ * operator in the child's subtree that reads a Comet operator would do so
through
+ * `CometExec.doExecute`, Spark's interpreted columnar-to-row path. So the
subtree gets its
+ * transitions now, from Spark's own rule, and
`EliminateRedundantTransitions` later replaces
+ * each one over a Comet child with Comet's own. Spark's rule leaves
existing transitions alone,
+ * which matters because this rule runs over the same plan twice under AQE.
+ */
+ private def convertShuffleInput(s: ShuffleExchangeExec): SparkPlan = {
+ val child =
+ ApplyColumnarRulesAndInsertTransitions(Seq.empty, outputsColumnar =
false).apply(s.child)
Review Comment:
Confirmed. With `transitionRevert.enabled=true` and `maxTransitions=0` over
the typed `map` case, the task failed casting Spark's `OffHeapColumnVector` to
`CometVector` in `reconcileStreamSchema`. 73c727c3b reverts the exchange with
its stage: when `RevertNativeForTransitionHeavyStages` reverts the stage below
a native shuffle whose child is the conversion, the shuffle goes back to the
JVM columnar shuffle that the conversion replaced, and
`EliminateRedundantTransitions` drops the `RowToColumnarExec` under it as it
does for any JVM columnar shuffle. The gate now only converts a shuffle that
`shuffleSupported` would give the JVM columnar shuffle, so that shuffle is
always available there. Your test is in the suite, with AQE on and off, and
fails without the change.
##########
spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometShuffleExchangeExec.scala:
##########
@@ -524,6 +548,58 @@ object CometShuffleExchangeExec
None
}
+ /**
+ * Whether a shuffle whose child is a Spark row-based plan can convert the
child's rows to Arrow
+ * with `CometSparkToColumnarExec` and use native shuffle, instead of the
JVM columnar shuffle.
+ * See [[CometConf.COMET_CONVERT_FROM_SHUFFLE_INPUT_ENABLED]]. The checks
are those that
+ * [[shuffleSupported]] would make for the converted plan, so
`CometExecRule` only converts a
+ * shuffle that will become native. Pure: does not tag the node.
+ */
+ def convertsInputForNativeShuffle(s: ShuffleExchangeExec): Boolean = {
Review Comment:
Updated all three in ef1d8ec60. The first two now give
`spark.comet.convert.shuffleInput.enabled` as the exception, and the plan
construction step says that a `CometNativeArrowSource` child exports its own
stream, which `prepareArrowStreamShuffleDependency` hands to native, while any
other input goes through the `prepareShuffleDependency(rdd, ...)` overload. The
tuning guide still quoted the revert's old log message, which this PR changed,
so that's fixed too.
##########
spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometShuffleExchangeExec.scala:
##########
@@ -524,6 +548,58 @@ object CometShuffleExchangeExec
None
}
+ /**
+ * Whether a shuffle whose child is a Spark row-based plan can convert the
child's rows to Arrow
+ * with `CometSparkToColumnarExec` and use native shuffle, instead of the
JVM columnar shuffle.
+ * See [[CometConf.COMET_CONVERT_FROM_SHUFFLE_INPUT_ENABLED]]. The checks
are those that
+ * [[shuffleSupported]] would make for the converted plan, so
`CometExecRule` only converts a
+ * shuffle that will become native. Pure: does not tag the node.
+ */
+ def convertsInputForNativeShuffle(s: ShuffleExchangeExec): Boolean = {
+ val conf = s.conf
+ CometConf.COMET_CONVERT_FROM_SHUFFLE_INPUT_ENABLED.get(conf) &&
+ // A decision an earlier pass recorded stands, as in shuffleSupported.
+ !hasFallbackReason(s) &&
+ isCometShuffleEnabledReason(s).isEmpty &&
+ CometConf.COMET_SHUFFLE_CONVERT_FROM_SPARK_PLAN_ENABLED.get(conf) &&
+ // Not exercised with Celeborn, which has no JVM columnar shuffle to
replace.
+ !isCometCelebornShuffleManagerEnabled(conf) &&
+ !isCometPlan(s.child) &&
+ !s.child.supportsColumnar &&
+ !isShuffleOperator(s.child) &&
+ !stageContainsDPPScan(s) &&
+ CometSparkToColumnarExec.isSchemaSupported(s.child.schema,
ListBuffer.empty) &&
+ // Arrow holds the time part of an interval in nanoseconds, so the
conversion overflows on a
+ // calendar interval with more microseconds than that can hold. The JVM
columnar shuffle
+ // leaves calendar intervals to Spark's shuffle as well.
+ !s.child.schema.existsRecursively(_.isInstanceOf[CalendarIntervalType]) &&
+ nativeShuffleFailureReasons(s).isEmpty &&
+ !hashesDifferentlyFromSpark(s)
Review Comment:
Done in 21651aa1d, as you laid it out: the config, `!isCometPlan(s.child)`,
`shuffleSupported(s).contains(CometColumnarShuffle)`, the schema check,
`nativeShuffleFailureReasons(s).isEmpty` and `!hashesDifferentlyFromSpark(s)`.
The `native` mode case wasn't intended. The setting is only meant to take over
from the JVM columnar shuffle, so `native` mode now keeps Spark's shuffle for a
Spark child, and the shuffle mode test checks `native` as well as `jvm`. The
predicate now tags the node when `shuffleSupported` finds no Comet shuffle, but
the next case in `CometExecRule` would record the same reasons anyway.
##########
spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometShuffleExchangeExec.scala:
##########
@@ -227,12 +236,27 @@ case class CometShuffleExchangeExec(
ctx,
positionalRoundRobin))
case None =>
- CometShuffleExchangeExec.prepareShuffleDependency(
- inputRDD.asInstanceOf[RDD[ColumnarBatch]],
- child.output,
- outputPartitioning,
- serializer,
- metrics)
+ child match {
+ case _: CometNativeArrowSource =>
+ CometShuffleExchangeExec.prepareArrowStreamShuffleDependency(
+ inputRDD.asInstanceOf[RDD[ArrowArrayStream]],
+ // The range partitioner samples rows, so it needs them as
batches.
+ outputPartitioning match {
+ case _: RangePartitioning => Some(child.executeColumnar())
Review Comment:
Done in 54defa139. When the conversion's child is row-based, the sampling
job now reads those rows through `conversion.child.execute()` and projects only
the sort keys, as `prepareJVMShuffleDependency` does.
`prepareNativeShuffleDependency` takes its sampling input as rows, and the
other callers turn their batches into rows as it did before. The old sampling
job also ran through the conversion's metrics, so a range-partitioned shuffle
reported twice the rows it wrote. The range partitioning test now checks the
count (2d5e9cf45).
The benchmark has a range case now, over all five columns. Best times in ms
for Spark, Comet, Comet converted and the Comet repeat were 2955, 3296, 3133
and 3079, so range partitioning comes out about even with the JVM columnar
shuffle, where the hash cases gain 15 to 30%. The sampling change didn't make a
difference I could measure above the noise, which was around 5% on a busy
machine. The tuning guide now says that range partitioning comes out about even
(80253c8fe).
##########
spark/src/test/scala/org/apache/comet/exec/CometShuffleInputConversionSuite.scala:
##########
@@ -0,0 +1,310 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.comet.exec
+
+import org.apache.spark.sql.{CometTestBase, DataFrame, Row}
+import org.apache.spark.sql.comet.CometSparkToColumnarExec
+import org.apache.spark.sql.comet.execution.shuffle.{CometColumnarShuffle,
CometNativeShuffle, CometShuffleExchangeExec}
+import org.apache.spark.sql.execution.{ColumnarToRowExec,
ColumnarToRowTransition, SparkPlan}
+import org.apache.spark.sql.functions.{array, avg, col, count, hash, length,
max, min, size, sum}
+import org.apache.spark.sql.internal.SQLConf
+import org.apache.spark.sql.types.{ArrayType, BinaryType,
CalendarIntervalType, DataTypes, DecimalType, IntegerType, LongType,
StringType, StructType}
+import org.apache.spark.unsafe.types.CalendarInterval
+
+import org.apache.comet.{CometConf, ExtendedExplainInfo}
+
+// Top-level, so the encoder needs no outer pointer.
+case class ShuffleInputRec(a: Int, b: String)
+
+/** Tests for [[CometConf.COMET_CONVERT_FROM_SHUFFLE_INPUT_ENABLED]]. */
+class CometShuffleInputConversionSuite extends CometTestBase {
+
+ import testImplicits._
+
+ /**
+ * `CometTestBase` turns on the conversion of leaf operators such as RDD
scans, which would
+ * convert the scans below the shuffles here before this conversion sees
them. It is off by
+ * default.
+ */
+ private def withoutLeafConversion(f: => Unit): Unit =
+ withSQLConf(sparkToArrowConversionConfs(enabled = false): _*)(f)
+
+ /** Defines a test that runs with the shuffle input conversion enabled. */
+ private def convertTest(name: String)(f: => Unit): Unit =
+ test(name) {
+ withoutLeafConversion {
+ withSQLConf(CometConf.COMET_CONVERT_FROM_SHUFFLE_INPUT_ENABLED.key ->
"true")(f)
+ }
+ }
+
+ private val rowSchema = new StructType()
+ .add("k", IntegerType)
+ .add("l", LongType)
+ .add("s", DataTypes.StringType)
+ .add("d", DataTypes.DoubleType)
+ .add("m", DecimalType(18, 2))
+
+ /**
+ * `rows` rows from an RDD, so the shuffle above them reads a Spark
operator: the conversion of
+ * RDD scans, `spark.comet.convert.rdd.enabled`, is off by default.
+ */
+ private def rowsDf(rows: Int = 1000): DataFrame = {
+ val data = (0 until rows).map { i =>
+ Row(
+ i % 23,
+ i.toLong,
+ if (i % 11 == 0) null else s"value-$i",
+ i + 0.5d,
+ java.math.BigDecimal.valueOf(i * 7919L % 1000000L, 2))
+ }
+ spark.createDataFrame(spark.sparkContext.parallelize(data, 4), rowSchema)
+ }
+
+ /** Rows with an `array<int>` column, which `CometSparkToColumnarExec` does
not convert. */
+ private def arraysDf(rows: Int = 100): DataFrame = {
+ val schema = new StructType().add("k", IntegerType).add("xs",
ArrayType(IntegerType))
+ val data = (0 until rows).map(i => Row(i % 23, Seq(i, i + 1)))
+ spark.createDataFrame(spark.sparkContext.parallelize(data, 4), schema)
+ }
+
+ private def conversions(plan: SparkPlan): Seq[CometSparkToColumnarExec] =
+ collectWithSubqueries(plan) { case c: CometSparkToColumnarExec => c }
+
+ private def cometShuffles(plan: SparkPlan): Seq[CometShuffleExchangeExec] =
+ collectWithSubqueries(plan) { case s: CometShuffleExchangeExec => s }
+
+ /** The native shuffles that read rows `CometSparkToColumnarExec` converted.
*/
+ private def convertedShuffles(plan: SparkPlan):
Seq[CometShuffleExchangeExec] =
+ cometShuffles(plan).filter { s =>
+ s.shuffleType == CometNativeShuffle &&
s.child.isInstanceOf[CometSparkToColumnarExec]
+ }
+
+ /**
+ * Spark inserts no columnar transitions below a `CometSparkToColumnarExec`,
so the rule adds
+ * them for the operators under the conversion. Without one, an operator
reads its Comet child
+ * through `CometExec.doExecute`, Spark's interpreted columnar-to-row path,
which gives the
+ * right answer slowly, so only the plan shows it.
+ */
+ private def assertRowOperatorsReadThroughTransitions(plan: SparkPlan): Unit
= {
+ val converted = conversions(plan)
+ assert(converted.nonEmpty, plan)
+ converted.foreach { conversion =>
+ // `InputAdapter` and `WholeStageCodegenExec` report their child's
`supportsColumnar`.
+ val bare = conversion.child.collect {
+ case p
+ if !p.supportsColumnar && !p.isInstanceOf[ColumnarToRowTransition]
&&
+ p.children.exists(_.supportsColumnar) =>
+ p
+ }
+ assert(bare.isEmpty, s"row operators read a columnar child without a
transition:\n$plan")
+ assert(
+ conversion.child.collect { case c: ColumnarToRowExec => c }.isEmpty,
+ s"expected Comet's columnar-to-row transitions, not Spark's:\n$plan")
+ }
+ // The transitions are added on the first of the rule's passes under AQE.
A later pass must
+ // not report them as operators Comet failed to convert.
+ val reasons = new ExtendedExplainInfo().getFallbackReasons(plan)
+ assert(!reasons.exists(_.contains("ColumnarToRow")), reasons)
+ }
+
+ convertTest("the shuffle of a Spark operator's rows runs as native shuffle")
{
+ Seq("true", "false").foreach { aqe =>
+ withSQLConf(SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> aqe) {
+ val (_, plan) = checkSparkAnswer(
+ rowsDf()
+ .repartition(10, col("k"))
+ .groupBy("k")
+ .agg(sum("l"), sum(length(col("s"))), sum("d"), sum("m")))
+ assert(convertedShuffles(plan).length == 1, s"AQE $aqe:\n$plan")
+ assert(
+ !cometShuffles(plan).exists(_.shuffleType == CometColumnarShuffle),
+ s"AQE $aqe:\n$plan")
+ }
+ }
+ }
+
+ convertTest("Spark operators over Comet operators read them through
transitions") {
+ withParquetTable((0 until 200).map(i => (i, (i % 13).toString)), "tbl") {
+ Seq("true", "false").foreach { aqe =>
+ withSQLConf(SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> aqe) {
+ val ds = spark.sql("SELECT _1 AS a, _2 AS b FROM
tbl").as[ShuffleInputRec]
+ // The shuffle reads the typed operation's SerializeFromObject.
+ val (_, repartitioned) = checkSparkAnswer(
+ ds.map(r => ShuffleInputRec(r.a % 13, r.b))
+ .repartition(7, col("a"))
+ .groupBy("a")
+ .agg(count("b")))
+ assert(convertedShuffles(repartitioned).length == 1, s"AQE
$aqe:\n$repartitioned")
+ assertRowOperatorsReadThroughTransitions(repartitioned)
+ // The shuffle reads a Spark partial aggregate over the typed
operation. The final
+ // aggregate stays on Spark too, so the shuffle would go back to
Spark's own.
+
withSQLConf(CometConf.COMET_SHUFFLE_REVERT_REDUNDANT_COLUMNAR_ENABLED.key ->
"false") {
+ val (_, aggregated) =
+ checkSparkAnswer(ds.map(r => ShuffleInputRec(r.a % 7,
r.b)).groupBy("a").count())
+ assert(convertedShuffles(aggregated).length == 1, s"AQE
$aqe:\n$aggregated")
+ assertRowOperatorsReadThroughTransitions(aggregated)
+ }
+ }
+ }
+ }
+ }
+
+ convertTest("aggregates whose partial aggregate runs on Spark") {
Review Comment:
Added in ddb428bb3. The shuffle did convert here, as it turns out: `min(s)`
on a string makes Spark plan sort aggregates, and the revert in the next test
only matches hash aggregates, so the shuffle stayed native between two Spark
`SortAggregateExec`s. The test now asserts each plan for two shapes: hash
aggregates whose final aggregate runs natively over the Spark partial, and the
original aggregates on Spark sort aggregates.
##########
spark/src/test/scala/org/apache/comet/exec/CometShuffleInputConversionSuite.scala:
##########
@@ -0,0 +1,310 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.comet.exec
+
+import org.apache.spark.sql.{CometTestBase, DataFrame, Row}
+import org.apache.spark.sql.comet.CometSparkToColumnarExec
+import org.apache.spark.sql.comet.execution.shuffle.{CometColumnarShuffle,
CometNativeShuffle, CometShuffleExchangeExec}
+import org.apache.spark.sql.execution.{ColumnarToRowExec,
ColumnarToRowTransition, SparkPlan}
+import org.apache.spark.sql.functions.{array, avg, col, count, hash, length,
max, min, size, sum}
+import org.apache.spark.sql.internal.SQLConf
+import org.apache.spark.sql.types.{ArrayType, BinaryType,
CalendarIntervalType, DataTypes, DecimalType, IntegerType, LongType,
StringType, StructType}
+import org.apache.spark.unsafe.types.CalendarInterval
+
+import org.apache.comet.{CometConf, ExtendedExplainInfo}
+
+// Top-level, so the encoder needs no outer pointer.
+case class ShuffleInputRec(a: Int, b: String)
+
+/** Tests for [[CometConf.COMET_CONVERT_FROM_SHUFFLE_INPUT_ENABLED]]. */
+class CometShuffleInputConversionSuite extends CometTestBase {
+
+ import testImplicits._
+
+ /**
+ * `CometTestBase` turns on the conversion of leaf operators such as RDD
scans, which would
+ * convert the scans below the shuffles here before this conversion sees
them. It is off by
+ * default.
+ */
+ private def withoutLeafConversion(f: => Unit): Unit =
+ withSQLConf(sparkToArrowConversionConfs(enabled = false): _*)(f)
+
+ /** Defines a test that runs with the shuffle input conversion enabled. */
+ private def convertTest(name: String)(f: => Unit): Unit =
+ test(name) {
+ withoutLeafConversion {
+ withSQLConf(CometConf.COMET_CONVERT_FROM_SHUFFLE_INPUT_ENABLED.key ->
"true")(f)
+ }
+ }
+
+ private val rowSchema = new StructType()
+ .add("k", IntegerType)
+ .add("l", LongType)
+ .add("s", DataTypes.StringType)
+ .add("d", DataTypes.DoubleType)
+ .add("m", DecimalType(18, 2))
+
+ /**
+ * `rows` rows from an RDD, so the shuffle above them reads a Spark
operator: the conversion of
+ * RDD scans, `spark.comet.convert.rdd.enabled`, is off by default.
+ */
+ private def rowsDf(rows: Int = 1000): DataFrame = {
+ val data = (0 until rows).map { i =>
+ Row(
+ i % 23,
+ i.toLong,
+ if (i % 11 == 0) null else s"value-$i",
+ i + 0.5d,
+ java.math.BigDecimal.valueOf(i * 7919L % 1000000L, 2))
+ }
+ spark.createDataFrame(spark.sparkContext.parallelize(data, 4), rowSchema)
+ }
+
+ /** Rows with an `array<int>` column, which `CometSparkToColumnarExec` does
not convert. */
+ private def arraysDf(rows: Int = 100): DataFrame = {
+ val schema = new StructType().add("k", IntegerType).add("xs",
ArrayType(IntegerType))
+ val data = (0 until rows).map(i => Row(i % 23, Seq(i, i + 1)))
+ spark.createDataFrame(spark.sparkContext.parallelize(data, 4), schema)
+ }
+
+ private def conversions(plan: SparkPlan): Seq[CometSparkToColumnarExec] =
+ collectWithSubqueries(plan) { case c: CometSparkToColumnarExec => c }
+
+ private def cometShuffles(plan: SparkPlan): Seq[CometShuffleExchangeExec] =
+ collectWithSubqueries(plan) { case s: CometShuffleExchangeExec => s }
+
+ /** The native shuffles that read rows `CometSparkToColumnarExec` converted.
*/
+ private def convertedShuffles(plan: SparkPlan):
Seq[CometShuffleExchangeExec] =
+ cometShuffles(plan).filter { s =>
+ s.shuffleType == CometNativeShuffle &&
s.child.isInstanceOf[CometSparkToColumnarExec]
+ }
+
+ /**
+ * Spark inserts no columnar transitions below a `CometSparkToColumnarExec`,
so the rule adds
+ * them for the operators under the conversion. Without one, an operator
reads its Comet child
+ * through `CometExec.doExecute`, Spark's interpreted columnar-to-row path,
which gives the
+ * right answer slowly, so only the plan shows it.
+ */
+ private def assertRowOperatorsReadThroughTransitions(plan: SparkPlan): Unit
= {
+ val converted = conversions(plan)
+ assert(converted.nonEmpty, plan)
+ converted.foreach { conversion =>
+ // `InputAdapter` and `WholeStageCodegenExec` report their child's
`supportsColumnar`.
+ val bare = conversion.child.collect {
+ case p
+ if !p.supportsColumnar && !p.isInstanceOf[ColumnarToRowTransition]
&&
+ p.children.exists(_.supportsColumnar) =>
+ p
+ }
+ assert(bare.isEmpty, s"row operators read a columnar child without a
transition:\n$plan")
+ assert(
+ conversion.child.collect { case c: ColumnarToRowExec => c }.isEmpty,
+ s"expected Comet's columnar-to-row transitions, not Spark's:\n$plan")
+ }
+ // The transitions are added on the first of the rule's passes under AQE.
A later pass must
+ // not report them as operators Comet failed to convert.
+ val reasons = new ExtendedExplainInfo().getFallbackReasons(plan)
+ assert(!reasons.exists(_.contains("ColumnarToRow")), reasons)
+ }
+
+ convertTest("the shuffle of a Spark operator's rows runs as native shuffle")
{
+ Seq("true", "false").foreach { aqe =>
+ withSQLConf(SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> aqe) {
+ val (_, plan) = checkSparkAnswer(
+ rowsDf()
+ .repartition(10, col("k"))
+ .groupBy("k")
+ .agg(sum("l"), sum(length(col("s"))), sum("d"), sum("m")))
+ assert(convertedShuffles(plan).length == 1, s"AQE $aqe:\n$plan")
+ assert(
+ !cometShuffles(plan).exists(_.shuffleType == CometColumnarShuffle),
+ s"AQE $aqe:\n$plan")
+ }
+ }
+ }
+
+ convertTest("Spark operators over Comet operators read them through
transitions") {
+ withParquetTable((0 until 200).map(i => (i, (i % 13).toString)), "tbl") {
+ Seq("true", "false").foreach { aqe =>
+ withSQLConf(SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> aqe) {
+ val ds = spark.sql("SELECT _1 AS a, _2 AS b FROM
tbl").as[ShuffleInputRec]
+ // The shuffle reads the typed operation's SerializeFromObject.
+ val (_, repartitioned) = checkSparkAnswer(
+ ds.map(r => ShuffleInputRec(r.a % 13, r.b))
+ .repartition(7, col("a"))
+ .groupBy("a")
+ .agg(count("b")))
+ assert(convertedShuffles(repartitioned).length == 1, s"AQE
$aqe:\n$repartitioned")
+ assertRowOperatorsReadThroughTransitions(repartitioned)
+ // The shuffle reads a Spark partial aggregate over the typed
operation. The final
+ // aggregate stays on Spark too, so the shuffle would go back to
Spark's own.
+
withSQLConf(CometConf.COMET_SHUFFLE_REVERT_REDUNDANT_COLUMNAR_ENABLED.key ->
"false") {
+ val (_, aggregated) =
+ checkSparkAnswer(ds.map(r => ShuffleInputRec(r.a % 7,
r.b)).groupBy("a").count())
+ assert(convertedShuffles(aggregated).length == 1, s"AQE
$aqe:\n$aggregated")
+ assertRowOperatorsReadThroughTransitions(aggregated)
+ }
+ }
+ }
+ }
+ }
+
+ convertTest("aggregates whose partial aggregate runs on Spark") {
+ checkSparkAnswer(
+ rowsDf()
+ .groupBy((col("k") % 7).as("g"))
+ .agg(sum("l"), count("s"), avg("d"), sum("m"), max("l"), min("s")))
+ }
+
+ convertTest("a shuffle between two Spark aggregates goes back to Spark's
shuffle") {
+ withSQLConf(CometConf.COMET_EXEC_AGGREGATE_ENABLED.key -> "false") {
+ val (_, plan) = checkSparkAnswer(rowsDf().groupBy("k").agg(sum("l")))
+ assert(cometShuffles(plan).isEmpty, plan)
+
withSQLConf(CometConf.COMET_SHUFFLE_REVERT_REDUNDANT_COLUMNAR_ENABLED.key ->
"false") {
+ val (_, kept) = checkSparkAnswer(rowsDf().groupBy("k").agg(sum("l")))
+ assert(convertedShuffles(kept).length == 1, kept)
+ }
+ }
+ }
+
+ convertTest("a join with an input that stays on the JVM columnar shuffle") {
Review Comment:
Done in ddb428bb3. There's now a test for each key type the gate admits:
boolean, the four integer types, float and double with NaN and both zeros,
decimal(18, 2), date, timestamp, timestamp_ntz and binary, each with boundary
values and NULL keys. Each test first compares the `spark_partition_id()` of
every row after a `repartition` with Spark's, which covers the NULL keys that
the join filters out, and then runs the join with one input on each shuffle.
They share a `withShuffledJoins` helper with the wide decimal and string tests.
To check that the placement comparison would notice a difference, I took out
the wide decimal guard and added a `decimal(38, 0)` key, and it failed on the
partition ids.
The Celeborn case is in `CometCelebornShufflePlanningSuite` (21651aa1d), and
fails if the gate stops excluding Celeborn. The test over several batches now
has `array<string>` and `map<string,string>` columns with NULL elements, values
and rows. Without the stream read it fails on the map too (`no more field nodes
for field key`), not only on the struct.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [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]
