From d31fa03004632c2e70bad74e1e295d60582acf64 Mon Sep 17 00:00:00 2001 From: Palash Tyagi <23239946+Magnus167@users.noreply.github.com> Date: Sun, 13 Apr 2025 00:39:13 +0100 Subject: [PATCH] Refactor business date retrieval by removing redundant functions and utilizing get_bdates_list_with_freq for improved efficiency --- src/utils/dateutils.rs | 135 ++--------------------------------------- 1 file changed, 5 insertions(+), 130 deletions(-) diff --git a/src/utils/dateutils.rs b/src/utils/dateutils.rs index 2261c98..4c8c594 100644 --- a/src/utils/dateutils.rs +++ b/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, Box> { - 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> { - // use `from_str` to convert the string to a BDateFreq enum - Self::from_str(&freq) - } - pub fn from_str(freq: &str) -> Result> { - 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> { - let business_days = get_bdates_list(start_date, end_date)?; - let group_keys: Vec = business_days - .iter() - .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) + let bdates_list = get_bdates_list_with_freq(&start_date, &end_date, freq)?; + let bdates_series = Series::new("bdates".into(), bdates_list); + Ok(bdates_series) } /// Get the business dates from a date column in a DataFrame.