sunchao commented on code in PR #5533: URL: https://github.com/apache/datafusion-comet/pull/5533#discussion_r4178667039
########## spark/src/test/scala/org/apache/comet/serde/CometEvaluationMaskSuite.scala: ########## @@ -0,0 +1,648 @@ +/* + * 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. + */ + +package org.apache.comet.serde + +import scala.reflect.ClassTag + +import org.apache.spark.sql.{CometTestBase, DataFrame, Row} +import org.apache.spark.sql.catalyst.expressions._ +import org.apache.spark.sql.catalyst.optimizer.BuildRight +import org.apache.spark.sql.catalyst.plans.logical.LocalRelation +import org.apache.spark.sql.catalyst.plans.physical.SinglePartition +import org.apache.spark.sql.comet._ +import org.apache.spark.sql.execution._ +import org.apache.spark.sql.execution.adaptive.{AdaptiveSparkPlanExec, LogicalQueryStage} +import org.apache.spark.sql.execution.aggregate.ObjectHashAggregateExec +import org.apache.spark.sql.execution.exchange.{BroadcastExchangeLike, ShuffleExchangeExec, ShuffleExchangeLike} +import org.apache.spark.sql.execution.joins._ +import org.apache.spark.sql.internal.SQLConf +import org.apache.spark.sql.types.{IntegerType, StringType} + +import org.apache.comet.CometConf +import org.apache.comet.CometSparkSessionExtensions.isSpark35Plus +import org.apache.comet.rules.CometExecRule +import org.apache.comet.shims.ShimCometWindowGroupLimit + +class CometEvaluationMaskSuite extends CometTestBase { + private val limitReason = "unbase64 requires Spark evaluation below LIMIT" + private val joinReason = "unbase64 requires Spark evaluation in first-match join conditions" + private val decodeMessage = "Last unit does not have enough valid bits" + + private def withModes(f: (Boolean, Boolean) => Unit): Unit = { + for (aqe <- Seq(false, true); dispatch <- Seq(false, true)) { + withSQLConf( + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> aqe.toString, + CometConf.COMET_SCALA_UDF_CODEGEN_ENABLED.key -> dispatch.toString) { + withClue(s"AQE=$aqe, dispatcher=$dispatch: ")(f(aqe, dispatch)) + } + } + } + + private def withInputs(inputs: (String, String)*)(f: => Unit): Unit = { + withTempPath { dir => + withTempView(inputs.map(_._1): _*) { + withSQLConf(CometConf.COMET_ENABLED.key -> "false") { + inputs.foreach { case (name, query) => + val path = s"${dir.getCanonicalPath}/$name" + sql(query).coalesce(1).write.parquet(path) + val input = spark.read.parquet(path) + assert(input.inputFiles.length == 1) + input.createOrReplaceTempView(name) + } + } + f + } + } + } + + private def sparkPlan(query: String): SparkPlan = { + var plan: SparkPlan = null + withSQLConf( + CometConf.COMET_ENABLED.key -> "false", + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false") { + plan = sql(query).queryExecution.executedPlan + } + plan + } + + private def applyRule(plan: SparkPlan): SparkPlan = + CometExecRule(spark).apply(stripAQEPlan(plan)) + private def count[T <: SparkPlan: ClassTag](plan: SparkPlan): Int = + collect(plan) { case node: T => node }.size + private def aggregateCounts(plan: SparkPlan): (Int, Int) = + (count[CometHashAggregateExec](plan), count[ObjectHashAggregateExec](plan)) + private def nativeScans(plan: SparkPlan): Int = collect(plan) { + case _: CometScanExec | _: CometNativeScanExec => true + }.size + private def original(plan: SparkPlan): SparkPlan = plan match { + case comet: CometExec => comet.originalPlan + case other => other + } + private def hasDecoder(expr: Expression): Boolean = expr.exists(_.isInstanceOf[UnBase64]) + private def malformed(error: Throwable): Boolean = + causeChain(error).exists(e => Option(e.getMessage).exists(_.contains(decodeMessage))) + private def decodeError(body: => Any): Unit = assert(malformed(intercept[Exception](body))) + private def decodeErrors(query: String): Unit = { + val (sparkError, cometError) = checkSparkAnswerMaybeThrows(sql(query)) + Seq(sparkError, cometError).foreach(error => assert(error.exists(malformed), query)) + } + + test("evaluation-mask protection can be disabled for known-valid input") { + val preserve = CometConf.COMET_EXEC_PRESERVE_EVALUATION_MASKS_ENABLED + assert(preserve.defaultValue.contains(true)) + withSQLConf( + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "true", + SQLConf.SHUFFLE_PARTITIONS.key -> "1", + CometConf.COMET_SCALA_UDF_CODEGEN_ENABLED.key -> "true") { + withInputs( + "opt_out_valid" -> "SELECT * FROM VALUES (1, 'YWJj'), (1, 'YWFh') AS t(k, bad)", + "opt_out_left" -> "SELECT * FROM VALUES (1, X'616262'), (2, X'616262') AS t(k, expected)") { + for (enabled <- Seq(true, false)) { + withSQLConf(preserve.key -> enabled.toString) { + val projection = "SELECT hex(unbase64(bad)) FROM opt_out_valid LIMIT 1" + val compound = "SELECT hex(unbase64(concat(bad, ''))) FROM opt_out_valid LIMIT 1" + if (enabled) { + checkSparkAnswerAndFallbackReason(projection, limitReason) + checkSparkAnswerAndFallbackReason(compound, limitReason) + } else { + checkSparkAnswerAndImpl(sql(projection), native = Seq("unbase64")) + checkSparkAnswerAndImpl(sql(compound), dispatched = Seq("unbase64")) + } + for (kind <- Seq("SEMI", "ANTI")) { + val query = s"""SELECT /*+ BROADCAST(r) */ l.k + |FROM opt_out_left l LEFT $kind JOIN opt_out_valid r + |ON l.k = r.k AND unbase64(r.bad) > l.expected""".stripMargin + if (enabled) { + checkSparkAnswerAndFallbackReason(query, joinReason) + } else { + val (_, plan) = checkSparkAnswerAndImpl(sql(query), native = Seq("unbase64")) + assert(count[CometBroadcastHashJoinExec](plan) == 1, plan.toString) + } + } + // AQE can remove the sort after Partial materializes. The opt-out keeps both + // aggregate halves native instead of preemptively restoring Spark's buffers. + withSQLConf( + SQLConf.WHOLESTAGE_CODEGEN_ENABLED.key -> "false", + CometConf.COMET_EXEC_LOCAL_TABLE_SCAN_ENABLED.key -> "true") { + val aggregate = sparkPlan( + "SELECT unbase64(bad) AS decoded, collect_list(k) FROM " + + "VALUES (1, 'YWJj'), (1, 'YWFh') AS t(k, bad) GROUP BY bad") + assert(aggregateCounts(aggregate) == (0, 2)) + val sorted = SortExec( + Seq(SortOrder(aggregate.output.head, Ascending)), + global = false, + aggregate) + val plan = applyRule(CollectLimitExec(1, sorted)) + assert(aggregateCounts(plan) == (if (enabled) (0, 2) else (2, 0)), plan.toString) + } + } + } + } + } + } + + test("LIMIT masks follow the decoder policy with ANSI, strict and compound inputs") { + assert(QueryPlanSerde.exprSerdeMap.collect { case (cls, _: RequiresSparkEvaluationMask[_]) => + cls + }.toSet == Set(classOf[UnBase64])) + val input = AttributeReference("encoded", StringType)() + withInputs( + "masked" -> "SELECT * FROM VALUES (1, 'YWJj'), (2, 'A') AS t(k, bad)", + "valid" -> "SELECT * FROM VALUES (1, 'YWJj'), (1, 'YWFh'), (1, NULL) AS t(k, bad)", + "inner_left" -> "SELECT 1 AS k, X'616262' AS expected", + "one_group" -> "SELECT 1 AS k, 'YWJj' AS bad") { + withModes { (_, _) => + for (ansi <- Seq(false, true)) { + withSQLConf(SQLConf.ANSI_ENABLED.key -> ansi.toString) { + for (child <- Seq(input, Literal("YWJj"), Concat(Seq(input, Literal("")))); + strict <- Seq(false, true)) { + val expressions = Seq(UnBase64(child, strict), new LimitUnBase64(child, strict)) + assert( + expressions.forall(QueryPlanSerde.evaluationMaskName(_).contains("unbase64"))) + } + // Existing unbase64/to_binary SQL tests exercise native and dispatched serialization. + // Other throwing expressions: https://github.com/apache/datafusion-comet/issues/6006. + Seq( + Add(Literal(Int.MaxValue), Literal(1)), + Cast(input, IntegerType), + Concat(Seq(input, Literal(""))), + FormatString(Literal("%s"), input), + Like(input, Literal("%"), escapeChar = '\\')).foreach { expr => + assert(QueryPlanSerde.evaluationMaskName(expr).isEmpty) + } + Seq( + "SELECT hex(unbase64(bad)) FROM masked LIMIT 1", + "SELECT bad FROM masked WHERE unbase64(bad) <=> X'616263' LIMIT 1", + "SELECT hex(unbase64(concat(bad, ''))) FROM masked LIMIT 1", + "SELECT hex(to_binary(bad, 'base64')) FROM masked LIMIT 1").foreach { query => + checkSparkAnswerAndFallbackReason(query, limitReason) + } + decodeErrors("SELECT hex(unbase64(bad)) FROM masked") + decodeErrors("SELECT hex(unbase64(bad)) FROM masked WHERE bad = 'A' LIMIT 1") + checkSparkAnswerAndOperator(sql("SELECT hex(unbase64(bad)) FROM valid")) + checkSparkAnswerAndOperator( + sql("SELECT /*+ BROADCAST(r) */ l.* FROM inner_left l " + + "INNER JOIN valid r ON l.k = r.k AND unbase64(r.bad) > l.expected")) + // AQE can remove this sort after Partial materializes; offset-only collection stays native. + checkSparkAnswerAndOperator( + sql("SELECT unbase64(bad) AS decoded, collect_list(k) FROM one_group " + + "GROUP BY bad ORDER BY decoded OFFSET 1")) + } + } + } + } + } + + test("whole-tree JVM dispatch preserves an unregistered decoder subclass below LIMIT") { + val name = "comet_test_unbase64" + spark.sessionState.functionRegistry.createOrReplaceTempFunction( + name, + children => new LimitUnBase64(children.head), + "scala_udf") + try { + withSQLConf( + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false", + CometConf.COMET_SCALA_UDF_CODEGEN_ENABLED.key -> "true") { + withInputs("subclass_input" -> "SELECT * FROM VALUES ('YQ=='), ('A') AS t(encoded)") { + val query = s"SELECT format_string('%d', length($name(encoded))) FROM subclass_input" + val native = sql(query).queryExecution.executedPlan.collect { + case p: CometProjectExec => p + } + assert( + native.size == 1 && native.head.nativeOp.getProjection + .getProjectList(0) + .hasJvmScalarUdf) + val df = sql(s"$query LIMIT 1") + val (_, plan) = checkSparkAnswerAndFallbackReason(df, limitReason) + assert(count[ProjectExec](plan) == 1 && count[CometProjectExec](plan) == 0) + assert(nativeScans(plan) == 1) + checkAnswer(df, Seq(Row("1"))) + decodeErrors(s"$query LIMIT 2") + } + } + } finally { + spark.sessionState.catalog.dropTempFunction(name, ignoreIfNotExists = true) + } + } + + test("unbase64 LIMIT masks survive physical boundaries and native reuse") { + withSQLConf( + SQLConf.WHOLESTAGE_CODEGEN_ENABLED.key -> "false", + CometConf.COMET_EXEC_LOCAL_TABLE_SCAN_ENABLED.key -> "true") { + withModes { (_, _) => + val input = sparkPlan("SELECT * FROM VALUES (1, 'YWJj'), (2, 'A') AS t(id, bad)") + def project = ProjectExec( + Seq(input.output.head, Alias(UnBase64(input.output(1)), "decoded")()), + input) + val staged = applyRule( + GlobalLimitExec(1, ShuffleExchangeExec(SinglePartition, LocalLimitExec(1, project)))) + assert(count[CometGlobalLimitExec](staged) == 1) + assert(count[CometLocalTableScanExec](staged) == 1) + val exchange = collect(staged) { case s: ShuffleExchangeLike => s }.head + for (plan <- Seq(applyRule(staged), applyRule(exchange))) { + assert(count[LocalLimitExec](plan) == 1) + assert(count[ProjectExec](plan) == 1) + assert(count[CometProjectExec](plan) == 0) + } + val barriers: Seq[SparkPlan => SparkPlan] = Seq( + child => SortExec(Seq(SortOrder(child.output.head, Ascending)), global = false, child), + child => ShuffleExchangeExec(SinglePartition, child)) + barriers.foreach { barrier => + val blocking = barrier(project) + // Materialization is safe for both fresh plans and native subtrees reused by AQE. + for (child <- Seq(blocking, applyRule(blocking))) { + val plan = applyRule(LocalLimitExec(1, child)) + withClue(blocking.nodeName) { + assert(count[CometProjectExec](plan) == 1) + assert(count[ProjectExec](plan) == 0) + } + } + } + val bridges: Seq[SparkPlan => SparkPlan] = Seq( + child => child, + child => RowToColumnarExec(ColumnarToRowExec(child)), + child => CometSparkToColumnarExec(CometColumnarToRowExec(child)), + child => CometSparkToColumnarExec(CometNativeColumnarToRowExec(child))) + bridges.foreach { bridge => + val decoded = project + val native = applyRule(ProjectExec(decoded.output, decoded)) + assert(count[CometProjectExec](native) == 2) + val logicalStage = LogicalQueryStage(LocalRelation(native.output), native) + native.setLogicalLink(logicalStage) + val transformed = applyRule(GlobalLimitExec(1, bridge(native))) + for (plan <- Seq(transformed, applyRule(transformed))) { + assert(count[ProjectExec](plan) == 2) + assert(count[CometProjectExec](plan) == 0) + assert(collect(plan) { case r: RowToColumnarTransition => r }.isEmpty) + val link = plan.children.head.getTagValue(SparkPlan.LOGICAL_PLAN_TAG) + assert(link.exists(_ eq logicalStage)) + } + } + } + } + } + + test("unbase64 chooses compatible aggregate buffers before AQE materialization") { + withSQLConf( + SQLConf.WHOLESTAGE_CODEGEN_ENABLED.key -> "false", + CometConf.COMET_EXEC_LOCAL_TABLE_SCAN_ENABLED.key -> "true") { + withModes { (aqe, _) => + for { + decoderInResult <- Seq(false, true) + boundary <- Seq("direct", "sort", "shuffle", "top-k", "offset") + } { + val result = if (decoderInResult) { + "unbase64(bad) AS decoded, collect_list(k) AS collected" + } else { + "bad, collect_list(unbase64(bad)) AS collected" + } + val aggregate = sparkPlan( + s"SELECT $result FROM " + + "VALUES (1, 'YWJj'), (2, 'YWJj') AS t(k, bad) GROUP BY bad") + assert(count[ObjectHashAggregateExec](aggregate) == 2) + val native = applyRule(aggregate) + assert(count[CometHashAggregateExec](native) == 2) + for (child <- Seq(aggregate, native)) { + val order = Seq(SortOrder(child.output.head, Ascending)) + val plan = boundary match { + case "direct" => CollectLimitExec(1, child) + case "sort" => CollectLimitExec(1, SortExec(order, global = false, child)) + case "shuffle" => CollectLimitExec(1, ShuffleExchangeExec(SinglePartition, child)) + case "top-k" => TakeOrderedAndProjectExec(1, order, child.output, child) + case "offset" => + CollectLimitExec(-1, SortExec(order, global = false, child), offset = 1) + } + // Final does not reevaluate decoder inputs, and offset-only never stops early. + val sparkBuffers = decoderInResult && + (boundary == "direct" || (aqe && boundary != "offset")) + val transformed = applyRule(plan) + withClue(s"decoderInResult=$decoderInResult, boundary=$boundary: ") { + for (replanned <- Seq(transformed, applyRule(transformed))) { + assert(aggregateCounts(replanned) == (if (sparkBuffers) (0, 2) else (2, 0))) + } + } + } + } + } + } + } + + test("unbase64 keeps aggregate buffers compatible through an actual AQE join change") { + withSQLConf( + SQLConf.WHOLESTAGE_CODEGEN_ENABLED.key -> "true", + SQLConf.SHUFFLE_PARTITIONS.key -> "1", + SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1", + "spark.sql.adaptive.autoBroadcastJoinThreshold" -> "10485760", + "spark.sql.join.preferSortMergeJoin" -> "true", + "spark.sql.adaptive.maxShuffledHashJoinLocalMapThreshold" -> "0", + CometConf.COMET_FORCE_SHJ.key -> "false") { + // One-file inputs keep the malformed suffix last; two bad groups retain LIMIT. + withInputs( + "aqe_valid" -> """SELECT 1 AS k, base64(cast(cast(id AS STRING) AS BINARY)) AS bad + |FROM range(0, 12, 1, 1)""".stripMargin, + "aqe_bad_last" -> + "SELECT k, CASE WHEN bad='MTE=' THEN 'A' ELSE bad END AS bad FROM aqe_valid", + "aqe_bad_only" -> """SELECT 1 AS k, CASE WHEN id=0 THEN 'A' ELSE 'B' END AS bad + |FROM range(0, 2, 1, 1)""".stripMargin, + "aqe_right" -> "SELECT 'MA==' AS bad") { + withModes { (aqe, _) => + for (join <- Seq("INNER", "LEFT SEMI")) { + withClue(s"join=$join: ") { + def query(input: String, decoder: String = "bad", agg: String = "collect_list(k)") + : String = + s"""SELECT a.decoded, a.collected FROM ( + | SELECT bad, unbase64($decoder) AS decoded, $agg AS collected + | FROM $input GROUP BY bad + |) a $join JOIN aqe_right b ON a.bad = b.bad LIMIT 1""".stripMargin + def checkSuccess(input: String, agg: String = "collect_list(k)"): SparkPlan = { + val text = query(input, agg = agg) + var expected = Seq.empty[Row] + withSQLConf(CometConf.COMET_ENABLED.key -> "false") { + expected = sql(text).collect().toSeq + } + val df = sql(text) + val initial = df.queryExecution.executedPlan + val initialBuffers = aggregateCounts(initial) + assert( + collect(initial) { + case p if original(p).isInstanceOf[SortMergeJoinExec] => p + }.size == 1, + initial.toString) + checkAnswer(df, expected) + if (aqe && agg == "collect_list(k)") { + assert(initialBuffers == (0, 2), initial.toString) + } + val plan = df.queryExecution.executedPlan + if (aqe) { + assert(plan.asInstanceOf[AdaptiveSparkPlanExec].isFinalPlan) + assert( + collect(plan) { + case j: BroadcastHashJoinExec => j.buildSide + case j: CometBroadcastHashJoinExec => j.buildSide + } == Seq(BuildRight), Review Comment: Fixed in `1ad78542da`: INNER now requires a single broadcast join while LEFT SEMI additionally requires BuildRight. The complete nine-test evaluation-mask suite passes. ########## spark/src/test/scala/org/apache/comet/serde/CometEvaluationMaskSuite.scala: ########## @@ -0,0 +1,648 @@ +/* + * 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. + */ + +package org.apache.comet.serde + +import scala.reflect.ClassTag + +import org.apache.spark.sql.{CometTestBase, DataFrame, Row} +import org.apache.spark.sql.catalyst.expressions._ +import org.apache.spark.sql.catalyst.optimizer.BuildRight +import org.apache.spark.sql.catalyst.plans.logical.LocalRelation +import org.apache.spark.sql.catalyst.plans.physical.SinglePartition +import org.apache.spark.sql.comet._ +import org.apache.spark.sql.execution._ +import org.apache.spark.sql.execution.adaptive.{AdaptiveSparkPlanExec, LogicalQueryStage} +import org.apache.spark.sql.execution.aggregate.ObjectHashAggregateExec +import org.apache.spark.sql.execution.exchange.{BroadcastExchangeLike, ShuffleExchangeExec, ShuffleExchangeLike} +import org.apache.spark.sql.execution.joins._ +import org.apache.spark.sql.internal.SQLConf +import org.apache.spark.sql.types.{IntegerType, StringType} + +import org.apache.comet.CometConf +import org.apache.comet.CometSparkSessionExtensions.isSpark35Plus +import org.apache.comet.rules.CometExecRule +import org.apache.comet.shims.ShimCometWindowGroupLimit + +class CometEvaluationMaskSuite extends CometTestBase { + private val limitReason = "unbase64 requires Spark evaluation below LIMIT" + private val joinReason = "unbase64 requires Spark evaluation in first-match join conditions" + private val decodeMessage = "Last unit does not have enough valid bits" + + private def withModes(f: (Boolean, Boolean) => Unit): Unit = { + for (aqe <- Seq(false, true); dispatch <- Seq(false, true)) { + withSQLConf( + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> aqe.toString, + CometConf.COMET_SCALA_UDF_CODEGEN_ENABLED.key -> dispatch.toString) { + withClue(s"AQE=$aqe, dispatcher=$dispatch: ")(f(aqe, dispatch)) + } + } + } + + private def withInputs(inputs: (String, String)*)(f: => Unit): Unit = { + withTempPath { dir => + withTempView(inputs.map(_._1): _*) { + withSQLConf(CometConf.COMET_ENABLED.key -> "false") { + inputs.foreach { case (name, query) => + val path = s"${dir.getCanonicalPath}/$name" + sql(query).coalesce(1).write.parquet(path) + val input = spark.read.parquet(path) + assert(input.inputFiles.length == 1) + input.createOrReplaceTempView(name) + } + } + f + } + } + } + + private def sparkPlan(query: String): SparkPlan = { + var plan: SparkPlan = null + withSQLConf( + CometConf.COMET_ENABLED.key -> "false", + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false") { + plan = sql(query).queryExecution.executedPlan + } + plan + } + + private def applyRule(plan: SparkPlan): SparkPlan = + CometExecRule(spark).apply(stripAQEPlan(plan)) + private def count[T <: SparkPlan: ClassTag](plan: SparkPlan): Int = + collect(plan) { case node: T => node }.size + private def aggregateCounts(plan: SparkPlan): (Int, Int) = + (count[CometHashAggregateExec](plan), count[ObjectHashAggregateExec](plan)) + private def nativeScans(plan: SparkPlan): Int = collect(plan) { + case _: CometScanExec | _: CometNativeScanExec => true + }.size + private def original(plan: SparkPlan): SparkPlan = plan match { + case comet: CometExec => comet.originalPlan + case other => other + } + private def hasDecoder(expr: Expression): Boolean = expr.exists(_.isInstanceOf[UnBase64]) + private def malformed(error: Throwable): Boolean = + causeChain(error).exists(e => Option(e.getMessage).exists(_.contains(decodeMessage))) + private def decodeError(body: => Any): Unit = assert(malformed(intercept[Exception](body))) + private def decodeErrors(query: String): Unit = { + val (sparkError, cometError) = checkSparkAnswerMaybeThrows(sql(query)) + Seq(sparkError, cometError).foreach(error => assert(error.exists(malformed), query)) + } + + test("evaluation-mask protection can be disabled for known-valid input") { + val preserve = CometConf.COMET_EXEC_PRESERVE_EVALUATION_MASKS_ENABLED + assert(preserve.defaultValue.contains(true)) + withSQLConf( + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "true", + SQLConf.SHUFFLE_PARTITIONS.key -> "1", + CometConf.COMET_SCALA_UDF_CODEGEN_ENABLED.key -> "true") { + withInputs( + "opt_out_valid" -> "SELECT * FROM VALUES (1, 'YWJj'), (1, 'YWFh') AS t(k, bad)", + "opt_out_left" -> "SELECT * FROM VALUES (1, X'616262'), (2, X'616262') AS t(k, expected)") { + for (enabled <- Seq(true, false)) { + withSQLConf(preserve.key -> enabled.toString) { + val projection = "SELECT hex(unbase64(bad)) FROM opt_out_valid LIMIT 1" + val compound = "SELECT hex(unbase64(concat(bad, ''))) FROM opt_out_valid LIMIT 1" + if (enabled) { + checkSparkAnswerAndFallbackReason(projection, limitReason) + checkSparkAnswerAndFallbackReason(compound, limitReason) + } else { + checkSparkAnswerAndImpl(sql(projection), native = Seq("unbase64")) + checkSparkAnswerAndImpl(sql(compound), dispatched = Seq("unbase64")) + } + for (kind <- Seq("SEMI", "ANTI")) { + val query = s"""SELECT /*+ BROADCAST(r) */ l.k + |FROM opt_out_left l LEFT $kind JOIN opt_out_valid r + |ON l.k = r.k AND unbase64(r.bad) > l.expected""".stripMargin + if (enabled) { + checkSparkAnswerAndFallbackReason(query, joinReason) + } else { + val (_, plan) = checkSparkAnswerAndImpl(sql(query), native = Seq("unbase64")) + assert(count[CometBroadcastHashJoinExec](plan) == 1, plan.toString) + } + } + // AQE can remove the sort after Partial materializes. The opt-out keeps both + // aggregate halves native instead of preemptively restoring Spark's buffers. + withSQLConf( + SQLConf.WHOLESTAGE_CODEGEN_ENABLED.key -> "false", + CometConf.COMET_EXEC_LOCAL_TABLE_SCAN_ENABLED.key -> "true") { + val aggregate = sparkPlan( + "SELECT unbase64(bad) AS decoded, collect_list(k) FROM " + + "VALUES (1, 'YWJj'), (1, 'YWFh') AS t(k, bad) GROUP BY bad") + assert(aggregateCounts(aggregate) == (0, 2)) + val sorted = SortExec( + Seq(SortOrder(aggregate.output.head, Ascending)), + global = false, + aggregate) + val plan = applyRule(CollectLimitExec(1, sorted)) + assert(aggregateCounts(plan) == (if (enabled) (0, 2) else (2, 0)), plan.toString) + } + } + } + } + } + } + + test("LIMIT masks follow the decoder policy with ANSI, strict and compound inputs") { + assert(QueryPlanSerde.exprSerdeMap.collect { case (cls, _: RequiresSparkEvaluationMask[_]) => + cls + }.toSet == Set(classOf[UnBase64])) + val input = AttributeReference("encoded", StringType)() + withInputs( + "masked" -> "SELECT * FROM VALUES (1, 'YWJj'), (2, 'A') AS t(k, bad)", + "valid" -> "SELECT * FROM VALUES (1, 'YWJj'), (1, 'YWFh'), (1, NULL) AS t(k, bad)", + "inner_left" -> "SELECT 1 AS k, X'616262' AS expected", + "one_group" -> "SELECT 1 AS k, 'YWJj' AS bad") { + withModes { (_, _) => + for (ansi <- Seq(false, true)) { + withSQLConf(SQLConf.ANSI_ENABLED.key -> ansi.toString) { + for (child <- Seq(input, Literal("YWJj"), Concat(Seq(input, Literal("")))); + strict <- Seq(false, true)) { + val expressions = Seq(UnBase64(child, strict), new LimitUnBase64(child, strict)) + assert( + expressions.forall(QueryPlanSerde.evaluationMaskName(_).contains("unbase64"))) + } + // Existing unbase64/to_binary SQL tests exercise native and dispatched serialization. + // Other throwing expressions: https://github.com/apache/datafusion-comet/issues/6006. + Seq( + Add(Literal(Int.MaxValue), Literal(1)), + Cast(input, IntegerType), + Concat(Seq(input, Literal(""))), + FormatString(Literal("%s"), input), + Like(input, Literal("%"), escapeChar = '\\')).foreach { expr => + assert(QueryPlanSerde.evaluationMaskName(expr).isEmpty) + } + Seq( + "SELECT hex(unbase64(bad)) FROM masked LIMIT 1", + "SELECT bad FROM masked WHERE unbase64(bad) <=> X'616263' LIMIT 1", + "SELECT hex(unbase64(concat(bad, ''))) FROM masked LIMIT 1", + "SELECT hex(to_binary(bad, 'base64')) FROM masked LIMIT 1").foreach { query => + checkSparkAnswerAndFallbackReason(query, limitReason) + } + decodeErrors("SELECT hex(unbase64(bad)) FROM masked") + decodeErrors("SELECT hex(unbase64(bad)) FROM masked WHERE bad = 'A' LIMIT 1") + checkSparkAnswerAndOperator(sql("SELECT hex(unbase64(bad)) FROM valid")) + checkSparkAnswerAndOperator( + sql("SELECT /*+ BROADCAST(r) */ l.* FROM inner_left l " + + "INNER JOIN valid r ON l.k = r.k AND unbase64(r.bad) > l.expected")) + // AQE can remove this sort after Partial materializes; offset-only collection stays native. + checkSparkAnswerAndOperator( + sql("SELECT unbase64(bad) AS decoded, collect_list(k) FROM one_group " + + "GROUP BY bad ORDER BY decoded OFFSET 1")) + } + } + } + } + } + + test("whole-tree JVM dispatch preserves an unregistered decoder subclass below LIMIT") { + val name = "comet_test_unbase64" + spark.sessionState.functionRegistry.createOrReplaceTempFunction( + name, + children => new LimitUnBase64(children.head), + "scala_udf") + try { + withSQLConf( + SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false", + CometConf.COMET_SCALA_UDF_CODEGEN_ENABLED.key -> "true") { + withInputs("subclass_input" -> "SELECT * FROM VALUES ('YQ=='), ('A') AS t(encoded)") { + val query = s"SELECT format_string('%d', length($name(encoded))) FROM subclass_input" + val native = sql(query).queryExecution.executedPlan.collect { + case p: CometProjectExec => p + } + assert( + native.size == 1 && native.head.nativeOp.getProjection + .getProjectList(0) + .hasJvmScalarUdf) + val df = sql(s"$query LIMIT 1") + val (_, plan) = checkSparkAnswerAndFallbackReason(df, limitReason) + assert(count[ProjectExec](plan) == 1 && count[CometProjectExec](plan) == 0) + assert(nativeScans(plan) == 1) + checkAnswer(df, Seq(Row("1"))) + decodeErrors(s"$query LIMIT 2") + } + } + } finally { + spark.sessionState.catalog.dropTempFunction(name, ignoreIfNotExists = true) + } + } + + test("unbase64 LIMIT masks survive physical boundaries and native reuse") { + withSQLConf( + SQLConf.WHOLESTAGE_CODEGEN_ENABLED.key -> "false", + CometConf.COMET_EXEC_LOCAL_TABLE_SCAN_ENABLED.key -> "true") { + withModes { (_, _) => + val input = sparkPlan("SELECT * FROM VALUES (1, 'YWJj'), (2, 'A') AS t(id, bad)") + def project = ProjectExec( + Seq(input.output.head, Alias(UnBase64(input.output(1)), "decoded")()), + input) + val staged = applyRule( + GlobalLimitExec(1, ShuffleExchangeExec(SinglePartition, LocalLimitExec(1, project)))) + assert(count[CometGlobalLimitExec](staged) == 1) + assert(count[CometLocalTableScanExec](staged) == 1) + val exchange = collect(staged) { case s: ShuffleExchangeLike => s }.head + for (plan <- Seq(applyRule(staged), applyRule(exchange))) { + assert(count[LocalLimitExec](plan) == 1) + assert(count[ProjectExec](plan) == 1) + assert(count[CometProjectExec](plan) == 0) + } + val barriers: Seq[SparkPlan => SparkPlan] = Seq( + child => SortExec(Seq(SortOrder(child.output.head, Ascending)), global = false, child), + child => ShuffleExchangeExec(SinglePartition, child)) + barriers.foreach { barrier => + val blocking = barrier(project) + // Materialization is safe for both fresh plans and native subtrees reused by AQE. + for (child <- Seq(blocking, applyRule(blocking))) { + val plan = applyRule(LocalLimitExec(1, child)) + withClue(blocking.nodeName) { + assert(count[CometProjectExec](plan) == 1) + assert(count[ProjectExec](plan) == 0) + } + } + } + val bridges: Seq[SparkPlan => SparkPlan] = Seq( + child => child, + child => RowToColumnarExec(ColumnarToRowExec(child)), + child => CometSparkToColumnarExec(CometColumnarToRowExec(child)), + child => CometSparkToColumnarExec(CometNativeColumnarToRowExec(child))) + bridges.foreach { bridge => + val decoded = project + val native = applyRule(ProjectExec(decoded.output, decoded)) + assert(count[CometProjectExec](native) == 2) + val logicalStage = LogicalQueryStage(LocalRelation(native.output), native) + native.setLogicalLink(logicalStage) + val transformed = applyRule(GlobalLimitExec(1, bridge(native))) + for (plan <- Seq(transformed, applyRule(transformed))) { + assert(count[ProjectExec](plan) == 2) + assert(count[CometProjectExec](plan) == 0) + assert(collect(plan) { case r: RowToColumnarTransition => r }.isEmpty) + val link = plan.children.head.getTagValue(SparkPlan.LOGICAL_PLAN_TAG) + assert(link.exists(_ eq logicalStage)) + } + } + } + } + } + + test("unbase64 chooses compatible aggregate buffers before AQE materialization") { + withSQLConf( + SQLConf.WHOLESTAGE_CODEGEN_ENABLED.key -> "false", + CometConf.COMET_EXEC_LOCAL_TABLE_SCAN_ENABLED.key -> "true") { + withModes { (aqe, _) => + for { + decoderInResult <- Seq(false, true) + boundary <- Seq("direct", "sort", "shuffle", "top-k", "offset") + } { + val result = if (decoderInResult) { + "unbase64(bad) AS decoded, collect_list(k) AS collected" + } else { + "bad, collect_list(unbase64(bad)) AS collected" + } + val aggregate = sparkPlan( + s"SELECT $result FROM " + + "VALUES (1, 'YWJj'), (2, 'YWJj') AS t(k, bad) GROUP BY bad") + assert(count[ObjectHashAggregateExec](aggregate) == 2) + val native = applyRule(aggregate) + assert(count[CometHashAggregateExec](native) == 2) + for (child <- Seq(aggregate, native)) { + val order = Seq(SortOrder(child.output.head, Ascending)) + val plan = boundary match { + case "direct" => CollectLimitExec(1, child) + case "sort" => CollectLimitExec(1, SortExec(order, global = false, child)) + case "shuffle" => CollectLimitExec(1, ShuffleExchangeExec(SinglePartition, child)) + case "top-k" => TakeOrderedAndProjectExec(1, order, child.output, child) + case "offset" => + CollectLimitExec(-1, SortExec(order, global = false, child), offset = 1) + } + // Final does not reevaluate decoder inputs, and offset-only never stops early. + val sparkBuffers = decoderInResult && + (boundary == "direct" || (aqe && boundary != "offset")) + val transformed = applyRule(plan) + withClue(s"decoderInResult=$decoderInResult, boundary=$boundary: ") { + for (replanned <- Seq(transformed, applyRule(transformed))) { + assert(aggregateCounts(replanned) == (if (sparkBuffers) (0, 2) else (2, 0))) + } + } + } + } + } + } + } + + test("unbase64 keeps aggregate buffers compatible through an actual AQE join change") { + withSQLConf( + SQLConf.WHOLESTAGE_CODEGEN_ENABLED.key -> "true", + SQLConf.SHUFFLE_PARTITIONS.key -> "1", + SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1", + "spark.sql.adaptive.autoBroadcastJoinThreshold" -> "10485760", + "spark.sql.join.preferSortMergeJoin" -> "true", + "spark.sql.adaptive.maxShuffledHashJoinLocalMapThreshold" -> "0", + CometConf.COMET_FORCE_SHJ.key -> "false") { + // One-file inputs keep the malformed suffix last; two bad groups retain LIMIT. + withInputs( + "aqe_valid" -> """SELECT 1 AS k, base64(cast(cast(id AS STRING) AS BINARY)) AS bad + |FROM range(0, 12, 1, 1)""".stripMargin, + "aqe_bad_last" -> + "SELECT k, CASE WHEN bad='MTE=' THEN 'A' ELSE bad END AS bad FROM aqe_valid", + "aqe_bad_only" -> """SELECT 1 AS k, CASE WHEN id=0 THEN 'A' ELSE 'B' END AS bad + |FROM range(0, 2, 1, 1)""".stripMargin, + "aqe_right" -> "SELECT 'MA==' AS bad") { + withModes { (aqe, _) => + for (join <- Seq("INNER", "LEFT SEMI")) { + withClue(s"join=$join: ") { + def query(input: String, decoder: String = "bad", agg: String = "collect_list(k)") + : String = + s"""SELECT a.decoded, a.collected FROM ( + | SELECT bad, unbase64($decoder) AS decoded, $agg AS collected + | FROM $input GROUP BY bad + |) a $join JOIN aqe_right b ON a.bad = b.bad LIMIT 1""".stripMargin + def checkSuccess(input: String, agg: String = "collect_list(k)"): SparkPlan = { + val text = query(input, agg = agg) + var expected = Seq.empty[Row] + withSQLConf(CometConf.COMET_ENABLED.key -> "false") { + expected = sql(text).collect().toSeq + } + val df = sql(text) + val initial = df.queryExecution.executedPlan + val initialBuffers = aggregateCounts(initial) + assert( + collect(initial) { + case p if original(p).isInstanceOf[SortMergeJoinExec] => p + }.size == 1, + initial.toString) + checkAnswer(df, expected) + if (aqe && agg == "collect_list(k)") { + assert(initialBuffers == (0, 2), initial.toString) + } + val plan = df.queryExecution.executedPlan + if (aqe) { + assert(plan.asInstanceOf[AdaptiveSparkPlanExec].isFinalPlan) + assert( + collect(plan) { + case j: BroadcastHashJoinExec => j.buildSide + case j: CometBroadcastHashJoinExec => j.buildSide + } == Seq(BuildRight), + plan.toString) + } + plan + } + val valid = checkSuccess("aqe_valid") + assert(aggregateCounts(valid) == (if (aqe) (0, 2) else (2, 0))) + // SEMI streams the left side; INNER may materialize a left broadcast candidate. + if (aqe && join == "LEFT SEMI") checkSuccess("aqe_bad_last") + else if (!aqe) decodeErrors(query("aqe_bad_last")) + decodeErrors(query("aqe_valid", "CASE WHEN bad='MA==' THEN 'A' ELSE bad END")) + // ObjectHash evaluates results even before a nonmatching join discards them. + decodeErrors(query("aqe_bad_only")) + val compatible = checkSuccess("aqe_valid", "max(k)") + assert(count[CometHashAggregateExec](compatible) > 0) Review Comment: Fixed in `1ad78542da`: AQE `max(k)` now requires initial aggregate counts `(2, 0)`, beside the incompatible collect_list control. Both AQE regressions pass in the full nine-test suite. -- 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]
