viirya commented on code in PR #6035:
URL: https://github.com/apache/datafusion-comet/pull/6035#discussion_r4321727663
##########
spark/src/main/scala/org/apache/comet/rules/CometExecRule.scala:
##########
@@ -724,6 +724,54 @@ case class CometExecRule(session: SparkSession,
queryStagePrep: Boolean = false)
}
}
+ /**
+ * Spark's collect limit consumes at most `limit` rows per partition, but a
native projection
+ * evaluates its whole input batch before that limit can stop reading. Limit
the projection's
+ * input first so an expression on an unconsumed row cannot raise an error.
Keep the outer
+ * limit: it still combines partitions and drops OFFSET rows, which Spark
evaluates before
+ * dropping.
+ *
+ * Only handle an immediately adjacent native projection. A filter or join
between the limit and
+ * projection can require more input rows, and a blocking operator must
consume its input.
+ */
Review Comment:
This is the part I would most like your read on.
`CollectLimitExec` is not the only operator that stops consuming before its
child finishes a batch. `BaseLimitExec.doExecute` is just as lazy:
```scala
protected override def doExecute(): RDD[InternalRow] =
child.execute().mapPartitionsInternal {
iter => iter.take(limit)
}
```
and `LocalLimitExec` / `GlobalLimitExec` both have their own Comet serdes
(`classOf[LocalLimitExec] -> CometLocalLimitExec`, `classOf[GlobalLimitExec] ->
CometGlobalLimitExec`), so the shape is reachable. Spark plans a non-collect
limit as `GlobalLimitExec(limit, LocalLimitExec(limit, child), offset)`
(`SparkStrategies.scala:1213`), which puts a `LocalLimitExec -> ProjectExec`
edge in plans like `SELECT CAST(IF(...) AS INT) FROM t LIMIT 1` used as a
subquery or insert source rather than a driver collect. `CollectTailExec` and
`TakeOrderedAndProjectExec` are in the same family.
The scaladoc says "Only handle an immediately adjacent native projection",
which explains the filter/join exclusion well, but the commit message's "Other
plan shapes are outside this PR" is the part that needs a reason rather than a
boundary. Is a `LocalLimit -> CometProject` edge safe for a reason I am missing
— because the native limit is itself pushed into the native plan and bounds the
scan upstream, say — or does it have the same defect?
If it is safe, that reason belongs here in the scaladoc, so the next person
does not have to re-derive it. If it is not, it is the same bug under a
different operator and I would rather see an issue tracking it than silence,
even if the fix stays out of this PR.
##########
spark/src/main/scala/org/apache/comet/rules/CometExecRule.scala:
##########
@@ -724,6 +724,54 @@ case class CometExecRule(session: SparkSession,
queryStagePrep: Boolean = false)
}
}
+ /**
+ * Spark's collect limit consumes at most `limit` rows per partition, but a
native projection
+ * evaluates its whole input batch before that limit can stop reading. Limit
the projection's
+ * input first so an expression on an unconsumed row cannot raise an error.
Keep the outer
+ * limit: it still combines partitions and drops OFFSET rows, which Spark
evaluates before
+ * dropping.
+ *
+ * Only handle an immediately adjacent native projection. A filter or join
between the limit and
+ * projection can require more input rows, and a blocking operator must
consume its input.
+ */
+ private def limitCollectProjectionInput(op: SparkPlan): SparkPlan = op match
{
+ case limit: CollectLimitExec if limit.limit >= 0 =>
+ limit.child match {
+ case project: CometProjectExec =>
+ // Selecting or renaming columns cannot throw. Keep Project ->
Filter adjacent so
+ // native planning can prune columns used only by the filter's
predicate.
+ val onlyColumns = project.projectList.forall {
+ case _: Attribute | Alias(_: Attribute, _) => true
+ case _ => false
+ }
+ if (onlyColumns) return op
Review Comment:
Minor style point: the two early `return`s inside the `case project:
CometProjectExec` block make the control flow harder to follow than the
`for`-comprehension below, especially since `return` inside a nested pattern
match in Scala is easy to misread as returning from the `match` rather than the
method.
A guard on the case, or folding both conditions into one `if (onlyColumns ||
alreadyCapped) op else { ... }`, would read more directly. Not worth a reroll
on its own.
##########
spark/src/test/scala/org/apache/comet/CometNativeCastSuite.scala:
##########
@@ -65,6 +66,112 @@ class CometNativeCastSuite extends CometTestBase with
AdaptiveSparkPlanHelper {
import testImplicits._
+ test("LIMIT skips a failing conditional literal cast and retains Comet
projection") {
+ withTempPath { path =>
+ withSQLConf(CometConf.COMET_ENABLED.key -> "false") {
+ spark
+ .range(5)
+ .selectExpr("id", "CAST(id AS STRING) AS value")
+ .coalesce(1)
+ .write
+ .parquet(path.getCanonicalPath)
+ }
+ withSQLConf(
+ SQLConf.ANSI_ENABLED.key -> "true",
+ CometConf.COMET_SCALA_UDF_CODEGEN_ENABLED.key -> "true",
+ CometConf.COMET_BATCH_SIZE.key -> "8192") {
+ withParquetTable(path.getCanonicalPath, "cast_branch_rows") {
+ // Both rows share a batch. The native projection must receive only
row 0: masking
+ // the cast by its IF condition alone would still evaluate the
failing row 1.
+ val cast = "CAST(IF(id = 1, 'bad', value) AS INT)"
+ val masked = s"SELECT $cast AS parsed FROM cast_branch_rows LIMIT 1"
+ // ConstantFolding must leave the throwing literal cast inside the
conditional.
+ val badCasts = sql(masked).queryExecution.optimizedPlan
+ .flatMap(_.expressions)
+ .flatMap(_.collect {
+ case c: Cast
+ if c.child.isInstanceOf[Literal] &&
+
Option(c.child.asInstanceOf[Literal].value).exists(_.toString == "bad") =>
+ c
+ })
+ assert(badCasts.nonEmpty)
+ val (_, plan) = checkSparkAnswerAndOperator(
+ sql(masked),
+ includeClasses = Seq(classOf[CometProjectExec]))
+ assert(new
ExtendedExplainInfo().getCodegenDispatchExpressions(plan).contains("cast"))
+ checkSparkError(
+ sql(s"SELECT $cast FROM cast_branch_rows LIMIT 2"),
+ "CAST_INVALID_INPUT")
+
+ // CollectLimit evaluates the prefix before dropping OFFSET rows.
Moving an offset
+ // below the projection would incorrectly hide this error in the
skipped first row.
+ checkSparkError(
+ sql(
+ "SELECT CAST(IF(id = 0, 'bad', value) AS INT) " +
+ "FROM cast_branch_rows LIMIT 1 OFFSET 1"),
+ "CAST_INVALID_INPUT")
+
+ withSQLConf(CometConf.COMET_BATCH_SIZE.key -> "2") {
Review Comment:
This cross-batch case with `COMET_BATCH_SIZE = 2` is the most valuable test
in the set — it is what proves the cap is cumulative across batches rather than
per batch, which a native `LocalLimit` could plausibly have got wrong. Same for
the per-partition test above that reads through `executeColumnar` rather than
`executeCollect`, since `executeCollect` can stop after the first partition and
would hide a missing cap on the second.
No change requested; flagging that these two are the ones worth keeping if
the suite ever gets trimmed for runtime.
##########
spark/src/main/scala/org/apache/comet/rules/CometExecRule.scala:
##########
@@ -724,6 +724,54 @@ case class CometExecRule(session: SparkSession,
queryStagePrep: Boolean = false)
}
}
+ /**
+ * Spark's collect limit consumes at most `limit` rows per partition, but a
native projection
+ * evaluates its whole input batch before that limit can stop reading. Limit
the projection's
+ * input first so an expression on an unconsumed row cannot raise an error.
Keep the outer
+ * limit: it still combines partitions and drops OFFSET rows, which Spark
evaluates before
+ * dropping.
+ *
+ * Only handle an immediately adjacent native projection. A filter or join
between the limit and
+ * projection can require more input rows, and a blocking operator must
consume its input.
+ */
+ private def limitCollectProjectionInput(op: SparkPlan): SparkPlan = op match
{
+ case limit: CollectLimitExec if limit.limit >= 0 =>
+ limit.child match {
+ case project: CometProjectExec =>
+ // Selecting or renaming columns cannot throw. Keep Project ->
Filter adjacent so
+ // native planning can prune columns used only by the filter's
predicate.
+ val onlyColumns = project.projectList.forall {
+ case _: Attribute | Alias(_: Attribute, _) => true
+ case _ => false
+ }
+ if (onlyColumns) return op
+ project.child match {
+ case cap: CometLocalLimitExec if cap.limit <= limit.limit =>
return op
+ case _ =>
+ }
+ val bounded = for {
+ cap <- convertToComet(LocalLimitExec(limit.limit, project.child),
CometLocalLimitExec)
Review Comment:
`convertToComet` is doing double duty here. Beyond returning the converted
plan, on failure it runs `rollUpFallbackReasons(op)` and
`reportUnexplainedFallback(op)` against the node passed in — and the node
passed in here is a synthetic `LocalLimitExec` that exists only inside this
method and is discarded when conversion fails.
I checked and this is currently harmless: every decline path in
`isOperatorEnabled` tags a reason, so the strict-mode `IllegalStateException`
cannot fire, and the discarded node's tag never reaches EXPLAIN. So this is not
a bug today.
It is worth a brief note in the comment, though, that the synthetic node is
deliberately routed through `convertToComet` for its config and serde gating,
and that its fallback tags are discarded with it. Otherwise the coupling is
invisible, and a future change to `reportUnexplainedFallback`'s contract — e.g.
making it strict whenever the tag is absent — would fail here for a reason that
has nothing to do with the user's plan.
--
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]
