magnus167/msyrs
1
0
Fork 0
mirror of https://github.com/Magnus167/msyrs.git synced 2025-08-20 04:20:00 +00:00
Code Issues Packages Projects Releases Wiki Activity

feat: refactor create_blacklist_from_qdf to accept multiple metrics and improve null handling

Browse Source
This commit is contained in:
Palash Tyagi 2025-04-16 19:27:35 +01:00
parent f0a9242d10
commit 0443906cc9
1 changed files with 43 additions and 17 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

60
src/utils/qdf/blacklist.rs
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),
@ -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!(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())
);
}
}
}