brijrajk commented on code in PR #12151: URL: https://github.com/apache/gluten/pull/12151#discussion_r3575606259
########## 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: @philo-he Great question. I verified this at three levels (code inspection, Spark-source trace, and runtime experiments in a container) rather than answering from theory. Full disclosure below, including one pre-existing gap. ## Direct answer: operator-level mixes are safe by construction This rule runs at pre-transform and rewrites the *expressions*. Whether the containing operator later offloads or falls back (validation failure, `columnar.filter=false`, etc.), the expressions remain `VeloxBloomFilterAggregate` / `VeloxBloomFilterMightContain`, and both execute on the JVM via JNI producing/consuming the same Velox (version=1) byte format. The byte format travels with the expression class, not with where the operator runs, so every producer/consumer offload combination is consistent. Runtime evidence (Spark 4.0, Velox backend; a probe suite ran the same runtime-filter join under different configs): | Probe | Non-default configs | Result | | --- | --- | --- | | A | none | SUCCESS, correct rows | | B | `columnar.filter.enabled=false` (consumer operator on JVM) | SUCCESS | | C | `columnar.hashagg.enabled=false` (producer operator on JVM) | SUCCESS | ## Full disclosure: whole-stage fallback *reversion* is a real, pre-existing gap There is a deeper case than operator-level fallback. `HeuristicApplier.makeRule` captures `originalPlan` *before* the pre-transform rules run, and `ExpandFallbackPolicy.fallbackToRowBasedPlan` reverts a fallen-back stage to that plan. The reversion therefore strips this rule's expression rewrite from that stage only, while the subquery producer is compiled by its own independent preparation pass (`PlanSubqueries` calls `QueryExecution.prepareExecutedPlan`). If the whole-stage fallback thresholds are enabled, the two sides can end up on different byte formats: | Probe | Non-default configs | Result | | --- | --- | --- | | D | `columnar.filter.enabled=false` + `wholeStage.fallback.threshold=1` | CRASH: Velox `(1 vs. 0) Unsupported BloomFilter version: 0` (producer stage reverted to vanilla, Velox side read Spark bytes) | | E | `wholeStage.fallback.threshold=1` only | CRASH: `java.io.IOException: Unexpected Bloom filter version number (16777218)` (consumer stage reverted to vanilla, read Velox bytes) | Caveats that bound this gap: 1. **Unreachable with default configs.** Both `wholeStage.fallback.threshold` and `query.fallback.threshold` default to `-1` (disabled). 2. **Pre-existing, not introduced by this PR.** I temporarily restored main's bloom-filter code (physical joint rewrite + `CallerInfo`) in the same environment and reran all five probes: main crashes identically on D and E. This PR neither introduces nor widens the gap. 3. **It fails loudly.** Both directions raise a version-check error (Spark `BloomFilterImpl.readFrom` / Velox `BloomFilterView`); there is no silent wrong result. 4. **The user-facing population is already reversion-safe in this PR.** The logical rule bakes the rewrite into the plan before physical planning, so `originalPlan` itself contains the Velox variants and reversion keeps both sides consistent. That is exactly the GLUTEN-12013 fix, and the threshold=1/2 tests in `GlutenBloomFilterFallbackSuite` prove it. ## Why runtime filters cannot get the same treatment in this PR Closing the gap the same way requires a logical rewrite that runs *after* `InjectRuntimeFilter`, but Spark's extension API has no injection point between that batch and physical planning (`injectOptimizerRule` lands in the earlier Operator Optimization batch). The alternative, re-patching expressions on the already-fallen-back plan, is the "patcher" approach we removed earlier in this PR at review feedback, and it has known misfire hazards when one side is legitimately on Spark-format bytes. ## Proposal Handle it as a follow-up issue rather than in this PR, with two candidate directions: 1. Propose an upstream Spark extension point for post-`InjectRuntimeFilter` logical rules, then move the runtime-filter rewrite there (making it reversion-safe identically to the user-facing case). 2. Alternatively, a narrowly-guarded joint re-rewrite in the `final` rule phase (which also applies to fallen-back plans), restricted to the `xxhash64` runtime-filter shape on both sides so it cannot misfire on Spark-format populations. I can file the follow-up issue with the reproduction (probe D/E configs) documented. Would you and @zhztheplayer prefer that, or do you consider the reversion case in scope for this PR? -- 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]
