Skip to content

Commit 8639584

Browse files
committed
[GLUTEN-12013][VL] Fix bloom-filter bytes corruption on whole-stage AQE fallback
`BloomFilterMightContainJointRewriteRule` previously rewrote every `BloomFilterAggregate` it encountered, including standalone usages such as `DataFrame.stat.bloomFilter()`. That API collects the aggregate output bytes and passes them directly to `BloomFilter.readFrom()`, which expects Spark-native format; receiving Velox-format bytes caused `java.io.IOException: Unexpected Bloom filter version number` (surfaced as a CI failure in `GlutenDataFrameStatSuite - Bloom filter`). Fix: only rewrite `BloomFilterAggregate` when it appears inside the `ScalarSubquery` of a `BloomFilterMightContain`. Standalone aggregates are left untouched so that collected bytes remain in Spark-native format. Add a regression test (`GlutenBloomFilterFallbackSuite`) to guard against reintroducing this regression. Local test results (Spark 4.0, Velox backend): - GlutenDataFrameStatSuite : 25/25 passed (was failing) - GlutenBloomFilterFallbackSuite : 4/4 passed - GlutenBloomFilterAggregateQuerySuite: 14/14 passed - GlutenInjectRuntimeFilterSuite : 13/13 passed
1 parent 7cc6faf commit 8639584

2 files changed

Lines changed: 34 additions & 8 deletions

File tree

backends-velox/src/main/scala/org/apache/gluten/extension/BloomFilterMightContainJointRewriteRule.scala

Lines changed: 17 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -21,7 +21,7 @@ import org.apache.gluten.expression.VeloxBloomFilterMightContain
2121
import org.apache.gluten.expression.aggregate.VeloxBloomFilterAggregate
2222

2323
import org.apache.spark.sql.SparkSession
24-
import org.apache.spark.sql.catalyst.expressions.BloomFilterMightContain
24+
import org.apache.spark.sql.catalyst.expressions.{BloomFilterMightContain, ScalarSubquery}
2525
import org.apache.spark.sql.catalyst.expressions.aggregate.{AggregateExpression, BloomFilterAggregate}
2626
import org.apache.spark.sql.catalyst.plans.logical.LogicalPlan
2727
import org.apache.spark.sql.catalyst.rules.Rule
@@ -35,6 +35,11 @@ import org.apache.spark.sql.catalyst.rules.Rule
3535
* an individual stage fallback to a whole-stage AQE fallback. This guarantees that both sides of
3636
* the bloom-filter pair always produce and consume the same byte format, regardless of whether
3737
* stages fall back to JVM execution after AQE re-planning.
38+
*
39+
* `BloomFilterAggregate` is only rewritten when it appears inside the [[ScalarSubquery]] of a
40+
* [[BloomFilterMightContain]]. Standalone usages (in particular `DataFrame.stat.bloomFilter()`,
41+
* which collects bloom filter bytes and passes them to `BloomFilter.readFrom()`) are intentionally
42+
* left untouched so that the returned bytes remain in Spark-native format.
3843
*/
3944
case class BloomFilterMightContainJointRewriteRule(spark: SparkSession)
4045
extends Rule[LogicalPlan] {
@@ -44,13 +49,17 @@ case class BloomFilterMightContainJointRewriteRule(spark: SparkSession)
4449
return plan
4550
}
4651
plan.transformAllExpressions {
47-
case aggExpr @ AggregateExpression(b: BloomFilterAggregate, _, _, _, _) =>
48-
aggExpr.copy(aggregateFunction = VeloxBloomFilterAggregate(
49-
b.child,
50-
b.estimatedNumItemsExpression,
51-
b.numBitsExpression,
52-
b.mutableAggBufferOffset,
53-
b.inputAggBufferOffset))
52+
case BloomFilterMightContain(subq: ScalarSubquery, v) =>
53+
val rewrittenPlan = subq.plan.transformAllExpressions {
54+
case ae @ AggregateExpression(b: BloomFilterAggregate, _, _, _, _) =>
55+
ae.copy(aggregateFunction = VeloxBloomFilterAggregate(
56+
b.child,
57+
b.estimatedNumItemsExpression,
58+
b.numBitsExpression,
59+
b.mutableAggBufferOffset,
60+
b.inputAggBufferOffset))
61+
}
62+
VeloxBloomFilterMightContain(subq.withNewPlan(rewrittenPlan), v)
5463
case BloomFilterMightContain(bf, v) =>
5564
VeloxBloomFilterMightContain(bf, v)
5665
}

gluten-ut/test/src/test/scala/org/apache/gluten/sql/GlutenBloomFilterFallbackSuite.scala

Lines changed: 17 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -155,6 +155,23 @@ class GlutenBloomFilterFallbackSuite extends WholeStageTransformerSuite {
155155
}
156156
}
157157

158+
// GLUTEN-12013: DataFrame.stat.bloomFilter() must not be affected by the optimizer rule.
159+
// The rule must only rewrite BloomFilterAggregate inside a BloomFilterMightContain subquery.
160+
// A standalone BloomFilterAggregate (as used here) must remain vanilla so that the collected
161+
// bytes are in Spark-native format and BloomFilter.readFrom() succeeds.
162+
test("GLUTEN-12013: DataFrame.stat.bloomFilter() produces Spark-readable bytes") {
163+
if (BackendsApiManager.getSettings.requireBloomFilterAggMightContainJointFallback()) {
164+
val table = "bloom_filter_stat_test"
165+
withTempView(table) {
166+
(1L to 1000L).toDF("col").createOrReplaceTempView(table)
167+
// Must not throw: java.io.IOException: Unexpected Bloom filter version number
168+
val bf = spark.table(table).stat.bloomFilter("col", 1000L, 0.01)
169+
// Bloom filters have no false negatives: every inserted value must be present.
170+
assert(bf.mightContainLong(500L), "Expected 500 to be in bloom filter")
171+
}
172+
}
173+
}
174+
158175
// GLUTEN-12013: native bloom filter disabled -- early-exit path of the optimizer rule.
159176
// When spark.gluten.sql.native.bloomFilter=false the rule returns the plan unchanged.
160177
// BloomFilterAggregate / BloomFilterMightContain remain as vanilla Spark expressions and

0 commit comments

Comments
 (0)