sunchao commented on code in PR #6035:
URL: https://github.com/apache/datafusion-comet/pull/6035#discussion_r4340653485
##########
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:
Confirmed: the non-collect limit plan has the same preexisting defect. I
reproduced `CometGlobalLimit -> exchange -> CometLocalLimit -> CometProject` on
Spark 4.0.4: Spark returns `0`, while Comet raises `CAST_INVALID_INPUT`. The
native local limit slices the projection's output after the projection has
evaluated the batch, so it does not protect that projection's input.
Filed #7147 with that reproducer and the throwing-filter case, under the
broader tracking issue #6006. The rule's scaladoc and compatibility guide now
link it. The issue also tracks direct GlobalLimit coverage and already-ordered
`TakeOrderedAndProject` inputs for investigation. It distinguishes these from
blocking operators and `CollectTail`, where Spark itself may need to evaluate
discarded rows. This PR keeps the adjacent collect-limit fix scoped as
discussed.
Updated in `eea8a6b00`. The focused Spark 4.0 run passed all 116 tests.
##########
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:
Added the comment beside the synthetic conversion. It explains that
`convertToComet` deliberately applies the normal configuration and serde
checks, and that a declined synthetic limit is discarded together with its
fallback tags while the original plan is retained.
Updated in `eea8a6b00`. The focused Spark 4.0 run passed all 116 tests.
##########
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:
Replaced both early returns with `if (onlyColumns || alreadyCapped) op else
{ ... }`. The existing `cap.limit <= limit.limit` check and column-pruning
guard are preserved.
Updated in `eea8a6b00`. The focused Spark 4.0 run passed all 116 tests.
##########
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:
Kept both tests unchanged: the batch-size-2 case checks the cumulative cap,
and the per-partition case still uses `executeColumnar` so the second partition
is exercised.
Updated in `eea8a6b00`. The focused Spark 4.0 run passed all 116 tests.
--
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]
