diff --git a/src/utils/dateutils.rs b/src/utils/dateutils.rs index 9f4c388..4f1697f 100644 --- a/src/utils/dateutils.rs +++ b/src/utils/dateutils.rs @@ -57,11 +57,125 @@ pub fn get_bdates_list( Ok(business_days) } +#[derive(Debug, Clone, Copy)] +pub enum BDateFreq { + Daily, + WeeklyMonday, + MonthStart, + QuarterStart, + YearStart, + MonthEnd, + QuarterEnd, + WeeklyFriday, + YearEnd, +} + +impl BDateFreq { + 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, + freq: Option, +) -> Result> { + let freq = freq.unwrap_or_else(|| "D".to_string()); + let freq = BDateFreq::from_str(&freq)?; + get_bdates_series(start_date, end_date, freq) +} + /// Get the business dates between two dates as a Series. -pub fn get_bdates_series(start_date: String, end_date: String) -> Result> { +pub fn get_bdates_series( + start_date: String, + end_date: String, + freq: BDateFreq, +) -> Result> { let business_days = get_bdates_list(start_date, end_date)?; - let series = Series::new("business_dates".into(), business_days); - Ok(series) + 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) } /// Get the business dates from a date column in a DataFrame.