andygrove commented on code in PR #6607:
URL: https://github.com/apache/datafusion-comet/pull/6607#discussion_r4184364143
##########
spark/src/main/scala/org/apache/spark/sql/comet/execution/shuffle/CometShuffleExchangeExec.scala:
##########
@@ -524,6 +525,45 @@ object CometShuffleExchangeExec
None
}
+ /**
+ * Whether a shuffle whose child is a Spark row-based plan can convert the
child's rows to Arrow
+ * with `CometSparkToColumnarExec` and use native shuffle, instead of the
JVM columnar shuffle.
+ * See [[CometConf.COMET_CONVERT_FROM_SHUFFLE_INPUT_ENABLED]]. The checks
are those that
+ * [[shuffleSupported]] would make for the converted plan, so
`CometExecRule` only converts a
+ * shuffle that will become native. Pure: does not tag the node.
+ */
+ def convertsInputForNativeShuffle(s: ShuffleExchangeExec): Boolean = {
+ val conf = s.conf
+ CometConf.COMET_CONVERT_FROM_SHUFFLE_INPUT_ENABLED.get(conf) &&
+ // A decision an earlier pass recorded stands, as in shuffleSupported.
+ !hasFallbackReason(s) &&
+ isCometShuffleEnabledReason(s).isEmpty &&
+ CometConf.COMET_SHUFFLE_CONVERT_FROM_SPARK_PLAN_ENABLED.get(conf) &&
+ // Not exercised with Celeborn, which has no JVM columnar shuffle to
replace.
+ !isCometCelebornShuffleManagerEnabled(conf) &&
+ !isCometPlan(s.child) &&
+ !s.child.supportsColumnar &&
+ !isShuffleOperator(s.child) &&
+ !stageContainsDPPScan(s) &&
+ CometSparkToColumnarExec.isSchemaSupported(s.child.schema,
ListBuffer.empty) &&
+ nativeShuffleFailureReasons(s).isEmpty &&
+ !hashesWideDecimal(s)
Review Comment:
Confirmed with your follow-up case: with the conversion on, the join on
`hash(lk) === hash(rk)` returned 0 rows to Spark's 800. a20a4a5e9c has the
guard look for a string anywhere in a hash key's expression rather than only at
the key's type, so a key computed from a string stays on the JVM columnar
shuffle too. A wide decimal still counts only as the key's own type, since
Comet's `hash` of a wide decimal already falls back to Spark.
The string join test now also joins on the hashes of the strings, and fails
without the change. It uses your two keys instead of four. Both shuffles write
the strings with invalid UTF-8 replaced (#4764), and with four keys two of the
hashes share a partition, where the join then matches them. That happens with
the conversion off as well.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [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]
