polars_plan/plans/python/
predicate.rs

1use polars_core::prelude::PolarsResult;
2use polars_utils::pl_serialize;
3
4use crate::prelude::*;
5
6fn accept_as_io_predicate(e: &Expr) -> bool {
7    const LIMIT: usize = 1 << 16;
8    match e {
9        Expr::Literal(lv) => match lv {
10            LiteralValue::Binary(v) => v.len() <= LIMIT,
11            LiteralValue::String(v) => v.len() <= LIMIT,
12            LiteralValue::Series(s) => s.estimated_size() < LIMIT,
13            // Don't accept dynamic types
14            LiteralValue::Int(_) => false,
15            LiteralValue::Float(_) => false,
16            _ => true,
17        },
18        Expr::Wildcard | Expr::Column(_) => true,
19        Expr::BinaryExpr { left, right, .. } => {
20            accept_as_io_predicate(left) && accept_as_io_predicate(right)
21        },
22        Expr::Ternary {
23            truthy,
24            falsy,
25            predicate,
26        } => {
27            accept_as_io_predicate(truthy)
28                && accept_as_io_predicate(falsy)
29                && accept_as_io_predicate(predicate)
30        },
31        Expr::Alias(_, _) => true,
32        Expr::Function {
33            function, input, ..
34        } => {
35            match function {
36                // we already checked if streaming, so we can all functions
37                FunctionExpr::Boolean(_) | FunctionExpr::BinaryExpr(_) | FunctionExpr::Coalesce => {
38                },
39                #[cfg(feature = "log")]
40                FunctionExpr::Entropy { .. }
41                | FunctionExpr::Log { .. }
42                | FunctionExpr::Log1p { .. }
43                | FunctionExpr::Exp { .. } => {},
44                #[cfg(feature = "abs")]
45                FunctionExpr::Abs => {},
46                #[cfg(feature = "trigonometry")]
47                FunctionExpr::Atan2 => {},
48                #[cfg(feature = "round_series")]
49                FunctionExpr::Clip { .. } => {},
50                #[cfg(feature = "fused")]
51                FunctionExpr::Fused(_) => {},
52                _ => return false,
53            }
54            input.iter().all(accept_as_io_predicate)
55        },
56        _ => false,
57    }
58}
59
60pub fn serialize(expr: &Expr) -> PolarsResult<Option<Vec<u8>>> {
61    if !accept_as_io_predicate(expr) {
62        return Ok(None);
63    }
64    let mut buf = vec![];
65    pl_serialize::serialize_into_writer(&mut buf, expr)?;
66
67    Ok(Some(buf))
68}