From 53510eceb19aa6602def51be46e03e119a640438 Mon Sep 17 00:00:00 2001 From: rich7420 Date: Sat, 24 Oct 2026 15:40:57 +0800 Subject: [PATCH] fix: cap projection input beneath local and global limits Part of #7147. --- .../user-guide/latest/compatibility/index.md | 13 +- .../apache/comet/rules/CometExecRule.scala | 41 ++-- .../apache/comet/CometNativeCastSuite.scala | 77 ++++++- .../comet/rules/CometExecRuleSuite.scala | 197 +++++++++++++----- 4 files changed, 249 insertions(+), 79 deletions(-) diff --git a/docs/source/user-guide/latest/compatibility/index.md b/docs/source/user-guide/latest/compatibility/index.md index 6fcafbf71c2..d66525ad273 100644 --- a/docs/source/user-guide/latest/compatibility/index.md +++ b/docs/source/user-guide/latest/compatibility/index.md @@ -154,15 +154,18 @@ Spark returns `'0'` after the first row matches. Comet can evaluate the filter p second row in the same batch and raise `CAST_INVALID_INPUT`. SQL without `ORDER BY` does not guarantee this input order; the example describes a physical scan order that exposes the difference. -Comet caps projection input when a collect limit directly consumes the projection. This protects -that projection's expressions, but does not bound evaluation inside its child operators. Throwing -filter predicates and projections beneath non-collect `LocalLimit` / `GlobalLimit` plans, such as -a `LIMIT` inside a subquery, remain affected. These existing limitations and the investigation of +With native local limits enabled, Comet caps projection input when `CollectLimit`, `LocalLimit`, +or `GlobalLimit` directly consumes the projection. This protects the projection's expressions +from evaluating rows beyond the limit, but does not bound evaluation inside its child operators. +Throwing filter predicates remain affected. These limitations and the investigation of already-ordered `TakeOrderedAndProject` inputs are tracked in [#7147](https://github.com/apache/datafusion-comet/issues/7147). An error on a row absent from the final result is not always a mismatch: Spark can also need to -evaluate that row, for example to sort the input or process an `OFFSET` prefix. +evaluate that row, for example to sort the input or process an `OFFSET` prefix. The projection +input cap includes that prefix. When `GlobalLimit` and the projection share a whole-stage codegen +stage, Spark can skip evaluating the prefix, while Comet still evaluates it. Comet can therefore +raise an error that Spark does not. The input cap does not address this existing difference. ## ANSI-mode error classes and messages diff --git a/spark/src/main/scala/org/apache/comet/rules/CometExecRule.scala b/spark/src/main/scala/org/apache/comet/rules/CometExecRule.scala index d79927f0c04..9ac5e5e3455 100644 --- a/spark/src/main/scala/org/apache/comet/rules/CometExecRule.scala +++ b/spark/src/main/scala/org/apache/comet/rules/CometExecRule.scala @@ -727,7 +727,7 @@ case class CometExecRule(session: SparkSession, queryStagePrep: Boolean = false) } plan.transformUp { case op => - val converted = convertNode(limitCollectProjectionInput(refreshStaleShuffleScans(op))) + val converted = convertNode(limitProjectionInput(refreshStaleShuffleScans(op))) // Replace SubqueryBroadcastExec with CometSubqueryBroadcastExec in DPP expressions // when the broadcast child has a Comet plan underneath. This enables exchange reuse // between the DPP subquery and the join's CometBroadcastExchangeExec because both @@ -737,22 +737,29 @@ 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. + * Spark's limits stop consuming rows, but a native limit slices a projection's output after the + * whole input batch has been evaluated. Cap the projection's input first so expressions on + * unconsumed rows cannot raise an error. The cap includes the OFFSET prefix; the outer limit + * still applies the offset and, for collect/global limits, preserves the single-partition + * result. Keep local limits too, preserving their plan shape and AQE links. * * 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. - * Non-collect LocalLimit/GlobalLimit plans can still evaluate a projection before limiting its - * output. That preexisting gap and throwing filter predicates are tracked separately in + * Throwing filter predicates and other partial-consumption gaps remain tracked in * https://github.com/apache/datafusion-comet/issues/7147. */ - private def limitCollectProjectionInput(op: SparkPlan): SparkPlan = op match { - case limit: CollectLimitExec if limit.limit >= 0 => - limit.child match { - case project: CometProjectExec => + private def limitProjectionInput(op: SparkPlan): SparkPlan = { + val rows = op match { + case limit: CollectLimitExec => limit.limit + case limit: LocalLimitExec => limit.limit + case limit: GlobalLimitExec => limit.limit + case _ => -1 + } + if (rows < 0) { + op + } else { + op.children match { + case Seq(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 { @@ -760,7 +767,7 @@ case class CometExecRule(session: SparkSession, queryStagePrep: Boolean = false) case _ => false } val alreadyCapped = project.child match { - case cap: CometLocalLimitExec => cap.limit <= limit.limit + case cap: CometLocalLimitExec => cap.limit <= rows case _ => false } if (onlyColumns || alreadyCapped) { @@ -769,9 +776,7 @@ case class CometExecRule(session: SparkSession, queryStagePrep: Boolean = false) // Use normal config and serde gating for the synthetic limit. If conversion fails, // its fallback tags are discarded with it and the original plan is retained. val bounded = for { - cap <- convertToComet( - LocalLimitExec(limit.limit, project.child), - CometLocalLimitExec) + cap <- convertToComet(LocalLimitExec(rows, project.child), CometLocalLimitExec) rebuilt <- convertToComet( project.originalPlan.withNewChildren(Seq(cap)), CometProjectExec) @@ -785,13 +790,13 @@ case class CometExecRule(session: SparkSession, queryStagePrep: Boolean = false) project.getTagValue(SparkPlan.LOGICAL_PLAN_INHERITED_TAG).foreach { link => rebuilt.setTagValue(SparkPlan.LOGICAL_PLAN_INHERITED_TAG, link) } - limit.withNewChildren(Seq(rebuilt)) + op.withNewChildren(Seq(rebuilt)) } bounded.getOrElse(op) } case _ => op } - case _ => op + } } /** diff --git a/spark/src/test/scala/org/apache/comet/CometNativeCastSuite.scala b/spark/src/test/scala/org/apache/comet/CometNativeCastSuite.scala index f16d993b0d1..82762b6c282 100644 --- a/spark/src/test/scala/org/apache/comet/CometNativeCastSuite.scala +++ b/spark/src/test/scala/org/apache/comet/CometNativeCastSuite.scala @@ -31,7 +31,7 @@ import org.apache.spark.SparkConf import org.apache.spark.sql.{CometTestBase, DataFrame, Row, SaveMode} import org.apache.spark.sql.catalyst.expressions.{Cast, Literal} import org.apache.spark.sql.catalyst.parser.ParseException -import org.apache.spark.sql.comet.{CometCollectLimitExec, CometProjectExec} +import org.apache.spark.sql.comet.{CometCollectLimitExec, CometGlobalLimitExec, CometLocalLimitExec, CometProjectExec} import org.apache.spark.sql.execution.adaptive.AdaptiveSparkPlanHelper import org.apache.spark.sql.functions.{col, monotonically_increasing_id} import org.apache.spark.sql.internal.SQLConf @@ -128,6 +128,81 @@ class CometNativeCastSuite } } + test("non-collect LIMIT skips a failing cast beyond its projection input cap") { + 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) + } + for (aqe <- Seq("false", "true"); + (limit, batchSize, offset) <- Seq((1, 8192, 0), (3, 2, 0), (2, 8192, 1))) { + withSQLConf( + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> aqe, + SQLConf.ANSI_ENABLED.key -> "true", + CometConf.COMET_SCALA_UDF_CODEGEN_ENABLED.key -> "true", + CometConf.COMET_BATCH_SIZE.key -> batchSize.toString) { + withParquetTable(path.getCanonicalPath, "non_collect_cast_rows") { + def query(n: Int, value: String, badId: Int = limit + offset): String = + s"""SELECT parsed FROM ( + | SELECT CAST(IF(id = $badId, '$value', value) AS INT) AS parsed + | FROM non_collect_cast_rows LIMIT $n OFFSET $offset + |) q WHERE parsed >= 0""".stripMargin + + val (_, plan) = checkSparkAnswerAndOperator( + sql(query(limit, "bad")), + includeClasses = Seq(classOf[CometLocalLimitExec], classOf[CometProjectExec])) + // AQE can remove the global limit after the shuffle materializes enough rows. + if (aqe == "false") { + assert(collect(plan) { case p: CometGlobalLimitExec => p }.nonEmpty) + } + val project = collect(plan) { case p: CometProjectExec => p }.find { + _.child.isInstanceOf[CometLocalLimitExec] + }.get + assert(project.child.asInstanceOf[CometLocalLimitExec].limit == limit + offset) + assert(new ExtendedExplainInfo().getCodegenDispatchExpressions(plan).contains("cast")) + checkSparkError(sql(query(limit + 1, "bad")), "CAST_INVALID_INPUT") + if (offset > 0) { + // The projection precedes the shuffle, so Spark also evaluates the offset prefix. + checkSparkError(sql(query(limit, "bad", badId = 0)), "CAST_INVALID_INPUT") + } + checkSparkAnswerAndOperator( + sql(query(limit + 1, "1")), + includeClasses = Seq(classOf[CometProjectExec])) + } + } + } + } + } + + test("non-collect LIMIT caps projection input in each partition") { + withSQLConf( + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false", + SQLConf.ANSI_ENABLED.key -> "true", + CometConf.COMET_SCALA_UDF_CODEGEN_ENABLED.key -> "true", + CometConf.COMET_EXEC_RANGE_ENABLED.key -> "true") { + withTempView("non_collect_partition_rows") { + spark.range(0, 4, 1, 2).createOrReplaceTempView("non_collect_partition_rows") + val (_, plan) = checkSparkAnswerAndOperator( + sql("""SELECT parsed FROM ( + | SELECT CAST(IF(id % 2 = 1, 'bad', '0') AS INT) AS parsed + | FROM non_collect_partition_rows LIMIT 1 + |) q WHERE parsed >= 0""".stripMargin), + includeClasses = Seq( + classOf[CometLocalLimitExec], + classOf[CometGlobalLimitExec], + classOf[CometProjectExec])) + val localLimit = collect(plan) { + case p: CometLocalLimitExec if p.child.isInstanceOf[CometProjectExec] => p + }.head + assert(localLimit.executeColumnar().getNumPartitions == 2) + } + } + } + test("LIMIT caps each partition before a failing conditional literal cast") { withSQLConf( SQLConf.ANSI_ENABLED.key -> "true", diff --git a/spark/src/test/scala/org/apache/comet/rules/CometExecRuleSuite.scala b/spark/src/test/scala/org/apache/comet/rules/CometExecRuleSuite.scala index bf0bacf0d9e..f8ae085ad26 100644 --- a/spark/src/test/scala/org/apache/comet/rules/CometExecRuleSuite.scala +++ b/spark/src/test/scala/org/apache/comet/rules/CometExecRuleSuite.scala @@ -305,6 +305,57 @@ class CometExecRuleSuite extends CometTestBase { } } + for ((name, makeLimit) <- Seq[(String, (Int, SparkPlan) => SparkPlan)]( + "local" -> ((n, child) => LocalLimitExec(n, child)), + "global" -> ((n, child) => GlobalLimitExec(n, child)))) { + test(s"$name limit caps a directly adjacent throwing projection") { + withSQLConf( + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false", + SQLConf.WHOLESTAGE_CODEGEN_ENABLED.key -> "false", + SQLConf.ANSI_ENABLED.key -> "true", + CometConf.COMET_CONVERT_FROM_RANGE_ENABLED.key -> "true", + CometConf.COMET_SCALA_UDF_CODEGEN_ENABLED.key -> "true") { + val query = "SELECT CAST(IF(id = 1, 'bad', '0') AS INT) AS parsed " + + "FROM range(0, 2, 1, 1)" + val expected = makeLimit(1, createSparkPlan(spark, query)) + .executeCollect() + .map(_.getInt(0)) + .toSeq + assert(expected == Seq(0)) + val result = applyCometExecRule(makeLimit(1, createSparkPlan(spark, query))) + assert(result.executeCollect().map(_.getInt(0)).toSeq == expected) + val project = result.collectFirst { case p: CometProjectExec => p }.get + assert(project.child.asInstanceOf[CometLocalLimitExec].limit == 1) + assert(project.nativeOp.getChildren(0).getLimit.getLimit == 1) + val empty = applyCometExecRule(makeLimit(0, createSparkPlan(spark, query))) + assert(empty.executeCollect().isEmpty) + } + } + } + + test("global limit keeps the offset prefix inside the projection input cap") { + withSQLConf( + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false", + SQLConf.WHOLESTAGE_CODEGEN_ENABLED.key -> "false", + SQLConf.ANSI_ENABLED.key -> "true", + CometConf.COMET_CONVERT_FROM_RANGE_ENABLED.key -> "true", + CometConf.COMET_SCALA_UDF_CODEGEN_ENABLED.key -> "true") { + val query = "SELECT CAST(IF(id = 3, 'bad', CAST(id AS STRING)) AS INT) AS parsed " + + "FROM range(0, 5, 1, 1)" + val expected = GlobalLimitExec(3, createSparkPlan(spark, query), offset = 1) + .executeCollect() + .map(_.getInt(0)) + .toSeq + assert(expected == Seq(1, 2)) + val result = + applyCometExecRule(GlobalLimitExec(3, createSparkPlan(spark, query), offset = 1)) + assert(result.executeCollect().map(_.getInt(0)).toSeq == expected) + val project = result.collectFirst { case p: CometProjectExec => p }.get + assert(project.child.asInstanceOf[CometLocalLimitExec].limit == 3) + assert(result.asInstanceOf[CometGlobalLimitExec].offset == 1) + } + } + test("collect limit preserves filter column pruning for column-only projections") { withSQLConf(SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false") { withParquetTable(Seq((0, "a", 1), (1, "b", 2), (2, "c", 3)), "limit_columns") { @@ -322,75 +373,111 @@ class CometExecRuleSuite extends CometTestBase { val serialized = OperatorOuterClass.Operator.parseFrom(project.serializedPlanOpt.plan.get) assert(serialized.getChildren(0).hasFilter) + Seq(LocalLimitExec(1, project), GlobalLimitExec(1, project)).foreach { limit => + val bounded = applyCometExecRule(limit).collectFirst { case p: CometProjectExec => + p + }.get + assert(bounded.child.isInstanceOf[CometFilterExec]) + assert(bounded.nativeOp.getChildren(0).hasFilter) + } } } } } - test("collect limit projection input cap respects disabled native local limits") { + test("projection input caps leave tail and offset-only limits unchanged") { withSQLConf( SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false", SQLConf.WHOLESTAGE_CODEGEN_ENABLED.key -> "false", - CometConf.COMET_CONVERT_FROM_RANGE_ENABLED.key -> "true", - CometConf.COMET_EXEC_LOCAL_LIMIT_ENABLED.key -> "false") { - val original = - createSparkPlan(spark, "SELECT id + 1 AS value FROM range(0, 10, 1, 1) LIMIT 2") - val result = applyCometExecRule(original) - assert(result.collect { case p: CometProjectExec => p }.size == 1) - assert(!result.exists(_.isInstanceOf[ProjectExec])) - assert(!result.exists(_.isInstanceOf[CometLocalLimitExec])) - assert(result.executeCollect().map(_.getLong(0)).toSeq == Seq(1L, 2L)) + CometConf.COMET_CONVERT_FROM_RANGE_ENABLED.key -> "true") { + val query = "SELECT id + 1 AS value FROM range(0, 5, 1, 1)" + val tail = applyCometExecRule(CollectTailExec(1, createSparkPlan(spark, query))) + assert(!tail.exists(_.isInstanceOf[CometLocalLimitExec])) + assert(tail.executeCollect().map(_.getLong(0)).toSeq == Seq(5L)) + + val offset = + applyCometExecRule(GlobalLimitExec(-1, createSparkPlan(spark, query), offset = 2)) + assert(!offset.exists(_.isInstanceOf[CometLocalLimitExec])) + assert(offset.executeCollect().map(_.getLong(0)).toSeq == Seq(3L, 4L, 5L)) } } - test("collect limit projection input cap preserves AQE links and repeated planning") { - withSQLConf( - SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false", - SQLConf.WHOLESTAGE_CODEGEN_ENABLED.key -> "false", - CometConf.COMET_CONVERT_FROM_RANGE_ENABLED.key -> "true") { - val original = createSparkPlan(spark, "SELECT id + 1 AS value FROM range(0, 10, 1, 1)") - val project = applyCometExecRule(original).asInstanceOf[CometProjectExec] - val stage = LogicalQueryStage(project.originalPlan.logicalLink.get, project) - val reused = spark.sessionState.planner.plan(stage).next() - assert(reused eq project) - - val result = applyCometExecRule(CollectLimitExec(2, reused)) - val bounded = result.collectFirst { case p: CometProjectExec => p }.get - assert(bounded.getTagValue(SparkPlan.LOGICAL_PLAN_TAG).exists(_ eq stage)) - assert(bounded.child.asInstanceOf[CometLocalLimitExec].limit == 2) - assert(bounded.nativeOp.getChildren(0).getLimit.getLimit == 2) - val serialized = OperatorOuterClass.Operator.parseFrom(bounded.serializedPlanOpt.plan.get) - assert(serialized.getChildren(0).getLimit.getLimit == 2) - - // AQE may wrap a previously converted projection in a newly planned collect limit. - val repeated = applyCometExecRule(CollectLimitExec(2, bounded)) - val again = repeated.collectFirst { case p: CometProjectExec => p }.get - assert(again.nativeOp eq bounded.nativeOp) - assert(again.getTagValue(SparkPlan.LOGICAL_PLAN_TAG).exists(_ eq stage)) - assert(repeated.collect { case cap: CometLocalLimitExec => cap }.size == 1) - } - } - - test("collect limit projection input cap does not cross a filter or sort") { - withSQLConf( - SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false", - SQLConf.WHOLESTAGE_CODEGEN_ENABLED.key -> "false", - CometConf.COMET_CONVERT_FROM_RANGE_ENABLED.key -> "true") { - Seq(false, true).foreach { sorted => + for ((name, makeLimit) <- Seq[(String, (Int, SparkPlan) => SparkPlan)]( + "collect" -> ((n, child) => CollectLimitExec(n, child)), + "local" -> ((n, child) => LocalLimitExec(n, child)), + "global" -> ((n, child) => GlobalLimitExec(n, child)))) { + test(s"$name limit projection input cap respects disabled native local limits") { + withSQLConf( + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false", + SQLConf.WHOLESTAGE_CODEGEN_ENABLED.key -> "false", + CometConf.COMET_CONVERT_FROM_RANGE_ENABLED.key -> "true", + CometConf.COMET_EXEC_LOCAL_LIMIT_ENABLED.key -> "false") { + val original = + makeLimit(2, createSparkPlan(spark, "SELECT id + 1 AS value FROM range(0, 10, 1, 1)")) + val result = applyCometExecRule(original) + assert(result.collect { case p: CometProjectExec => p }.size == 1) + assert(!result.exists(_.isInstanceOf[ProjectExec])) + assert(!result.exists(_.isInstanceOf[CometLocalLimitExec])) + assert(result.executeCollect().map(_.getLong(0)).toSeq == Seq(1L, 2L)) + } + } + + test(s"$name limit projection input cap preserves AQE links and repeated planning") { + withSQLConf( + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false", + SQLConf.WHOLESTAGE_CODEGEN_ENABLED.key -> "false", + CometConf.COMET_CONVERT_FROM_RANGE_ENABLED.key -> "true") { val original = createSparkPlan(spark, "SELECT id + 1 AS value FROM range(0, 10, 1, 1)") val project = applyCometExecRule(original).asInstanceOf[CometProjectExec] - val child = if (sorted) { - SortExec( - Seq(SortOrder(project.output.head, Descending)), - global = false, - child = project) - } else { - FilterExec(GreaterThan(project.output.head, Literal(2L)), project) + val stage = LogicalQueryStage(project.originalPlan.logicalLink.get, project) + val reused = spark.sessionState.planner.plan(stage).next() + assert(reused eq project) + + val result = applyCometExecRule(makeLimit(2, reused)) + val bounded = result.collectFirst { case p: CometProjectExec => p }.get + assert(bounded.getTagValue(SparkPlan.LOGICAL_PLAN_TAG).exists(_ eq stage)) + assert(bounded.child.asInstanceOf[CometLocalLimitExec].limit == 2) + assert(bounded.nativeOp.getChildren(0).getLimit.getLimit == 2) + val block = result.collectFirst { + case p: CometNativeExec if p.serializedPlanOpt.plan.isDefined => p + }.get + val serialized = OperatorOuterClass.Operator.parseFrom(block.serializedPlanOpt.plan.get) + val serializedProject = + if (serialized.hasProjection) serialized else serialized.getChildren(0) + assert(serializedProject.hasProjection) + assert(serializedProject.getChildren(0).getLimit.getLimit == 2) + + // AQE may wrap a previously converted projection in a newly planned limit. + val repeated = applyCometExecRule(makeLimit(2, bounded)) + val again = repeated.collectFirst { case p: CometProjectExec => p }.get + assert(again.nativeOp eq bounded.nativeOp) + assert(again.getTagValue(SparkPlan.LOGICAL_PLAN_TAG).exists(_ eq stage)) + assert(again.collect { case cap: CometLocalLimitExec => cap }.size == 1) + } + } + + test(s"$name limit projection input cap does not cross a filter or sort") { + withSQLConf( + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false", + SQLConf.WHOLESTAGE_CODEGEN_ENABLED.key -> "false", + CometConf.COMET_CONVERT_FROM_RANGE_ENABLED.key -> "true") { + Seq(false, true).foreach { sorted => + val original = createSparkPlan(spark, "SELECT id + 1 AS value FROM range(0, 10, 1, 1)") + val project = applyCometExecRule(original).asInstanceOf[CometProjectExec] + val child = if (sorted) { + SortExec( + Seq(SortOrder(project.output.head, Descending)), + global = false, + child = project) + } else { + FilterExec(GreaterThan(project.output.head, Literal(2L)), project) + } + val result = applyCometExecRule(makeLimit(1, child)) + val expectedLimits = if (name == "local") 1 else 0 + assert(result.collect { case cap: CometLocalLimitExec => cap }.size == expectedLimits) + val expected = if (sorted) 10L else 3L + assert(result.executeCollect().map(_.getLong(0)).toSeq == Seq(expected)) } - val result = applyCometExecRule(CollectLimitExec(1, child)) - assert(!result.exists(_.isInstanceOf[CometLocalLimitExec])) - val expected = if (sorted) 10L else 3L - assert(result.executeCollect().map(_.getLong(0)).toSeq == Seq(expected)) } } }