sunchao commented on code in PR #6607:
URL: https://github.com/apache/datafusion-comet/pull/6607#discussion_r4178941107
##########
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:
[P2] Keep the converted input’s struct vectors valid across batches. This
new plan shape sends `CometSparkToColumnarExec` through `executeColumnar()` and
then `ColumnarBatchArrowReader`, which closes batches backed by
`RowArrowReader`’s reusable vectors. A struct payload succeeds for one batch
but fails when the next batch is loaded. This affects ordinary struct columns
even when the partition key is an integer. Consume the conversion’s Arrow
stream directly or repair the bridge’s ownership handling before admitting
these inputs.
Evidence: On Spark 4.1.3, create 8,200 RDD rows in one input partition with
schema `k: int, payload: struct<v: long>`, then run `repartition(3,
col("k")).collect()` with leaf conversion disabled and shuffle-input conversion
enabled. At batch size 8,192, the job fails with `C Data interface error:
java.lang.IllegalArgumentException: no more field nodes for field v: Int(64,
true) and vector []`. Spark, the supplied base, and conversion-disabled head
return all 8,200 rows. A 24-row reproduction also passes with one batch and
fails with batch size 7.
##########
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:
[P2] Retain fallback for calendar intervals whose range the Arrow conversion
cannot represent. `isSchemaSupported` accepts `CalendarIntervalType`, and the
native shuffle checks admit it, so this predicate also converts inputs that
previously stayed on Spark shuffle. Spark stores a signed 64-bit microsecond
component, but `IntervalMonthDayNanoWriter` multiplies it by 1,000 using
`Math.multiplyExact`. Valid Spark intervals therefore become query failures.
Reject calendar intervals recursively for this conversion until the full Spark
range is supported.
Evidence: On Spark 4.1.3, create one RDD row with schema `k: int, i:
CalendarIntervalType` and value `Row(1, new CalendarInterval(0, 0,
10800000000000000L))`; execute `repartition(2, col("k")).collect()`. Spark, the
supplied base, and conversion-disabled head return the row. Enabling
`spark.comet.convert.shuffleInput.enabled` fails with `ArithmeticException:
long overflow` at `IntervalMonthDayNanoWriter.setValueUnsafe`,
`ArrowWriters.scala:1623`. The base plan retains Spark’s exchange.
##########
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:
[P1] Preserve Spark’s raw-byte hash when converting string partition keys.
With `spark.comet.convert.shuffleInput.enabled=true`, a string-key join can
convert one input while an `array<int>` payload keeps the other on JVM shuffle.
Spark hashes `UTF8String` bytes directly, whereas the converted input passes
through `ScanExec.import_column`, which replaces invalid UTF-8 before native
hashing. Equal Spark keys then reach different partitions and matching rows
disappear. Preserve the original partition hash before sanitization, or retain
JVM shuffle for affected inputs.
Evidence: Spark 4.1.3 reproduction: create 40 RDD rows `(Array((0x80 + i %
4).toByte), i.toLong)`, cast the binary column to string on both join inputs,
and retain an `array<int>` payload only on the right. Disable leaf conversion,
AQE, and broadcast joins; use 10 shuffle partitions. Spark, the supplied base,
and head with conversion disabled return 400 rows. Head with conversion enabled
returns 0. Reproduced by `FullShuffleReviewSuite` in
`/tmp/pr6607-review-full-l2ehudzv/head-tests.log`; the executed plan contains
one native and one JVM columnar exchange.
--
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]
