andygrove commented on code in PR #5533:
URL: https://github.com/apache/datafusion-comet/pull/5533#discussion_r3884708892
##########
spark/src/main/scala/org/apache/comet/rules/CometExecRule.scala:
##########
@@ -714,6 +722,132 @@ case class CometExecRule(session: SparkSession)
}
}
+ /** Keep operators that can skip malformed unbase64 inputs in Spark's row
pipeline. */
+ private def preserveUnbase64EvaluationMasks(plan: SparkPlan): SparkPlan = {
+ val limitReason = "unbase64 requires Spark evaluation below LIMIT"
+ val joinReason = "unbase64 requires Spark evaluation in first-match join
conditions"
+
+ def containsUnbase64(expr: Expression): Boolean =
+ expr.exists(_.isInstanceOf[UnBase64])
Review Comment:
This is the only `unbase64`-specific thing in the whole rule. Everything
else in `preserveUnbase64EvaluationMasks` is a general mechanism, but the
underlying divergence is not specific to `unbase64` either. Any expression that
throws on data has it, under these same operators. `SELECT a + b FROM
ansi_int_overflow LIMIT 1` with ANSI on fails in Comet and succeeds in Spark
for exactly the same reason, and so do ANSI cast, decimal divide-by-zero, and
`element_at`. So we would be shipping a fair amount of planner machinery that
fixes one member of the family.
Could the trigger come from the serde instead of a hard-coded class check?
Something like a marker trait on `CometExpressionSerde` meaning "this
expression can raise a data-dependent error", with `CometUnBase64` mixing it
in, so `containsUnbase64` becomes a registry lookup. Then ANSI arithmetic can
opt in later without anyone touching `CometExecRule` again.
I want to flag the flip side honestly, because I think it decides the shape
of this PR. Once the trigger widens to ANSI `Add` or `Cast`, the fallback
swallows a large fraction of real queries. That makes me wonder whether the
right call is the one we have made for the other divergences in this family,
which is to document it under "Known result-value divergences" in
`docs/source/user-guide/latest/compatibility/index.md` and file the general
problem, rather than special-case it in the planner. I would rather we pick one
of those two directions deliberately than land the narrow version by default.
There is also a shape this approach cannot reach even for `unbase64` alone.
`WHERE k = 1 AND unbase64(bad) = X'616263'` short-circuits per row in Spark and
evaluates the whole batch natively in Comet, with no `LIMIT` involved. Same for
`CASE WHEN k = 1 THEN unbase64(bad) END`. Is that worth a tracking issue
alongside this, so we are not implicitly claiming the class is closed?
##########
spark/src/main/scala/org/apache/comet/serde/strings.scala:
##########
@@ -733,8 +735,12 @@ object CometUnBase64 extends
CometExpressionSerde[UnBase64] with CodegenDispatch
"unbase64 with a non-trivial child expression uses the JVM codegen
dispatcher to preserve" +
" Spark's short-circuit evaluation (native path is limited to column and
literal children)"
+ private val maskedEvaluationReason =
+ "unbase64 below LIMIT or in first-match join conditions falls back to
Spark so malformed" +
+ " input is not decoded on rows Spark skips"
+
override def getUnsupportedReasons(): Seq[String] =
- Seq(failOnErrorReason, nonTrivialChildReason)
+ Seq(failOnErrorReason, nonTrivialChildReason, maskedEvaluationReason)
Review Comment:
The wording here reads well, and I know the contract in
`adding_a_new_expression.md` allows these to differ from what `getSupportLevel`
returns, so I am not asking you to move it. Worth noting for anyone reading
later that this reason is never produced by the serde: the decision is made by
`CometExecRule` at the operator level.
The part I would like changed is the committed docs, which `GenerateDocs`
will not touch. `docs/source/user-guide/latest/expressions.md:618` still has an
empty Notes cell for `unbase64`, and the audit entry at
`docs/source/contributor-guide/expression-audits/string_funcs.md:248` still
lists `failOnError` and non-trivial children as the only restrictions. Could
you mirror this reason into both?
##########
spark/src/main/scala/org/apache/comet/rules/CometExecRule.scala:
##########
@@ -714,6 +722,132 @@ case class CometExecRule(session: SparkSession)
}
}
+ /** Keep operators that can skip malformed unbase64 inputs in Spark's row
pipeline. */
+ private def preserveUnbase64EvaluationMasks(plan: SparkPlan): SparkPlan = {
+ val limitReason = "unbase64 requires Spark evaluation below LIMIT"
+ val joinReason = "unbase64 requires Spark evaluation in first-match join
conditions"
+
+ def containsUnbase64(expr: Expression): Boolean =
+ expr.exists(_.isInstanceOf[UnBase64])
+
+ def firstMatch(joinType: JoinType): Boolean = joinType match {
+ case LeftSemi | LeftAnti => true
+ case _ => false
+ }
+
+ // A reused native Final may need to fall back along with its incompatible
Partial buffer.
+ // Restore only that buffer-producing chain so tagUnsafePartialAggregates
can protect it.
+ // Do not cross a materialized query stage or descend below the input of a
pure Partial.
+ def restoreNativeAggregateBuffers(node: SparkPlan): Option[SparkPlan] = {
+ val original = node match {
+ case agg: CometHashAggregateExec => agg.originalPlan
+ case shuffle: CometShuffleExchangeExec => shuffle.originalPlan
+ case _ => node
+ }
+ def restore(children: Seq[SparkPlan]): SparkPlan = {
+ val restored = original.withNewChildren(children)
+
node.getTagValue(SparkPlan.LOGICAL_PLAN_TAG).foreach(restored.setLogicalLink)
+ restored
+ }
+ original match {
+ case agg: BaseAggregateExec
+ if agg.aggregateExpressions.nonEmpty &&
+ agg.aggregateExpressions.forall(_.mode == Partial) =>
+ if (original ne node) Some(restore(node.children)) else None
+ case agg: BaseAggregateExec
+ if agg.aggregateExpressions.forall(e =>
+ e.mode == Partial || e.mode == PartialMerge) =>
+ restoreNativeAggregateBuffers(node.children.head).map(child =>
restore(Seq(child)))
+ case _: ShuffleExchangeLike =>
+ restoreNativeAggregateBuffers(node.children.head).map(child =>
restore(Seq(child)))
+ case _ => None
+ }
+ }
+
+ def protect(node: SparkPlan, belowLimit: Boolean): (SparkPlan,
Option[String]) = {
+ val original = node match {
+ case scan: CometScanExec =>
+ scan.wrapped
+ .copy(partitionFilters = scan.partitionFilters, dataFilters =
scan.dataFilters)
+ case comet: CometExec => comet.originalPlan
+ case shuffle: CometShuffleExchangeExec => shuffle.originalPlan
+ case broadcast: CometBroadcastExchangeExec => broadcast.originalPlan
+ case _ => node
+ }
+ val startsLimit = original match {
+ case _: CollectLimitExec | _: LocalLimitExec | _: GlobalLimitExec =>
true
+ case topK: TakeOrderedAndProjectExec =>
+ SortOrder.orderingSatisfies(node.children.head.outputOrdering,
topK.sortOrder)
+ case windowLimit
+ if ShimCometWindowGroupLimit.windowGroupLimitClass.exists(
+ _.isInstance(windowLimit)) =>
+ SortOrder.orderingSatisfies(
+ node.children.head.outputOrdering,
+ windowLimit.requiredChildOrdering.head)
+ case _ => false
+ }
+ // These operators consume their input before yielding rows. Still visit
their children:
+ // an inner LocalLimit below an exchange must establish its own
evaluation boundary.
+ val materializesInput = original match {
+ case _: SortExec | _: HashAggregateExec | _: ObjectHashAggregateExec |
+ _: ShuffleExchangeLike | _: BroadcastExchangeLike | _:
QueryStageExec |
+ _: ReusedExchangeExec =>
+ true
+ case _ => false
+ }
+ val protectedChildren = node.children.map { child =>
+ protect(child, startsLimit || (belowLimit && !materializesInput))
+ }
+ val childReason = protectedChildren.flatMap(_._2).headOption
+ val condition = original match {
+ case join: HashJoin if firstMatch(join.joinType) => join.condition
+ case join: SortMergeJoinExec if firstMatch(join.joinType) =>
join.condition
+ case join: BroadcastNestedLoopJoinExec if firstMatch(join.joinType) =>
join.condition
+ case _ => None
+ }
+ val ownReason = if (belowLimit &&
original.expressions.exists(containsUnbase64)) {
+ Some(limitReason)
+ } else if (condition.exists(containsUnbase64)) {
+ Some(joinReason)
+ } else {
+ node.getTagValue(CometExecRule.UNSAFE_UNBASE64_EVALUATION)
+ }
+ // Exchanges can restart native execution after consuming a Spark row
pipeline.
+ val restartsNative = original.isInstanceOf[ShuffleExchangeLike] ||
+ original.isInstanceOf[BroadcastExchangeLike]
+ val reason = ownReason.orElse(if (restartsNative) None else childReason)
+ val children = protectedChildren.map(_._1).map { child =>
+ original match {
+ case agg: BaseAggregateExec
+ if reason.isDefined &&
agg.aggregateExpressions.map(_.mode).distinct == Seq(
+ Final) &&
+
!QueryPlanSerde.allAggsSupportMixedExecution(agg.aggregateExpressions) =>
+ restoreNativeAggregateBuffers(child).getOrElse(child)
+ case _ => child
+ }
+ }
+ val prepared = node match {
+ // Do not refill a batch between the row decoder and its
short-circuiting consumer.
+ case _: RowToColumnarExec | _: CometSparkToColumnarExec if
childReason.isDefined =>
+ children.head
+ case _: ColumnarToRowExec | _: CometColumnarToRowExec | _:
CometNativeColumnarToRowExec
+ if childReason.isDefined && !children.head.supportsColumnar =>
+ children.head
+ case _ if (original ne node) && (reason.isDefined || children !=
node.children) =>
+ // AQE can reuse an existing native subtree. Rebuild affected
ancestors as well so
+ // their serialized native plans do not retain the decoder that just
fell back.
+ val restored = original.withNewChildren(children)
+
node.getTagValue(SparkPlan.LOGICAL_PLAN_TAG).foreach(restored.setLogicalLink)
+ restored
+ case _ => node.withNewChildren(children)
+ }
+
reason.foreach(prepared.setTagValue(CometExecRule.UNSAFE_UNBASE64_EVALUATION,
_))
+ (prepared, reason)
+ }
+
+ protect(plan, belowLimit = false)._1
Review Comment:
This walks every operator's expression tree on every plan on every rule
application, and for the overwhelming majority of queries there is no
`UnBase64` to find. Given #5199 is open on planner rule cost, could we bail out
up front with something like `if
(!plan.exists(_.expressions.exists(containsUnbase64))) return plan`? The
sticky-tag behaviour stays correct, because a node that carries the tag still
contains the decoder that earned it.
##########
spark/src/main/scala/org/apache/comet/rules/CometExecRule.scala:
##########
@@ -714,6 +722,132 @@ case class CometExecRule(session: SparkSession)
}
}
+ /** Keep operators that can skip malformed unbase64 inputs in Spark's row
pipeline. */
+ private def preserveUnbase64EvaluationMasks(plan: SparkPlan): SparkPlan = {
+ val limitReason = "unbase64 requires Spark evaluation below LIMIT"
+ val joinReason = "unbase64 requires Spark evaluation in first-match join
conditions"
+
+ def containsUnbase64(expr: Expression): Boolean =
+ expr.exists(_.isInstanceOf[UnBase64])
+
+ def firstMatch(joinType: JoinType): Boolean = joinType match {
+ case LeftSemi | LeftAnti => true
+ case _ => false
+ }
+
+ // A reused native Final may need to fall back along with its incompatible
Partial buffer.
+ // Restore only that buffer-producing chain so tagUnsafePartialAggregates
can protect it.
+ // Do not cross a materialized query stage or descend below the input of a
pure Partial.
+ def restoreNativeAggregateBuffers(node: SparkPlan): Option[SparkPlan] = {
+ val original = node match {
+ case agg: CometHashAggregateExec => agg.originalPlan
+ case shuffle: CometShuffleExchangeExec => shuffle.originalPlan
+ case _ => node
+ }
+ def restore(children: Seq[SparkPlan]): SparkPlan = {
+ val restored = original.withNewChildren(children)
+
node.getTagValue(SparkPlan.LOGICAL_PLAN_TAG).foreach(restored.setLogicalLink)
+ restored
+ }
+ original match {
+ case agg: BaseAggregateExec
+ if agg.aggregateExpressions.nonEmpty &&
+ agg.aggregateExpressions.forall(_.mode == Partial) =>
+ if (original ne node) Some(restore(node.children)) else None
+ case agg: BaseAggregateExec
+ if agg.aggregateExpressions.forall(e =>
+ e.mode == Partial || e.mode == PartialMerge) =>
+ restoreNativeAggregateBuffers(node.children.head).map(child =>
restore(Seq(child)))
+ case _: ShuffleExchangeLike =>
+ restoreNativeAggregateBuffers(node.children.head).map(child =>
restore(Seq(child)))
+ case _ => None
+ }
+ }
+
+ def protect(node: SparkPlan, belowLimit: Boolean): (SparkPlan,
Option[String]) = {
+ val original = node match {
+ case scan: CometScanExec =>
+ scan.wrapped
+ .copy(partitionFilters = scan.partitionFilters, dataFilters =
scan.dataFilters)
+ case comet: CometExec => comet.originalPlan
+ case shuffle: CometShuffleExchangeExec => shuffle.originalPlan
+ case broadcast: CometBroadcastExchangeExec => broadcast.originalPlan
+ case _ => node
+ }
+ val startsLimit = original match {
+ case _: CollectLimitExec | _: LocalLimitExec | _: GlobalLimitExec =>
true
+ case topK: TakeOrderedAndProjectExec =>
+ SortOrder.orderingSatisfies(node.children.head.outputOrdering,
topK.sortOrder)
+ case windowLimit
+ if ShimCometWindowGroupLimit.windowGroupLimitClass.exists(
+ _.isInstance(windowLimit)) =>
+ SortOrder.orderingSatisfies(
+ node.children.head.outputOrdering,
+ windowLimit.requiredChildOrdering.head)
Review Comment:
This branch looks like it only holds without a `PARTITION BY`.
`WindowGroupLimitEvaluatorFactory` picks the plain limit iterator when
`partitionSpec.isEmpty`, and that one stops pulling once `rank >= limit`. With
a partition spec it uses `GroupedLimitIterator`, whose `skipRemainingRows()`
calls `fetchNextRow()` in a loop, so Spark pulls and evaluates every input row
anyway. In that case we fall back and gain nothing.
Since `EnsureRequirements` guarantees the child satisfies
`requiredChildOrdering`, this check is effectively always true, so the
partitioned case is reachable whenever the child happens to be sorted already
and no `SortExec` gets inserted to stop the propagation. Should this be
restricted to `partitionSpec.isEmpty`? The new fixture only covers the
unpartitioned shape, so a partitioned case would be worth adding either way.
##########
spark/src/main/scala/org/apache/comet/rules/CometExecRule.scala:
##########
@@ -714,6 +722,132 @@ case class CometExecRule(session: SparkSession)
}
}
+ /** Keep operators that can skip malformed unbase64 inputs in Spark's row
pipeline. */
+ private def preserveUnbase64EvaluationMasks(plan: SparkPlan): SparkPlan = {
Review Comment:
Related to my other comment, this changes performance for queries that were
never broken. The reason propagates to every ancestor up to the nearest
exchange, so any projection with `unbase64` under a `LIMIT` leaves native
execution for the whole chain, and `CollectLimitExec` shows up on every
`df.show()` and every BI-tool `LIMIT`. That gives back the 2.4x to 3.5x the
native kernel bought in #5451, for data that is almost always well-formed, and
there is no way to turn it off. Could we gate this behind a `CometConf` key so
a user who knows their base64 is clean can keep the native path?
##########
spark/src/test/resources/sql-tests/expressions/misc/unbase64_operator_masks.sql:
##########
@@ -0,0 +1,118 @@
+-- Licensed to the Apache Software Foundation (ASF) under one
+-- or more contributor license agreements. See the NOTICE file
+-- distributed with this work for additional information
+-- regarding copyright ownership. The ASF licenses this file
+-- to you under the Apache License, Version 2.0 (the
+-- "License"); you may not use this file except in compliance
+-- with the License. You may obtain a copy of the License at
+--
+-- http://www.apache.org/licenses/LICENSE-2.0
+--
+-- Unless required by applicable law or agreed to in writing,
+-- software distributed under the License is distributed on an
+-- "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+-- KIND, either express or implied. See the License for the
+-- specific language governing permissions and limitations
+-- under the License.
+
+-- LIMIT and first-match joins can skip malformed input even when their
expressions would
+-- throw if evaluated. Native batch evaluation and the JVM dispatcher must
preserve those
+-- operator boundaries, including after AQE replans a join. Malformed terminal
Base64 units
+-- throw regardless of ANSI mode, and disabling JVM dispatch still permits
native decoding.
+-- ConfigMatrix: spark.comet.exec.scalaUDF.codegen.enabled=false,true
+-- ConfigMatrix: spark.sql.adaptive.enabled=false,true
+-- ConfigMatrix: spark.sql.ansi.enabled=false,true
+-- Config: spark.sql.shuffle.partitions=1
+
+-- Parquet keeps the decoder inputs as attributes. A single file fixes the
physical row order
+-- so LIMIT consumes the valid first row and never reaches the malformed
second row.
+statement
+CREATE TABLE test_unbase64_operator_limit USING parquet AS
+SELECT /*+ COALESCE(1) */ bad FROM VALUES ('YWJj'), ('A') AS t(bad)
+
+query expect_fallback(unbase64 requires Spark evaluation below LIMIT)
+SELECT hex(unbase64(bad)) FROM test_unbase64_operator_limit LIMIT 1
+
+query expect_fallback(unbase64 requires Spark evaluation below LIMIT)
+SELECT bad FROM test_unbase64_operator_limit
+WHERE unbase64(bad) <=> X'616263' LIMIT 1
+
+-- When enabled, the JVM dispatcher evaluates this compound child over whole
batches too.
+query expect_fallback(unbase64 requires Spark evaluation below LIMIT)
+SELECT hex(unbase64(concat(bad, ''))) FROM test_unbase64_operator_limit LIMIT 1
+
+-- Strict decoding also uses UnBase64 through the JVM dispatcher when enabled.
+query expect_fallback(unbase64 requires Spark evaluation below LIMIT)
+SELECT hex(to_binary(bad, 'base64')) FROM test_unbase64_operator_limit LIMIT 1
+
+-- Removing the mask, or making the malformed row the consumed row, must still
throw.
+query expect_error(Last unit does not have enough valid bits)
+SELECT hex(unbase64(bad)) FROM test_unbase64_operator_limit
+
+query expect_error(Last unit does not have enough valid bits)
+SELECT hex(unbase64(bad)) FROM test_unbase64_operator_limit WHERE bad = 'A'
LIMIT 1
+
+statement
+CREATE TABLE test_unbase64_operator_left USING parquet AS
+SELECT 1 AS k, X'616262' AS expected
+
+-- Spark visits the last inserted build row first. Its successful match must
prevent the
+-- malformed candidate from being decoded, for both semi and anti joins.
+statement
+CREATE TABLE test_unbase64_operator_right USING parquet AS
+SELECT /*+ COALESCE(1) */ k, bad FROM VALUES (1, 'A'), (1, 'YWJj') AS t(k, bad)
+
+query expect_fallback(unbase64 requires Spark evaluation in first-match join
conditions)
+SELECT /*+ BROADCAST(r) */ l.* FROM test_unbase64_operator_left l
+LEFT SEMI JOIN test_unbase64_operator_right r
+ON l.k = r.k AND unbase64(r.bad) > l.expected
+
+query expect_fallback(unbase64 requires Spark evaluation in first-match join
conditions)
+SELECT /*+ BROADCAST(r) */ l.* FROM test_unbase64_operator_left l
+LEFT ANTI JOIN test_unbase64_operator_right r
+ON l.k = r.k AND unbase64(r.bad) > l.expected
+
+-- Reversing the build rows removes the first-match mask and exposes the
malformed input.
+statement
+CREATE TABLE test_unbase64_operator_reverse USING parquet AS
+SELECT /*+ COALESCE(1) */ k, bad FROM VALUES (1, 'YWJj'), (1, 'A') AS t(k, bad)
+
+query expect_error(Last unit does not have enough valid bits)
+SELECT /*+ BROADCAST(r) */ l.* FROM test_unbase64_operator_left l
+LEFT SEMI JOIN test_unbase64_operator_reverse r
+ON l.k = r.k AND unbase64(r.bad) > l.expected
+
+query expect_error(Last unit does not have enough valid bits)
+SELECT /*+ BROADCAST(r) */ l.* FROM test_unbase64_operator_left l
+LEFT ANTI JOIN test_unbase64_operator_reverse r
+ON l.k = r.k AND unbase64(r.bad) > l.expected
+
+-- Safe projections and inner joins must retain native decoding, including
nullable input.
+statement
+CREATE TABLE test_unbase64_operator_valid USING parquet AS
+SELECT /*+ COALESCE(1) */ k, bad FROM VALUES (1, 'YWJj'), (1, 'YWFh'), (1,
NULL) AS t(k, bad)
+
+-- Falling back for a decoded group key must also keep the partial
collect_list in Spark,
+-- because its intermediate buffer cannot be shared with a native partial
aggregate.
+-- Select one valid group so LIMIT cannot choose different groups in Spark and
Comet.
+query expect_fallback(unbase64 requires Spark evaluation below LIMIT)
+SELECT hex(unbase64(bad)), sort_array(collect_list(k))
+FROM test_unbase64_operator_valid WHERE bad = 'YWJj'
+GROUP BY bad LIMIT 1
+
+query
+SELECT hex(unbase64(bad)) FROM test_unbase64_operator_valid
+
+query
+SELECT /*+ BROADCAST(r) */ l.* FROM test_unbase64_operator_left l
+INNER JOIN test_unbase64_operator_valid r
+ON l.k = r.k AND unbase64(r.bad) > l.expected
+
+-- Semi/anti joins without a throwing decoder remain native.
+query
+SELECT /*+ BROADCAST(r) */ l.* FROM test_unbase64_operator_left l
+LEFT SEMI JOIN test_unbase64_operator_right r ON l.k = r.k
+
+query
+SELECT /*+ BROADCAST(r) */ l.* FROM test_unbase64_operator_left l
+LEFT ANTI JOIN test_unbase64_operator_right r ON l.k = r.k
Review Comment:
The fixtures cover the limit and broadcast-join shapes nicely, and I like
that the native-execution controls are here rather than only in the rule suite.
Two shapes only exist as hand-built plans in `CometExecRuleSuite` and never
run end to end. One is `TakeOrderedAndProject` with pre-satisfied child
ordering, which is the branch whose guard depends on what `outputOrdering`
actually reports after conversion. The other is the shuffled-hash and
sort-merge semi/anti strategies, since this file only exercises broadcast.
Would you mind adding an `ORDER BY ... LIMIT n` case and a case with
`spark.sql.autoBroadcastJoinThreshold=-1`? Those are the two places where a
synthetic plan and a real one are most likely to disagree.
--
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]
