Compare commits

...

2 Commits

2 changed files with 46 additions and 21 deletions

View File

@ -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| {

View File

@ -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 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 +30,19 @@ 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)?;
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()),
)
// .filter(&null_mask)
// .filter(
// col(metric.as_str())
// .is_null()
// .or(col(metric.as_str()).is_nan()),
// )
.sort(
["cid", "xcat"],
SortMultipleOptions::default().with_maintain_order(true),
@ -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!(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())
);
}
}