linkrs_query/optimizer/heuristic/join_optimization/
push_project_down_join.rs

1//! Rules for pushing projection operations down to JOIN
2//!
3//! This rule pushes the projection operation down to both sides of the JOIN, thereby reducing the amount of intermediate data.
4//!
5//! # Conversion example
6//!
7//! Before:
8//! ```text
9//!   Project(col1, col2, col3)
10//!           |
11//!   InnerJoin
12//!   /          \
13//! Left        Right
14//! (col1-col5) (col2-col6)
15//! ```
16//!
17//! After:
18//! ```text
19//!   InnerJoin
20//!   /          \
21//! Project     Project
22//! (col1)      (col2, col3)
23//!   |            |
24//! Left        Right
25//! ```
26//!
27//! # Applicable Conditions
28//!
29//! The Project node is located above the JOIN node.
30//! The projection columns can be separated into left and right sides.
31//! The JOIN keys are retained in the projection.
32
33use crate::optimizer::heuristic::context::RewriteContext;
34use crate::optimizer::heuristic::pattern::Pattern;
35use crate::optimizer::heuristic::result::{RewriteError, RewriteResult, TransformResult};
36use crate::optimizer::heuristic::rule::{PushDownRule, RewriteRule};
37use crate::planning::plan::core::nodes::base::plan_node_traits::SingleInputNode;
38use crate::planning::plan::core::nodes::join::join_node::{InnerJoinNode, LeftJoinNode};
39use crate::planning::plan::core::nodes::operation::project_node::ProjectNode;
40use crate::planning::plan::PlanNodeEnum;
41use linkrs_core::types::expr::contextual::ContextualExpression;
42use linkrs_core::types::expr::expression_context::ExpressionAnalysisContext;
43use linkrs_core::types::expr::visitor::ExpressionVisitor;
44use linkrs_core::types::expr::visitor_collectors::VariableCollector;
45use linkrs_core::types::YieldColumn;
46use linkrs_core::Expression;
47use std::sync::Arc;
48
49/// Rules for pushing projection operations down to JOIN
50///
51/// Push the projection operation down to both sides of the JOIN to reduce intermediate results.
52#[derive(Debug)]
53pub struct PushProjectDownJoinRule;
54
55impl PushProjectDownJoinRule {
56    pub fn new() -> Self {
57        Self
58    }
59
60    fn collect_variables(&self, expr: &Expression) -> Vec<String> {
61        let mut collector = VariableCollector::new();
62        collector.visit(expr);
63        collector.variables
64    }
65
66    fn get_columns_from_expr(&self, expr: &ContextualExpression) -> Vec<String> {
67        if let Some(expr_meta) = expr.expression() {
68            self.collect_variables(expr_meta.inner())
69        } else {
70            Vec::new()
71        }
72    }
73
74    fn get_columns_from_yield_column(&self, col: &YieldColumn) -> Vec<String> {
75        let mut vars = self.get_columns_from_expr(&col.expression);
76        vars.push(col.alias.clone());
77        vars
78    }
79
80    fn create_project_node(
81        &self,
82        input: PlanNodeEnum,
83        columns: Vec<String>,
84        ctx: Arc<ExpressionAnalysisContext>,
85    ) -> Result<PlanNodeEnum, RewriteError> {
86        let yield_columns: Vec<YieldColumn> = columns
87            .into_iter()
88            .map(|col| {
89                let expr = Expression::variable(&col);
90                let meta = linkrs_core::types::expr::ExpressionMeta::new(expr);
91                let id = ctx.register_expression(meta);
92                let ctx_expr = ContextualExpression::new(id, ctx.clone());
93                YieldColumn {
94                    expression: ctx_expr,
95                    alias: col,
96                }
97            })
98            .collect();
99
100        if yield_columns.is_empty() {
101            return Ok(input);
102        }
103
104        ProjectNode::new(input, yield_columns)
105            .map(PlanNodeEnum::Project)
106            .map_err(|e| {
107                RewriteError::rewrite_failed(format!("Failed to create ProjectNode: {:?}", e))
108            })
109    }
110
111    fn split_columns_for_join(
112        &self,
113        project_columns: &[YieldColumn],
114        left_col_names: &[String],
115        right_col_names: &[String],
116        hash_keys: &[ContextualExpression],
117        probe_keys: &[ContextualExpression],
118    ) -> (Vec<String>, Vec<String>) {
119        let mut left_needed: Vec<String> = Vec::new();
120        let mut right_needed: Vec<String> = Vec::new();
121
122        for col in project_columns {
123            let vars = self.get_columns_from_yield_column(col);
124            for var in vars {
125                if left_col_names.contains(&var) && !left_needed.contains(&var) {
126                    left_needed.push(var.clone());
127                }
128                if right_col_names.contains(&var) && !right_needed.contains(&var) {
129                    right_needed.push(var);
130                }
131            }
132        }
133
134        for key in hash_keys {
135            let vars = self.get_columns_from_expr(key);
136            for var in vars {
137                if left_col_names.contains(&var) && !left_needed.contains(&var) {
138                    left_needed.push(var);
139                }
140            }
141        }
142
143        for key in probe_keys {
144            let vars = self.get_columns_from_expr(key);
145            for var in vars {
146                if right_col_names.contains(&var) && !right_needed.contains(&var) {
147                    right_needed.push(var);
148                }
149            }
150        }
151
152        (left_needed, right_needed)
153    }
154
155    fn apply_to_inner_join(
156        &self,
157        project: &ProjectNode,
158        join: &InnerJoinNode,
159        ctx: &RewriteContext,
160    ) -> RewriteResult<Option<TransformResult>> {
161        let left_col_names = join.left_input().col_names().to_vec();
162        let right_col_names = join.right_input().col_names().to_vec();
163
164        let (left_needed, right_needed) = self.split_columns_for_join(
165            project.columns(),
166            &left_col_names,
167            &right_col_names,
168            join.hash_keys(),
169            join.probe_keys(),
170        );
171
172        if left_needed.is_empty() && right_needed.is_empty() {
173            return Ok(None);
174        }
175
176        if left_needed.len() == left_col_names.len() && right_needed.len() == right_col_names.len()
177        {
178            return Ok(None);
179        }
180
181        let expr_ctx = ctx.expr_context();
182        let mut new_left = join.left_input().clone();
183        let mut new_right = join.right_input().clone();
184
185        if !left_needed.is_empty() && left_needed.len() < left_col_names.len() {
186            new_left = self.create_project_node(new_left, left_needed, expr_ctx.clone())?;
187        }
188
189        if !right_needed.is_empty() && right_needed.len() < right_col_names.len() {
190            new_right = self.create_project_node(new_right, right_needed, expr_ctx)?;
191        }
192
193        let new_join = InnerJoinNode::new(
194            new_left,
195            new_right,
196            join.hash_keys().to_vec(),
197            join.probe_keys().to_vec(),
198        )
199        .map_err(|e| {
200            RewriteError::rewrite_failed(format!("Failed to create InnerJoinNode: {:?}", e))
201        })?;
202
203        let mut result = TransformResult::new();
204        result.erase_curr = true;
205        result.add_new_node(PlanNodeEnum::InnerJoin(new_join));
206
207        Ok(Some(result))
208    }
209
210    fn apply_to_left_join(
211        &self,
212        project: &ProjectNode,
213        join: &LeftJoinNode,
214        ctx: &RewriteContext,
215    ) -> RewriteResult<Option<TransformResult>> {
216        let left_col_names = join.left_input().col_names().to_vec();
217        let right_col_names = join.right_input().col_names().to_vec();
218
219        let (left_needed, right_needed) = self.split_columns_for_join(
220            project.columns(),
221            &left_col_names,
222            &right_col_names,
223            join.hash_keys(),
224            join.probe_keys(),
225        );
226
227        if left_needed.is_empty() && right_needed.is_empty() {
228            return Ok(None);
229        }
230
231        if left_needed.len() == left_col_names.len() && right_needed.len() == right_col_names.len()
232        {
233            return Ok(None);
234        }
235
236        let expr_ctx = ctx.expr_context();
237        let mut new_left = join.left_input().clone();
238        let mut new_right = join.right_input().clone();
239
240        if !left_needed.is_empty() && left_needed.len() < left_col_names.len() {
241            new_left = self.create_project_node(new_left, left_needed, expr_ctx.clone())?;
242        }
243
244        if !right_needed.is_empty() && right_needed.len() < right_col_names.len() {
245            new_right = self.create_project_node(new_right, right_needed, expr_ctx)?;
246        }
247
248        let new_join = LeftJoinNode::new(
249            new_left,
250            new_right,
251            join.hash_keys().to_vec(),
252            join.probe_keys().to_vec(),
253        )
254        .map_err(|e| {
255            RewriteError::rewrite_failed(format!("Failed to create LeftJoinNode: {:?}", e))
256        })?;
257
258        let mut result = TransformResult::new();
259        result.erase_curr = true;
260        result.add_new_node(PlanNodeEnum::LeftJoin(new_join));
261
262        Ok(Some(result))
263    }
264}
265
266impl Default for PushProjectDownJoinRule {
267    fn default() -> Self {
268        Self::new()
269    }
270}
271
272impl RewriteRule for PushProjectDownJoinRule {
273    fn name(&self) -> &'static str {
274        "PushProjectDownJoinRule"
275    }
276
277    fn pattern(&self) -> Pattern {
278        Pattern::new_with_name("Project")
279    }
280
281    fn apply(
282        &self,
283        ctx: &mut RewriteContext,
284        node: &PlanNodeEnum,
285    ) -> RewriteResult<Option<TransformResult>> {
286        let project = match node {
287            PlanNodeEnum::Project(n) => n,
288            _ => return Ok(None),
289        };
290
291        let input = project.input();
292
293        match input {
294            PlanNodeEnum::InnerJoin(join) => self.apply_to_inner_join(project, join, ctx),
295            PlanNodeEnum::LeftJoin(join) => self.apply_to_left_join(project, join, ctx),
296            _ => Ok(None),
297        }
298    }
299}
300
301impl PushDownRule for PushProjectDownJoinRule {
302    fn can_push_down(&self, node: &PlanNodeEnum, target: &PlanNodeEnum) -> bool {
303        matches!(node, PlanNodeEnum::Project(_))
304            && matches!(
305                target,
306                PlanNodeEnum::InnerJoin(_) | PlanNodeEnum::LeftJoin(_)
307            )
308    }
309
310    fn push_down(
311        &self,
312        ctx: &mut RewriteContext,
313        node: &PlanNodeEnum,
314        _target: &PlanNodeEnum,
315    ) -> RewriteResult<Option<TransformResult>> {
316        self.apply(ctx, node)
317    }
318}
319
320#[cfg(test)]
321mod tests {
322    use super::*;
323
324    #[test]
325    fn test_rule_name() {
326        let rule = PushProjectDownJoinRule::new();
327        assert_eq!(rule.name(), "PushProjectDownJoinRule");
328    }
329
330    #[test]
331    fn test_rule_pattern() {
332        let rule = PushProjectDownJoinRule::new();
333        let pattern = rule.pattern();
334        assert!(pattern.node.is_some());
335    }
336}
linkrs_query/optimizer/heuristic/join_optimization/push_project_down_join.rs

linkrs_query/optimizer/heuristic/join_optimization/
push_project_down_join.rs
