Compare commits

...

2 Commits

2 changed files with 8 additions and 132 deletions

View File

@ -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 group_keys: Vec<String> = 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.

View File

@ -1,3 +1,4 @@
pub mod qdf;
pub mod bdates;
pub mod dateutils;
pub mod misc;
pub mod dateutils;
pub mod qdf;