update create_blacklist_from_qdf function signature to accept metrics as a vector

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,16 @@ use polars::prelude::*;
 use std::collections::{BTreeMap, HashMap};
 use std::error::Error;
 
+use super::get_unique_metrics;
+
 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 metrics = metrics.unwrap_or_else(|| get_unique_metrics(df).unwrap());
-    let metric = metric.unwrap_or_else(|| "value".into());
     let blacklist_name = blacklist_name.unwrap_or_else(|| "BLACKLIST".into());
     let group_by_cid = group_by_cid.unwrap_or(true);
 
@@ -29,19 +30,19 @@ pub fn create_blacklist_from_qdf(
         BDateFreq::Daily,
     )?;
 
-    // filter df
+    let null_mask = get_nan_mask(df, metrics)?;
-    let null_mask = df.column(metric.as_str())?.is_null();
+
-    let nan_mask = df.column(metric.as_str())?.is_nan()?;
+    let df = df.filter(&null_mask)?.clone();
-    let null_mask = null_mask | nan_mask;
-    let df = df.filter(&null_mask)?;
 
     let df = df
+        .clone()
         .lazy()
-        .filter(
+        // .filter(&null_mask)
-            col(metric.as_str())
+        // .filter(
-                .is_null()
+        //     col(metric.as_str())
-                .or(col(metric.as_str()).is_nan()),
+        //         .is_null()
-        )
+        //         .or(col(metric.as_str()).is_nan()),
+        // )
         .sort(
             ["cid", "xcat"],
             SortMultipleOptions::default().with_maintain_order(true),
@@ -85,7 +86,7 @@ pub fn create_blacklist_from_qdf(
     for (_key, vals) in blk.iter_mut() {
         // order is important - dedup depends on the vec being sorted
         vals.sort();
-        vals.dedup(); 
+        vals.dedup();
     }
 
     let all_bdates_strs = all_bdates
@@ -111,6 +112,25 @@ pub fn create_blacklist_from_qdf(
     Ok(btree_map)
 }
 
+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 +295,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!(
-        assert_eq!(result["1"], ("2023-01-05".to_string(), "2023-01-05".to_string()));
+            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())
+        );
     }
-}
+}