[GLUTEN-12013][VL] Fix bloom-filter bytes corruption on whole-stage AQE fallback

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.

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(...), <non-literal>)`: 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(...), <literal>)` (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)
     }
   }
 }

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(<scalar subquery>, <value>)` 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)
+        }
+    }
+  }
+}