brijrajk commented on code in PR #12151: URL: https://github.com/apache/gluten/pull/12151#discussion_r3619723780
########## 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 Good catch to press on this -- you're right that `injectFinal` always executes and the expression rewrite itself does work. The confusion is that "rewrite succeeds" and "result is correct" are separate claims. Two different problems, and `injectFinal` only fixes one: | Problem | Fixed by injectFinal? | Why | | --- | --- | --- | | Byte **format** (velox vs. spark) | Yes | The rewrite does run on the reverted plan; probe E's aggregate shows `ObjectHashAggregateExec[VeloxBloomFilterAggregate:Final]`, not vanilla `BloomFilterAggregate`. No version-mismatch crash. | | Bloom **capacity** (bit array size) | No | `VeloxBloomFilterAggregate.createAggregationBuffer()` always calls `VeloxBloomFilter.empty(estimatedNumItems)` (line 102) -- reading the capacity baked into the expression tree by `InjectRuntimeFilter` (100 in my probe), independent of who executes it or when it was rewritten. | So in probe E: the producer's **partial** stage ran natively and used the session-config capacity (1M items -- native `bloom_filter_agg` ignores the expression's own args). The **final** stage reverted, and `injectFinal` correctly re-rewrote its expression to `VeloxBloomFilterAggregate` -- but that class builds its buffer from the expression's `estimatedNumItems` (100), not from the incoming bytes' actual size. Merging a 1M-capacity filter into a 100-capacity buffer hits Velox's `BloomFilter::merge`, whose size guard is a `VELOX_DCHECK_EQ` (compiled out in release) -- so it silently ORs mismatched ranges instead of failing. In short: `injectFinal` guarantees both sides are the *same expression class*. It does not guarantee they agree on *capacity*, because capacity is sourced differently by the native and JVM code paths and the rewrite never touches that. That's the prerequisite I flagged in my last comment -- fixing it means making the JVM class size its buffer from the same source the native side uses (or adopt the incoming filter's size on first merge), independent of whether we end up using `injectFinal`, the fallback-policy block, or the upstream extension point. -- 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]
