From 46d705a5f6e256ba45c1364fa8549bba4a825a1f Mon Sep 17 00:00:00 2001
From: Palash Tyagi <23239946+Magnus167@users.noreply.github.com>
Date: Sun, 6 Apr 2025 00:57:51 +0100
Subject: [PATCH] Add function to retrieve business dates from a date column in
 a DataFrame

---
 src/utils/misc.rs | 49 +++++++++++++++++++++++++++++++++++++++++++++++
 1 file changed, 49 insertions(+)

diff --git a/src/utils/misc.rs b/src/utils/misc.rs
index 9bca1d1..1e28163 100644
--- a/src/utils/misc.rs
+++ b/src/utils/misc.rs
@@ -1,4 +1,5 @@
 use chrono::NaiveDate;
+use chrono::{Datelike, Weekday};
 use polars::prelude::*;
 use std::collections::HashMap;
 use std::error::Error;
@@ -43,6 +44,54 @@ pub fn get_min_max_real_dates(
     }
 }
 
+/// Get the business dates from a date column in a DataFrame.
+pub fn get_bdates_from_col(date_col: &Series, freq: &str) -> Result<Series, Box<dyn Error>> {
+    let dates = date_col
+        .date()?
+        .into_iter()
+        .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
+        .iter()
+        .map(|date| match freq {
+            "D" => *date,
+            "W" => {
+                let weekday = date.weekday().num_days_from_monday(); // Get the weekday directly
+                *date - chrono::Duration::days(weekday as i64)
+            }
+            "M" => NaiveDate::from_ymd_opt(date.year(), date.month(), 1)
+                .unwrap_or_else(|| NaiveDate::from_ymd_opt(1970, 1, 1).unwrap()),
+            "Q" => {
+                let quarter = (date.month0() / 3) + 1;
+                let month = match quarter {
+                    1 => 1,
+                    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();
+
+    Ok(DateChunked::from_naive_date(date_col.name().clone(), aligned_dates).into_series())
+}
+
 /// Get the `cid` from a ticker string.
 #[allow(dead_code)]
 pub fn get_cid(ticker: String) -> Result<String, Box<dyn Error>> {