diff --git a/backends-velox/src/main/scala/org/apache/gluten/backendsapi/velox/VeloxRuleApi.scala b/backends-velox/src/main/scala/org/apache/gluten/backendsapi/velox/VeloxRuleApi.scala index d63928527df..067f620d676 100644 --- a/backends-velox/src/main/scala/org/apache/gluten/backendsapi/velox/VeloxRuleApi.scala +++ b/backends-velox/src/main/scala/org/apache/gluten/backendsapi/velox/VeloxRuleApi.scala @@ -53,6 +53,7 @@ object VeloxRuleApi { private def injectSpark(injector: SparkInjector): Unit = { // Inject the regular Spark rules directly. injector.injectOptimizerRule(CollectRewriteRule.apply) + injector.injectOptimizerRule(BloomFilterMightContainJointRewriteRule.apply) injector.injectOptimizerRule(HLLRewriteRule.apply) injector.injectOptimizerRule(CollapseGetJsonObjectExpressionRule.apply) injector.injectOptimizerRule(RewriteCastFromArray.apply) @@ -81,11 +82,7 @@ object VeloxRuleApi { injector.injectPreTransform(c => FallbackMultiCodegens.apply(c.session)) injector.injectPreTransform(c => MergeTwoPhasesHashBaseAggregate(c.session)) injector.injectPreTransform(_ => RewriteSubqueryBroadcast()) - injector.injectPreTransform( - c => - BloomFilterMightContainJointRewriteRule.apply( - c.session, - c.caller.isBloomFilterStatFunction())) + injector.injectPreTransform(c => RuntimeBloomFilterRewriteRule(c.session)) injector.injectPreTransform(_ => EliminateRedundantGetTimestamp) // Legacy: The legacy transform rule. diff --git a/backends-velox/src/main/scala/org/apache/gluten/extension/BloomFilterMightContainJointRewriteRule.scala b/backends-velox/src/main/scala/org/apache/gluten/extension/BloomFilterMightContainJointRewriteRule.scala index efbedc6ca00..c46c5df12b9 100644 --- a/backends-velox/src/main/scala/org/apache/gluten/extension/BloomFilterMightContainJointRewriteRule.scala +++ b/backends-velox/src/main/scala/org/apache/gluten/extension/BloomFilterMightContainJointRewriteRule.scala @@ -21,63 +21,71 @@ import org.apache.gluten.expression.VeloxBloomFilterMightContain import org.apache.gluten.expression.aggregate.VeloxBloomFilterAggregate import org.apache.spark.sql.SparkSession -import org.apache.spark.sql.catalyst.expressions.{BinaryExpression, BloomFilterMightContain, Expression} -import org.apache.spark.sql.catalyst.expressions.aggregate.{BloomFilterAggregate, TypedImperativeAggregate} +import org.apache.spark.sql.catalyst.expressions.{BloomFilterMightContain, Literal, ScalarSubquery} +import org.apache.spark.sql.catalyst.expressions.aggregate.{AggregateExpression, BloomFilterAggregate} +import org.apache.spark.sql.catalyst.plans.logical.LogicalPlan import org.apache.spark.sql.catalyst.rules.Rule -import org.apache.spark.sql.execution.SparkPlan -case class BloomFilterMightContainJointRewriteRule( - spark: SparkSession, - isBloomFilterStatFunction: Boolean) - extends Rule[SparkPlan] { - override def apply(plan: SparkPlan): SparkPlan = { - if (isBloomFilterStatFunction || !GlutenConfig.get.enableNativeBloomFilter) { +/** + * Optimizer rule that rewrites `BloomFilterAggregate` -> `VeloxBloomFilterAggregate` and + * `BloomFilterMightContain` -> `VeloxBloomFilterMightContain` at the logical plan level. + * + * Running as an optimizer rule ensures the substitution is captured in the `originalPlan` snapshot + * that [[org.apache.gluten.extension.columnar.heuristic.ExpandFallbackPolicy]] uses when promoting + * an individual stage fallback to a whole-stage AQE fallback. This guarantees that both sides of + * the bloom-filter pair always produce and consume the same byte format, regardless of whether + * stages fall back to JVM execution after AQE re-planning. + * + * This rule only covers bloom filters that are present in the plan during Spark's Operator + * Optimization batch, i.e. the ones written by the user. Runtime bloom filters are injected by + * `InjectRuntimeFilter`, which runs in a later batch of `SparkOptimizer`, after this rule has + * already fired -- they are handled by the physical-level [[RuntimeBloomFilterRewriteRule]] + * instead. The aggregate (producer) and the might-contain (consumer) are always rewritten as a + * pair, or not at all, so they never end up on different serialized byte formats: + * - `might_contain(ScalarSubquery(...), )`: rewrite both sides to their Velox forms + * (version=1). This is the path that GLUTEN-12013 protects across whole-stage AQE fallback. + * - `might_contain(ScalarSubquery(...), )` (as in SPARK-54336): leave both vanilla + * (version=0). Rewriting only the outer side to Velox while the inner aggregate stayed vanilla + * `bloom_filter_agg` is exactly what caused the `kBloomFilterV1 == version` (1 vs. 0) crash. + * Leaving both vanilla also preserves vanilla's NULL-on-empty-input semantics. + * + * Standalone `BloomFilterAggregate` (e.g., `DataFrame.stat.bloomFilter()`) is never matched, so its + * bytes stay in Spark-native format. + */ +case class BloomFilterMightContainJointRewriteRule(spark: SparkSession) + extends Rule[LogicalPlan] { + + override def apply(plan: LogicalPlan): LogicalPlan = { + if (!GlutenConfig.get.enableNativeBloomFilter) { return plan } - val out = plan.transformWithSubqueries { - case p => - applyForNode(p) - } - out - } - - private def replaceBloomFilterAggregate[T]( - expr: Expression, - bloomFilterAggReplacer: ( - Expression, - Expression, - Expression, - Int, - Int) => TypedImperativeAggregate[T]): Expression = expr match { - case BloomFilterAggregate( - child, - estimatedNumItemsExpression, - numBitsExpression, - mutableAggBufferOffset, - inputAggBufferOffset) => - bloomFilterAggReplacer( - child, - estimatedNumItemsExpression, - numBitsExpression, - mutableAggBufferOffset, - inputAggBufferOffset) - case other => other - } - - private def replaceMightContain[T]( - expr: Expression, - mightContainReplacer: (Expression, Expression) => BinaryExpression): Expression = expr match { - case BloomFilterMightContain(bloomFilterExpression, valueExpression) => - mightContainReplacer(bloomFilterExpression, valueExpression) - case other => other - } - - private def applyForNode(p: SparkPlan) = { - p.transformExpressions { - case e => - replaceMightContain( - replaceBloomFilterAggregate(e, VeloxBloomFilterAggregate.apply), - VeloxBloomFilterMightContain.apply) + plan.transformAllExpressions { + case BloomFilterMightContain(subq: ScalarSubquery, v) if !v.isInstanceOf[Literal] => + // Non-literal value: rewrite both the outer might-contain and the inner aggregate to + // Velox format so that both sides always produce/consume the same byte layout (version=1), + // even when stages fall back to JVM execution via ExpandFallbackPolicy. + val rewrittenPlan = subq.plan.transformAllExpressions { + case ae @ AggregateExpression(b: BloomFilterAggregate, _, _, _, _) => + ae.copy(aggregateFunction = VeloxBloomFilterAggregate( + b.child, + b.estimatedNumItemsExpression, + b.numBitsExpression, + b.mutableAggBufferOffset, + b.inputAggBufferOffset)) + } + VeloxBloomFilterMightContain(subq.withNewPlan(rewrittenPlan), v) + case bfmc @ BloomFilterMightContain(_: ScalarSubquery, _) => + // Literal value (SPARK-54336: `might_contain((SELECT bloom_filter_agg(col) FROM t), 0L)`). + // Leave BOTH the inner aggregate and the outer might-contain as vanilla Spark expressions + // so they stay on the same Spark-native (version=0) byte format, and so an empty aggregate + // input still yields vanilla's NULL bloom filter. Rewriting only the outer side to Velox + // (which expects version=1) while the inner aggregate stays vanilla `bloom_filter_agg` (no + // Substrait mapping -> JVM -> version=0) is what caused the `kBloomFilterV1 == version` + // (1 vs. 0) crash. + bfmc + case BloomFilterMightContain(bf, v) => + // Pre-computed literal bloom filter bytes -- rewrite to consume Velox-format bytes. + VeloxBloomFilterMightContain(bf, v) } } } diff --git a/backends-velox/src/main/scala/org/apache/gluten/extension/RuntimeBloomFilterRewriteRule.scala b/backends-velox/src/main/scala/org/apache/gluten/extension/RuntimeBloomFilterRewriteRule.scala new file mode 100644 index 00000000000..f3d81b53402 --- /dev/null +++ b/backends-velox/src/main/scala/org/apache/gluten/extension/RuntimeBloomFilterRewriteRule.scala @@ -0,0 +1,78 @@ +/* + * 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.gluten.extension + +import org.apache.gluten.config.GlutenConfig +import org.apache.gluten.expression.VeloxBloomFilterMightContain +import org.apache.gluten.expression.aggregate.VeloxBloomFilterAggregate + +import org.apache.spark.sql.SparkSession +import org.apache.spark.sql.catalyst.expressions.{BloomFilterMightContain, XxHash64} +import org.apache.spark.sql.catalyst.expressions.aggregate.{AggregateExpression, BloomFilterAggregate} +import org.apache.spark.sql.catalyst.rules.Rule +import org.apache.spark.sql.execution.SparkPlan + +/** + * Physical pre-transform rule that rewrites runtime-filter bloom filters (the ones injected by + * Spark's `InjectRuntimeFilter` optimizer rule) to their Velox variants so they offload natively. + * + * Runtime bloom filters cannot be handled by [[BloomFilterMightContainJointRewriteRule]]: that rule + * is registered via `injectOptimizerRule`, which lands in Spark's Operator Optimization batch, + * while `InjectRuntimeFilter` runs in a later batch of `SparkOptimizer`. The runtime-filter + * expressions therefore do not exist yet when the logical rule fires, and a physical-level rewrite + * (as was always done before the logical rule was introduced) is required to keep + * `FilterExecTransformer` and the bloom-filter aggregate native. + * + * The rewrite is restricted to `InjectRuntimeFilter`'s exact expression shapes, which always wrap + * the key in [[XxHash64]] on both the producer and the consumer side: + * - producer: `bloom_filter_agg(xxhash64(key), ...)` -> `velox_bloom_filter_agg(...)` + * - consumer: `might_contain(bf, xxhash64(key))` -> `velox_might_contain(...)` + * + * Because each side is identifiable on its own, both are rewritten consistently to the Velox byte + * format (version=1) even when AQE compiles the bloom-filter subquery separately from the consuming + * filter stage. The `XxHash64` fingerprint also guarantees the other bloom-filter populations are + * never touched: + * - `DataFrame.stat.bloomFilter()` builds `bloom_filter_agg(col, ...)` on the raw column (no + * `XxHash64` wrapper) and deserializes the result with Spark's `BloomFilter.readFrom`, so its + * bytes must stay in Spark-native format. + * - User-facing `might_contain(, )` pairs are already rewritten at the + * logical level by [[BloomFilterMightContainJointRewriteRule]] (the GLUTEN-12013 fix), making + * this rule a no-op for them. + * - Literal-value pairs (SPARK-54336) contain no `XxHash64` and stay fully vanilla. + */ +case class RuntimeBloomFilterRewriteRule(spark: SparkSession) extends Rule[SparkPlan] { + override def apply(plan: SparkPlan): SparkPlan = { + if (!GlutenConfig.get.enableNativeBloomFilter) { + return plan + } + plan.transformWithSubqueries { + case p => + p.transformExpressions { + case ae @ AggregateExpression(b: BloomFilterAggregate, _, _, _, _) + if b.child.isInstanceOf[XxHash64] => + ae.copy(aggregateFunction = VeloxBloomFilterAggregate( + b.child, + b.estimatedNumItemsExpression, + b.numBitsExpression, + b.mutableAggBufferOffset, + b.inputAggBufferOffset)) + case BloomFilterMightContain(bf, value: XxHash64) => + VeloxBloomFilterMightContain(bf, value) + } + } + } +} diff --git a/backends-velox/src/test/scala/org/apache/gluten/sql/GlutenBloomFilterFallbackSuite.scala b/backends-velox/src/test/scala/org/apache/gluten/sql/GlutenBloomFilterFallbackSuite.scala new file mode 100644 index 00000000000..49a7ba86eac --- /dev/null +++ b/backends-velox/src/test/scala/org/apache/gluten/sql/GlutenBloomFilterFallbackSuite.scala @@ -0,0 +1,398 @@ +/* + * 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.gluten.sql + +import org.apache.gluten.backendsapi.BackendsApiManager +import org.apache.gluten.config.GlutenConfig +import org.apache.gluten.execution.{FilterExecTransformerBase, WholeStageTransformerSuite} +import org.apache.gluten.expression.VeloxBloomFilterMightContain +import org.apache.gluten.expression.aggregate.VeloxBloomFilterAggregate + +import org.apache.spark.sql.Row +import org.apache.spark.sql.catalyst.FunctionIdentifier +import org.apache.spark.sql.catalyst.expressions.BloomFilterMightContain +import org.apache.spark.sql.catalyst.expressions.ExpressionInfo +import org.apache.spark.sql.catalyst.expressions.aggregate.BloomFilterAggregate +import org.apache.spark.sql.execution.aggregate.{BaseAggregateExec, ObjectHashAggregateExec} +import org.apache.spark.sql.internal.SQLConf + +/** + * Regression tests for https://github.com/apache/gluten/issues/12013. + * + * Verifies that the bloom-filter producer (`bloom_filter_agg`) and consumer (`might_contain`) + * always stay on a consistent serialized byte format: + * - user-facing pairs are rewritten to their Velox forms at the logical level by + * `BloomFilterMightContainJointRewriteRule`, surviving whole-stage AQE fallback; + * - literal-valued pairs (SPARK-54336) are left fully vanilla; + * - runtime-filter pairs injected by `InjectRuntimeFilter` are rewritten at the physical level by + * `RuntimeBloomFilterRewriteRule` and must offload natively. + */ +class GlutenBloomFilterFallbackSuite extends WholeStageTransformerSuite { + protected val resourcePath: String = null + protected val fileFormat: String = null + + import testImplicits._ + + private val funcIdBloomFilterAgg = FunctionIdentifier("bloom_filter_agg") + private val funcIdMightContain = FunctionIdentifier("might_contain") + + override def beforeAll(): Unit = { + super.beforeAll() + spark.sessionState.functionRegistry.registerFunction( + funcIdBloomFilterAgg, + new ExpressionInfo(classOf[BloomFilterAggregate].getName, "bloom_filter_agg"), + args => + args.size match { + case 1 => new BloomFilterAggregate(args(0)) + case 2 => new BloomFilterAggregate(args(0), args(1)) + case 3 => new BloomFilterAggregate(args(0), args(1), args(2)) + case _ => throw new IllegalArgumentException("bloom_filter_agg requires 1-3 arguments") + } + ) + spark.sessionState.functionRegistry.registerFunction( + funcIdMightContain, + new ExpressionInfo(classOf[BloomFilterMightContain].getName, "might_contain"), + args => BloomFilterMightContain(args(0), args(1))) + } + + override def afterAll(): Unit = { + spark.sessionState.functionRegistry.dropFunction(funcIdBloomFilterAgg) + spark.sessionState.functionRegistry.dropFunction(funcIdMightContain) + super.afterAll() + } + + private val veloxBloomFilterMaxNumBits = 4194304L + + // GLUTEN-12013: only filter stage falls back (threshold=2). + // bloom_filter_agg subquery runs natively and produces Velox-format bytes; the filter stage + // falls back via ExpandFallbackPolicy. The optimizer-level substitution ensures the fallback + // plan still uses VeloxBloomFilterMightContain so the JVM filter reads Velox-format bytes. + test("GLUTEN-12013: bloom_filter_agg whole-stage fallback does not corrupt bloom filter bytes") { + if (BackendsApiManager.getSettings.requireBloomFilterAggMightContainJointFallback()) { + val table = "bloom_filter_test" + val numEstimatedItems = 5000000L + val sqlString = + s""" + |SELECT col positive_membership_test + |FROM $table + |WHERE might_contain( + | (SELECT bloom_filter_agg(col, + | cast($numEstimatedItems as long), + | cast($veloxBloomFilterMaxNumBits as long)) + | FROM $table), col) + |""".stripMargin + withTempView(table) { + (Seq(Long.MinValue, 0, Long.MaxValue) ++ (1L to 200000L)) + .toDF("col") + .createOrReplaceTempView(table) + // Threshold=2: FilterExec fallback cost=2 triggers whole-stage fallback; agg cost=1 + // does not, so Stage 0 runs natively. ANSI off keeps agg cost at 1 on Spark 4.0+. + withSQLConf( + GlutenConfig.COLUMNAR_FILTER_ENABLED.key -> "false", + GlutenConfig.COLUMNAR_WHOLESTAGE_FALLBACK_THRESHOLD.key -> "2", + SQLConf.ANSI_ENABLED.key -> "false" + ) { + val df = spark.sql(sqlString) + // Must not throw: java.io.IOException: Unexpected Bloom filter version number. + assert(df.collect().length == 200003) + // Verify the optimizer rule ran: VeloxBloomFilterMightContain must be present even + // though Stage 1 executes inside a FallbackNode. + assert( + df.queryExecution.optimizedPlan.toString.contains("velox_might_contain"), + "Expected velox_might_contain in optimized plan -- optimizer rule may not have run" + ) + } + } + } + } + + // GLUTEN-12013: both stages fall back (threshold=1). + // Stage 0's inherent transition cost of 1 meets the threshold so ExpandFallbackPolicy + // promotes it to a whole-stage fallback too. The optimizer rule has already rewritten both + // sides to Velox variants before ExpandFallbackPolicy captures its snapshot. Even in JVM + // row-mode, VeloxBloomFilterAggregate produces Velox-format bytes (via JNI) and + // VeloxBloomFilterMightContain consumes them -- both sides are consistent. + test("GLUTEN-12013: bloom_filter_agg whole-stage fallback when both stages fall back") { + if (BackendsApiManager.getSettings.requireBloomFilterAggMightContainJointFallback()) { + val table = "bloom_filter_test" + val numEstimatedItems = 5000000L + val sqlString = + s""" + |SELECT col positive_membership_test + |FROM $table + |WHERE might_contain( + | (SELECT bloom_filter_agg(col, + | cast($numEstimatedItems as long), + | cast($veloxBloomFilterMaxNumBits as long)) + | FROM $table), col) + |""".stripMargin + withTempView(table) { + (Seq(Long.MinValue, 0, Long.MaxValue) ++ (1L to 200000L)) + .toDF("col") + .createOrReplaceTempView(table) + // Threshold=1: both stages fall back; both use Velox variants via JNI. + withSQLConf( + GlutenConfig.COLUMNAR_FILTER_ENABLED.key -> "false", + GlutenConfig.COLUMNAR_WHOLESTAGE_FALLBACK_THRESHOLD.key -> "1", + SQLConf.ANSI_ENABLED.key -> "false" + ) { + val df = spark.sql(sqlString) + // Must not throw: java.io.IOException: Unexpected Bloom filter version number. + assert(df.collect().length == 200003) + // Verify the optimizer rule ran on both sides. + assert( + df.queryExecution.optimizedPlan.toString.contains("velox_might_contain"), + "Expected velox_might_contain in optimized plan -- optimizer rule may not have run" + ) + } + } + } + } + + // GLUTEN-12013: DataFrame.stat.bloomFilter() must not be affected by the optimizer rule. + // The rule must only rewrite BloomFilterAggregate inside a BloomFilterMightContain subquery. + // A standalone BloomFilterAggregate (as used here) must remain vanilla so that the collected + // bytes are in Spark-native format and BloomFilter.readFrom() succeeds. + test("GLUTEN-12013: DataFrame.stat.bloomFilter() produces Spark-readable bytes") { + if (BackendsApiManager.getSettings.requireBloomFilterAggMightContainJointFallback()) { + val table = "bloom_filter_stat_test" + withTempView(table) { + (1L to 1000L).toDF("col").createOrReplaceTempView(table) + // Must not throw: java.io.IOException: Unexpected Bloom filter version number + val bf = spark.table(table).stat.bloomFilter("col", 1000L, 0.01) + // Bloom filters have no false negatives: every inserted value must be present. + assert(bf.mightContainLong(500L), "Expected 500 to be in bloom filter") + } + } + } + + // GLUTEN-12013: native bloom filter disabled -- early-exit path of the optimizer rule. + // When spark.gluten.sql.native.bloomFilter=false the rule returns the plan unchanged. + // BloomFilterAggregate / BloomFilterMightContain remain as vanilla Spark expressions and + // produce/consume consistent Spark-format bytes. + test( + "GLUTEN-12013: native bloom filter disabled skips rewrite and produces correct results") { + if (BackendsApiManager.getSettings.requireBloomFilterAggMightContainJointFallback()) { + val table = "bloom_filter_test" + val numEstimatedItems = 5000000L + val sqlString = + s""" + |SELECT col positive_membership_test + |FROM $table + |WHERE might_contain( + | (SELECT bloom_filter_agg(col, + | cast($numEstimatedItems as long), + | cast($veloxBloomFilterMaxNumBits as long)) + | FROM $table), col) + |""".stripMargin + withTempView(table) { + (Seq(Long.MinValue, 0, Long.MaxValue) ++ (1L to 200000L)) + .toDF("col") + .createOrReplaceTempView(table) + withSQLConf( + GlutenConfig.COLUMNAR_NATIVE_BLOOMFILTER_ENABLED.key -> "false", + SQLConf.ANSI_ENABLED.key -> "false" + ) { + val df = spark.sql(sqlString) + assert(df.collect().length == 200003) + // Verify the rule early-exited: plan must NOT contain Velox variants. + assert( + !df.queryExecution.optimizedPlan.toString.contains("velox_might_contain"), + "Expected vanilla BloomFilterMightContain when native bloom filter is disabled" + ) + } + } + } + } + + // GLUTEN-12013: verify that the bloom-filter subquery uses VeloxBloomFilterAggregate even when + // the aggregate node executes in JVM mode (ObjectHashAggregateExec). + // + // When hash-aggregate offloading is disabled, the subquery runs as + // ObjectHashAggregateExec(VeloxBloomFilterAggregate). VeloxBloomFilterAggregate.eval() calls + // serialize(buffer) directly without a cardinality guard, so it always produces Velox-format + // (version=1) bytes. VeloxBloomFilterMightContain on the outer side reads those bytes correctly. + test("GLUTEN-12013: VeloxBloomFilterAggregate in JVM subquery produces correct Velox bytes") { + if (BackendsApiManager.getSettings.requireBloomFilterAggMightContainJointFallback()) { + val table = "bloom_filter_test" + val numEstimatedItems = 5000000L + val sqlString = + s""" + |SELECT col positive_membership_test + |FROM $table + |WHERE might_contain( + | (SELECT bloom_filter_agg(col, + | cast($numEstimatedItems as long), + | cast($veloxBloomFilterMaxNumBits as long)) + | FROM $table), col) + |""".stripMargin + withTempView(table) { + (Seq(Long.MinValue, 0, Long.MaxValue) ++ (1L to 200000L)) + .toDF("col") + .createOrReplaceTempView(table) + // Disable hash-aggregate offloading: the bloom_filter_agg subquery executes in JVM mode as + // ObjectHashAggregateExec(VeloxBloomFilterAggregate). This mirrors the q59 golden shape. + withSQLConf( + GlutenConfig.COLUMNAR_HASHAGG_ENABLED.key -> "false", + SQLConf.ANSI_ENABLED.key -> "false" + ) { + val df = spark.sql(sqlString) + // Must not throw: java.io.IOException: Unexpected Bloom filter version number. + // VeloxBloomFilterAggregate.eval() produces Velox-format bytes even in JVM mode. + val result = df.collect() + assert(result.length == 200003, s"Expected 200003 rows, got ${result.length}") + + // Directly verify the subquery's aggregate function class at runtime. + // ObjectHashAggregateExec(VeloxBloomFilterAggregate) must be present -- NOT vanilla + // BloomFilterAggregate -- so we know the physical rewrite actually happened. + val subqueryVeloxAggs = collectWithSubqueries(df.queryExecution.executedPlan) { + case agg: ObjectHashAggregateExec + if agg.aggregateExpressions.exists( + _.aggregateFunction.isInstanceOf[VeloxBloomFilterAggregate]) => + agg + } + assert( + subqueryVeloxAggs.nonEmpty, + "Expected ObjectHashAggregateExec(VeloxBloomFilterAggregate) in the bloom-filter " + + "subquery. Actual subquery aggs: " + + collectWithSubqueries(df.queryExecution.executedPlan) { + case agg: ObjectHashAggregateExec => agg + }.map( + a => + a.aggregateExpressions + .map(_.aggregateFunction.getClass.getSimpleName) + .mkString(",")) + .mkString("; ") + ) + } + } + } + } + + // GLUTEN-12013 follow-up: runtime bloom filters injected by Spark's InjectRuntimeFilter must + // stay native. InjectRuntimeFilter runs in a later SparkOptimizer batch than the Operator + // Optimization batch that hosts BloomFilterMightContainJointRewriteRule, so the logical rule + // never sees these expressions -- RuntimeBloomFilterRewriteRule rewrites them at the physical + // level instead. Without that rewrite, vanilla might_contain/bloom_filter_agg have no Substrait + // mapping and the consuming FilterExec plus the producing aggregate fall back to the JVM with + // R2C/C2R transitions (the regression originally visible in the TPC-DS q59 golden). + test("GLUTEN-12013: runtime bloom filter keeps FilterExecTransformer native") { + withTable("bf_fact", "bf_dim") { + spark + .range(0, 10000) + .selectExpr("id as key", "id % 100 as payload") + .write + .saveAsTable("bf_fact") + spark + .range(0, 100) + .selectExpr("id as key", "id % 10 as f") + .write + .saveAsTable("bf_dim") + withSQLConf( + SQLConf.RUNTIME_BLOOM_FILTER_ENABLED.key -> "true", + SQLConf.RUNTIME_BLOOM_FILTER_APPLICATION_SIDE_SCAN_SIZE_THRESHOLD.key -> "3000", + SQLConf.AUTO_BROADCASTJOIN_THRESHOLD.key -> "-1", + SQLConf.ANSI_ENABLED.key -> "false" + ) { + val df = spark.sql( + "SELECT * FROM bf_fact JOIN bf_dim ON bf_fact.key = bf_dim.key WHERE bf_dim.f = 5") + assert( + df.queryExecution.optimizedPlan.toString.contains("might_contain"), + "Precondition failed: InjectRuntimeFilter did not inject a bloom filter" + ) + assert(df.collect().length == 10) + // The consumer side must be a native FilterExecTransformer evaluating + // velox_might_contain -- not a fallen-back JVM FilterExec. + val nativeBloomFilters = collectWithSubqueries(df.queryExecution.executedPlan) { + case f: FilterExecTransformerBase + if f.cond.exists(_.isInstanceOf[VeloxBloomFilterMightContain]) => + f + } + assert( + nativeBloomFilters.nonEmpty, + "Expected a native FilterExecTransformer with velox_might_contain; the runtime " + + s"bloom filter fell back to JVM. Plan:\n${df.queryExecution.executedPlan}" + ) + // The producer side must use VeloxBloomFilterAggregate so the bytes are version=1. + val veloxAggs = collectWithSubqueries(df.queryExecution.executedPlan) { + case agg: BaseAggregateExec + if agg.aggregateExpressions.exists( + _.aggregateFunction.isInstanceOf[VeloxBloomFilterAggregate]) => + agg + } + assert( + veloxAggs.nonEmpty, + "Expected VeloxBloomFilterAggregate in the runtime-filter subquery. Plan:\n" + + s"${df.queryExecution.executedPlan}" + ) + } + } + } + + // SPARK-54336 / GLUTEN-12013: a might_contain whose value argument is a literal (not a column), + // fed by a nested scalar subquery. This mirrors Spark's upstream + // `BloomFilterAggregateQuerySuite."SPARK-54336"`. For a non-column value the rule must leave + // BOTH the inner bloom_filter_agg and the outer might_contain vanilla so they stay on the same + // Spark-native (version=0) byte format. Previously only the outer side was rewritten to Velox + // (version=1), while the inner aggregate -- having no Substrait mapping -- fell back to JVM and + // emitted version=0 bytes, so velox_might_contain failed to deserialize them + // (kBloomFilterV1 == version, 1 vs. 0). + // + // Gated to Spark 4.0+: the query exercises the exact `MergeScalarSubqueries` path that Spark's + // own SPARK-54336 fixes, and that fix only exists in Spark 4.0.2+/4.1. On earlier Spark the + // analyzer/optimizer throws `UnresolvedException` (dataType on an unresolved ScalarSubquery) + // before Gluten's rule runs, so there is nothing for this test to verify there. + testWithMinSparkVersion( + "GLUTEN-12013: might_contain with a literal value keeps both sides vanilla (SPARK-54336)", + "4.0") { + if (BackendsApiManager.getSettings.requireBloomFilterAggMightContainJointFallback()) { + val table = "bloom_filter_lit_test" + withTempView(table) { + // Single non-null row [0]; the bloom filter therefore contains 0L. + Seq(0L).toDF("col").createOrReplaceTempView(table) + val sqlString = + s""" + |SELECT + | (SELECT + | first(might_contain( + | (SELECT bloom_filter_agg(col) FROM $table), + | 0L + | )) + | FROM $table) + |FROM $table + |""".stripMargin + // Codegen-off, matching the upstream BloomFilterAggregateQuerySuiteCGOff variant that + // originally surfaced this crash. + withSQLConf( + SQLConf.WHOLESTAGE_CODEGEN_ENABLED.key -> "false", + SQLConf.CODEGEN_FACTORY_MODE.key -> "NO_CODEGEN", + SQLConf.ANSI_ENABLED.key -> "false" + ) { + val df = spark.sql(sqlString) + // For a literal value, neither side is offloaded to Velox: both stay vanilla so the + // byte formats are consistent (version=0). + assert( + !df.queryExecution.optimizedPlan.toString.contains("velox_might_contain"), + "Expected the literal-valued might_contain to stay vanilla, not velox_might_contain" + ) + // Must not throw kBloomFilterV1 == version (1 vs. 0); 0L was inserted, so it is present. + checkAnswer(df, Row(true)) + } + } + } + } +} diff --git a/gluten-core/src/main/scala/org/apache/gluten/extension/caller/CallerInfo.scala b/gluten-core/src/main/scala/org/apache/gluten/extension/caller/CallerInfo.scala index a65c79da7a7..f7bd42916bf 100644 --- a/gluten-core/src/main/scala/org/apache/gluten/extension/caller/CallerInfo.scala +++ b/gluten-core/src/main/scala/org/apache/gluten/extension/caller/CallerInfo.scala @@ -30,7 +30,6 @@ trait CallerInfo { def isAqe(): Boolean def isCache(): Boolean def isStreaming(): Boolean - def isBloomFilterStatFunction(): Boolean } object CallerInfo { @@ -42,8 +41,7 @@ object CallerInfo { private class Impl( override val isAqe: Boolean, override val isCache: Boolean, - override val isStreaming: Boolean, - override val isBloomFilterStatFunction: Boolean + override val isStreaming: Boolean ) extends CallerInfo /* @@ -57,8 +55,7 @@ object CallerInfo { new Impl( isAqe = inAqeCall(stack), isCache = inCacheCall(stack), - isStreaming = inStreamingCall(stack), - isBloomFilterStatFunction = inBloomFilterStatFunctionCall(stack)) + isStreaming = inStreamingCall(stack)) } private def inAqeCall(stack: Seq[StackTraceElement]): Boolean = { @@ -78,25 +75,13 @@ object CallerInfo { stack.exists(_.getClassName.equals(streamName)) } - private def inBloomFilterStatFunctionCall(stack: Seq[StackTraceElement]): Boolean = { - // Two independent existence checks over the stack. The previous shape used a nested - // `stack.exists(_.getMethodName.equals(...))` inside the outer predicate, which re-walked - // the whole stack for every candidate frame - O(n^2) in stack depth. Splitting the checks - // makes each a single pass. `&&` short-circuits, so the bloomFilter scan only runs when a - // DataFrameStatFunctions frame is present. - val hasStatFunctionsClass = - stack.exists(_.getClassName.equals("org.apache.spark.sql.DataFrameStatFunctions")) - hasStatFunctionsClass && stack.exists(_.getMethodName.equals("bloomFilter")) - } - /** For testing only. */ def withLocalValue[T]( isAqe: Boolean, isCache: Boolean, - isStreaming: Boolean = false, - isBloomFilterStatFunction: Boolean = false)(body: => T): T = { + isStreaming: Boolean = false)(body: => T): T = { val prevValue = localStorage.get() - val newValue = new Impl(isAqe, isCache, isStreaming, isBloomFilterStatFunction) + val newValue = new Impl(isAqe, isCache, isStreaming) localStorage.set(Some(newValue)) try { body