Skip to main content

akar_processor/processor/mapper/
map_projection.rs

1use super::ExecutionContext;
2use crate::expression_evaluator::ExpressionEvaluator;
3use crate::physical_operator::*;
4use crate::processor::projection_helper::resolve_projection_column_index;
5use akar_common::error::ProcessorError;
6use akar_common::vector::DataChunk;
7use akar_parser::ast::Expression;
8use akar_planner::logical_operator::LogicalOperator;
9use std::sync::{Arc, Mutex};
10
11fn projection_needs_expression_eval(expr: &Expression) -> bool {
12    matches!(
13        expr,
14        Expression::FunctionCall(_, _)
15            | Expression::Constant(_)
16            | Expression::BinaryOp(_, _, _)
17            | Expression::UnaryOp(_, _)
18            | Expression::List(_)
19            | Expression::Map(_)
20            | Expression::Parameter(_)
21            | Expression::ExistsSubquery(_)
22            | Expression::ListPredicate { .. }
23    )
24}
25
26/// Derive the output column name of a projected expression, so downstream
27/// operators (ORDER BY / TOP-K) can resolve sort keys by name rather than by
28/// position (P52.1).
29fn expression_field_name(expr: &Expression) -> String {
30    match expr {
31        Expression::PropertyAccess(obj, prop) => {
32            if let Expression::Variable(var) = &**obj {
33                format!("{var}.{prop}")
34            } else {
35                prop.clone()
36            }
37        }
38        Expression::Variable(name) => name.clone(),
39        _ => String::new(),
40    }
41}
42
43/// Resolve ORDER BY / TOP-K sort keys to column indices.
44///
45/// Sort keys are expressions (e.g. `p.age`); they must be mapped to the actual
46/// output column they refer to. Positional mapping sorts by the i-th key's
47/// position, which is wrong whenever ORDER BY references a non-first column
48/// (e.g. `RETURN p.name, p.age ORDER BY p.age` sorted by name).
49fn resolve_sort_keys(sort_keys: &[(Expression, bool)], input: &[DataChunk]) -> Vec<(u32, bool)> {
50    sort_keys
51        .iter()
52        .enumerate()
53        .map(|(i, (expr, asc))| {
54            let col = input.first().and_then(|c| resolve_projection_column_index(expr, c));
55            (col.unwrap_or(i) as u32, *asc)
56        })
57        .collect()
58}
59
60pub fn map_and_execute_projection(
61    op: &LogicalOperator,
62    current_input: Vec<DataChunk>,
63    ctx: &mut ExecutionContext,
64) -> Result<Vec<DataChunk>, ProcessorError> {
65    match op {
66        LogicalOperator::Projection(p) => {
67            let input = if p.children.is_empty() {
68                current_input
69            } else {
70                ctx.execute_children(&p.children)?
71            };
72
73            let result = if p.expressions.is_empty() {
74                input
75            } else {
76                let needs_eval = p
77                    .expressions
78                    .iter()
79                    .any(|be| projection_needs_expression_eval(&be.expression));
80
81                if needs_eval {
82                    let registry = ctx
83                        .function_registry
84                        .clone()
85                        .ok_or_else(|| "No function registry available for expression projection".to_string())?;
86
87                    let mut eval = ExpressionEvaluator::new(registry);
88                    if let Some(ref seq_fn) = ctx.sequence_fn {
89                        eval = eval.with_sequence_fn(seq_fn.clone());
90                    }
91                    if let Some(ref subquery_fn) = ctx.subquery_fn {
92                        eval = eval.with_subquery_fn(subquery_fn.clone());
93                    }
94
95                    let mut output = Vec::with_capacity(input.len());
96                    for chunk in input {
97                        let mut fields = Vec::with_capacity(p.expressions.len());
98                        let mut field_types = Vec::with_capacity(p.expressions.len());
99                        for be in &p.expressions {
100                            let result_vec = eval.evaluate(&be.expression, &chunk)?;
101                            let pt = result_vec.physical_type();
102                            let arr = akar_common::arrow_vector::ArrowVector::from_legacy(&result_vec).array;
103                            fields.push(arr);
104                            field_types.push(pt);
105                        }
106                        let size = fields.first().map(|f| f.len()).unwrap_or(chunk.size);
107                        let field_names = p
108                            .expressions
109                            .iter()
110                            .map(|be| expression_field_name(&be.expression))
111                            .collect();
112                        output.push(DataChunk {
113                            fields,
114                            field_types,
115                            size,
116                            field_names,
117                            sel_vector: None,
118                        });
119                    }
120                    output
121                } else {
122                    let column_indices: Vec<usize> = if let Some(first_chunk) = input.first() {
123                        p.expressions
124                            .iter()
125                            .filter_map(|be| resolve_projection_column_index(&be.expression, first_chunk))
126                            .collect()
127                    } else {
128                        Vec::new()
129                    };
130                    let column_indices = if column_indices.len() == p.expressions.len() {
131                        column_indices
132                    } else {
133                        (0..p.expressions.len()).collect()
134                    };
135                    let proj = PhysicalProjection { column_indices };
136                    proj.execute(input)?
137                }
138            };
139            Ok(result)
140        }
141        LogicalOperator::Filter(f) => {
142            let evaluator = ctx.function_registry.clone().map(|reg| {
143                let mut eval = ExpressionEvaluator::new(reg);
144                if let Some(ref seq_fn) = ctx.sequence_fn {
145                    eval = eval.with_sequence_fn(seq_fn.clone());
146                }
147                if let Some(ref subquery_fn) = ctx.subquery_fn {
148                    eval = eval.with_subquery_fn(subquery_fn.clone());
149                }
150                Arc::new(Mutex::new(eval))
151            });
152            let filter = if let Some(eval) = evaluator {
153                PhysicalFilter::with_evaluator(f.expression.clone(), eval)
154            } else {
155                PhysicalFilter::new(f.expression.clone())
156            };
157            let result = filter.execute(current_input)?;
158            Ok(result)
159        }
160        LogicalOperator::Limit(l) => {
161            let limit = PhysicalLimit {
162                limit: l.limit,
163                offset: l.offset,
164            };
165            let result = limit.execute(current_input)?;
166            Ok(result)
167        }
168        LogicalOperator::TopK(tk) => {
169            let sort_keys = resolve_sort_keys(&tk.sort_keys, &current_input);
170            let topk = PhysicalTopK {
171                sort_keys,
172                limit: tk.limit,
173                offset: tk.offset,
174            };
175            let result = topk.execute(current_input)?;
176            Ok(result)
177        }
178        LogicalOperator::OrderBy(o) => {
179            let sort_keys = resolve_sort_keys(&o.sort_keys, &current_input);
180            let order = PhysicalOrderBy { sort_keys };
181            let result = order.execute(current_input)?;
182            Ok(result)
183        }
184        LogicalOperator::Flatten(f) => {
185            let input = if f.children.is_empty() {
186                current_input
187            } else {
188                ctx.execute_children(&f.children)?
189            };
190            let flatten = PhysicalFlatten::new(f.group_pos);
191            flatten.execute(input)
192        }
193        LogicalOperator::Unwind(uw) => {
194            let unwind = PhysicalUnwind {
195                expression: uw.expression.clone(),
196                variable: uw.variable.clone(),
197            };
198            let result = unwind.execute(current_input)?;
199            Ok(result)
200        }
201        LogicalOperator::Partitioner(p) => {
202            const MORSEL_SIZE: usize = 1024;
203            let partitioner = Partitioner::new(MORSEL_SIZE);
204            let morsels = partitioner.execute(current_input)?;
205
206            let mut results = Vec::new();
207            for _morsel in morsels {
208                let child_result = ctx.execute_children(&p.children)?;
209                results.extend(child_result);
210            }
211            Ok(results)
212        }
213        _ => Err(format!("Not a projection/filter operator: {:?}", op).into()),
214    }
215}