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

[brijrajk]

What changes are proposed in this pull request?

Fixes #12013, and the SPARK-54336 regression that the issue's query shape exposes.

Background

BloomFilterMightContainJointRewriteRule rewrites a bloom-filter producer
(bloom_filter_agg) and its consumer (might_contain) to their Velox variants so Velox
evaluates them natively. Previously it was a physical pre-transform rule; this PR splits the
rewrite into two rules according to where the expressions become visible in the optimizer
pipeline.

1. User-facing bloom filters: logical rule (the GLUTEN-12013 fix)

BloomFilterMightContainJointRewriteRule is now a Rule[LogicalPlan] registered via
injectOptimizerRule, which lands in Spark's Operator Optimization batch. Running there
ensures the substitution is baked into the originalPlan snapshot that
ExpandFallbackPolicy captures when it promotes an individual-stage fallback to a
whole-stage AQE fallback, so both sides keep the same serialized byte format (version=1)
even if a stage reverts to JVM execution. That is the original GLUTEN-12013 crash
(java.io.IOException: Unexpected Bloom filter version number).

The producer and consumer are always rewritten as a pair, or not at all:

• might_contain(ScalarSubquery(...), <non-literal>): rewrite both sides to Velox forms.
• might_contain(ScalarSubquery(...), <literal>) (the SPARK-54336 shape,
  might_contain((SELECT bloom_filter_agg(col) FROM t), 0L)): leave BOTH sides vanilla.
  Rewriting only the outer side to Velox (version=1) while the inner vanilla
  bloom_filter_agg has no Substrait mapping and emits version=0 bytes is what caused the
  kBloomFilterV1 == version (1 vs. 0) crash. Leaving both vanilla also preserves
  vanilla's NULL-on-empty-input semantics.

2. Runtime bloom filters: physical rule

injectOptimizerRule registers into extendedOperatorOptimizationRules, which run inside
the Operator Optimization batch, while InjectRuntimeFilter runs in a later batch of
SparkOptimizer. The logical rule therefore never sees runtime-filter bloom expressions:
they do not exist yet when it fires. Without a rewrite they have no Substrait mapping, and
both the consuming FilterExec and the producing aggregate would fall back to the JVM with
R2C/C2R transitions (a serious performance regression on TPC-DS q59 and other
runtime-filter queries, caught by reviewers on an earlier revision of this PR).

A new physical pre-transform rule, RuntimeBloomFilterRewriteRule, handles them. It is
restricted to InjectRuntimeFilter's exact expression shape, which always wraps the key in
xxhash64(...) on both sides:

• producer: bloom_filter_agg(xxhash64(key), ...) to velox_bloom_filter_agg(...)
• consumer: might_contain(bf, xxhash64(key)) to velox_might_contain(...)

Each side is identifiable on its own, so the rewrite stays consistent even when AQE
compiles the bloom-filter subquery separately from the consuming filter stage. This keeps
FilterExecTransformer and the bloom-filter aggregate native, matching the behavior before
this PR; no TPC-DS plan-stability goldens change.

The xxhash64 fingerprint also means DataFrame.stat.bloomFilter() (which builds
bloom_filter_agg on the raw column and deserializes the bytes with Spark's
BloomFilter.readFrom) is structurally never matched, so the
CallerInfo.isBloomFilterStatFunction stack-walk hack is removed.

3. Whole-stage reversion safety for runtime bloom filters (injectFinal)

ExpandFallbackPolicy's whole-stage fallback can revert a runtime bloom filter's partial
and final aggregation stages independently, since they run as separate physical operators
across a shuffle boundary. If only one side reverts, a native-Velox stage would end up
paired with a vanilla-reverted stage.

RuntimeBloomFilterRewriteRule is now also registered at injectFinal, which runs
regardless of whether ExpandFallbackPolicy reverted the plan. Re-applying the same
rewrite there restores both sides to the Velox form after any such reversion. The rule is
idempotent for already-rewritten expressions, so this is a no-op in the common case.

This is only safe because of GLUTEN-12613 (#12614, merged): VeloxBloomFilterAggregate's
JVM-side buffer capacity now agrees with the native bloom_filter_agg aggregate's. Before
that fix, a reverted stage merging with a still-native stage could silently corrupt the
filter instead of crashing on a version mismatch.

How was this patch tested?

New GlutenBloomFilterFallbackSuite in backends-velox covering:

• whole-stage AQE fallback at thresholds 1 and 2 (the GLUTEN-12013 scenarios)
• the SPARK-54336 literal-value path (both sides stay vanilla)
• JVM-mode subquery aggregation (VeloxBloomFilterAggregate inside ObjectHashAggregateExec)
• DataFrame.stat.bloomFilter() staying on Spark-native bytes
• spark.gluten.sql.native.bloomFilter=false early exit
• runtime bloom filters keeping a native FilterExecTransformer with velox_might_contain
  and a VeloxBloomFilterAggregate producer

All green locally (Spark 4.0, Velox backend), plus GlutenBloomFilterAggregateQuerySuite
(14/14, including SPARK-54336) and all 7 TPC-DS/TPC-H plan-stability suites against the
unchanged goldens (322/322).

[github-actions]

Run Gluten Clickhouse CI on x86

[Copilot]

Copilot encountered an error and was unable to review this pull request. You can try again by re-requesting a review.

[philo-he]

Gentle ping for a maintainer review. The link-referenced-issues CI check that initially failed has since re-run successfully — all checks are green.

Also re-raising: could a maintainer remove the CORE label? The three changed files are all Velox-backend-specific (backends-velox/ and gluten-ut/spark40/) — no common core code is touched, so VELOX label only is correct.

@brijrajk, thanks for the PR. Could you rebase the code to see if the CI failures go away?

[brijrajk]

Done — rebased onto current main and force-pushed. Fresh CI triggered.

[Copilot]

Pull request overview

Copilot reviewed 3 out of 3 changed files in this pull request and generated 2 comments.

[philo-he]

@brijrajk, could you first check if Copilot's comments make sense?

[brijrajk]

Thanks for flagging this, @philo-he!

Both of Copilot's comments were valid:

1. Patcher active when native bloom filter is disabled

When spark.gluten.sql.native.bloomFilter=false, Stage 0 falls back to Spark and produces Spark-format bytes. The joint-fallback rule still wraps Stage 1 in a FallbackNode, so the patcher was incorrectly rewriting it to VeloxBloomFilterMightContain — which would cause the same IOException the patcher was introduced to fix, just from the opposite trigger.

Added a second guard: if (!GlutenConfig.get.enableNativeBloomFilter) return plan. This mirrors the existing guard already in BloomFilterMightContainJointRewriteRule.

2. df.collect + df.count() runs the query twice

Combined into assert(df.collect().length == 200003L) — single execution, same failure signal if the IOException is thrown.

[philo-he]

@brijrajk, thanks for the update. Could you check if my following understanding is correct?

Besides the spark.gluten.sql.native.bloomFilter=false setting, which makes the bloom filter fall back in stage 0, there's another case: the fallback policy can also cause it to fall back. In that case, if we rely solely on checking that config, could it lead to an incompatibility issue in stage 1?

[brijrajk]

@philo-he You are absolutely right. We confirmed it with a test case.

How threshold and cost work

ExpandFallbackPolicy counts the number of columnar↔row conversion boundaries inside a stage. If that count (cost) meets COLUMNAR_WHOLESTAGE_FALLBACK_THRESHOLD, the entire stage is wrapped in a FallbackNode and runs as plain Spark.

Scenario	Threshold	Stage 0 cost	Stage 1 cost	Outcome
Original fix (PR as-is)	2	1 → native ✓	2 → whole-stage fallback	Stage 0 Velox bytes, Stage 1 JVM — patcher correct
Your scenario	1	1 → whole-stage fallback	≥ 1 → whole-stage fallback	Stage 0 Spark bytes, Stage 1 JVM — patcher misfires

Test case confirming the failure

testGluten(
  "Test bloom_filter_agg whole-stage fallback when both stages fall back",
  Issue12013) {
  ...
  if (BackendsApiManager.getSettings.requireBloomFilterAggMightContainJointFallback()) {
    // threshold=1: Stage 0's inherent transition cost of 1 meets the threshold, so
    // ExpandFallbackPolicy promotes Stage 0 to a whole-stage fallback as well.
    // Stage 0 runs as Spark and produces Spark-format bytes. Stage 1 also falls back.
    // The patcher must NOT rewrite BloomFilterMightContain -> VeloxBloomFilterMightContain
    // in this case.
    withSQLConf(
      GlutenConfig.COLUMNAR_FILTER_ENABLED.key -> "false",
      GlutenConfig.COLUMNAR_WHOLESTAGE_FALLBACK_THRESHOLD.key -> "1",
      SQLConf.ANSI_ENABLED.key -> "false"
    ) {
      val df = spark.sql(sqlString)
      assert(df.collect().length == 200003L)
    }
  }
}

Output

- Gluten - Test bloom_filter_agg whole-stage fallback when both stages fall back *** FAILED ***
  org.apache.spark.SparkException: Job aborted due to stage failure: Task 0 in stage 7.0 failed 1 times,
  most recent failure: Lost task 0.0 in stage 7.0: org.apache.gluten.exception.GlutenException:
  Exception: VeloxUserError
  Error Source: USER
  Error Code: INVALID_ARGUMENT
  Reason: (1 vs. 0)
  Retriable: False
  Expression: kBloomFilterV1 == version
  Function: mayContain
  File: velox/common/base/BloomFilter.h
  Line: 70

    at org.apache.gluten.utils.VeloxBloomFilterJniWrapper.mightContainLongOnSerializedBloom(Native Method)
    at org.apache.gluten.utils.VeloxBloomFilter.mightContainLongOnSerializedBloom(VeloxBloomFilter.java:163)
    ...

Tests: succeeded 1, failed 1

kBloomFilterV1 == version failing with (1 vs. 0) is the exact version-byte mismatch: Velox's reader expected its own format (1) but got Spark's format (0).

Proposed fix

The root cause is that enableNativeBloomFilter answers "is native bloom filter on in config?" but the right question is "did Stage 0 actually run natively?" The fix is to make the guard structural: inside patchBloomFilterMightContain, before rewriting, inspect the physical plan referenced by bloomFilterExpression. If Stage 0's plan is itself a FallbackNode, it will produce Spark-format bytes and Stage 1 must be left with the vanilla BloomFilterMightContain.

Do you see any concerns with this approach, or is there a cleaner way you would handle it?

[zhztheplayer]

@brijrajk thanks.

[zhztheplayer]

I don't recall why BloomFilterMightContainJointRewriteRule was made a physical rule, but can you try turning it to a logical rule anyway? So such a patcher rule can be avoided?

[brijrajk]

Done — BloomFilterMightContainJointRewriteRule is now a Rule[LogicalPlan] registered via injectOptimizerRule, modelled after CollectRewriteRule. The patcher is gone. Running as an optimizer rule ensures both substitutions (BloomFilterAggregate → VeloxBloomFilterAggregate and BloomFilterMightContain → VeloxBloomFilterMightContain) are captured in the originalPlan snapshot before ExpandFallbackPolicy takes it, so the byte format stays consistent regardless of which stages fall back. This also fixes the threshold=1 case where Stage 0 itself falls back (the patcher would incorrectly rewrite the filter side while Stage 0 was producing Spark-format bytes).

[rdtr]

could CallerInfo.isBloomFilterStatFunction() and inBloomFilterStatFunctionCall() be removed now with this PR?

[rdtr]

I think Patcher is now gone so this comment is outdated?

[brijrajk]

Fixed — updated the comment to describe the optimizer rule approach instead.

[github-actions]

Run Gluten Clickhouse CI on x86

[github-actions]

Run Gluten Clickhouse CI on x86

[github-actions]

Run Gluten Clickhouse CI on x86

[brijrajk]

Quick update: rebased twice more since the above to stay current with upstream (no code changes -- same commit, same test results as listed above).

You may see red X's in the CI history from a couple of the reruns. I looked into each one and they're unrelated infra flakiness, not caused by this PR:

• A hung spark-test-spark41 job (unrelated test suite spinning on a temp-directory warning loop, later fixed by an unrelated CI test-splitting change)
• A native crash in SparkScriptTransformationExec's subprocess-failure test (pre-existing Velox memory-manager teardown bug, unrelated to bloom filters)
• A transient Maven dependency-resolution failure fetching avalon-framework in an unrelated integration-test module

None reproduce locally and none touch the bloom-filter code path. Current CI run should be clean; happy to dig into any specific failure if it looks otherwise.

[zhztheplayer]

@brijrajk did I miss something, why 3 rules in total are still needed for the bloom filter offloading? I was thinking we only need one single injectFinal now?

[zhztheplayer]

@brijrajk The bottom line is, we just ensure all Spark bloomfilter_agg and might_contain are replaced with velox_bloomfilter_agg (both partial and final) and velox_might_contain, with no regard to the operator-wise fallbacks. A vanilla (fallen back) subquery including velox_bloomfilter_agg should be compatible with an offloaded scan including velox_might_contain, and vice versa. Anything beyond that in the rules may over-complicate the approach.

[brijrajk]

@zhztheplayer Good question, and the underlying insight is right: byte-format compatibility only depends on which expression class runs (VeloxBloomFilterAggregate/VeloxBloomFilterMightContain vs. the vanilla ones), not on whether the surrounding operator ended up native or JVM-fallback. That's exactly why the injectFinal re-application is safe to add on top of the existing rules.

I built the literal version of the proposal to check it properly before replying: one rule, unconditionally rewriting every bloom_filter_agg/might_contain to Velox form (no literal/non-literal distinction, no xxhash64 restriction), registered only at injectFinal. Three real regressions showed up:

1. DataFrame.stat.bloomFilter() breaks. Its standalone bloom_filter_agg (not paired with any might_contain) gets swept up by an unconditional match and rewritten to Velox format, and Spark's own BloomFilter.readFrom() can't parse it:

java.io.IOException: Unexpected Bloom filter version number (16777472)

2. Native offload is lost for runtime bloom filters. Confirmed via plan diff: the producing aggregate ran as vanilla ObjectHashAggregate and the consuming filter as plain Filter wrapped in RowToVeloxColumnar, even though nothing structurally prevented native execution. Cause: HeuristicTransform's offload validator runs before injectFinal, so it sees the still-vanilla expression and rejects native offload for that operator -- by the time injectFinal swaps in the Velox class, the physical operator type is already locked in. This is the same R2C/C2R regression that was the reason the physical injectPreTransform registration got added for runtime filters earlier in this PR's review. TPCH q19's plan-stability golden also diverged under this build, so it's not an isolated case.

3. SPARK-54336's NULL-on-empty-input semantics silently break.

== Correct Answer ==   == Gluten Answer ==
[null]                 [false]

Not a crash -- a query that should return null returns false instead, because the literal-value pair no longer stays vanilla.

So I don't think we can collapse to a single late rule: fixing correctness on reversion (what injectFinal is for) and getting native offload in the common case (which needs the rewrite to exist before HeuristicTransform validates) are two different requirements, and a single late hook can only satisfy one of them -- the offload decision is irreversible by the time injectFinal runs.

Separately, unconditional matching (dropping the pairing/fingerprint check) isn't safe regardless of timing, since nothing then distinguishes a bloom_filter_agg feeding a might_contain from a standalone stat.bloomFilter() call.

Happy to explore a narrower version, though -- e.g. dropping the literal-value special case in the logical rule so both literal and non-literal always rewrite to Velox, if we first close the NULL-on-empty-input gap in VeloxBloomFilterMightContain/VeloxBloomFilterAggregate. That would remove one axis of complexity without touching the injection points. Let me know if that's the kind of simplification you had in mind, or if I'm missing something in what you were picturing.

[zhztheplayer]

@brijrajk Thanks.

Then my suggestion is to add a new API injectPre, which modifies the original plan and won't be reverted when whole stage fallback is induced. E.g,:

  injector.injectPre(c =>
    BloomFilterMightContainJointRewriteRule(
      c.session,
      c.caller.isBloomFilterStatFunction()))

and

 val rewrittenPlan =
    transformPlan("pre", preRules(call), inputPlan)

  val suggestedPlan =
    transformPlan("transform", transformRules(call), rewrittenPlan)

  val finalPlan =
    transformPlan(
      "fallback",
      fallbackPolicies(call).map(_(rewrittenPlan)),
      suggestedPlan)

What do you think?

[brijrajk]

@zhztheplayer Took your injectPre idea and actually built it, to check it properly before replying.

What I implemented:

• A new injectPre hook on GlutenInjector.LegacyInjector, backed by an empty-by-default builder list (mirrors the existing hooks -- additive only, no changes needed anywhere that doesn't opt in).
• HeuristicApplier runs a new "pre" phase before "transform", and fallbackPolicies now closes over the post-"pre" plan instead of the raw input -- so a whole-stage revert can no longer strip away anything registered there.
• Switched RuntimeBloomFilterRewriteRule to register once, at injectPre, dropping both current registrations (injectPreTransform + injectFinal).

Result on Velox: it works, cleanly. Full suite passes with the collapsed single registration:

• GlutenBloomFilterFallbackSuite: 8/8, including both whole-stage-reversion tests and the native-offload test
• GlutenBloomFilterAggregateQuerySuite: 28/30 (2 known pre-existing/unrelated failures)
• GlutenInjectRuntimeFilterSuite: 13/13
• Full TPC-DS/TPC-H plan-stability: 322/322

So the mechanism is sound, and it does genuinely simplify things for the runtime-filter rule specifically.

Checked ClickHouse too, since this touches shared gluten-core: CHRuleApi.scala needs zero changes and builds, links, and packages successfully against the modified core. That's expected since it never calls injectPre (the new builder list stays empty for CH), but wanted it verified against a real build rather than just argued.

Why I'm not just pushing this into #12151: it's a gluten-core/HeuristicApplier change, shared across every backend, not scoped to bloom filters at all. Given how this PR has already gone once around on scope (the capacity fix got split into #12614), I'd rather do the same here: propose injectPre as its own follow-up PR once we're aligned on the API shape, and keep #12151 as the three-rule version it already is, which is correct and fully tested. Happy to open that PR, or if you'd rather drive the core API design yourself given you proposed it, that works too. Let me know which you'd prefer.

[github-actions]

Run Gluten Clickhouse CI on x86

[zhztheplayer]

@brijrajk

Does it work for you, if you can open a refactor PR to add injectPre which can be reviewed and merged first, then we rebase this PR to move to injectPre + one rule solution? Thanks!

[brijrajk]

Yes, that works -- opened it: #12669 (tracked as GLUTEN-12668). Verified it there (322/322 plan-stability regression clean on Velox, backends-clickhouse builds unmodified). Once it merges, I'll rebase this PR to switch RuntimeBloomFilterRewriteRule to the single injectPre registration.

[brijrajk]

injectPre (#12669) merged, and this PR is now rebased onto it. RuntimeBloomFilterRewriteRule is registered once now, at injectPre, dropping the previous injectPreTransform + injectFinal pair.

Full suite re-verified against the actual merged hook (not just the earlier prototype):

• GlutenBloomFilterFallbackSuite: 8/8, including both whole-stage-reversion scenarios and native-offload preservation
• GlutenBloomFilterAggregateQuerySuite: 28/30 (2 known pre-existing/unrelated failures)
• GlutenInjectRuntimeFilterSuite: 13/13
• Full TPC-DS/TPC-H plan-stability: 322/322

Ready for another look.

[github-actions]

Run Gluten Clickhouse CI on x86

[zhztheplayer]

@brijrajk Can you check whether the optimizer logical rule BloomFilterMightContainJointRewriteRule is still needed?

[github-actions]

Run Gluten Clickhouse CI on x86

[brijrajk]

@zhztheplayer Confirmed: the logical rule is not needed. I have rewritten the PR to the one-rule injectPre design. Rebased onto current main.

The whole PR is now:

• BloomFilterMightContainJointRewriteRule stays exactly as it is on main: a physical rule, unconditional, guarded by c.caller.isBloomFilterStatFunction(). The earlier conversion to a logical optimizer rule is reverted, and the CallerInfo flag is restored.
• Its registration moves from injectPreTransform to injectPre. That is the entire production change.
• RuntimeBloomFilterRewriteRule is removed. It turns out to be unnecessary: the unconditional rewrite already covers InjectRuntimeFilter's expressions, and the runtime-filter reversion test passes without it.
• The regression suite stays.

Production diff is now 12 added / 5 removed lines in VeloxRuleApi.scala, versus 629/77 across 5 files before.

Worth recording why the earlier iterations pointed the other way, since it was a measurement artifact rather than a real constraint. What was being measured was the conditional rule ported to a physical hook. That variant has to reach into the scalar subquery's plan to rewrite the producer, and it cannot: AdaptiveSparkPlanExec is a LeafExecNode, so traversal stops there. Instrumenting it shows aggsReachedByTransform=0, planUnchanged=true while the aggregate is plainly visible in the subquery's treeString. The consumer still gets rewritten, producing exactly the mismatch we are trying to avoid.

The unconditional rewrite never needs that reach-in. The producer is rewritten when Gluten's rules run on the subquery's own AdaptiveSparkPlanExec, so both sides agree by construction and nothing crosses the AQE boundary. That also removes the need for the SPARK-54336 literal special case, which only existed as a consequence of operating at the logical level.

Measured on Spark 4.0 / Scala 2.13:

Configuration	Result
current main	bug reproduces: SparkException in ScalarSubquery.updateResult, Unsupported BloomFilter version: 0
main + registration at injectPre	fallback suite 8/8, aggregate + runtime filter + stat + subquery 270/270, plan stability 322/322, no crashes

Why the hook matters: injectPreTransform runs after HeuristicApplier captures the plan ExpandFallbackPolicy reverts to, so a whole-stage revert strips the rewrite from one stage while another keeps it. injectPre runs before that capture, so the rewrite is part of the revert target. The byte-format mismatch is concrete: Velox writes the version as a single int8 of 1, while Spark's BloomFilterImpl.writeTo writes a 4-byte big-endian int, so vanilla bytes start with 0x00 and the native reader reports version 0.

Two notes on the tests:

• The two whole-stage-fallback tests previously asserted optimizedPlan contained velox_might_contain. That only holds for a logical rule, so they now assert on executedPlan, which is what actually determines whether the rewrite survived the fallback.
• The SPARK-54336 test no longer asserts that both sides stay vanilla. It asserts the invariant that matters: never velox_might_contain without a matching velox_bloom_filter_agg.

This also relies on VeloxBloomFilterAggregate's JVM-side buffer sizing agreeing with the native aggregate's (GLUTEN-12613, merged), otherwise a reverted stage merging with a native one corrupts the filter silently instead of failing loudly.