polars_plan/plans/python/
predicate.rs

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
use polars_core::error::polars_err;
use polars_core::prelude::PolarsResult;

use crate::prelude::*;

fn accept_as_io_predicate(e: &Expr) -> bool {
    const LIMIT: usize = 1 << 16;
    match e {
        Expr::Literal(lv) => match lv {
            LiteralValue::Binary(v) => v.len() <= LIMIT,
            LiteralValue::String(v) => v.len() <= LIMIT,
            LiteralValue::Series(s) => s.estimated_size() < LIMIT,
            // Don't accept dynamic types
            LiteralValue::Int(_) => false,
            LiteralValue::Float(_) => false,
            _ => true,
        },
        Expr::Wildcard | Expr::Column(_) => true,
        Expr::BinaryExpr { left, right, .. } => {
            accept_as_io_predicate(left) && accept_as_io_predicate(right)
        },
        Expr::Ternary {
            truthy,
            falsy,
            predicate,
        } => {
            accept_as_io_predicate(truthy)
                && accept_as_io_predicate(falsy)
                && accept_as_io_predicate(predicate)
        },
        Expr::Alias(_, _) => true,
        Expr::Function {
            function, input, ..
        } => {
            match function {
                // we already checked if streaming, so we can all functions
                FunctionExpr::Boolean(_) | FunctionExpr::BinaryExpr(_) | FunctionExpr::Coalesce => {
                },
                #[cfg(feature = "log")]
                FunctionExpr::Entropy { .. }
                | FunctionExpr::Log { .. }
                | FunctionExpr::Log1p { .. }
                | FunctionExpr::Exp { .. } => {},
                #[cfg(feature = "abs")]
                FunctionExpr::Abs => {},
                #[cfg(feature = "trigonometry")]
                FunctionExpr::Atan2 => {},
                #[cfg(feature = "round_series")]
                FunctionExpr::Clip { .. } => {},
                #[cfg(feature = "fused")]
                FunctionExpr::Fused(_) => {},
                _ => return false,
            }
            input.iter().all(accept_as_io_predicate)
        },
        _ => false,
    }
}

pub fn serialize(expr: &Expr) -> PolarsResult<Option<Vec<u8>>> {
    if !accept_as_io_predicate(expr) {
        return Ok(None);
    }
    let mut buf = vec![];
    ciborium::into_writer(expr, &mut buf)
        .map_err(|_| polars_err!(ComputeError: "could not serialize: {}", expr))?;

    Ok(Some(buf))
}