alopex-dataframe 0.7.2

/// Expression AST used by `DataFrame` and `LazyFrame`.
#[derive(Debug, Clone, PartialEq)]
pub enum Expr {
    /// Column reference.
    Column(String),
    /// Literal scalar value.
    Literal(Scalar),
    /// Binary operator expression.
    BinaryOp {
        left: Box<Expr>,
        op: Operator,
        right: Box<Expr>,
    },
    /// Unary operator expression.
    UnaryOp { op: UnaryOperator, expr: Box<Expr> },
    /// Aggregation expression (only valid under `group_by().agg()`).
    Agg { func: AggFunc, expr: Box<Expr> },
    /// Namespace function expression.
    Function {
        input: Box<Expr>,
        function: ExprFunction,
    },
    /// Expression alias (renames the resulting column).
    Alias { expr: Box<Expr>, name: String },
    /// Wildcard (`*`) that expands to all columns in projections.
    Wildcard,
}

/// Supported namespace expression functions.
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum ExprFunction {
    /// String namespace function.
    String(StringFunction),
    /// Datetime namespace function.
    Datetime(DatetimeFunction),
    /// List namespace function.
    List(ListFunction),
}

/// Supported `str.*` functions.
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum StringFunction {
    /// Convert UTF-8 strings to lowercase.
    ToLowercase,
    /// Convert UTF-8 strings to uppercase.
    ToUppercase,
    /// Regex contains.
    Contains { pattern: String },
    /// Regex replacement.
    Replace {
        pattern: String,
        replacement: String,
    },
    /// Strip whitespace or the provided characters.
    StripChars { chars: Option<String> },
    /// Split by a literal separator.
    Split { separator: String },
    /// Count Unicode scalar values.
    LenChars,
    /// Extract a regex capture group.
    Extract {
        pattern: String,
        capture_group: usize,
    },
}

/// Supported `dt.*` functions.
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum DatetimeFunction {
    /// Extract UTC year.
    Year,
    /// Extract UTC month.
    Month,
    /// Extract UTC day of month.
    Day,
    /// Extract ISO weekday.
    Weekday,
    /// Format as UTC text.
    ToString,
    /// Convert between fixed-offset time zones.
    ConvertTimeZone {
        from_offset: String,
        to_offset: String,
    },
}

/// Supported `list.*` functions.
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum ListFunction {
    /// Join string list elements.
    Join {
        separator: String,
        null_value: Option<String>,
    },
    /// Return list length.
    Len,
    /// Test whether a list contains a non-null string value.
    Contains { value: String },
}

/// Supported binary operators.
#[derive(Debug, Copy, Clone, PartialEq, Eq)]
pub enum Operator {
    /// Addition.
    Add,
    /// Subtraction.
    Sub,
    /// Multiplication.
    Mul,
    /// Division.
    Div,
    /// Equality.
    Eq,
    /// Inequality.
    Neq,
    /// Greater-than.
    Gt,
    /// Less-than.
    Lt,
    /// Greater-than-or-equal.
    Ge,
    /// Less-than-or-equal.
    Le,
    /// Boolean AND.
    And,
    /// Boolean OR.
    Or,
}

/// Supported unary operators.
#[derive(Debug, Copy, Clone, PartialEq, Eq)]
pub enum UnaryOperator {
    /// Boolean NOT.
    Not,
}

/// Supported aggregation functions.
#[derive(Debug, Copy, Clone, PartialEq, Eq)]
pub enum AggFunc {
    /// Sum of non-null values.
    Sum,
    /// Mean of non-null values.
    Mean,
    /// Count of non-null values.
    Count,
    /// Minimum of non-null values.
    Min,
    /// Maximum of non-null values.
    Max,
}

/// Scalar literal values.
#[derive(Debug, Clone, PartialEq)]
pub enum Scalar {
    /// Null literal.
    Null,
    /// Boolean literal.
    Boolean(bool),
    /// 64-bit integer literal.
    Int64(i64),
    /// 64-bit float literal.
    Float64(f64),
    /// UTF-8 string literal.
    Utf8(String),
}

impl From<()> for Scalar {
    fn from(_: ()) -> Self {
        Scalar::Null
    }
}

impl From<bool> for Scalar {
    fn from(v: bool) -> Self {
        Scalar::Boolean(v)
    }
}

impl From<i64> for Scalar {
    fn from(v: i64) -> Self {
        Scalar::Int64(v)
    }
}

impl From<f64> for Scalar {
    fn from(v: f64) -> Self {
        Scalar::Float64(v)
    }
}

impl From<String> for Scalar {
    fn from(v: String) -> Self {
        Scalar::Utf8(v)
    }
}

impl From<&str> for Scalar {
    fn from(v: &str) -> Self {
        Scalar::Utf8(v.to_string())
    }
}

impl Expr {
    /// Alias this expression (used to name output columns).
    pub fn alias(self, name: impl Into<String>) -> Expr {
        Expr::Alias {
            expr: Box::new(self),
            name: name.into(),
        }
    }

    /// Build an addition expression.
    #[allow(clippy::should_implement_trait)]
    pub fn add(self, rhs: Expr) -> Expr {
        Expr::BinaryOp {
            left: Box::new(self),
            op: Operator::Add,
            right: Box::new(rhs),
        }
    }

    /// Build a subtraction expression.
    #[allow(clippy::should_implement_trait)]
    pub fn sub(self, rhs: Expr) -> Expr {
        Expr::BinaryOp {
            left: Box::new(self),
            op: Operator::Sub,
            right: Box::new(rhs),
        }
    }

    /// Build a multiplication expression.
    #[allow(clippy::should_implement_trait)]
    pub fn mul(self, rhs: Expr) -> Expr {
        Expr::BinaryOp {
            left: Box::new(self),
            op: Operator::Mul,
            right: Box::new(rhs),
        }
    }

    /// Build a division expression.
    #[allow(clippy::should_implement_trait)]
    pub fn div(self, rhs: Expr) -> Expr {
        Expr::BinaryOp {
            left: Box::new(self),
            op: Operator::Div,
            right: Box::new(rhs),
        }
    }

    /// Build an equality predicate.
    pub fn eq(self, rhs: Expr) -> Expr {
        Expr::BinaryOp {
            left: Box::new(self),
            op: Operator::Eq,
            right: Box::new(rhs),
        }
    }

    /// Build an inequality predicate.
    pub fn neq(self, rhs: Expr) -> Expr {
        Expr::BinaryOp {
            left: Box::new(self),
            op: Operator::Neq,
            right: Box::new(rhs),
        }
    }

    /// Build a greater-than predicate.
    pub fn gt(self, rhs: Expr) -> Expr {
        Expr::BinaryOp {
            left: Box::new(self),
            op: Operator::Gt,
            right: Box::new(rhs),
        }
    }

    /// Build a less-than predicate.
    pub fn lt(self, rhs: Expr) -> Expr {
        Expr::BinaryOp {
            left: Box::new(self),
            op: Operator::Lt,
            right: Box::new(rhs),
        }
    }

    /// Build a greater-than-or-equal predicate.
    pub fn ge(self, rhs: Expr) -> Expr {
        Expr::BinaryOp {
            left: Box::new(self),
            op: Operator::Ge,
            right: Box::new(rhs),
        }
    }

    /// Build a less-than-or-equal predicate.
    pub fn le(self, rhs: Expr) -> Expr {
        Expr::BinaryOp {
            left: Box::new(self),
            op: Operator::Le,
            right: Box::new(rhs),
        }
    }

    /// Build a boolean AND predicate.
    pub fn and_(self, rhs: Expr) -> Expr {
        Expr::BinaryOp {
            left: Box::new(self),
            op: Operator::And,
            right: Box::new(rhs),
        }
    }

    /// Build a boolean OR predicate.
    pub fn or_(self, rhs: Expr) -> Expr {
        Expr::BinaryOp {
            left: Box::new(self),
            op: Operator::Or,
            right: Box::new(rhs),
        }
    }

    /// Build a boolean NOT predicate.
    pub fn not_(self) -> Expr {
        Expr::UnaryOp {
            op: UnaryOperator::Not,
            expr: Box::new(self),
        }
    }

    /// Build a `sum` aggregation.
    pub fn sum(self) -> Expr {
        Expr::Agg {
            func: AggFunc::Sum,
            expr: Box::new(self),
        }
    }

    /// Build a `mean` aggregation.
    pub fn mean(self) -> Expr {
        Expr::Agg {
            func: AggFunc::Mean,
            expr: Box::new(self),
        }
    }

    /// Build a `count` aggregation (nulls excluded).
    pub fn count(self) -> Expr {
        Expr::Agg {
            func: AggFunc::Count,
            expr: Box::new(self),
        }
    }

    /// Build a `min` aggregation.
    pub fn min(self) -> Expr {
        Expr::Agg {
            func: AggFunc::Min,
            expr: Box::new(self),
        }
    }

    /// Build a `max` aggregation.
    pub fn max(self) -> Expr {
        Expr::Agg {
            func: AggFunc::Max,
            expr: Box::new(self),
        }
    }

    /// Enter the `str.*` expression namespace.
    pub fn str(self) -> StringExpr {
        StringExpr { input: self }
    }

    /// Enter the `dt.*` expression namespace.
    pub fn dt(self) -> DatetimeExpr {
        DatetimeExpr { input: self }
    }

    /// Enter the `list.*` expression namespace.
    pub fn list(self) -> ListExpr {
        ListExpr { input: self }
    }
}

/// Builder for `str.*` expression functions.
#[derive(Debug, Clone, PartialEq)]
pub struct StringExpr {
    input: Expr,
}

impl StringExpr {
    fn function(self, function: StringFunction) -> Expr {
        Expr::Function {
            input: Box::new(self.input),
            function: ExprFunction::String(function),
        }
    }

    /// Convert UTF-8 strings to lowercase.
    pub fn to_lowercase(self) -> Expr {
        self.function(StringFunction::ToLowercase)
    }

    /// Convert UTF-8 strings to uppercase.
    pub fn to_uppercase(self) -> Expr {
        self.function(StringFunction::ToUppercase)
    }

    /// Return whether each string matches a regex pattern.
    pub fn contains(self, pattern: impl Into<String>) -> Expr {
        self.function(StringFunction::Contains {
            pattern: pattern.into(),
        })
    }

    /// Replace regex matches.
    pub fn replace(self, pattern: impl Into<String>, replacement: impl Into<String>) -> Expr {
        self.function(StringFunction::Replace {
            pattern: pattern.into(),
            replacement: replacement.into(),
        })
    }

    /// Strip whitespace from both ends.
    pub fn strip_chars(self, chars: Option<impl Into<String>>) -> Expr {
        self.function(StringFunction::StripChars {
            chars: chars.map(Into::into),
        })
    }

    /// Split by a literal separator.
    pub fn split(self, separator: impl Into<String>) -> Expr {
        self.function(StringFunction::Split {
            separator: separator.into(),
        })
    }

    /// Count Unicode scalar values.
    pub fn len_chars(self) -> Expr {
        self.function(StringFunction::LenChars)
    }

    /// Extract a regex capture group.
    pub fn extract(self, pattern: impl Into<String>, capture_group: usize) -> Expr {
        self.function(StringFunction::Extract {
            pattern: pattern.into(),
            capture_group,
        })
    }
}

/// Builder for `dt.*` expression functions.
#[derive(Debug, Clone, PartialEq)]
pub struct DatetimeExpr {
    input: Expr,
}

impl DatetimeExpr {
    fn function(self, function: DatetimeFunction) -> Expr {
        Expr::Function {
            input: Box::new(self.input),
            function: ExprFunction::Datetime(function),
        }
    }

    /// Extract UTC year.
    pub fn year(self) -> Expr {
        self.function(DatetimeFunction::Year)
    }

    /// Extract UTC month.
    pub fn month(self) -> Expr {
        self.function(DatetimeFunction::Month)
    }

    /// Extract UTC day of month.
    pub fn day(self) -> Expr {
        self.function(DatetimeFunction::Day)
    }

    /// Extract ISO weekday.
    pub fn weekday(self) -> Expr {
        self.function(DatetimeFunction::Weekday)
    }

    /// Format timestamps as UTC text.
    #[allow(clippy::inherent_to_string)]
    pub fn to_string(self) -> Expr {
        self.function(DatetimeFunction::ToString)
    }

    /// Convert between fixed-offset time zones.
    pub fn convert_time_zone(
        self,
        from_offset: impl Into<String>,
        to_offset: impl Into<String>,
    ) -> Expr {
        self.function(DatetimeFunction::ConvertTimeZone {
            from_offset: from_offset.into(),
            to_offset: to_offset.into(),
        })
    }
}

/// Builder for `list.*` expression functions.
#[derive(Debug, Clone, PartialEq)]
pub struct ListExpr {
    input: Expr,
}

impl ListExpr {
    fn function(self, function: ListFunction) -> Expr {
        Expr::Function {
            input: Box::new(self.input),
            function: ExprFunction::List(function),
        }
    }

    /// Join string list elements with a separator.
    pub fn join(self, separator: impl Into<String>, null_value: Option<impl Into<String>>) -> Expr {
        self.function(ListFunction::Join {
            separator: separator.into(),
            null_value: null_value.map(Into::into),
        })
    }

    /// Return list lengths.
    pub fn len(self) -> Expr {
        self.function(ListFunction::Len)
    }

    /// Return whether a string list contains `value`.
    pub fn contains(self, value: impl Into<String>) -> Expr {
        self.function(ListFunction::Contains {
            value: value.into(),
        })
    }
}

#[cfg(test)]
mod tests {
    use super::{AggFunc, Expr, ExprFunction, Operator, Scalar, StringFunction, UnaryOperator};
    use crate::expr::{col, lit};

    #[test]
    fn builder_and_chaining_works() {
        let expr = col("a").add(lit(1_i64)).alias("b");
        assert_eq!(
            expr,
            Expr::Alias {
                expr: Box::new(Expr::BinaryOp {
                    left: Box::new(Expr::Column("a".to_string())),
                    op: Operator::Add,
                    right: Box::new(Expr::Literal(Scalar::Int64(1))),
                }),
                name: "b".to_string(),
            }
        );
    }

    #[test]
    fn logical_and_agg_works() {
        let expr = col("x")
            .gt(lit(1_i64))
            .and_(col("y").lt(lit(10_i64)).not_())
            .alias("p");

        assert!(matches!(
            expr,
            Expr::Alias {
                expr: _,
                name
            } if name == "p"
        ));

        let agg = col("v").sum();
        assert_eq!(
            agg,
            Expr::Agg {
                func: AggFunc::Sum,
                expr: Box::new(Expr::Column("v".to_string()))
            }
        );

        let u = Expr::Column("a".to_string()).not_();
        assert_eq!(
            u,
            Expr::UnaryOp {
                op: UnaryOperator::Not,
                expr: Box::new(Expr::Column("a".to_string()))
            }
        );
    }

    #[test]
    fn namespace_builders_create_function_exprs() {
        let expr = col("name").str().to_lowercase().alias("name_lower");
        assert_eq!(
            expr,
            Expr::Alias {
                expr: Box::new(Expr::Function {
                    input: Box::new(Expr::Column("name".to_string())),
                    function: ExprFunction::String(StringFunction::ToLowercase),
                }),
                name: "name_lower".to_string(),
            }
        );
    }
}