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). An `AS alias` overrides the derived name (P53.16, G5).
29fn expression_field_name(alias: Option<&str>, expr: &Expression) -> String {
30    if let Some(a) = alias {
31        return a.to_string();
32    }
33    match expr {
34        Expression::PropertyAccess(obj, prop) => {
35            if let Expression::Variable(var) = &**obj {
36                format!("{var}.{prop}")
37            } else {
38                prop.clone()
39            }
40        }
41        Expression::Variable(name) => name.clone(),
42        _ => String::new(),
43    }
44}
45
46/// Resolve ORDER BY / TOP-K sort keys to column indices.
47///
48/// Sort keys are expressions (e.g. `p.age`); they must be mapped to the actual
49/// output column they refer to. Positional mapping sorts by the i-th key's
50/// position, which is wrong whenever ORDER BY references a non-first column
51/// (e.g. `RETURN p.name, p.age ORDER BY p.age` sorted by name).
52///
53/// Sort keys that are computed expressions (`array_cosine_similarity(...)`,
54/// `a + b`, ...) cannot be mapped to a column: `resolve_projection_column_index`
55/// only handles PropertyAccess/Variable. They are evaluated per row via the
56/// `ExpressionEvaluator` and appended to a copy of the input as synthetic
57/// trailing columns (P53.23). The returned `Vec<DataChunk>` is the (possibly
58/// augmented) input to sort on, and `usize` is the number of appended columns
59/// that must be stripped from the operator output.
60fn resolve_sort_keys(
61    sort_keys: &[(Expression, bool)],
62    input: &[DataChunk],
63    ctx: &mut ExecutionContext,
64) -> Result<(Vec<(u32, bool)>, Vec<DataChunk>, usize), ProcessorError> {
65    let base_cols = input.first().map(|c| c.num_fields()).unwrap_or(0);
66    let mut resolved = Vec::with_capacity(sort_keys.len());
67    let mut computed: Vec<Expression> = Vec::new();
68    for (expr, asc) in sort_keys.iter() {
69        let col = input.first().and_then(|c| resolve_projection_column_index(expr, c));
70        match col {
71            Some(idx) => resolved.push((idx as u32, *asc)),
72            None => {
73                resolved.push(((base_cols + computed.len()) as u32, *asc));
74                computed.push(expr.clone());
75            }
76        }
77    }
78    if computed.is_empty() {
79        return Ok((resolved, input.to_vec(), 0));
80    }
81    let registry = ctx
82        .function_registry
83        .clone()
84        .ok_or_else(|| "No function registry available for computed ORDER BY key".to_string())?;
85    let mut eval = ExpressionEvaluator::new(registry);
86    if let Some(ref seq_fn) = ctx.sequence_fn {
87        eval = eval.with_sequence_fn(seq_fn.clone());
88    }
89    if let Some(ref subquery_fn) = ctx.subquery_fn {
90        eval = eval.with_subquery_fn(subquery_fn.clone());
91    }
92    let mut augmented = Vec::with_capacity(input.len());
93    for chunk in input {
94        let mut fields = chunk.fields.clone();
95        let mut field_types = chunk.field_types.clone();
96        for expr in &computed {
97            let vv = eval.evaluate_arrow(expr, chunk)?;
98            fields.push(vv.array);
99            field_types.push(vv.physical_type);
100        }
101        augmented.push(DataChunk {
102            fields,
103            field_types,
104            size: chunk.size,
105            field_names: chunk.field_names.clone(),
106            sel_vector: None,
107        });
108    }
109    Ok((resolved, augmented, computed.len()))
110}
111
112/// Drop the synthetic computed sort-key columns appended by `resolve_sort_keys`
113/// (P53.23) from the operator output chunks.
114fn strip_sort_columns(chunks: &mut [DataChunk], extra: usize) {
115    if extra == 0 {
116        return;
117    }
118    for chunk in chunks {
119        let keep = chunk.fields.len().saturating_sub(extra);
120        chunk.fields.truncate(keep);
121        chunk.field_types.truncate(keep);
122        chunk.field_names.truncate(keep);
123    }
124}
125
126/// Resolve a projection expression to one or more input column indices.
127///
128/// Property accesses resolve to a single column; a bare variable either matches
129/// an exact input column or — when it names a node/relationship variable —
130/// expands to every input column prefixed `{var}.` (P53.25, so `WITH a, b, row`
131/// carries all of `a.id/a.content/a._id/...` forward). Returns `None` when the
132/// expression cannot be mapped to columns (caller falls back to positional
133/// indices).
134fn resolve_projection_column_expand(expr: &Expression, chunk: &DataChunk) -> Option<Vec<usize>> {
135    match expr {
136        Expression::PropertyAccess(obj, prop) => {
137            let col_name = if let Expression::Variable(var) = &**obj {
138                format!("{}.{}", var, prop)
139            } else {
140                prop.clone()
141            };
142            if !chunk.field_names.is_empty() {
143                if let Some(idx) = chunk.field_names.iter().position(|n| n == &col_name || n == prop) {
144                    return Some(vec![idx]);
145                }
146            }
147            if let Expression::Variable(var) = &**obj
148                && let Ok(idx) = var.parse::<usize>()
149            {
150                return Some(vec![idx]);
151            }
152            None
153        }
154        Expression::Variable(name) => {
155            if let Ok(idx) = name.parse::<usize>() {
156                return Some(vec![idx]);
157            }
158            if !chunk.field_names.is_empty() {
159                if let Some(idx) = chunk.field_names.iter().position(|n| n == name) {
160                    return Some(vec![idx]);
161                }
162                let prefix = format!("{}.", name);
163                let idxs: Vec<usize> = chunk
164                    .field_names
165                    .iter()
166                    .enumerate()
167                    .filter(|(_, n)| n.starts_with(&prefix))
168                    .map(|(i, _)| i)
169                    .collect();
170                if !idxs.is_empty() {
171                    return Some(idxs);
172                }
173            }
174            None
175        }
176        _ => None,
177    }
178}
179
180/// True when a property access's base is a chunk column (a map/struct value,
181/// e.g. an UNWIND row variable `e`), so `e.src` must be extracted per row
182/// instead of resolving to a qualified column (P53.26/P53.37a).
183fn property_base_is_chunk_column(expr: &Expression, chunk: &DataChunk) -> bool {
184    if let Expression::PropertyAccess(base, _) = expr
185        && let Expression::Variable(var) = &**base
186        && !chunk.field_names.is_empty()
187    {
188        return chunk.field_names.iter().any(|n| n == var);
189    }
190    false
191}
192
193pub fn map_and_execute_projection(
194    op: &LogicalOperator,
195    current_input: Vec<DataChunk>,
196    ctx: &mut ExecutionContext,
197) -> Result<Vec<DataChunk>, ProcessorError> {
198    match op {
199        LogicalOperator::Projection(p) => {
200            let input = if p.children.is_empty() {
201                current_input
202            } else {
203                ctx.execute_children(&p.children)?
204            };
205
206            let result = if p.expressions.is_empty() {
207                input
208            } else {
209                // Expressions that cannot map to a plain column — computed exprs
210                // (function calls, arithmetic) OR property accesses on a map/
211                // UNWIND variable (`e.src`: base `e` is a chunk column but no
212                // qualified `e.src` column exists) — must go through the per-row
213                // evaluator (P53.37a). Other unresolvable accesses keep the
214                // positional fallback so plans whose scan/optional-merge dropped
215                // columns (masked pre-existing gaps) keep their old behavior.
216                let needs_eval = p
217                    .expressions
218                    .iter()
219                    .any(|be| projection_needs_expression_eval(&be.expression))
220                    || input.first().is_some_and(|chunk| {
221                        p.expressions.iter().any(|be| {
222                            matches!(
223                                &be.expression,
224                                Expression::PropertyAccess(_, _) | Expression::Variable(_)
225                            ) && resolve_projection_column_expand(&be.expression, chunk).is_none()
226                                && property_base_is_chunk_column(&be.expression, chunk)
227                        })
228                    });
229
230                if needs_eval {
231                    let registry = ctx
232                        .function_registry
233                        .clone()
234                        .ok_or_else(|| "No function registry available for expression projection".to_string())?;
235
236                    let mut eval = ExpressionEvaluator::new(registry);
237                    if let Some(ref seq_fn) = ctx.sequence_fn {
238                        eval = eval.with_sequence_fn(seq_fn.clone());
239                    }
240                    if let Some(ref subquery_fn) = ctx.subquery_fn {
241                        eval = eval.with_subquery_fn(subquery_fn.clone());
242                    }
243
244                    let mut output = Vec::with_capacity(input.len());
245                    for chunk in input {
246                        let mut fields = Vec::with_capacity(p.expressions.len());
247                        let mut field_types = Vec::with_capacity(p.expressions.len());
248                        for be in &p.expressions {
249                            // Use the Arrow-native evaluator so complex-typed
250                            // results (List/Struct columns and literals) round-trip
251                            // without a ValueVector (which has no side-storage).
252                            let result_vec = eval.evaluate_arrow(&be.expression, &chunk)?;
253                            let pt = result_vec.physical_type;
254                            let arr = result_vec.array.clone();
255                            fields.push(arr);
256                            field_types.push(pt);
257                        }
258                        let size = fields.first().map(|f| f.len()).unwrap_or(chunk.size);
259                        let field_names = p
260                            .expressions
261                            .iter()
262                            .map(|be| expression_field_name(be.alias.as_deref(), &be.expression))
263                            .collect();
264                        output.push(DataChunk {
265                            fields,
266                            field_types,
267                            size,
268                            field_names,
269                            sel_vector: None,
270                        });
271                    }
272                    output
273                } else {
274                    let column_indices: Vec<usize> = if let Some(first_chunk) = input.first() {
275                        let mut all: Option<Vec<usize>> = Some(Vec::new());
276                        for be in &p.expressions {
277                            match resolve_projection_column_expand(&be.expression, first_chunk) {
278                                Some(idxs) => all.as_mut().expect("all is Some").extend(idxs),
279                                None => {
280                                    all = None;
281                                    break;
282                                }
283                            }
284                        }
285                        all.unwrap_or_default()
286                    } else {
287                        Vec::new()
288                    };
289                    let column_indices = if !column_indices.is_empty() {
290                        column_indices
291                    } else {
292                        (0..p.expressions.len()).collect()
293                    };
294                    let rename = column_indices.len() == p.expressions.len() && !p.expressions.is_empty();
295                    let proj = PhysicalProjection { column_indices };
296                    let mut result = proj.execute(input)?;
297                    // Rename output columns to alias-aware names (P53.16): the
298                    // plain-column path copies input field_names, so `RETURN
299                    // m.name AS nm` would otherwise keep `m.name` as the label.
300                    if rename {
301                        let names: Vec<String> = p
302                            .expressions
303                            .iter()
304                            .map(|be| expression_field_name(be.alias.as_deref(), &be.expression))
305                            .collect();
306                        for chunk in &mut result {
307                            if chunk.fields.len() == names.len() {
308                                chunk.field_names = names.clone();
309                            }
310                        }
311                    }
312                    result
313                }
314            };
315            Ok(result)
316        }
317        LogicalOperator::Filter(f) => {
318            let evaluator = ctx.function_registry.clone().map(|reg| {
319                let mut eval = ExpressionEvaluator::new(reg);
320                if let Some(ref seq_fn) = ctx.sequence_fn {
321                    eval = eval.with_sequence_fn(seq_fn.clone());
322                }
323                if let Some(ref subquery_fn) = ctx.subquery_fn {
324                    eval = eval.with_subquery_fn(subquery_fn.clone());
325                }
326                Arc::new(Mutex::new(eval))
327            });
328            let filter = if let Some(eval) = evaluator {
329                PhysicalFilter::with_evaluator(f.expression.clone(), eval)
330            } else {
331                PhysicalFilter::new(f.expression.clone())
332            };
333            let result = filter.execute(current_input)?;
334            Ok(result)
335        }
336        LogicalOperator::Limit(l) => {
337            let limit = PhysicalLimit {
338                limit: l.limit,
339                offset: l.offset,
340            };
341            let result = limit.execute(current_input)?;
342            Ok(result)
343        }
344        LogicalOperator::TopK(tk) => {
345            let (sort_keys, augmented, extra) = resolve_sort_keys(&tk.sort_keys, &current_input, ctx)?;
346            let topk = PhysicalTopK {
347                sort_keys,
348                limit: tk.limit,
349                offset: tk.offset,
350            };
351            let mut result = topk.execute(augmented)?;
352            strip_sort_columns(&mut result, extra);
353            Ok(result)
354        }
355        LogicalOperator::OrderBy(o) => {
356            let (sort_keys, augmented, extra) = resolve_sort_keys(&o.sort_keys, &current_input, ctx)?;
357            let order = PhysicalOrderBy { sort_keys };
358            let mut result = order.execute(augmented)?;
359            strip_sort_columns(&mut result, extra);
360            Ok(result)
361        }
362        LogicalOperator::Flatten(f) => {
363            let input = if f.children.is_empty() {
364                current_input
365            } else {
366                ctx.execute_children(&f.children)?
367            };
368            let flatten = PhysicalFlatten::new(f.group_pos);
369            flatten.execute(input)
370        }
371        LogicalOperator::Unwind(uw) => {
372            let unwind = PhysicalUnwind {
373                expression: uw.expression.clone(),
374                variable: uw.variable.clone(),
375            };
376            let result = unwind.execute(current_input)?;
377            Ok(result)
378        }
379        LogicalOperator::Partitioner(p) => {
380            const MORSEL_SIZE: usize = 1024;
381            let partitioner = Partitioner::new(MORSEL_SIZE);
382            let morsels = partitioner.execute(current_input)?;
383
384            let mut results = Vec::new();
385            for _morsel in morsels {
386                let child_result = ctx.execute_children(&p.children)?;
387                results.extend(child_result);
388            }
389            Ok(results)
390        }
391        _ => Err(format!("Not a projection/filter operator: {:?}", op).into()),
392    }
393}
akar_processor/processor/mapper/map_projection.rs

akar_processor/processor/mapper/
map_projection.rs
