src/utils/qdf/blacklist.rs

@@ -7,45 +7,8 @@ use polars::prelude::*;
 use std::collections::{BTreeMap, HashMap};
 use std::error::Error;
 
-use crate::utils::qdf::get_unique_metrics;
+use super::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>,
@@ -67,26 +30,19 @@ pub fn create_blacklist_from_qdf(
         BDateFreq::Daily,
     )?;
 
-    // if none of the metrics are null or NaN, return an empty blacklist
+    let null_mask = get_nan_mask(df, metrics)?;
-    if !metrics.iter().any(|metric| {
-        df.column(metric)
-            .map(|col| col.is_null().any())
-            .unwrap_or(false)
-    }) {
-        return Ok(BTreeMap::new());
-    }
 
-    // let null_mask = get_nan_mask(df, metrics)?;
+    let df = df.filter(&null_mask)?.clone();
-    // let df = df.filter(&null_mask)?.clone();
 
     let df = df
         .clone()
         .lazy()
-        .with_columns([
+        // .filter(&null_mask)
-            (cols(metrics.clone()).is_null().or(cols(metrics).is_nan())).alias("null_mask")
+        // .filter(
-        ])
+        //     col(metric.as_str())
-        .filter(col("null_mask"))
+        //         .is_null()
-        // if is now empty, return an empty blacklist
+        //         .or(col(metric.as_str()).is_nan()),
+        // )
         .sort(
             ["cid", "xcat"],
             SortMultipleOptions::default().with_maintain_order(true),
@@ -156,8 +112,6 @@ 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>,