sunchao commented on code in PR #5533:
URL: https://github.com/apache/datafusion-comet/pull/5533#discussion_r3885125232
##########
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:
Fixed in 28074b67. A known partitioned `WindowGroupLimit` no longer starts
the fallback itself. It drains its groups when fully consumed, but an outer
`LIMIT` can still stop that input, so the outer protection continues through
it. Unknown shim extractions retain the conservative behavior.
I expanded the planner matrix across partitioned/unpartitioned input, all
three ranks, partial/final modes, outer limits, dispatcher settings, native
child reuse, and repeated planning. The new execution test keeps valid input
native, requires the consumed malformed suffix to throw, and checks that `LIMIT
1` can skip that suffix. These pass on Spark 3.5.9 and 4.1.3.
The updated description also records a separate existing empty-build join
difference that the earlier broad fallback happened to hide. The
partitioned-window change does not claim complete error-evaluation parity.
##########
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:
Updated both in 28074b67: the `unbase64` Notes cell in `expressions.md` and
the entry in `expression-audits/string_funcs.md` now describe the
operator-level fallback. The audit also distinguishes partitioned windows,
inherited outer-limit protection, and materialization barriers. I left the
generated Implementation cell unchanged.
##########
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:
Added executed coverage in 28074b67. The top-K tests use two ordered
partitions, assert there is no intervening sort or exchange, and exercise both
`collect()` and RDD collection with the decoder in a child projection, child
filter, or output projection.
The join tests disable both static and adaptive broadcast thresholds, force
`SHUFFLE_HASH` or `MERGE` for both semi and anti joins, and inspect the
executed plans to confirm the strategy and decoder residual. Both groups run
with AQE and JVM dispatch on/off, require errors when malformed input is
actually consumed, and include native-execution controls. All seven new
top-K/join tests pass on Spark 3.5.9 and 4.1.3.
--
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]
