polars_plan/dsl/
arity.rs

1use super::*;
2
3/// Utility struct for the `when-then-otherwise` expression.
4///
5/// Represents the state of the expression after [when] is called.
6///
7/// In this state, `then` must be called to continue to finish the expression.
8#[derive(Clone)]
9pub struct When {
10    condition: Expr,
11}
12
13/// Utility struct for the `when-then-otherwise` expression.
14///
15/// Represents the state of the expression after `when(...).then(...)` is called.
16#[derive(Clone)]
17pub struct Then {
18    condition: Expr,
19    statement: Expr,
20}
21
22/// Utility struct for the `when-then-otherwise` expression.
23///
24/// Represents the state of the expression after an additional `when` is called.
25///
26/// In this state, `then` must be called to continue to finish the expression.
27#[derive(Clone)]
28pub struct ChainedWhen {
29    conditions: Vec<Expr>,
30    statements: Vec<Expr>,
31}
32
33/// Utility struct for the `when-then-otherwise` expression.
34///
35/// Represents the state of the expression after an additional `then` is called.
36#[derive(Clone)]
37pub struct ChainedThen {
38    conditions: Vec<Expr>,
39    statements: Vec<Expr>,
40}
41
42impl When {
43    /// Add a condition to the `when-then-otherwise` expression.
44    pub fn then<E: Into<Expr>>(self, expr: E) -> Then {
45        Then {
46            condition: self.condition,
47            statement: expr.into(),
48        }
49    }
50}
51
52impl Then {
53    /// Attach a statement to the corresponding condition.
54    pub fn when<E: Into<Expr>>(self, condition: E) -> ChainedWhen {
55        ChainedWhen {
56            conditions: vec![self.condition, condition.into()],
57            statements: vec![self.statement],
58        }
59    }
60
61    /// Define a default for the `when-then-otherwise` expression.
62    pub fn otherwise<E: Into<Expr>>(self, statement: E) -> Expr {
63        ternary_expr(self.condition, self.statement, statement.into())
64    }
65}
66
67impl ChainedWhen {
68    pub fn then<E: Into<Expr>>(mut self, statement: E) -> ChainedThen {
69        self.statements.push(statement.into());
70        ChainedThen {
71            conditions: self.conditions,
72            statements: self.statements,
73        }
74    }
75}
76
77impl ChainedThen {
78    /// Add another condition to the `when-then-otherwise` expression.
79    pub fn when<E: Into<Expr>>(mut self, condition: E) -> ChainedWhen {
80        self.conditions.push(condition.into());
81
82        ChainedWhen {
83            conditions: self.conditions,
84            statements: self.statements,
85        }
86    }
87
88    /// Define a default for the `when-then-otherwise` expression.
89    pub fn otherwise<E: Into<Expr>>(self, expr: E) -> Expr {
90        // we iterate the preds/ exprs last in first out
91        // and nest them.
92        //
93        // // this expr:
94        //   when((col('x') == 'a')).then(1)
95        //         .when(col('x') == 'b').then(2)
96        //         .when(col('x') == 'c').then(3)
97        //         .otherwise(4)
98        //
99        // needs to become:
100        //       when((col('x') == 'a')).then(1)                        -
101        //         .otherwise(                                           |
102        //             when(col('x') == 'b').then(2)            -        |
103        //             .otherwise(                               |       |
104        //                 pl.when(col('x') == 'c').then(3)      |       |
105        //                 .otherwise(4)                         | inner | outer
106        //             )                                         |       |
107        //         )                                            _|      _|
108        //
109        // by iterating LIFO we first create
110        // `inner` and then assign that to `otherwise`,
111        // which will be used in the next layer `outer`
112        //
113
114        let conditions_iter = self.conditions.into_iter().rev();
115        let mut statements_iter = self.statements.into_iter().rev();
116
117        let mut otherwise = expr.into();
118
119        for e in conditions_iter {
120            otherwise = ternary_expr(
121                e,
122                statements_iter
123                    .next()
124                    .expect("expr expected, did you call when().then().otherwise?"),
125                otherwise,
126            );
127        }
128
129        otherwise
130    }
131}
132
133/// Start a `when-then-otherwise` expression.
134pub fn when<E: Into<Expr>>(condition: E) -> When {
135    When {
136        condition: condition.into(),
137    }
138}
139
140pub fn ternary_expr(predicate: Expr, truthy: Expr, falsy: Expr) -> Expr {
141    Expr::Ternary {
142        predicate: Arc::new(predicate),
143        truthy: Arc::new(truthy),
144        falsy: Arc::new(falsy),
145    }
146}
147
148/// Compute `op(l, r)` (or equivalently `l op r`). `l` and `r` must have types compatible with the Operator.
149pub fn binary_expr(l: Expr, op: Operator, r: Expr) -> Expr {
150    Expr::BinaryExpr {
151        left: Arc::new(l),
152        op,
153        right: Arc::new(r),
154    }
155}
polars_plan/dsl/arity.rs

polars_plan/dsl/
arity.rs
