mirror of
https://github.com/Magnus167/msyrs.git
synced 2025-08-20 07:40:00 +00:00
Compare commits
2 Commits
cfbd54be7a
...
8f5f276efd
Author | SHA1 | Date | |
---|---|---|---|
![]() |
8f5f276efd | ||
![]() |
d31fa03004 |
@ -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.
|
||||
|
@ -1,3 +1,4 @@
|
||||
pub mod qdf;
|
||||
pub mod bdates;
|
||||
pub mod dateutils;
|
||||
pub mod misc;
|
||||
pub mod dateutils;
|
||||
pub mod qdf;
|
||||
|
Loading…
x
Reference in New Issue
Block a user