andygrove commented on code in PR #6607:
URL: https://github.com/apache/datafusion-comet/pull/6607#discussion_r4179215168
##########
spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometShuffleExchangeExec.scala:
##########
@@ -524,6 +525,45 @@ object CometShuffleExchangeExec
None
}
+ /**
+ * Whether a shuffle whose child is a Spark row-based plan can convert the
child's rows to Arrow
+ * with `CometSparkToColumnarExec` and use native shuffle, instead of the
JVM columnar shuffle.
+ * See [[CometConf.COMET_CONVERT_FROM_SHUFFLE_INPUT_ENABLED]]. The checks
are those that
+ * [[shuffleSupported]] would make for the converted plan, so
`CometExecRule` only converts a
+ * shuffle that will become native. Pure: does not tag the node.
+ */
+ def convertsInputForNativeShuffle(s: ShuffleExchangeExec): Boolean = {
+ val conf = s.conf
+ CometConf.COMET_CONVERT_FROM_SHUFFLE_INPUT_ENABLED.get(conf) &&
+ // A decision an earlier pass recorded stands, as in shuffleSupported.
+ !hasFallbackReason(s) &&
+ isCometShuffleEnabledReason(s).isEmpty &&
+ CometConf.COMET_SHUFFLE_CONVERT_FROM_SPARK_PLAN_ENABLED.get(conf) &&
+ // Not exercised with Celeborn, which has no JVM columnar shuffle to
replace.
+ !isCometCelebornShuffleManagerEnabled(conf) &&
+ !isCometPlan(s.child) &&
+ !s.child.supportsColumnar &&
+ !isShuffleOperator(s.child) &&
+ !stageContainsDPPScan(s) &&
+ CometSparkToColumnarExec.isSchemaSupported(s.child.schema,
ListBuffer.empty) &&
+ nativeShuffleFailureReasons(s).isEmpty &&
+ !hashesWideDecimal(s)
Review Comment:
Confirmed with your case: with the conversion on, the join returned 0 rows
to Spark's 400. 9795b61f52 keeps a shuffle that hashes a string, at any nesting
depth, on the JVM columnar shuffle, the way the wide decimal guard already
worked. Preserving the raw-byte hash would mean native shuffle hashing the
bytes before the import replaces invalid UTF-8, which is more than this PR
should take on, so I went with the fallback. Your join is now a test in
`CometShuffleInputConversionSuite`, and the config doc and the tuning guide say
that a string key stays on columnar shuffle. The typed `map` test and the
benchmark case that hashed a string now hash a number.
##########
spark/src/main/scala/org/apache/comet/rules/CometExecRule.scala:
##########
@@ -1256,6 +1273,30 @@ case class CometExecRule(session: SparkSession)
private def hasEnabledHandler(op: SparkPlan): Boolean =
allExecs.get(op.getClass).exists(_.enabledConfig.forall(_.get(op.conf)))
+ /**
+ * Converts the rows a Spark operator feeds to a shuffle to Arrow, so that
the shuffle runs as
+ * native shuffle instead of the JVM columnar shuffle. See
+ * [[CometConf.COMET_CONVERT_FROM_SHUFFLE_INPUT_ENABLED]], and
+ * [[CometShuffleExchangeExec.convertsInputForNativeShuffle]] for when it
applies.
+ *
+ * Spark inserts the columnar transitions after this rule, but it does not
look below a
+ * `RowToColumnarTransition` such as `CometSparkToColumnarExec`. Without a
transition, a Spark
+ * operator in the child's subtree that reads a Comet operator would do so
through
+ * `CometExec.doExecute`, Spark's interpreted columnar-to-row path. So the
subtree gets its
+ * transitions now, from Spark's own rule, and
`EliminateRedundantTransitions` later replaces
+ * each one over a Comet child with Comet's own. Spark's rule leaves
existing transitions alone,
+ * which matters because this rule runs over the same plan twice under AQE.
+ */
+ private def convertShuffleInput(s: ShuffleExchangeExec): SparkPlan = {
+ val child =
+ ApplyColumnarRulesAndInsertTransitions(Seq.empty, outputsColumnar =
false).apply(s.child)
+ convertToComet(child, CometSparkToColumnarExec)
Review Comment:
Confirmed: your 24-row case failed with `no more field nodes for field v` at
batch size 7. 16a1f73aa4 has native shuffle read a `CometNativeArrowSource`
child through `doExecuteAsArrowStream`, as a native operator already does with
such an input, so its batches no longer go through `ColumnarBatchArrowReader`.
Range partitioning still samples from `executeColumnar`, but the sampling
doesn't close the batches. The new test runs hash and range partitioning over a
struct column that fills several batches, with leaf conversion on and off. With
leaf conversion on, the same shuffle over a converted RDD scan failed the same
way before this change, so `spark.comet.sparkToColumnar.enabled` gets the fix
too.
##########
spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometShuffleExchangeExec.scala:
##########
@@ -524,6 +525,45 @@ object CometShuffleExchangeExec
None
}
+ /**
+ * Whether a shuffle whose child is a Spark row-based plan can convert the
child's rows to Arrow
+ * with `CometSparkToColumnarExec` and use native shuffle, instead of the
JVM columnar shuffle.
+ * See [[CometConf.COMET_CONVERT_FROM_SHUFFLE_INPUT_ENABLED]]. The checks
are those that
+ * [[shuffleSupported]] would make for the converted plan, so
`CometExecRule` only converts a
+ * shuffle that will become native. Pure: does not tag the node.
+ */
+ def convertsInputForNativeShuffle(s: ShuffleExchangeExec): Boolean = {
+ val conf = s.conf
+ CometConf.COMET_CONVERT_FROM_SHUFFLE_INPUT_ENABLED.get(conf) &&
+ // A decision an earlier pass recorded stands, as in shuffleSupported.
+ !hasFallbackReason(s) &&
+ isCometShuffleEnabledReason(s).isEmpty &&
+ CometConf.COMET_SHUFFLE_CONVERT_FROM_SPARK_PLAN_ENABLED.get(conf) &&
+ // Not exercised with Celeborn, which has no JVM columnar shuffle to
replace.
+ !isCometCelebornShuffleManagerEnabled(conf) &&
+ !isCometPlan(s.child) &&
+ !s.child.supportsColumnar &&
+ !isShuffleOperator(s.child) &&
+ !stageContainsDPPScan(s) &&
+ CometSparkToColumnarExec.isSchemaSupported(s.child.schema,
ListBuffer.empty) &&
Review Comment:
Confirmed, your row failed with `long overflow` in
`IntervalMonthDayNanoWriter`. d8d7ce4b41 leaves an input with a calendar
interval anywhere in its schema out of the conversion, so that shuffle stays on
Spark's shuffle as it did before. There's a test with your value, and the
config doc and the tuning guide list calendar intervals among the columns the
conversion rules out.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [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]
