alopex_dataframe/expr/
expr.rs

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

alopex_dataframe/expr/
expr.rs
