alopex_dataframe/expr/
expr.rs

1/// Expression AST used by `DataFrame` and `LazyFrame`.
2#[derive(Debug, Clone, PartialEq)]
3pub enum Expr {
4    /// Column reference.
5    Column(String),
6    /// Literal scalar value.
7    Literal(Scalar),
8    /// Binary operator expression.
9    BinaryOp {
10        left: Box<Expr>,
11        op: Operator,
12        right: Box<Expr>,
13    },
14    /// Unary operator expression.
15    UnaryOp { op: UnaryOperator, expr: Box<Expr> },
16    /// Aggregation expression (only valid under `group_by().agg()`).
17    Agg { func: AggFunc, expr: Box<Expr> },
18    /// Namespace function expression.
19    Function {
20        input: Box<Expr>,
21        function: ExprFunction,
22    },
23    /// Expression alias (renames the resulting column).
24    Alias { expr: Box<Expr>, name: String },
25    /// Wildcard (`*`) that expands to all columns in projections.
26    Wildcard,
27}
28
29/// Supported namespace expression functions.
30#[derive(Debug, Clone, PartialEq, Eq)]
31pub enum ExprFunction {
32    /// String namespace function.
33    String(StringFunction),
34    /// Datetime namespace function.
35    Datetime(DatetimeFunction),
36    /// List namespace function.
37    List(ListFunction),
38}
39
40/// Supported `str.*` functions.
41#[derive(Debug, Clone, PartialEq, Eq)]
42pub enum StringFunction {
43    /// Convert UTF-8 strings to lowercase.
44    ToLowercase,
45    /// Convert UTF-8 strings to uppercase.
46    ToUppercase,
47    /// Regex contains.
48    Contains { pattern: String },
49    /// Regex replacement.
50    Replace {
51        pattern: String,
52        replacement: String,
53    },
54    /// Strip whitespace or the provided characters.
55    StripChars { chars: Option<String> },
56    /// Split by a literal separator.
57    Split { separator: String },
58    /// Count Unicode scalar values.
59    LenChars,
60    /// Extract a regex capture group.
61    Extract {
62        pattern: String,
63        capture_group: usize,
64    },
65}
66
67/// Supported `dt.*` functions.
68#[derive(Debug, Clone, PartialEq, Eq)]
69pub enum DatetimeFunction {
70    /// Extract UTC year.
71    Year,
72    /// Extract UTC month.
73    Month,
74    /// Extract UTC day of month.
75    Day,
76    /// Extract ISO weekday.
77    Weekday,
78    /// Format as UTC text.
79    ToString,
80    /// Convert between fixed-offset time zones.
81    ConvertTimeZone {
82        from_offset: String,
83        to_offset: String,
84    },
85}
86
87/// Supported `list.*` functions.
88#[derive(Debug, Clone, PartialEq, Eq)]
89pub enum ListFunction {
90    /// Join string list elements.
91    Join {
92        separator: String,
93        null_value: Option<String>,
94    },
95    /// Return list length.
96    Len,
97    /// Test whether a list contains a non-null string value.
98    Contains { value: String },
99}
100
101/// Supported binary operators.
102#[derive(Debug, Copy, Clone, PartialEq, Eq)]
103pub enum Operator {
104    /// Addition.
105    Add,
106    /// Subtraction.
107    Sub,
108    /// Multiplication.
109    Mul,
110    /// Division.
111    Div,
112    /// Equality.
113    Eq,
114    /// Inequality.
115    Neq,
116    /// Greater-than.
117    Gt,
118    /// Less-than.
119    Lt,
120    /// Greater-than-or-equal.
121    Ge,
122    /// Less-than-or-equal.
123    Le,
124    /// Boolean AND.
125    And,
126    /// Boolean OR.
127    Or,
128}
129
130/// Supported unary operators.
131#[derive(Debug, Copy, Clone, PartialEq, Eq)]
132pub enum UnaryOperator {
133    /// Boolean NOT.
134    Not,
135}
136
137/// Supported aggregation functions.
138#[derive(Debug, Copy, Clone, PartialEq, Eq)]
139pub enum AggFunc {
140    /// Sum of non-null values.
141    Sum,
142    /// Mean of non-null values.
143    Mean,
144    /// Count of non-null values.
145    Count,
146    /// Minimum of non-null values.
147    Min,
148    /// Maximum of non-null values.
149    Max,
150}
151
152/// Scalar literal values.
153#[derive(Debug, Clone, PartialEq)]
154pub enum Scalar {
155    /// Null literal.
156    Null,
157    /// Boolean literal.
158    Boolean(bool),
159    /// 64-bit integer literal.
160    Int64(i64),
161    /// 64-bit float literal.
162    Float64(f64),
163    /// UTF-8 string literal.
164    Utf8(String),
165}
166
167impl From<()> for Scalar {
168    fn from(_: ()) -> Self {
169        Scalar::Null
170    }
171}
172
173impl From<bool> for Scalar {
174    fn from(v: bool) -> Self {
175        Scalar::Boolean(v)
176    }
177}
178
179impl From<i64> for Scalar {
180    fn from(v: i64) -> Self {
181        Scalar::Int64(v)
182    }
183}
184
185impl From<f64> for Scalar {
186    fn from(v: f64) -> Self {
187        Scalar::Float64(v)
188    }
189}
190
191impl From<String> for Scalar {
192    fn from(v: String) -> Self {
193        Scalar::Utf8(v)
194    }
195}
196
197impl From<&str> for Scalar {
198    fn from(v: &str) -> Self {
199        Scalar::Utf8(v.to_string())
200    }
201}
202
203impl Expr {
204    /// Alias this expression (used to name output columns).
205    pub fn alias(self, name: impl Into<String>) -> Expr {
206        Expr::Alias {
207            expr: Box::new(self),
208            name: name.into(),
209        }
210    }
211
212    /// Build an addition expression.
213    #[allow(clippy::should_implement_trait)]
214    pub fn add(self, rhs: Expr) -> Expr {
215        Expr::BinaryOp {
216            left: Box::new(self),
217            op: Operator::Add,
218            right: Box::new(rhs),
219        }
220    }
221
222    /// Build a subtraction expression.
223    #[allow(clippy::should_implement_trait)]
224    pub fn sub(self, rhs: Expr) -> Expr {
225        Expr::BinaryOp {
226            left: Box::new(self),
227            op: Operator::Sub,
228            right: Box::new(rhs),
229        }
230    }
231
232    /// Build a multiplication expression.
233    #[allow(clippy::should_implement_trait)]
234    pub fn mul(self, rhs: Expr) -> Expr {
235        Expr::BinaryOp {
236            left: Box::new(self),
237            op: Operator::Mul,
238            right: Box::new(rhs),
239        }
240    }
241
242    /// Build a division expression.
243    #[allow(clippy::should_implement_trait)]
244    pub fn div(self, rhs: Expr) -> Expr {
245        Expr::BinaryOp {
246            left: Box::new(self),
247            op: Operator::Div,
248            right: Box::new(rhs),
249        }
250    }
251
252    /// Build an equality predicate.
253    pub fn eq(self, rhs: Expr) -> Expr {
254        Expr::BinaryOp {
255            left: Box::new(self),
256            op: Operator::Eq,
257            right: Box::new(rhs),
258        }
259    }
260
261    /// Build an inequality predicate.
262    pub fn neq(self, rhs: Expr) -> Expr {
263        Expr::BinaryOp {
264            left: Box::new(self),
265            op: Operator::Neq,
266            right: Box::new(rhs),
267        }
268    }
269
270    /// Build a greater-than predicate.
271    pub fn gt(self, rhs: Expr) -> Expr {
272        Expr::BinaryOp {
273            left: Box::new(self),
274            op: Operator::Gt,
275            right: Box::new(rhs),
276        }
277    }
278
279    /// Build a less-than predicate.
280    pub fn lt(self, rhs: Expr) -> Expr {
281        Expr::BinaryOp {
282            left: Box::new(self),
283            op: Operator::Lt,
284            right: Box::new(rhs),
285        }
286    }
287
288    /// Build a greater-than-or-equal predicate.
289    pub fn ge(self, rhs: Expr) -> Expr {
290        Expr::BinaryOp {
291            left: Box::new(self),
292            op: Operator::Ge,
293            right: Box::new(rhs),
294        }
295    }
296
297    /// Build a less-than-or-equal predicate.
298    pub fn le(self, rhs: Expr) -> Expr {
299        Expr::BinaryOp {
300            left: Box::new(self),
301            op: Operator::Le,
302            right: Box::new(rhs),
303        }
304    }
305
306    /// Build a boolean AND predicate.
307    pub fn and_(self, rhs: Expr) -> Expr {
308        Expr::BinaryOp {
309            left: Box::new(self),
310            op: Operator::And,
311            right: Box::new(rhs),
312        }
313    }
314
315    /// Build a boolean OR predicate.
316    pub fn or_(self, rhs: Expr) -> Expr {
317        Expr::BinaryOp {
318            left: Box::new(self),
319            op: Operator::Or,
320            right: Box::new(rhs),
321        }
322    }
323
324    /// Build a boolean NOT predicate.
325    pub fn not_(self) -> Expr {
326        Expr::UnaryOp {
327            op: UnaryOperator::Not,
328            expr: Box::new(self),
329        }
330    }
331
332    /// Build a `sum` aggregation.
333    pub fn sum(self) -> Expr {
334        Expr::Agg {
335            func: AggFunc::Sum,
336            expr: Box::new(self),
337        }
338    }
339
340    /// Build a `mean` aggregation.
341    pub fn mean(self) -> Expr {
342        Expr::Agg {
343            func: AggFunc::Mean,
344            expr: Box::new(self),
345        }
346    }
347
348    /// Build a `count` aggregation (nulls excluded).
349    pub fn count(self) -> Expr {
350        Expr::Agg {
351            func: AggFunc::Count,
352            expr: Box::new(self),
353        }
354    }
355
356    /// Build a `min` aggregation.
357    pub fn min(self) -> Expr {
358        Expr::Agg {
359            func: AggFunc::Min,
360            expr: Box::new(self),
361        }
362    }
363
364    /// Build a `max` aggregation.
365    pub fn max(self) -> Expr {
366        Expr::Agg {
367            func: AggFunc::Max,
368            expr: Box::new(self),
369        }
370    }
371
372    /// Enter the `str.*` expression namespace.
373    pub fn str(self) -> StringExpr {
374        StringExpr { input: self }
375    }
376
377    /// Enter the `dt.*` expression namespace.
378    pub fn dt(self) -> DatetimeExpr {
379        DatetimeExpr { input: self }
380    }
381
382    /// Enter the `list.*` expression namespace.
383    pub fn list(self) -> ListExpr {
384        ListExpr { input: self }
385    }
386}
387
388/// Builder for `str.*` expression functions.
389#[derive(Debug, Clone, PartialEq)]
390pub struct StringExpr {
391    input: Expr,
392}
393
394impl StringExpr {
395    fn function(self, function: StringFunction) -> Expr {
396        Expr::Function {
397            input: Box::new(self.input),
398            function: ExprFunction::String(function),
399        }
400    }
401
402    /// Convert UTF-8 strings to lowercase.
403    pub fn to_lowercase(self) -> Expr {
404        self.function(StringFunction::ToLowercase)
405    }
406
407    /// Convert UTF-8 strings to uppercase.
408    pub fn to_uppercase(self) -> Expr {
409        self.function(StringFunction::ToUppercase)
410    }
411
412    /// Return whether each string matches a regex pattern.
413    pub fn contains(self, pattern: impl Into<String>) -> Expr {
414        self.function(StringFunction::Contains {
415            pattern: pattern.into(),
416        })
417    }
418
419    /// Replace regex matches.
420    pub fn replace(self, pattern: impl Into<String>, replacement: impl Into<String>) -> Expr {
421        self.function(StringFunction::Replace {
422            pattern: pattern.into(),
423            replacement: replacement.into(),
424        })
425    }
426
427    /// Strip whitespace from both ends.
428    pub fn strip_chars(self, chars: Option<impl Into<String>>) -> Expr {
429        self.function(StringFunction::StripChars {
430            chars: chars.map(Into::into),
431        })
432    }
433
434    /// Split by a literal separator.
435    pub fn split(self, separator: impl Into<String>) -> Expr {
436        self.function(StringFunction::Split {
437            separator: separator.into(),
438        })
439    }
440
441    /// Count Unicode scalar values.
442    pub fn len_chars(self) -> Expr {
443        self.function(StringFunction::LenChars)
444    }
445
446    /// Extract a regex capture group.
447    pub fn extract(self, pattern: impl Into<String>, capture_group: usize) -> Expr {
448        self.function(StringFunction::Extract {
449            pattern: pattern.into(),
450            capture_group,
451        })
452    }
453}
454
455/// Builder for `dt.*` expression functions.
456#[derive(Debug, Clone, PartialEq)]
457pub struct DatetimeExpr {
458    input: Expr,
459}
460
461impl DatetimeExpr {
462    fn function(self, function: DatetimeFunction) -> Expr {
463        Expr::Function {
464            input: Box::new(self.input),
465            function: ExprFunction::Datetime(function),
466        }
467    }
468
469    /// Extract UTC year.
470    pub fn year(self) -> Expr {
471        self.function(DatetimeFunction::Year)
472    }
473
474    /// Extract UTC month.
475    pub fn month(self) -> Expr {
476        self.function(DatetimeFunction::Month)
477    }
478
479    /// Extract UTC day of month.
480    pub fn day(self) -> Expr {
481        self.function(DatetimeFunction::Day)
482    }
483
484    /// Extract ISO weekday.
485    pub fn weekday(self) -> Expr {
486        self.function(DatetimeFunction::Weekday)
487    }
488
489    /// Format timestamps as UTC text.
490    #[allow(clippy::inherent_to_string)]
491    pub fn to_string(self) -> Expr {
492        self.function(DatetimeFunction::ToString)
493    }
494
495    /// Convert between fixed-offset time zones.
496    pub fn convert_time_zone(
497        self,
498        from_offset: impl Into<String>,
499        to_offset: impl Into<String>,
500    ) -> Expr {
501        self.function(DatetimeFunction::ConvertTimeZone {
502            from_offset: from_offset.into(),
503            to_offset: to_offset.into(),
504        })
505    }
506}
507
508/// Builder for `list.*` expression functions.
509#[derive(Debug, Clone, PartialEq)]
510pub struct ListExpr {
511    input: Expr,
512}
513
514impl ListExpr {
515    fn function(self, function: ListFunction) -> Expr {
516        Expr::Function {
517            input: Box::new(self.input),
518            function: ExprFunction::List(function),
519        }
520    }
521
522    /// Join string list elements with a separator.
523    pub fn join(self, separator: impl Into<String>, null_value: Option<impl Into<String>>) -> Expr {
524        self.function(ListFunction::Join {
525            separator: separator.into(),
526            null_value: null_value.map(Into::into),
527        })
528    }
529
530    /// Return list lengths.
531    pub fn len(self) -> Expr {
532        self.function(ListFunction::Len)
533    }
534
535    /// Return whether a string list contains `value`.
536    pub fn contains(self, value: impl Into<String>) -> Expr {
537        self.function(ListFunction::Contains {
538            value: value.into(),
539        })
540    }
541}
542
543#[cfg(test)]
544mod tests {
545    use super::{AggFunc, Expr, ExprFunction, Operator, Scalar, StringFunction, UnaryOperator};
546    use crate::expr::{col, lit};
547
548    #[test]
549    fn builder_and_chaining_works() {
550        let expr = col("a").add(lit(1_i64)).alias("b");
551        assert_eq!(
552            expr,
553            Expr::Alias {
554                expr: Box::new(Expr::BinaryOp {
555                    left: Box::new(Expr::Column("a".to_string())),
556                    op: Operator::Add,
557                    right: Box::new(Expr::Literal(Scalar::Int64(1))),
558                }),
559                name: "b".to_string(),
560            }
561        );
562    }
563
564    #[test]
565    fn logical_and_agg_works() {
566        let expr = col("x")
567            .gt(lit(1_i64))
568            .and_(col("y").lt(lit(10_i64)).not_())
569            .alias("p");
570
571        assert!(matches!(
572            expr,
573            Expr::Alias {
574                expr: _,
575                name
576            } if name == "p"
577        ));
578
579        let agg = col("v").sum();
580        assert_eq!(
581            agg,
582            Expr::Agg {
583                func: AggFunc::Sum,
584                expr: Box::new(Expr::Column("v".to_string()))
585            }
586        );
587
588        let u = Expr::Column("a".to_string()).not_();
589        assert_eq!(
590            u,
591            Expr::UnaryOp {
592                op: UnaryOperator::Not,
593                expr: Box::new(Expr::Column("a".to_string()))
594            }
595        );
596    }
597
598    #[test]
599    fn namespace_builders_create_function_exprs() {
600        let expr = col("name").str().to_lowercase().alias("name_lower");
601        assert_eq!(
602            expr,
603            Expr::Alias {
604                expr: Box::new(Expr::Function {
605                    input: Box::new(Expr::Column("name".to_string())),
606                    function: ExprFunction::String(StringFunction::ToLowercase),
607                }),
608                name: "name_lower".to_string(),
609            }
610        );
611    }
612}
alopex_dataframe/expr/expr.rs

alopex_dataframe/expr/
expr.rs
