Enhance dateutils with BDateFreq and AggregationType enums for improved date grouping and aggregation functionality

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<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,
+    freq: Option<String>,
+) -> Result<Series, Box<dyn Error>> {
+    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<Series, Box<dyn Error>> {
+pub fn get_bdates_series(
+    start_date: String,
+    end_date: String,
+    freq: BDateFreq,
+) -> Result<Series, Box<dyn Error>> {
     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<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)
 }
 
 /// Get the business dates from a date column in a DataFrame.