brijrajk commented on code in PR #12151: URL: https://github.com/apache/gluten/pull/12151#discussion_r3624938945
########## 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 Short answer: the capacity fix stands on its own and is worth doing regardless. Consolidating to one rule doesn't work as a simple move though. I went ahead and built it to check. ## What I built and ran Capacity fix plus one unified rule, registered at both preTransform and injectFinal, with the logical rule dropped entirely. Ran it against the existing regression suite (GlutenBloomFilterFallbackSuite). ## What failed 4 of the 7 tests crash with the same `(1 vs. 0)` version mismatch this PR exists to fix: | Test | Needs producer/consumer coordination? | Result | | --- | --- | --- | | threshold=2 fallback | Yes | FAILED, `(1 vs. 0)` | | threshold=1 fallback | Yes | FAILED, same crash | | JVM-mode subquery | Yes | FAILED, same crash | | SPARK-54336 literal | Yes | FAILED, same crash | | stat.bloomFilter | No | passed | | native.bloomFilter=false | No | passed | | runtime-filter-native | No, xxhash64 alone is enough signal | passed | The split is clean: anything needing "check the consumer's value, then update the producer to match" fails. Anything that doesn't need that coordination passes. ## Why it fails I logged the plan before and after collect() on the threshold=2 test to see what was actually happening: | | Filter (consumer) | Subquery aggregate (producer) | | --- | --- | --- | | Before collect() | vanilla | vanilla | | After collect() | rewritten correctly | still vanilla | Same rule, same match arm, same subq.withNewPlan(...) call. Only half of it takes effect at runtime. Here's why: `InsertAdaptiveSparkPlan.buildSubqueryMap` compiles and caches each subquery's physical plan once, keyed by exprId, before Gluten's rule ever runs. A physical rule that reaches into the subquery and rewrites it produces a perfectly valid new object, but the subquery executes off that earlier cached copy, not the new one. The consumer rewrite works because it's just a plain expression swap sitting in the outer plan. The producer rewrite needs the cache itself to change, and there's no way to reach it from a rule running this late. That also explains why the runtime-filter rule was never affected: it doesn't reach into the subquery manually at all, it matches the aggregate expression wherever the framework's own traversal already visits it, so it's wired correctly by construction. And it explains why the logical rule works: it rewrites the subquery's logical plan before the cache is even built, so the cache is already correct by the time AQE captures it. ## Suggestion Keep the two rules as they are for this PR. File the capacity fix as its own follow-up since it doesn't depend on any of this. And treat "can a physical rule ever be made reversion-safe for the coordinated case" as a separate question, one the Spark extension point idea might actually answer better than anything we can do rule-side. -- 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]
