polars-plan 0.54.1

Lazy query engine for the Polars DataFrame library
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47

use std::fmt::{Display, Formatter};

use polars_core::prelude::*;
use polars_ops::prelude::Roll;

use super::FunctionOptions;
use crate::plans::aexpr::function_expr::FieldsMapper;
use crate::prelude::FunctionFlags;

#[cfg_attr(feature = "ir_serde", derive(serde::Serialize, serde::Deserialize))]
#[derive(Clone, PartialEq, Debug, Eq, Hash)]
pub enum IRBusinessFunction {
    BusinessDayCount { week_mask: [bool; 7] },
    AddBusinessDay { week_mask: [bool; 7], roll: Roll },
    IsBusinessDay { week_mask: [bool; 7] },
}

impl IRBusinessFunction {
    pub fn get_field(&self, mapper: FieldsMapper) -> PolarsResult<Field> {
        match self {
            Self::BusinessDayCount { .. } => mapper.with_dtype(DataType::Int32),
            Self::AddBusinessDay { .. } => mapper.with_same_dtype(),
            Self::IsBusinessDay { .. } => mapper.with_dtype(DataType::Boolean),
        }
    }
    pub fn function_options(&self) -> FunctionOptions {
        use IRBusinessFunction as B;
        match self {
            B::BusinessDayCount { .. } => {
                FunctionOptions::elementwise().with_flags(|f| f | FunctionFlags::ALLOW_RENAME)
            },
            B::AddBusinessDay { .. } | B::IsBusinessDay { .. } => FunctionOptions::elementwise(),
        }
    }
}

impl Display for IRBusinessFunction {
    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
        use IRBusinessFunction::*;
        let s = match self {
            BusinessDayCount { .. } => "business_day_count",
            AddBusinessDay { .. } => "add_business_days",
            IsBusinessDay { .. } => "is_business_day",
        };
        write!(f, "{s}")
    }
}