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
26pub fn map_and_execute_projection(
27    op: &LogicalOperator,
28    current_input: Vec<DataChunk>,
29    ctx: &mut ExecutionContext,
30) -> Result<Vec<DataChunk>, ProcessorError> {
31    match op {
32        LogicalOperator::Projection(p) => {
33            let input = if p.children.is_empty() {
34                current_input
35            } else {
36                ctx.execute_children(&p.children)?
37            };
38
39            let result = if p.expressions.is_empty() {
40                input
41            } else {
42                let needs_eval = p
43                    .expressions
44                    .iter()
45                    .any(|be| projection_needs_expression_eval(&be.expression));
46
47                if needs_eval {
48                    let registry = ctx
49                        .function_registry
50                        .clone()
51                        .ok_or_else(|| "No function registry available for expression projection".to_string())?;
52
53                    let mut eval = ExpressionEvaluator::new(registry);
54                    if let Some(ref seq_fn) = ctx.sequence_fn {
55                        eval = eval.with_sequence_fn(seq_fn.clone());
56                    }
57                    if let Some(ref subquery_fn) = ctx.subquery_fn {
58                        eval = eval.with_subquery_fn(subquery_fn.clone());
59                    }
60
61                    let mut output = Vec::with_capacity(input.len());
62                    for chunk in input {
63                        let mut fields = Vec::with_capacity(p.expressions.len());
64                        let mut field_types = Vec::with_capacity(p.expressions.len());
65                        for be in &p.expressions {
66                            let result_vec = eval.evaluate(&be.expression, &chunk)?;
67                            let pt = result_vec.physical_type();
68                            let arr = akar_common::arrow_vector::ArrowVector::from_legacy(&result_vec).array;
69                            fields.push(arr);
70                            field_types.push(pt);
71                        }
72                        let size = fields.first().map(|f| f.len()).unwrap_or(chunk.size);
73                        output.push(DataChunk {
74                            fields,
75                            field_types,
76                            size,
77                            field_names: vec![],
78                            sel_vector: None,
79                        });
80                    }
81                    output
82                } else {
83                    let column_indices: Vec<usize> = if let Some(first_chunk) = input.first() {
84                        p.expressions
85                            .iter()
86                            .filter_map(|be| resolve_projection_column_index(&be.expression, first_chunk))
87                            .collect()
88                    } else {
89                        Vec::new()
90                    };
91                    let column_indices = if column_indices.len() == p.expressions.len() {
92                        column_indices
93                    } else {
94                        (0..p.expressions.len()).collect()
95                    };
96                    let proj = PhysicalProjection { column_indices };
97                    proj.execute(input)?
98                }
99            };
100            Ok(result)
101        }
102        LogicalOperator::Filter(f) => {
103            let evaluator = ctx.function_registry.clone().map(|reg| {
104                let mut eval = ExpressionEvaluator::new(reg);
105                if let Some(ref seq_fn) = ctx.sequence_fn {
106                    eval = eval.with_sequence_fn(seq_fn.clone());
107                }
108                if let Some(ref subquery_fn) = ctx.subquery_fn {
109                    eval = eval.with_subquery_fn(subquery_fn.clone());
110                }
111                Arc::new(Mutex::new(eval))
112            });
113            let filter = if let Some(eval) = evaluator {
114                PhysicalFilter::with_evaluator(f.expression.clone(), eval)
115            } else {
116                PhysicalFilter::new(f.expression.clone())
117            };
118            let result = filter.execute(current_input)?;
119            Ok(result)
120        }
121        LogicalOperator::Limit(l) => {
122            let limit = PhysicalLimit {
123                limit: l.limit,
124                offset: l.offset,
125            };
126            let result = limit.execute(current_input)?;
127            Ok(result)
128        }
129        LogicalOperator::TopK(tk) => {
130            let sort_keys: Vec<(u32, bool)> = tk
131                .sort_keys
132                .iter()
133                .enumerate()
134                .map(|(i, _s)| (i as u32, tk.sort_keys.get(i).map(|s| s.1).unwrap_or(true)))
135                .collect();
136            let topk = PhysicalTopK {
137                sort_keys,
138                limit: tk.limit,
139                offset: tk.offset,
140            };
141            let result = topk.execute(current_input)?;
142            Ok(result)
143        }
144        LogicalOperator::OrderBy(o) => {
145            let sort_keys: Vec<(u32, bool)> = o
146                .sort_keys
147                .iter()
148                .enumerate()
149                .map(|(i, _s)| (i as u32, o.sort_keys.get(i).map(|s| s.1).unwrap_or(true)))
150                .collect();
151            let order = PhysicalOrderBy { sort_keys };
152            let result = order.execute(current_input)?;
153            Ok(result)
154        }
155        LogicalOperator::Flatten(f) => {
156            let input = if f.children.is_empty() {
157                current_input
158            } else {
159                ctx.execute_children(&f.children)?
160            };
161            let flatten = PhysicalFlatten::new(f.group_pos);
162            flatten.execute(input)
163        }
164        LogicalOperator::Unwind(uw) => {
165            let unwind = PhysicalUnwind {
166                expression: uw.expression.clone(),
167                variable: uw.variable.clone(),
168            };
169            let result = unwind.execute(current_input)?;
170            Ok(result)
171        }
172        LogicalOperator::Partitioner(p) => {
173            const MORSEL_SIZE: usize = 1024;
174            let partitioner = Partitioner::new(MORSEL_SIZE);
175            let morsels = partitioner.execute(current_input)?;
176
177            let mut results = Vec::new();
178            for _morsel in morsels {
179                let child_result = ctx.execute_children(&p.children)?;
180                results.extend(child_result);
181            }
182            Ok(results)
183        }
184        _ => Err(format!("Not a projection/filter operator: {:?}", op).into()),
185    }
186}
akar_processor/processor/mapper/map_projection.rs

akar_processor/processor/mapper/
map_projection.rs
