Refactor get_bdates_from_col to enhance business day identification and bucket dates by period

This commit is contained in:
Palash Tyagi 2025-04-06 05:13:27 +01:00
parent 46d705a5f6
commit 3bca3a931f

View File

@ -43,53 +43,65 @@ pub fn get_min_max_real_dates(
))) )))
} }
} }
/// Get the business dates from a date column in a DataFrame. /// Get the business dates from a date column in a DataFrame.
/// Identify business days, bucket them by period, and pick the first available date from each period.
pub fn get_bdates_from_col(date_col: &Series, freq: &str) -> Result<Series, Box<dyn Error>> { pub fn get_bdates_from_col(date_col: &Series, freq: &str) -> Result<Series, Box<dyn Error>> {
let dates = date_col // Ensure the column is of Date type
.date()? if date_col.dtype() != &DataType::Date {
.into_iter() return Err("The column is not of Date type".into());
.filter_map(|opt| opt.map(|d| NaiveDate::from_num_days_from_ce_opt(d as i32))) }
.filter(|d| {
if let Some(date) = d {
let wd = date.weekday();
wd != Weekday::Sat && wd != Weekday::Sun
} else {
false
}
})
.filter_map(|opt| opt) // Filter out None and unwrap Some
.collect::<Vec<_>>();
let aligned_dates: Vec<NaiveDate> = dates // Step 1: Identify business days (exclude weekends)
.iter() let date_as_days = date_col.cast(&DataType::Int32)?;
.map(|date| match freq { let business_days: Vec<NaiveDate> = date_as_days
"D" => *date, .i32()?
"W" => { .into_iter()
let weekday = date.weekday().num_days_from_monday(); // Get the weekday directly .filter_map(|opt_days| {
*date - chrono::Duration::days(weekday as i64) opt_days.map(|days| {
} NaiveDate::from_ymd_opt(1970, 1, 1).unwrap() + chrono::Duration::days(days as i64)
"M" => NaiveDate::from_ymd_opt(date.year(), date.month(), 1) })
.unwrap_or_else(|| NaiveDate::from_ymd_opt(1970, 1, 1).unwrap()), })
"Q" => { .filter(|date| {
let quarter = (date.month0() / 3) + 1; // Exclude weekends (Saturday and Sunday)
let month = match quarter { let weekday = date.weekday();
1 => 1, weekday != Weekday::Sat && weekday != Weekday::Sun
2 => 4,
3 => 7,
4 => 10,
_ => unreachable!(),
};
NaiveDate::from_ymd_opt(date.year(), month, 1)
.unwrap_or_else(|| NaiveDate::from_ymd_opt(1970, 1, 1).unwrap())
}
"A" => NaiveDate::from_ymd_opt(date.year(), 1, 1)
.unwrap_or_else(|| NaiveDate::from_ymd_opt(1970, 1, 1).unwrap()),
_ => *date, // fallback
}) })
.collect(); .collect();
Ok(DateChunked::from_naive_date(date_col.name().clone(), aligned_dates).into_series()) // Step 2: Bucket dates by period
let mut buckets: HashMap<String, Vec<NaiveDate>> = HashMap::new();
for date in &business_days {
let bucket_key = match freq {
"D" => date.format("%Y-%m-%d").to_string(),
"W" => format!("{}-W{:02}", date.year(), date.iso_week().week()),
"M" => date.format("%Y-%m").to_string(),
"Q" => format!("{}-Q{}", date.year(), (date.month() - 1) / 3 + 1),
"A" => date.year().to_string(),
_ => return Err("Invalid frequency specified".into()),
};
buckets.entry(bucket_key).or_default().push(*date);
}
// Step 3: Pick the first available date from each bucket
let mut selected_dates: Vec<NaiveDate> = Vec::new();
for (_, mut dates) in buckets {
dates.sort(); // Ensure dates are sorted within the bucket
if let Some(first_date) = dates.first() {
selected_dates.push(*first_date);
}
}
// Step 4: Convert selected dates back to a Series of Date type
let bdates_series = Series::new(
"bdates".into(),
selected_dates
.into_iter()
.map(|date| date.format("%Y-%m-%d").to_string()) // Format as strings
.collect::<Vec<String>>(),
)
.cast(&DataType::Date)?; // Cast to Date type
Ok(bdates_series)
} }
/// Get the `cid` from a ticker string. /// Get the `cid` from a ticker string.