andygrove commented on code in PR #6564:
URL: https://github.com/apache/datafusion-comet/pull/6564#discussion_r4184778300
##########
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:
Confirmed: with the conversion on, `map(f).filter(...).map(g).rdd.take(1)`
failed on row 30, with `MapElements` at the root of the plan. 1e238811fe takes
any root that produces objects as code reading `Dataset.rdd`: an
`ObjectProducerExec`, or a filter or project over one. That covers the
`DeserializeToObjectExec` that `Dataset.rdd` adds, the typed operation left at
the root once `EliminateSerialization` drops it, and a typed filter Spark pulls
above that operation. A Dataset's own plan ends in rows, so no other plan has
such a root.
The RDD test now runs your query, and the same query with a typed filter at
the end, with AQE on and off. Your query fails with the old check, and the
typed filter one fails without the walk through the filter.
##########
spark/src/main/scala/org/apache/comet/rules/CometExecRule.scala:
##########
@@ -1256,6 +1268,35 @@ 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 typed Dataset operation produces to Arrow, so the
operators above it can
+ * run natively. See [[CometConf.COMET_CONVERT_FROM_TYPED_DATASET_ENABLED]].
+ *
+ * Spark inserts the columnar transitions after this rule, but it does not
look below a
+ * `RowToColumnarTransition` such as `CometSparkToColumnarExec`. That is
harmless above a leaf.
+ * Here the typed operation's own operators sit below the conversion, and
without a transition
+ * they would read a Comet child 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 convertTypedDatasetOutput(op: SerializeFromObjectExec):
SparkPlan = {
+ val unsupported = op.output.filterNot(a =>
+ CometSparkToColumnarExec.isTypeSupported(a.dataType, a.name,
ListBuffer.empty))
+ if (unsupported.nonEmpty) {
+ withFallbackReason(
+ op,
+ "Comet cannot convert the output of a typed Dataset operation to Arrow
because it does " +
+ "not support the type of these columns: " +
+ unsupported.map(a => s"${a.name}:
${a.dataType.simpleString}").mkString(", "))
+ } else {
+ val withTransitions =
+ ApplyColumnarRulesAndInsertTransitions(Seq.empty, outputsColumnar =
false).apply(op)
+ convertToComet(withTransitions,
CometSparkToColumnarExec).getOrElse(withTransitions)
Review Comment:
Fixed in 2737816792, and generalized in c202f51588. The conversion leaves a
typed operation's output unconverted when a limit above it can stop reading
early, unless an operator that reads all of its input, such as an exchange, a
sort or a hash aggregate, sits in between. Your query is the test `a limit does
not evaluate typed Dataset rows beyond the result`, with AQE on and off.
##########
spark/src/main/scala/org/apache/comet/rules/CometExecRule.scala:
##########
@@ -473,6 +498,23 @@ case class CometExecRule(session: SparkSession)
case op if shouldApplySparkToColumnar(conf, op) =>
convertToComet(op, CometSparkToColumnarExec).getOrElse(op)
+ // Typed Dataset operations (`map`, `flatMap`, `mapPartitions`,
`mapGroups`, ...) pass JVM
+ // objects between their operators, so those stay on Spark. Each of them
ends in
+ // `SerializeFromObjectExec`, though, whose output is ordinary rows, and
converting those to
+ // Arrow lets the operators above the typed operation run natively.
+ case op: SerializeFromObjectExec
+ if CometConf.COMET_CONVERT_FROM_TYPED_DATASET_ENABLED.get(conf) =>
+ if (op
+
.getTagValue(CometExecRule.SKIP_TYPED_DATASET_CONVERSION_UNDER_LIMIT)
+ .isDefined) {
+ withFallbackReason(
+ op,
+ "Comet does not convert the output of a typed Dataset operation
below a limit " +
+ "because Arrow batching could evaluate rows beyond Spark's
row-level limit")
+ } else {
+ convertTypedDatasetOutput(op)
Review Comment:
Fixed in c202f51588. A `MapPartitionsExec` above the typed operation now
counts as a reader that can stop early, as a limit does. Your query is the test
`a mapPartitions function does not evaluate typed Dataset rows it never reads`,
with AQE on and off.
--
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]
