polars_plan/plans/aexpr/
scalar.rs

1use recursive::recursive;
2
3use super::*;
4
5#[recursive]
6pub fn is_scalar_ae(node: Node, expr_arena: &Arena<AExpr>) -> bool {
7    match expr_arena.get(node) {
8        AExpr::Literal(lv) => lv.is_scalar(),
9        AExpr::Function { options, input, .. }
10        | AExpr::AnonymousFunction { options, input, .. } => {
11            if options.flags.contains(FunctionFlags::RETURNS_SCALAR) {
12                true
13            } else if options.is_elementwise()
14                || !options.flags.contains(FunctionFlags::CHANGES_LENGTH)
15            {
16                input.iter().all(|e| e.is_scalar(expr_arena))
17            } else {
18                false
19            }
20        },
21        AExpr::BinaryExpr { left, right, .. } => {
22            is_scalar_ae(*left, expr_arena) && is_scalar_ae(*right, expr_arena)
23        },
24        AExpr::Ternary { truthy, falsy, .. } => {
25            is_scalar_ae(*truthy, expr_arena) && is_scalar_ae(*falsy, expr_arena)
26        },
27        AExpr::Agg(_) | AExpr::Len => true,
28        AExpr::Cast { expr, .. } | AExpr::Alias(expr, _) => is_scalar_ae(*expr, expr_arena),
29        _ => false,
30    }
31}