feat: add zero-wait distributed runtime filtering to AQE

ballista/client/tests/runtime_filter.rs

@@ -0,0 +1,132 @@
+// 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.
+
+#![cfg(feature = "standalone")]
+
+use ballista::prelude::SessionContextExt;
+use datafusion::arrow::datatypes::{DataType, Field, Schema};
+use datafusion::arrow::util::pretty::pretty_format_batches;
+use datafusion::prelude::*;
+use std::sync::Arc;
+
+async fn run(ctx: &SessionContext, sql: &str) -> String {
+    let batches = tokio::time::timeout(std::time::Duration::from_secs(30), async {
+        ctx.sql(sql).await.unwrap().collect().await.unwrap()
+    })
+    .await
+    .expect("runtime filter must not deadlock the query");
+    pretty_format_batches(&batches).unwrap().to_string()
+}
+
+#[tokio::test]
+async fn distributed_runtime_filter_matches_unfiltered_joins() {
+    let dir = tempfile::tempdir().unwrap();
+    for name in ["build", "probe"] {
+        std::fs::create_dir(dir.path().join(name)).unwrap();
+    }
+    std::fs::write(dir.path().join("build/a.csv"), "k,v\n2,1\n,2\n").unwrap();
+    std::fs::write(dir.path().join("build/b.csv"), "k,v\n8,3\n2,4\n").unwrap();
+    for part in 0..4 {
+        let mut csv = "k,v\n,0\n".to_string();
+        for k in 0..25 {
+            csv.push_str(&format!("{},{}\n", part * 25 + k, k));
+        }
+        std::fs::write(dir.path().join(format!("probe/{part}.csv")), csv).unwrap();
+    }
+    let ctx = SessionContext::standalone().await.unwrap();
+    for setting in [
+        "ballista.planner.adaptive.enabled = true",
+        "datafusion.optimizer.prefer_hash_join = true",
+        "datafusion.execution.target_partitions = 4",
+    ] {
+        run(&ctx, &format!("SET {setting}")).await;
+    }
+    let schema = Arc::new(Schema::new(vec![
+        Field::new("k", DataType::Int64, true),
+        Field::new("v", DataType::Int64, true),
+    ]));
+    // Parquet provides build cardinality, allowing AQE to choose a broadcast
+    // stage whose completion naturally precedes the probe scan.
+    let build_parquet = dir.path().join("build_parquet");
+    let local =
+        SessionContext::new_with_config(SessionConfig::new().with_target_partitions(2));
+    local
+        .read_csv(
+            dir.path().join("build").to_str().unwrap(),
+            CsvReadOptions::new().schema(&schema),
+        )
+        .await
+        .unwrap()
+        .write_parquet(build_parquet.to_str().unwrap(), Default::default(), None)
+        .await
+        .unwrap();
+    ctx.register_parquet(
+        "build",
+        build_parquet.to_str().unwrap(),
+        ParquetReadOptions::default(),
+    )
+    .await
+    .unwrap();
+    ctx.register_csv(
+        "probe",
+        dir.path().join("probe").to_str().unwrap(),
+        CsvReadOptions::new().schema(&schema),
+    )
+    .await
+    .unwrap();
+    let queries = [
+        "SELECT b.k, b.v, p.v FROM build b JOIN probe p ON b.k = p.k WHERE b.v > 0 AND p.v >= 0 ORDER BY b.k, b.v, p.v",
+        "SELECT b.k, b.v, p.v FROM build b JOIN probe p ON b.k IS NOT DISTINCT FROM p.k WHERE b.v > 0 AND p.v >= 0 ORDER BY b.k, b.v, p.v",
+        "SELECT p.k FROM build b RIGHT JOIN probe p ON b.k = p.k ORDER BY p.k",
+        "SELECT b.k, p.k FROM build b JOIN (SELECT k FROM probe ORDER BY k LIMIT 3) p ON b.k = p.k ORDER BY b.k, p.k",
+    ];
+    for query in queries {
+        run(&ctx, "SET ballista.planner.runtime_filter.enabled = false").await;
+        let expected = run(&ctx, query).await;
+        run(&ctx, "SET ballista.planner.runtime_filter.enabled = true").await;
+        assert_eq!(run(&ctx, query).await, expected, "{query}");
+    }
+    let plan = run(&ctx, &format!("EXPLAIN ANALYZE {}", queries[0])).await;
+    assert!(plan.contains("RuntimeFilterBuildExec"), "{plan}");
+    assert!(plan.contains("RuntimeFilterApplyExec"), "{plan}");
+    assert!(plan.contains("ready=true"), "{plan}");
+    assert!(
+        plan.split("rows_pruned=").skip(1).any(|suffix| {
+            suffix
+                .chars()
+                .take_while(|c| c.is_ascii_digit())
+                .collect::<String>()
+                .parse::<usize>()
+                .is_ok_and(|rows| rows > 0)
+        }),
+        "runtime filter must actually prune rows: {plan}"
+    );
+    // Independently runnable partitioned stages must preserve results even
+    // when the producer has not reported before the probe is dispatched.
+    for setting in [
+        "ballista.optimizer.broadcast_join_threshold_bytes = 0",
+        "ballista.optimizer.broadcast_join_threshold_rows = 0",
+        "ballista.optimizer.stage_build_side = false",
+    ] {
+        run(&ctx, &format!("SET {setting}")).await;
+    }
+    let partitioned = run(&ctx, queries[0]).await;
+    let plan = run(&ctx, &format!("EXPLAIN ANALYZE {}", queries[0])).await;
+    assert!(plan.contains("mode=Partitioned"), "{plan}");
+    run(&ctx, "SET ballista.planner.runtime_filter.enabled = false").await;
+    assert_eq!(partitioned, run(&ctx, queries[0]).await);
+}

ballista/core/src/config.rs

@@ -182,6 +182,10 @@ pub const BALLISTA_SEMI_JOIN_REDUCTION: &str = "ballista.optimizer.semi_join_red
 pub const BALLISTA_BALANCE_SCAN_FILE_GROUPS: &str =
     "ballista.optimizer.balance_scan_file_groups";
 
+/// Enable cross-stage hash-join runtime filters (AQE only).
+pub const BALLISTA_RUNTIME_FILTER_ENABLED: &str =
+    "ballista.planner.runtime_filter.enabled";
+
 /// Configuration key to enable AQE coalesce-shuffle-partitions rule.
 /// Disabled by default — opt in when the workload benefits from larger
 /// downstream tasks more than from preserved parallelism.
@@ -509,6 +513,9 @@ static CONFIG_ENTRIES: LazyLock<HashMap<String, ConfigEntry>> = LazyLock::new(||
                          "Wait time in milliseconds between IO retries in the Ballista client.".to_string(),
                          DataType::UInt64,
                          Some(3000.to_string())),
+        ConfigEntry::new(BALLISTA_RUNTIME_FILTER_ENABLED.to_string(),
+                         "Enables bounded distributed runtime filters for AQE inner and semi hash joins.".to_string(),
+                         DataType::Boolean, Some(false.to_string())),
         ConfigEntry::new(BALLISTA_COALESCE_ENABLED.to_string(),
                          "Enables the AQE coalesce-shuffle-partitions rule. \
                           Disabled by default — opt in when fewer/larger \
@@ -1125,6 +1132,11 @@ impl BallistaConfig {
         self.get_bool_setting(BALLISTA_BALANCE_SCAN_FILE_GROUPS)
     }
 
+    /// Whether distributed runtime filtering is enabled.
+    pub fn runtime_filter_enabled(&self) -> bool {
+        self.get_bool_setting(BALLISTA_RUNTIME_FILTER_ENABLED)
+    }
+
     /// Returns whether the AQE coalesce-shuffle-partitions rule is enabled.
     pub fn coalesce_enabled(&self) -> bool {
         self.get_bool_setting(BALLISTA_COALESCE_ENABLED)

ballista/core/src/execution_plans/runtime_filter.rs

@@ -114,25 +114,46 @@ fn bloom_bits(hash: u64) -> impl Iterator<Item = usize> {
 }
 
 impl RuntimeFilter {
+    /// Whether the summary accepts every non-NULL key. Union cannot make a
+    /// saturated Bloom selective again; its range bounds may still be useful.
+    pub fn is_saturated(&self) -> bool {
+        self.in_hashes.is_none() && self.bloom.iter().all(|&byte| byte == u8::MAX)
+    }
+
+    fn is_nearly_saturated(&self) -> bool {
+        // With four probes, 99% occupancy implies about 96% false positives.
+        // Keep the complete payload and its bounds, but avoid probing this Bloom.
+        self.in_hashes.is_none()
+            && self
+                .bloom
+                .iter()
+                .map(|byte| byte.count_ones() as usize)
+                .sum::<usize>()
+                * 100
+                >= BLOOM_BYTES * 8 * 99
+    }
+
     fn update(&mut self, keys: &ArrayRef) -> Result<()> {
         self.has_null |= keys.null_count() > 0;
-        with_hashes([keys], &HASH_STATE, |hashes| {
-            for (row, &hash) in hashes.iter().enumerate() {
-                if keys.is_null(row) {
-                    continue;
-                }
-                for bit in bloom_bits(hash) {
-                    self.bloom[bit / 8] |= 1 << (bit % 8);
-                }
-                if let Some(values) = &mut self.in_hashes {
-                    values.insert(hash);
-                    if values.len() > MAX_VALUES {
-                        self.in_hashes = None;
+        if !self.is_saturated() {
+            with_hashes([keys], &HASH_STATE, |hashes| {
+                for (row, &hash) in hashes.iter().enumerate() {
+                    if keys.is_null(row) {
+                        continue;
+                    }
+                    for bit in bloom_bits(hash) {
+                        self.bloom[bit / 8] |= 1 << (bit % 8);
+                    }
+                    if let Some(values) = &mut self.in_hashes {
+                        values.insert(hash);
+                        if values.len() > MAX_VALUES {
+                            self.in_hashes = None;
+                        }
                     }
                 }
-            }
-            Ok(())
-        })?;
+                Ok(())
+            })?;
+        }
         if !self.bounds_disabled && keys.null_count() != keys.len() {
             let mut min = MinAccumulator::try_new(keys.data_type())?;
             let mut max = MaxAccumulator::try_new(keys.data_type())?;
@@ -406,6 +427,16 @@ impl RuntimeFilterApplyExec {
             metrics: ExecutionPlanMetricsSet::new(),
         })
     }
+
+    /// Bind a complete report while preserving the input and join key.
+    pub fn with_payload(&self, payload: RuntimeFilterPayload) -> Result<Self> {
+        Self::try_new(
+            self.input.clone(),
+            self.filter_id,
+            self.key.clone(),
+            Some(payload),
+        )
+    }
 }
 
 impl DisplayAs for RuntimeFilterApplyExec {
@@ -481,6 +512,12 @@ impl ExecutionPlan for RuntimeFilterApplyExec {
         let Some(filter) = self.filter.clone() else {
             return Ok(input);
         };
+        // Preserve the payload's bounds for scan pruning, but avoid hashing
+        // batches against a nearly full Bloom with little filtering capacity.
+        if filter.is_nearly_saturated() {
+            return Ok(input);
+        }
+
         let key = self.key.clone();
         let pruned = MetricBuilder::new(&self.metrics).counter("rows_pruned", partition);
         let stream = input.map(move |batch| {
@@ -674,4 +711,63 @@ mod tests {
         })
         .unwrap();
     }
+
+    #[test]
+    fn runtime_filter_saturation_preserves_later_bounds_and_nulls() {
+        let mut filter = RuntimeFilter::default();
+        filter
+            .update(&(Arc::new(Int64Array::from_iter_values(0..1_000_000)) as ArrayRef))
+            .unwrap();
+        assert!(filter.is_saturated());
+        filter
+            .update(
+                &(Arc::new(Int64Array::from(vec![Some(-1), Some(2_000_000), None]))
+                    as ArrayRef),
+            )
+            .unwrap();
+        let payload = filter.to_proto(1).unwrap();
+        assert!(payload.has_null);
+        assert_eq!(
+            ScalarValue::try_from(payload.min.as_ref().unwrap()).unwrap(),
+            ScalarValue::Int64(Some(-1))
+        );
+        assert_eq!(
+            ScalarValue::try_from(payload.max.as_ref().unwrap()).unwrap(),
+            ScalarValue::Int64(Some(2_000_000))
+        );
+    }
+
+    #[tokio::test]
+    async fn runtime_filter_nearly_full_bloom_bypasses_row_filtering() {
+        let mut filter = RuntimeFilter::default();
+        filter
+            .update(&(Arc::new(Int64Array::from_iter_values(0..120_000)) as ArrayRef))
+            .unwrap();
+        assert!(!filter.is_saturated());
+        let batch = RecordBatch::try_from_iter([(
+            "key",
+            Arc::new(Int64Array::from(vec![Some(2), None])) as ArrayRef,
+        )])
+        .unwrap();
+        let input = MemorySourceConfig::try_new_exec(
+            &[vec![batch.clone()]],
+            batch.schema(),
+            None,
+        )
+        .unwrap();
+        let apply: Arc<dyn ExecutionPlan> = Arc::new(
+            RuntimeFilterApplyExec::try_new(
+                input,
+                7,
+                Arc::new(Column::new("key", 0)),
+                Some(filter.to_proto(7).unwrap()),
+            )
+            .unwrap(),
+        );
+        let result = collect(apply, SessionContext::new().task_ctx())
+            .await
+            .unwrap();
+        // Bypass passes the batch through, including NULLs the join will check.
+        assert_eq!(result, vec![batch]);
+    }
 }

ballista/core/src/extension.rs

@@ -821,8 +821,9 @@ impl SessionConfigHelperExt for SessionConfig {
             // boundaries that run as independent tasks, so when the producing
             // operator and the consuming scan land in different stages the
             // filter is never populated across the boundary and the scan blocks
-            // forever. Disable dynamic filter pushdown until Ballista can carry
-            // dynamic filters across stage boundaries.
+            // forever. Keep these local dynamic expressions disabled. Ballista's
+            // opt-in runtime-filter hub transports independent completed join
+            // summaries; it does not populate DataFusion's dynamic expressions.
             //
             // See https://github.com/apache/datafusion-ballista/issues/1375
             .set_bool("datafusion.optimizer.enable_dynamic_filter_pushdown", false)