brijrajk commented on code in PR #12151: URL: https://github.com/apache/gluten/pull/12151#discussion_r3600614771
########## 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] { Review Comment: @zhztheplayer @philo-he thanks for both proposals. Rather than answer from theory, I prototyped the `injectFinal` idea and probed it at runtime. The results changed my recommendation, so full findings below, including one correction to my own earlier comment. ## zhztheplayer's questions **"Why is the logical rule still kept?"** The Operator Optimization batch is a valid injection point for the expressions it targets: user-written `might_contain` exists at analysis time; only `InjectRuntimeFilter`'s expressions are missing there (created in a later `SparkOptimizer` batch). It must stay for two reasons: 1. The literal-vs-non-literal pair decision (SPARK-54336 must stay vanilla; GLUTEN-12013 must become Velox) needs whole-plan visibility. After `PlanSubqueries` splits compilation, a producer-side `bloom_filter_agg(col)` alone is ambiguous between the two cases; there is no structural fingerprint like the runtime filters' `xxhash64` wrapper to disambiguate. 2. Being logical bakes the rewrite into the plan before physical planning, so `originalPlan` itself carries the Velox variants; the user-facing population is reversion-safe with zero patching (proven by the threshold=1/2 tests). **"Any risk if we always replace in the physical phase / move to `injectFinal`?"** Yes, two, and the second is empirical: 1. *Move* alone would regress offload: validation needs the Velox expression before `HeuristicTransform`, so it has to be keep-preTransform-plus-add-final. 2. I implemented exactly that (same `xxhash64` guards, registered additionally at `injectFinal`) and reran the probes. Probe D flipped from crash to correct, and the guards held (`stat.bloomFilter`, SPARK-54336, `native.bloomFilter=false` all fine). But probe E (`wholeStage.fallback.threshold=1` alone) flipped from a crash to a **silent wrong result: 6 of 10 rows, deterministically missing the same keys across reruns**. A bloom pre-filter can only produce false positives, so missing rows means the filter itself was corrupt. That is strictly worse than the crash, so I have not pushed it. ## Root cause of the silent corruption (new finding) Plan diagnostics for probe E showed a **phase split**: the producer's partial aggregate stayed native (`FlushableHashAggregateExecTransformer[VeloxBloomFilterAggregate:Partial]`) while the reverted final-aggregation stage ran on the JVM (`ObjectHashAggregateExec[VeloxBloomFilterAggregate:Final]`). Same expression class, same byte format, no crash. But the two implementations size their blooms from different sources: - JVM `VeloxBloomFilterAggregate.createAggregationBuffer()` uses the expression's `estimatedNumItems` (for runtime filters, `InjectRuntimeFilter` sets it to the creation-side row count: 100 here). - Native `bloom_filter_agg` ignores the expression arguments and uses the session confs Gluten forwards to the native query context (`spark.sql.optimizer.runtime.bloomFilter.*`, defaults 1M items / 8388608 bits). The JVM final agg then merges a large native partial into its small buffer through Velox's `BloomFilter::merge` (common/base/BloomFilter.h), whose size guard is `VELOX_DCHECK_EQ(bits_.size(), otherSize)`, compiled out in release builds, followed by an unconditional `bits::orBits(..., 64 * otherSize)`. Bit positions were inserted modulo the large capacity and queried modulo the small one: deterministic false negatives. This incompatibility is latent and independent of the PoC: any execution pairing a native-partial with a JVM-final `velox_bloom_filter_agg` under divergent capacity parameters is exposed. Today the known reversion paths surface it as a crash because the vanilla/velox format check fails first; the PoC removed the crash and unmasked the silent path. **Correction to my Jul 14 comment:** probe E's original crash was this same phase split, i.e. the vanilla final aggregate choking on native partial buffers, not the consumer filter reverting (the consumer stayed a native `FilterExecTransformer` throughout). Same class of failure, wrong operator attributed. ## philo-he's proposals **Spark extension API for post-`InjectRuntimeFilter` logical rules.** Agreed this is the cleanest long-term shape: the runtime-filter rewrite would land in the optimized logical plan, so `originalPlan` carries it into every stage and reversion can never strip it, exactly like the user-facing path today. The physical rule would then be deleted entirely. Two qualifications: (a) as you noted, upstream acceptance is uncertain and it is a multi-release timeline, so it cannot be this PR's blocker; (b) the phase-split finding above means even this does not fully close the problem by itself. A reverted final-agg stage still executes the JVM implementation against native partial buffers, so the capacity alignment is a prerequisite regardless of where the rewrite runs. I think it is still worth filing the SPARK JIRA; happy to draft it. **Block whole-stage fallback when a stage contains a bloom filter.** The phase-split finding actually strengthens this option: it is the only proposal that prevents reversion-induced splits entirely (both the producer/consumer split and the partial/final split), sidestepping the format problem and the capacity problem at once. On implementation, the layering concern is solvable without `ExpandFallbackPolicy` learning Velox types: the policy is constructed in `VeloxRuleApi` (`injectFallbackPolicy(c => p => ExpandFallbackPolicy(...))`), so backends-velox can pass a bloom-detecting predicate into it. The trade-off is that it overrides the user's explicit fallback decision for those stages; whether that is acceptable is a judgement call I would defer to you both. ## Where this leaves the three options - **`injectFinal` (zhztheplayer):** right structure, but only safe after the capacity divergence is fixed. Prerequisite: make JVM `VeloxBloomFilterAggregate` size its buffer from the same source as native (the forwarded session confs), or adopt the incoming filter's size on first merge (`bits_.empty()` branch already supports adoption). Also worth proposing to Velox that the `DCHECK` become a user check so a mismatch can never be silent. - **Fallback-policy block (philo-he):** most robust short-term; no capacity prerequisite; costs an override of the user's fallback intent. - **Upstream extension point (philo-he):** best long-term; needs capacity alignment anyway; not a blocker for this PR. ## Proposed sequencing 1. Merge this PR as-is: it fixes GLUTEN-12013 and SPARK-54336, restores native runtime filters, changes no goldens, and the residual reversion gap is default-off, fails loudly, and is byte-identical to main (verified by running the same probes against main's bloom code). 2. Follow-up issue 1: capacity alignment between JVM and native `velox_bloom_filter_agg` (prerequisite), plus a phase-split regression test. 3. Follow-up issue 2: close the reversion gap on top of (2), via `injectFinal` or the policy predicate, whichever you prefer. 4. Long-term: SPARK JIRA for a post-`InjectRuntimeFilter` extension point, which would let the physical rule collapse into the logical one. Happy to file the follow-up issues with the probe reproductions. Does this sequencing work for you both? -- 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]
