use std::sync::Arc;
use crate::dataframe::groupby::{AggFunc, CustomAggFn};
use crate::dataframe::hierarchical_groupby::HierarchicalAgg;
pub struct HierarchicalAggBuilder {
column: String,
level_functions: Vec<(usize, AggFunc, String)>,
custom_fn: Option<CustomAggFn>,
propagate_up: bool,
include_intermediate: bool,
}
impl HierarchicalAggBuilder {
pub fn new(column: String) -> Self {
Self {
column,
level_functions: Vec::new(),
custom_fn: None,
propagate_up: false,
include_intermediate: true,
}
}
pub fn at_level(mut self, level: usize, func: AggFunc, alias: String) -> Self {
self.level_functions.push((level, func, alias));
self
}
pub fn at_all_levels(mut self, func: AggFunc, base_alias: String) -> Self {
self.level_functions.push((0, func, base_alias));
self
}
pub fn with_propagation(mut self) -> Self {
self.propagate_up = true;
self
}
pub fn with_custom<F>(mut self, func: F) -> Self
where
F: Fn(&[f64]) -> f64 + Send + Sync + 'static,
{
self.custom_fn = Some(Arc::new(func));
self
}
pub fn build(self) -> HierarchicalAgg {
HierarchicalAgg {
column: self.column,
level_functions: self.level_functions,
custom_fn: self.custom_fn,
propagate_up: self.propagate_up,
include_intermediate: self.include_intermediate,
}
}
}
pub mod utils {
use super::*;
pub fn simple_hierarchical_agg(column: &str, func: AggFunc, level: usize) -> HierarchicalAgg {
HierarchicalAggBuilder::new(column.to_string())
.at_level(level, func, format!("{}_{}", func.as_str(), column))
.build()
}
pub fn all_levels_agg(column: &str, func: AggFunc, max_levels: usize) -> Vec<HierarchicalAgg> {
(0..max_levels)
.map(|level| simple_hierarchical_agg(column, func, level))
.collect()
}
pub fn comprehensive_agg(column: &str, level: usize) -> Vec<HierarchicalAgg> {
vec![
simple_hierarchical_agg(column, AggFunc::Sum, level),
simple_hierarchical_agg(column, AggFunc::Mean, level),
simple_hierarchical_agg(column, AggFunc::Count, level),
]
}
}