Refactor business date retrieval by removing redundant functions and utilizing get_bdates_list_with_freq for improved efficiency

src/utils/dateutils.rs

@@ -1,3 +1,5 @@
+use crate::utils::bdates::get_bdates_list_with_freq;
+use crate::utils::bdates::BDateFreq;
 use chrono::NaiveDate;
 use chrono::{Datelike, Weekday};
 use polars::prelude::*;
@@ -34,107 +36,6 @@ pub fn get_min_max_real_dates(
     }
 }
 
-/// Get the business dates between two dates.
-pub fn get_bdates_list(
-    start_date: String,
-    end_date: String,
-) -> Result<Vec<NaiveDate>, Box<dyn Error>> {
-    let start_date = NaiveDate::parse_from_str(&start_date, "%Y-%m-%d")?;
-    let end_date = NaiveDate::parse_from_str(&end_date, "%Y-%m-%d")?;
-
-    let mut business_days = Vec::new();
-    let mut current_date = start_date;
-    while current_date <= end_date {
-        // Check if the current date is a business day (not Saturday or Sunday)
-        if current_date.weekday() != Weekday::Sat && current_date.weekday() != Weekday::Sun {
-            business_days.push(current_date);
-        }
-        current_date = current_date.succ_opt().ok_or(format!(
-            "Failed to get the next day for : {:?}",
-            current_date
-        ))?;
-    }
-    Ok(business_days)
-}
-
-#[derive(Debug, Clone, Copy)]
-pub enum BDateFreq {
-    Daily,
-    WeeklyMonday,
-    MonthStart,
-    QuarterStart,
-    YearStart,
-    MonthEnd,
-    QuarterEnd,
-    WeeklyFriday,
-    YearEnd,
-}
-
-impl BDateFreq {
-    pub fn from_string(freq: String) -> Result<Self, Box<dyn Error>> {
-        // use `from_str` to convert the string to a BDateFreq enum
-        Self::from_str(&freq)
-    }
-    pub fn from_str(freq: &str) -> Result<Self, Box<dyn Error>> {
-        match freq {
-            "D" => Ok(BDateFreq::Daily),
-            "W" => Ok(BDateFreq::WeeklyMonday),
-            "M" => Ok(BDateFreq::MonthStart),
-            "Q" => Ok(BDateFreq::QuarterStart),
-            "A" => Ok(BDateFreq::YearStart),
-            "ME" => Ok(BDateFreq::MonthEnd),
-            "QE" => Ok(BDateFreq::QuarterEnd),
-            "WF" => Ok(BDateFreq::WeeklyFriday),
-            "YE" => Ok(BDateFreq::YearEnd),
-            _ => Err("Invalid frequency specified".into()),
-        }
-    }
-
-    pub fn agg_type(&self) -> AggregationType {
-        match self {
-            BDateFreq::Daily
-            | BDateFreq::WeeklyMonday
-            | BDateFreq::MonthStart
-            | BDateFreq::QuarterStart
-            | BDateFreq::YearStart => AggregationType::Start,
-            BDateFreq::WeeklyFriday
-            | BDateFreq::MonthEnd
-            | BDateFreq::QuarterEnd
-            | BDateFreq::YearEnd => AggregationType::End,
-        }
-    }
-}
-
-#[derive(Debug, Clone, Copy)]
-pub enum AggregationType {
-    Start, // Indicates picking the first date in a group.
-    End,   // Indicates picking the last date in a group.
-}
-
-// Map a BDateFreq to an AggregationType.
-fn compute_group_key(d: NaiveDate, freq: BDateFreq) -> String {
-    match freq {
-        // For Daily, each date is its own group.
-        BDateFreq::Daily => format!("{}", d),
-        // For weekly grouping, we use ISO week information.
-        BDateFreq::WeeklyMonday | BDateFreq::WeeklyFriday => {
-            let iso = d.iso_week();
-            format!("{}-W{:02}", iso.year(), iso.week())
-        }
-        // Group by Year-Month.
-        BDateFreq::MonthStart | BDateFreq::MonthEnd => {
-            format!("{}-M{:02}", d.year(), d.month())
-        }
-        // Group by Year-Quarter.
-        BDateFreq::QuarterStart | BDateFreq::QuarterEnd => {
-            let quarter = (d.month() - 1) / 3 + 1;
-            format!("{}-Q{}", d.year(), quarter)
-        }
-        // Group by Year.
-        BDateFreq::YearStart | BDateFreq::YearEnd => format!("{}", d.year()),
-    }
-}
-
 pub fn get_bdates_series_default(
     start_date: String,
     end_date: String,
@@ -151,35 +52,9 @@ pub fn get_bdates_series(
     end_date: String,
     freq: BDateFreq,
 ) -> Result<Series, Box<dyn Error>> {
-    let business_days = get_bdates_list(start_date, end_date)?;
+    let bdates_list = get_bdates_list_with_freq(&start_date, &end_date, freq)?;
-    let group_keys: Vec<String> = business_days
+    let bdates_series = Series::new("bdates".into(), bdates_list);
-        .iter()
+    Ok(bdates_series)
-        .map(|&d| compute_group_key(d, freq))
-        .collect();
-
-    let df = DataFrame::new(vec![
-        Column::new("bdates".into(), business_days),
-        Column::new("group".into(), group_keys),
-    ])?;
-    let gb = df.lazy().group_by(["group"]);
-    let aggx = match freq.agg_type() {
-        AggregationType::Start => gb.agg([col("bdates").first()]),
-        AggregationType::End => gb.agg([col("bdates").last()]),
-    };
-    let result = aggx.collect()?;
-    let result = result
-        .column("bdates")?
-        .as_series()
-        .ok_or("Column 'bdates' not found")?
-        .clone();
-    let result = result.sort(SortOptions {
-        descending: false,
-        nulls_last: false,
-        multithreaded: false,
-        maintain_order: false,
-    })?;
-
-    Ok(result)
 }
 
 /// Get the business dates from a date column in a DataFrame.