sunchao commented on code in PR #6564:
URL: https://github.com/apache/datafusion-comet/pull/6564#discussion_r4178950799
##########
spark/src/main/scala/org/apache/comet/rules/CometExecRule.scala:
##########
@@ -1256,6 +1288,78 @@ case class CometExecRule(session: SparkSession)
private def hasEnabledHandler(op: SparkPlan): Boolean =
allExecs.get(op.getClass).exists(_.enabledConfig.forall(_.get(op.conf)))
+ /**
+ * Tags each `SerializeFromObjectExec` whose output an operator above it can
stop reading early
+ * with that operator, so `transform` does not convert it. Spark computes
the rows of a typed
+ * Dataset operation one at a time, as the operator above reads them, while
the conversion fills
+ * a whole Arrow batch first. Below a limit, a `mapPartitions` function such
as `_.take(1)`, or
+ * code reading `Dataset.rdd`, the conversion would run the user function on
rows Spark never
+ * reaches, and a function that throws on one of them would fail a query
that succeeds in Spark.
+ * An operator that reads all of its input before it returns a row ends the
search, since Spark
+ * computes every row below it anyway: an exchange, a sort, a hash
aggregate, or a top-k over
+ * input that is not already sorted.
+ *
+ * Conversion is bottom-up, so this runs first. TreeNode tags survive the
child copies made
+ * during transformUp, while an identity set would not.
+ */
+ private def tagPartiallyReadTypedDatasetOutputs(plan: SparkPlan): Unit = {
+ def visit(op: SparkPlan, partialReader: Option[String]): Unit = {
+ val childReader = op match {
+ case serialize: SerializeFromObjectExec =>
+ partialReader.foreach(
+ serialize.setTagValue(CometExecRule.TYPED_DATASET_PARTIAL_READER,
_))
+ partialReader
+ case _: CollectLimitExec | _: LocalLimitExec | _: GlobalLimitExec =>
Some("a limit")
+ // A top-k reads only its first rows when its input is already sorted.
+ case topK: TakeOrderedAndProjectExec
+ if SortOrder.orderingSatisfies(topK.child.outputOrdering,
topK.sortOrder) =>
+ Some("a limit")
+ case _: MapPartitionsExec => Some("a mapPartitions function")
+ case _: Exchange | _: SortExec | _: HashAggregateExec | _:
ObjectHashAggregateExec |
+ _: TakeOrderedAndProjectExec =>
+ None
+ case _ => partialReader
+ }
+ op.children.foreach(visit(_, childReader))
+ }
+ // `Dataset.rdd` plans a `DeserializeToObjectExec` at the root, and the
RDD's own code
+ // decides how much of it to read, as `take(1)` does.
+ val rootReader = plan match {
+ case _: DeserializeToObjectExec => Some("code reading Dataset.rdd")
Review Comment:
[P2] Preserve RDD laziness after Spark eliminates serialization.
`Dataset.rdd` does not always retain a `DeserializeToObjectExec` root: when the
Dataset ends with another `map`, Spark’s `EliminateSerialization` removes that
deserializer and leaves `MapElementsExec`. Consequently,
`map(f).filter(...).map(g).rdd.take(1)` bypasses this guard and converts the
earlier serializer. If `f` throws on row 30, conversion-enabled Comet fails
although Spark and conversion-disabled Comet return the first result
successfully. Could the safeguard recognize these optimized RDD roots and
preserve row-level consumption? Add this case alongside the current RDD
regression test.
Evidence: Reproduced at this head on Spark 4.1.3/JDK 17: `spark.range(0,
100, 1, 1).map { i => if (i == 30L) throw new
IllegalArgumentException("unexpected evaluation of row 30"); i + 1L
}.filter(col("value") > 0L).map(i => i + 1L).rdd.take(1)`. With AQE both false
and true, Spark and Comet with `spark.comet.convert.typedDataset.enabled=false`
return `Array(2L)`. Enabling conversion throws `SparkException` caused by
`CometNativeException` wrapping that `IllegalArgumentException`, with
`RowArrowReader.loadNextBatch` in the stack. The RDD execution plan is rooted
at `MapElements`, above `DeserializeToObject -> CometColumnarToRow ->
CometFilter -> CometSparkRowToColumnar -> SerializeFromObject`. Spark’s
`EliminateSerialization` rule removes the outer deserialize/serialize pair
across all five supported versions. The six-configuration probe failed only in
the two conversion-enabled cases.
--
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]
