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 Good question. Verified at three levels: code inspection, Spark-source trace, and runtime probes in a container. Summary: **1. Operator-level mixes are safe by construction.** This rule rewrites the *expressions* at pre-transform. Wherever the operators land (native or JVM), the expressions stay `VeloxBloomFilterAggregate` / `VeloxBloomFilterMightContain`, which run on the JVM via JNI with the same version=1 byte format. Byte format travels with the expression class, not the operator. | 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 | **2. Full disclosure: whole-stage fallback *reversion* is a real gap.** `HeuristicApplier` captures `originalPlan` before pre-transform, so `ExpandFallbackPolicy` reversion strips the rewrite from that stage only, while the subquery is prepared independently (`PlanSubqueries`). With the fallback thresholds enabled, formats can split in both directions: | Probe | Non-default configs | Result | | --- | --- | --- | | D | `columnar.filter.enabled=false` + `wholeStage.fallback.threshold=1` | CRASH: Velox `(1 vs. 0) Unsupported BloomFilter version: 0` | | E | `wholeStage.fallback.threshold=1` only | CRASH: `IOException: Unexpected Bloom filter version number (16777218)` | **3. Caveats bounding the gap:** - Unreachable with default configs (both thresholds default to `-1`). - Pre-existing: I reran all five probes with main's bloom-filter code in the same environment; main crashes identically on D and E. This PR neither introduces nor widens it. - Fails loudly (version-check errors on both sides); never a silent wrong result. - The user-facing path is already reversion-safe in this PR: the logical rule bakes the rewrite into `originalPlan` itself (that is the GLUTEN-12013 fix, proven by the threshold=1/2 tests). **4. Why not fixable the same way here:** it needs a logical rewrite after `InjectRuntimeFilter`, and Spark's extension API has no hook between that batch and physical planning. Re-patching fallen-back plans is the "patcher" we removed earlier at review feedback. **5. Candidate fixes (follow-up):** - Propose an upstream Spark extension point for post-`InjectRuntimeFilter` logical rules, then move the runtime-filter rewrite there. - Or a narrowly-guarded joint re-rewrite in the `final` rule phase (applies to fallen-back plans too), restricted to the `xxhash64` shape on both sides. @philo-he @zhztheplayer what do you think the next steps should be for resolving this gap: a follow-up issue (I can file it with the probe D/E reproduction), or do you consider it 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]
