wip: blacklist util

src/_py/utils.rs

@@ -1,4 +1,5 @@
-use pyo3::prelude::*;
+use pyo3::types::IntoPyDict;
+use pyo3::{prelude::*, types::PyDict};
 use pyo3_polars::{PyDataFrame, PySeries};
 
 /// Python wrapper for [`crate::utils::qdf`] module.
@@ -7,6 +8,7 @@ use pyo3_polars::{PyDataFrame, PySeries};
 pub fn utils(_py: Python, m: &PyModule) -> PyResult<()> {
     m.add_function(wrap_pyfunction!(get_bdates_series_default_pl, m)?)?;
     m.add_function(wrap_pyfunction!(get_bdates_series_default_opt, m)?)?;
+    m.add_function(wrap_pyfunction!(create_blacklist_from_qdf, m)?)?;
     Ok(())
 }
 
@@ -33,3 +35,19 @@ pub fn get_bdates_series_default_opt(
             .map_err(|e| PyErr::new::<pyo3::exceptions::PyValueError, _>(format!("{}", e)))?,
     ))
 }
+
+#[allow(deprecated)]
+#[pyfunction(signature = (df, metric = None))]
+pub fn create_blacklist_from_qdf(df: PyDataFrame, metric: Option<String>) -> PyResult<PyObject> {
+    let result = crate::utils::qdf::blacklist::create_blacklist_from_qdf(&df.into(), metric)
+        .map_err(|e| PyErr::new::<pyo3::exceptions::PyValueError, _>(format!("{}", e)))?;
+    Python::with_gil(|py| {
+        let dict = PyDict::new(py);
+        // for (key, (start_date, end_date)) in result {
+        //     dict.set_item(key, (start_date, end_date))
+        for (key, dates) in result {
+            dict.set_item(key, dates).map_err(|e| PyErr::from(e))?;
+        }
+        Ok(dict.into())
+    })
+}

src/msyrs.pyi

@@ -56,3 +56,5 @@ class utils:
     def get_bdates_series_default_pl(*args, **kwargs) -> Series: ...
     @staticmethod
     def get_bdates_series_default_opt(*args, **kwargs) -> Series: ...
+    @staticmethod
+    def create_blacklist_from_qdf(*args, **kwargs) -> dict: ...

src/utils/qdf/blacklist.rs

@@ -1,27 +1,161 @@
+use crate::utils::bdates::{self, get_bdates_list_with_freq, BDateFreq};
 use crate::utils::dateutils::{get_bdates_series_default_opt, get_min_max_real_dates};
-use crate::utils::qdf::core::*;
-use chrono::{Duration, NaiveDate};
+use crate::utils::qdf::core::check_quantamental_dataframe;
+use chrono::NaiveDate;
 use polars::prelude::*;
 use std::collections::HashMap;
 use std::error::Error;
 
-use super::pivots::pivot_dataframe_by_ticker;
-
-/// The required columns for a Quantamental DataFrame.
-const QDF_INDEX_COLUMNS: [&str; 3] = ["real_date", "cid", "xcat"];
-
 pub fn create_blacklist_from_qdf(
     df: &DataFrame,
     metric: Option<String>,
-    start: Option<String>,
-    end: Option<String>,
+    // ) -> Result<HashMap<String, (String, String)>, Box<dyn Error>> {
 ) -> Result<HashMap<String, Vec<String>>, Box<dyn Error>> {
     // Verify that the DataFrame follows the Quantamental structure.
     check_quantamental_dataframe(df)?;
     let mut blacklist: HashMap<String, Vec<String>> = HashMap::new();
+    // let mut blacklist: HashMap<String, (String, String)> = HashMap::new();
+    let mut blk: HashMap<String, Vec<NaiveDate>> = HashMap::new();
 
     // Use the provided metric or default to "value".
     let metric = metric.unwrap_or_else(|| "value".into());
 
+    let (min_date, max_date) = get_min_max_real_dates(df, "real_date".into())?;
+    let min_date_str = min_date.format("%Y-%m-%d").to_string();
+    let max_date_str = max_date.format("%Y-%m-%d").to_string();
+    // let all_bdates = get_bdates_series_default_opt(min_date_str, max_date_str, None)?;
+    let all_bdates = get_bdates_list_with_freq(
+        min_date_str.clone().as_str(),
+        max_date_str.clone().as_str(),
+        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)?;
+
+    let df = df
+        .lazy()
+        .filter(
+            col(metric.as_str())
+                .is_null()
+                .or(col(metric.as_str()).is_nan()),
+        )
+        .sort(
+            ["cid", "xcat"],
+            SortMultipleOptions::default().with_maintain_order(true),
+        )
+        .group_by([col("cid"), col("xcat")])
+        // .agg([col("real_date").sort(SortOptions::default())])
+        .agg([col("real_date")
+            .dt()
+            .strftime("%Y-%m-%d")
+            .sort(SortOptions::default())])
+        .select([
+            concat_str([col("cid"), col("xcat")], "_", true).alias("ticker"),
+            col("real_date").alias("real_dates"),
+        ])
+        .collect()?;
+
+    // assert!(0 == 1, "{:?}", df);
+
+    let ticker_vec = df
+        .column("ticker")?
+        .str()?
+        .into_iter()
+        .filter_map(|opt| opt.map(|s| s.to_string()))
+        .collect::<Vec<String>>();
+
+    let rdt = get_vec_of_vec_of_dates_from_df(df)?;
+
+    // assert!(0 == 1, "{:?}", rdt);
+    for (tkr, dates) in ticker_vec.iter().zip(rdt.iter()) {
+        blacklist.insert(tkr.to_string(), dates.clone());
+    }
+
     Ok(blacklist)
 }
+
+fn convert_dates_list_to_date_ranges(
+    blacklist: HashMap<String, Vec<String>>,
+    all_bdates: Vec<NaiveDate>,
+) -> HashMap<String, (String, String)> {
+    let blk = HashMap::new();
+    let bdates = all_bdates
+        .iter()
+        .map(|date| date.format("%Y-%m-%d").to_string())
+        .collect::<Vec<String>>();
+
+    //
+
+    blk
+}
+
+fn get_vec_of_vec_of_dates_from_df(df: DataFrame) -> Result<Vec<Vec<String>>, Box<dyn Error>> {
+    let rdt = df
+        .column("real_dates")?
+        // .clone()
+        .as_series()
+        .unwrap()
+        .list()?
+        .into_iter()
+        .filter_map(|opt| opt)
+        .collect::<Vec<Series>>()
+        .iter()
+        .map(|s| {
+            s.str()
+                .unwrap()
+                .into_iter()
+                .filter_map(|opt| opt.map(|s| s.to_string()))
+                .collect::<Vec<String>>()
+        })
+        .collect::<Vec<Vec<String>>>();
+    Ok(rdt)
+}
+
+fn get_vec_of_vec_of_naivedates_from_df(
+    df: DataFrame,
+) -> Result<Vec<Vec<NaiveDate>>, Box<dyn Error>> {
+    let rdt = df
+        .column("real_dates")?
+        // .clone()
+        .as_series()
+        .unwrap()
+        .list()?
+        .into_iter()
+        .filter_map(|opt| opt)
+        .collect::<Vec<Series>>()
+        .iter()
+        .map(|s| {
+            s.date()
+                .unwrap()
+                .into_iter()
+                .filter_map(|opt| opt.and_then(|date| NaiveDate::from_num_days_from_ce_opt(date)))
+                .collect::<Vec<NaiveDate>>()
+        })
+        .collect::<Vec<Vec<NaiveDate>>>();
+    Ok(rdt)
+}
+
+// fn get_vec_of_vec_of_dates_from_df(df: DataFrame) -> Result<Vec<Vec<String>>, Box<dyn Error>> {
+//     let real_dates_column = df.column("real_dates")?.clone();
+//     let series = real_dates_column.as_series().unwrap().clone();
+//     let rdt = series.list()?.clone();
+//     let rdt = rdt
+//         .into_iter()
+//         .filter_map(|opt| opt)
+//         .collect::<Vec<Series>>();
+//     let rdt = rdt
+//         .iter()
+//         .map(|s| {
+//             s.str()
+//                 .unwrap()
+//                 .into_iter()
+//                 .filter_map(|opt| opt.map(|s| s.to_string()))
+//                 .collect::<Vec<String>>()
+//         })
+//         .collect::<Vec<Vec<String>>>();
+//     Ok(rdt)
+// }