polars_plan/dsl/function_expr/
business.rs

1use std::fmt::{Display, Formatter};
2
3use polars_core::prelude::*;
4use polars_ops::prelude::Roll;
5#[cfg(feature = "serde")]
6use serde::{Deserialize, Serialize};
7
8use super::FunctionOptions;
9use crate::dsl::{FieldsMapper, SpecialEq};
10use crate::map_as_slice;
11use crate::prelude::{ColumnsUdf, FunctionFlags};
12
13#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
14#[derive(Clone, PartialEq, Debug, Eq, Hash)]
15pub enum BusinessFunction {
16    BusinessDayCount {
17        week_mask: [bool; 7],
18        holidays: Vec<i32>,
19    },
20    AddBusinessDay {
21        week_mask: [bool; 7],
22        holidays: Vec<i32>,
23        roll: Roll,
24    },
25    IsBusinessDay {
26        week_mask: [bool; 7],
27        holidays: Vec<i32>,
28    },
29}
30
31impl BusinessFunction {
32    pub fn get_field(&self, mapper: FieldsMapper) -> PolarsResult<Field> {
33        match self {
34            Self::BusinessDayCount { .. } => mapper.with_dtype(DataType::Int32),
35            Self::AddBusinessDay { .. } => mapper.with_same_dtype(),
36            Self::IsBusinessDay { .. } => mapper.with_dtype(DataType::Boolean),
37        }
38    }
39    pub fn function_options(&self) -> FunctionOptions {
40        use BusinessFunction as B;
41        match self {
42            B::BusinessDayCount { .. } => {
43                FunctionOptions::elementwise().with_flags(|f| f | FunctionFlags::ALLOW_RENAME)
44            },
45            B::AddBusinessDay { .. } | B::IsBusinessDay { .. } => FunctionOptions::elementwise(),
46        }
47    }
48}
49
50impl Display for BusinessFunction {
51    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
52        use BusinessFunction::*;
53        let s = match self {
54            BusinessDayCount { .. } => "business_day_count",
55            AddBusinessDay { .. } => "add_business_days",
56            IsBusinessDay { .. } => "is_business_day",
57        };
58        write!(f, "{s}")
59    }
60}
61impl From<BusinessFunction> for SpecialEq<Arc<dyn ColumnsUdf>> {
62    fn from(func: BusinessFunction) -> Self {
63        use BusinessFunction::*;
64        match func {
65            BusinessDayCount {
66                week_mask,
67                holidays,
68            } => {
69                map_as_slice!(business_day_count, week_mask, &holidays)
70            },
71            AddBusinessDay {
72                week_mask,
73                holidays,
74                roll,
75            } => {
76                map_as_slice!(add_business_days, week_mask, &holidays, roll)
77            },
78            IsBusinessDay {
79                week_mask,
80                holidays,
81            } => {
82                map_as_slice!(is_business_day, week_mask, &holidays)
83            },
84        }
85    }
86}
87
88pub(super) fn business_day_count(
89    s: &[Column],
90    week_mask: [bool; 7],
91    holidays: &[i32],
92) -> PolarsResult<Column> {
93    let start = &s[0];
94    let end = &s[1];
95    polars_ops::prelude::business_day_count(
96        start.as_materialized_series(),
97        end.as_materialized_series(),
98        week_mask,
99        holidays,
100    )
101    .map(Column::from)
102}
103pub(super) fn add_business_days(
104    s: &[Column],
105    week_mask: [bool; 7],
106    holidays: &[i32],
107    roll: Roll,
108) -> PolarsResult<Column> {
109    let start = &s[0];
110    let n = &s[1];
111    polars_ops::prelude::add_business_days(
112        start.as_materialized_series(),
113        n.as_materialized_series(),
114        week_mask,
115        holidays,
116        roll,
117    )
118    .map(Column::from)
119}
120
121pub(super) fn is_business_day(
122    s: &[Column],
123    week_mask: [bool; 7],
124    holidays: &[i32],
125) -> PolarsResult<Column> {
126    let dates = &s[0];
127    polars_ops::prelude::is_business_day(dates.as_materialized_series(), week_mask, holidays)
128        .map(Column::from)
129}
polars_plan/dsl/function_expr/business.rs

polars_plan/dsl/function_expr/
business.rs
