refactor: rename apply_blacklist_lazy to apply_blacklist and update documentation

.gitattributes

@@ -0,0 +1 @@
+notebooks/** linguist-vendored

src/_py/utils.rs

@@ -1,7 +1,6 @@
 use pyo3::{prelude::*, types::PyDict};
 use pyo3_polars::{PyDataFrame, PySeries};
 
-
 /// Python wrapper for [`crate::utils::qdf`] module.
 #[allow(deprecated)]
 #[pymodule]
@@ -37,18 +36,18 @@ pub fn get_bdates_series_default_opt(
 }
 
 #[allow(deprecated)]
-#[pyfunction(signature = (df, group_by_cid=None, blacklist_name=None, metric=None))]
+#[pyfunction(signature = (df, group_by_cid=None, blacklist_name=None, metrics=None))]
 pub fn create_blacklist_from_qdf(
     df: PyDataFrame,
     group_by_cid: Option<bool>,
     blacklist_name: Option<String>,
-    metric: Option<String>,
+    metrics: Option<Vec<String>>,
 ) -> PyResult<PyObject> {
     let result = crate::utils::qdf::blacklist::create_blacklist_from_qdf(
         &df.into(),
         group_by_cid,
         blacklist_name,
-        metric,
+        metrics,
     )
     .map_err(|e| PyErr::new::<pyo3::exceptions::PyValueError, _>(format!("{}", e)))?;
     Python::with_gil(|py| {

src/utils/qdf/blacklist.rs

@@ -7,15 +7,73 @@ use polars::prelude::*;
 use std::collections::{BTreeMap, HashMap};
 use std::error::Error;
 
+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.
+///
+/// * `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.
+pub fn apply_blacklist(
+    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)?;
+    // dataframe is like:
+    // | cid | xcat | real_date | metric1 | metric2 |
+    // |-----|------|-----------|---------|---------|
+    // | A   | B    | 2023-01-01| 1.0     | 2.0     |
+    // | A   | B    | 2023-01-02| 1.0     | 2.0     |
+    // | A   | C    | 2023-01-01| 1.0     | 2.0     |
+    // | A   | C    | 2023-01-02| 1.0     | 2.0     |
+    // | D   | E    | 2023-01-01| 1.0     | 2.0     |
+    // | D   | E    | 2023-01-02| 1.0     | 2.0     |
+
+    // (real date column is Naive date)
+
+    // blacklist is like:
+    // {'A_B_1': ('2023-01-02', '2023-01-03'),
+    //  'A_B_2': ('2023-01-04', '2023-01-05'),
+    //  'A_C_1': ('2023-01-02', '2023-01-03'), }
+
+    // get_cid('A_B_1') = 'A'
+    // get_cid('A_B_2') = 'A'
+    // get_cid('D_E_1') = 'D'
+
+    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>,
     blacklist_name: Option<String>,
-    metric: Option<String>,
+    metrics: Option<Vec<String>>,
 ) -> Result<BTreeMap<String, (String, String)>, Box<dyn Error>> {
     check_quantamental_dataframe(df)?;
-
-    let metric = metric.unwrap_or_else(|| "value".into());
+    let metrics = metrics.unwrap_or_else(|| get_unique_metrics(df).unwrap());
     let blacklist_name = blacklist_name.unwrap_or_else(|| "BLACKLIST".into());
     let group_by_cid = group_by_cid.unwrap_or(true);
 
@@ -29,19 +87,26 @@ pub fn create_blacklist_from_qdf(
         BDateFreq::Daily,
     )?;
 
-    // filter df
-    let null_mask = df.column(metric.as_str())?.is_null();
-    let nan_mask = df.column(metric.as_str())?.is_nan()?;
-    let null_mask = null_mask | nan_mask;
-    let df = df.filter(&null_mask)?;
+    // 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 null_mask = get_nan_mask(df, metrics)?;
+    // let df = df.filter(&null_mask)?.clone();
 
     let df = df
+        .clone()
         .lazy()
-        .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),
@@ -111,6 +176,27 @@ 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>,
+) -> Result<ChunkedArray<BooleanType>, Box<dyn Error>> {
+    let null_masks: Vec<ChunkedArray<BooleanType>> = metrics
+        .iter()
+        .map(|metric| {
+            let null_mask = df.column(metric.as_str())?.is_null();
+            let nan_mask = df.column(metric.as_str())?.is_nan()?;
+            Ok(null_mask | nan_mask)
+        })
+        .collect::<Result<_, Box<dyn Error>>>()?;
+    let null_mask = null_masks
+        .into_iter()
+        .reduce(|acc, mask| acc | mask)
+        .unwrap_or_else(|| BooleanChunked::full_null("null_mask".into(), df.height()));
+    Ok(null_mask)
+}
+
 fn convert_dates_list_to_date_ranges(
     blacklist: Vec<String>,
     all_bdates_strs: Vec<String>,
@@ -275,7 +361,13 @@ mod tests {
         // Expect two ranges:
         //   range 0 => ("2023-01-02", "2023-01-03")
         //   range 1 => ("2023-01-05", "2023-01-05")
-        assert_eq!(result["0"], ("2023-01-02".to_string(), "2023-01-03".to_string()));
-        assert_eq!(result["1"], ("2023-01-05".to_string(), "2023-01-05".to_string()));
+        assert_eq!(
+            result["0"],
+            ("2023-01-02".to_string(), "2023-01-03".to_string())
+        );
+        assert_eq!(
+            result["1"],
+            ("2023-01-05".to_string(), "2023-01-05".to_string())
+        );
     }
 }