krishiv_plan/optimizer/
broadcast.rs

1//! Broadcast-auto logical optimizer rule.
2
3use crate::{LogicalPlan, NodeOp, PlanNode};
4
5use super::OptimizerRule;
6
7/// Default threshold for auto-broadcast: tables with estimated rows below
8/// this value are candidates for broadcast join.  ~1M rows ≈ 100 MiB at 100
9/// bytes/row.
10pub const DEFAULT_BROADCAST_THRESHOLD_ROWS: u64 = 1_000_000;
11
12/// Logical optimizer rule that marks small scan nodes as broadcast-eligible.
13///
14/// Scans the logical plan for `NodeOp::Scan` nodes whose `estimated_rows` is
15/// set and below the threshold.  Such nodes are annotated with
16/// `broadcast_eligible = true` so the lowering pass promotes their exchange
17/// to `Broadcast` partitioning.
18///
19/// The threshold is deliberately conservative (1M rows).  Without `estimated_rows`
20/// populated from source metadata (parquet footer, Kafka stats, etc.) the rule
21/// is a no-op.
22pub struct BroadcastAutoRule {
23    /// Max rows a table can have to be considered broadcast-eligible.
24    max_rows: u64,
25}
26
27impl BroadcastAutoRule {
28    /// Create a new rule with the given max row threshold.
29    pub fn new(max_rows: u64) -> Self {
30        Self { max_rows }
31    }
32}
33
34impl OptimizerRule for BroadcastAutoRule {
35    fn name(&self) -> &str {
36        "broadcast-auto"
37    }
38
39    fn apply(&self, plan: &LogicalPlan) -> Option<LogicalPlan> {
40        let nodes = plan.nodes();
41        let mut changed = false;
42        let mut new_nodes: Vec<PlanNode> = Vec::with_capacity(nodes.len());
43
44        for node in nodes {
45            let is_small_scan = matches!(node.op(), Some(NodeOp::Scan { .. }))
46                && node.estimated_rows().is_some_and(|r| r <= self.max_rows);
47
48            if is_small_scan && !node.broadcast_eligible() {
49                changed = true;
50                new_nodes.push(node.clone().with_broadcast_eligible(true));
51            } else {
52                new_nodes.push(node.clone());
53            }
54        }
55
56        if !changed {
57            return None;
58        }
59
60        let mut new_plan = LogicalPlan::new(plan.name(), plan.kind());
61        for n in new_nodes {
62            new_plan = new_plan.with_node(n);
63        }
64        Some(new_plan)
65    }
66}
krishiv_plan/optimizer/broadcast.rs

krishiv_plan/optimizer/
broadcast.rs
