Implement code changes to enhance functionality and improve performance

src/utils/qdf/blacklist.rs

@@ -7,8 +7,45 @@ use polars::prelude::*;
 use std::collections::{BTreeMap, HashMap};
 use std::error::Error;
 
-use super::get_unique_metrics;
+use crate::utils::qdf::get_unique_metrics;
 
+// struct Blacklist which is a wrapper around hashmap and btreemap
+#[derive(Debug, Clone)]
+pub struct Blacklist {
+    pub blacklist: BTreeMap<String, (String, String)>,
+}
+
+// impl hashmap into
+impl Blacklist {
+    pub fn into_hashmap(self) -> HashMap<String, (String, String)> {
+        self.blacklist.into_iter().collect()
+    }
+}
+
+/// Apply a blacklist to a Quantamental DataFrame with Lazy API.
+///
+/// * `blacklist` is a map from any “ticker‑like” key to a tuple of
+///   `(start_date, end_date)` in **inclusive** `"YYYY‑MM‑DD"` format.
+/// * `metrics` – if `None`, every metric from `get_unique_metrics(df)`
+///   is used.
+/// * `group_by_cid = Some(false)` is not implemented yet (parity with
+///   the eager version).
+pub fn apply_blacklist_lazy(
+    df: &mut DataFrame,
+    blacklist: &BTreeMap<String, (String, String)>,
+    metrics: Option<Vec<String>>,
+    group_by_cid: Option<bool>,
+) -> Result<DataFrame, Box<dyn std::error::Error>> {
+    check_quantamental_dataframe(df)?;
+    Ok(df.clone())
+}
+/// Create a blacklist from a Quantamental DataFrame.
+/// The blacklist is a mapping of tickers to date ranges where the specified metrics are null or NaN.
+/// # Arguments:
+/// * `df` - The Quantamental DataFrame.
+/// * `group_by_cid` - If true, group the blacklist by `cid`. Defaults to true.
+/// * `blacklist_name` - The name of the blacklist. Defaults to "BLACKLIST".
+/// * `metrics` - The metrics to check for null or NaN values. If None, all metrics are used.
 pub fn create_blacklist_from_qdf(
     df: &DataFrame,
     group_by_cid: Option<bool>,
@@ -30,19 +67,26 @@ pub fn create_blacklist_from_qdf(
         BDateFreq::Daily,
     )?;
 
-    let null_mask = get_nan_mask(df, metrics)?;
+    // if none of the metrics are null or NaN, return an empty blacklist
+    if !metrics.iter().any(|metric| {
+        df.column(metric)
+            .map(|col| col.is_null().any())
+            .unwrap_or(false)
+    }) {
+        return Ok(BTreeMap::new());
+    }
 
-    let df = df.filter(&null_mask)?.clone();
+    // let null_mask = get_nan_mask(df, metrics)?;
+    // let df = df.filter(&null_mask)?.clone();
 
     let df = df
         .clone()
         .lazy()
-        // .filter(&null_mask)
-        // .filter(
-        //     col(metric.as_str())
-        //         .is_null()
-        //         .or(col(metric.as_str()).is_nan()),
-        // )
+        .with_columns([
+            (cols(metrics.clone()).is_null().or(cols(metrics).is_nan())).alias("null_mask")
+        ])
+        .filter(col("null_mask"))
+        // if is now empty, return an empty blacklist
         .sort(
             ["cid", "xcat"],
             SortMultipleOptions::default().with_maintain_order(true),
@@ -112,6 +156,8 @@ pub fn create_blacklist_from_qdf(
     Ok(btree_map)
 }
 
+/// Get a mask of NaN values for the specified metrics in the DataFrame.
+#[allow(dead_code)]
 fn get_nan_mask(
     df: &DataFrame,
     metrics: Vec<String>,