brijrajk commented on code in PR #12151: URL: https://github.com/apache/gluten/pull/12151#discussion_r3619865949
########## 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: Correction to my last comment: I re-verified the root cause by building and running in the container with richer diagnostics, and the "capacity source" explanation I gave (JVM using the expression's item count vs. native using session confs) was wrong. Posting the real mechanism, confirmed with actual numbers. ## What I found instead `estimatedNumItems`/`numBits` are identical on both engines: `1000000`/`8388608` (Spark's plain defaults; my test table has no `ANALYZE` stats, so `InjectRuntimeFilter` never sees a real row count). There is no source divergence. The actual bug: JVM and native compute the buffer size from the **same two literals** using **different formulas**, each using only one of the two arguments. | Path | Formula | Result for (items=1000000, bits=8388608) | | --- | --- | --- | | JVM `VeloxBloomFilterAggregate.createAggregationBuffer()` (line 102) | `VeloxBloomFilter.empty(estimatedNumItems)` -> Velox `reset(1000000)` -> `nextPow2(1000000)/4` words. Ignores `numBits` entirely. | 16,777,216 bits | | Native `BloomFilterAggAggregate.cpp` | `capacity_ = numBits/16` -> `reset(capacity_)`. Ignores raw item count once `numBits` is present. | 8,388,608 bits | Exactly 2x, deterministic, and independent of table stats or session config -- reproducible for any runtime bloom filter through this code path, not specific to my probe's row counts. | Probe | Partial stage | Final stage | Sizes match? | Result | | --- | --- | --- | --- | --- | | A baseline | native | native | yes (both native formula) | 10/10 | | D `filter=false` + `threshold=1` | JVM | JVM | yes (both JVM formula) | 10/10 | | E `threshold=1` only | native | JVM | **no** | 6/10 | Probe D reverts the *whole* subquery stage (partial+final together), so both ends use the same formula and stay consistent. Probe E reverts only the final-aggregation stage while the partial stage stays native -- that's the actual dangerous case, and it produces silent row loss for exactly this reason: native partial sets bits modulo an 8,388,608-bit array; the JVM final buffer (sized for 16,777,216 bits, i.e. `VeloxBloomFilter.empty(1000000)`) merges those bits into the first half of a differently-sized array (unchecked, per the `DCHECK` noted earlier), and later queries hash modulo the *larger* array -- so previously-set bits land at different positions than where they were inserted. ## What doesn't change The conclusion is the same, just for the right reason: any fix needs the JVM aggregate's buffer sizing reconciled with the native formula (e.g. size from `numBits` the same way, or adopt the incoming filter's size on first merge), independent of whether we land on `injectFinal`, the fallback-policy block, or the upstream extension point. Apologies for the imprecise mechanism in my last comment -- appreciate you pushing on this, it's what surfaced the container run that caught it. -- 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]
