use std::collections::HashMap;
use crate::{CanonicalColumnName, SourceColumnName, request::AggregationType};
#[derive(Debug, Clone)]
pub enum CompositionStrategy {
Direct { source: String },
Stack {
sources: Vec<String>,
reductions: Vec<ComponentReduction>,
},
Incompatible { unit: String, reason: String },
}
impl CompositionStrategy {
pub fn is_direct(&self) -> bool {
matches!(self, Self::Direct { .. })
}
pub fn is_stack(&self) -> bool {
matches!(self, Self::Stack { .. })
}
pub fn is_incompatible(&self) -> bool {
matches!(self, Self::Incompatible { .. })
}
pub fn source_names(&self) -> Vec<&str> {
match self {
Self::Direct { source } => vec![source.as_str()],
Self::Stack { sources, .. } => sources.iter().map(|s| s.as_str()).collect(),
Self::Incompatible { .. } => vec![],
}
}
}
#[derive(Debug, Clone)]
pub struct ComponentReduction {
pub source: String,
pub reduce_components: Vec<SourceColumnName>,
pub aggregation: AggregationType,
}
#[derive(Debug, Clone)]
pub struct CompositionPlan {
pub unit_strategies: HashMap<CanonicalColumnName, CompositionStrategy>,
pub join_units: Vec<CanonicalColumnName>,
}
impl CompositionPlan {
pub fn new() -> Self {
Self {
unit_strategies: HashMap::new(),
join_units: Vec::new(),
}
}
pub fn is_simple(&self) -> bool {
self.unit_strategies.len() == 1 && self.unit_strategies.values().all(|s| s.is_direct())
}
pub fn requires_stacking(&self) -> bool {
self.unit_strategies.values().any(|s| s.is_stack())
}
pub fn requires_joining(&self) -> bool {
self.join_units.len() > 1
}
pub fn requires_reduction(&self) -> bool {
self.unit_strategies.values().any(
|s| matches!(s, CompositionStrategy::Stack { reductions, .. } if !reductions.is_empty()),
)
}
pub fn has_errors(&self) -> bool {
self.unit_strategies.values().any(|s| s.is_incompatible())
}
pub fn errors(&self) -> Vec<&str> {
self.unit_strategies
.values()
.filter_map(|s| {
if let CompositionStrategy::Incompatible { reason, .. } = s {
Some(reason.as_str())
} else {
None
}
})
.collect()
}
pub fn all_sources(&self) -> Vec<&str> {
let mut sources: Vec<&str> = self
.unit_strategies
.values()
.flat_map(|s| s.source_names())
.collect();
sources.sort();
sources.dedup();
sources
}
pub fn get_strategy(&self, unit_name: &CanonicalColumnName) -> Option<&CompositionStrategy> {
self.unit_strategies.get(unit_name)
}
}
impl Default for CompositionPlan {
fn default() -> Self {
Self::new()
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_composition_strategy_direct() {
let strategy = CompositionStrategy::Direct {
source: "source_a".into(),
};
assert!(strategy.is_direct());
assert!(!strategy.is_stack());
assert_eq!(strategy.source_names(), vec!["source_a"]);
}
#[test]
fn test_composition_strategy_stack() {
let strategy = CompositionStrategy::Stack {
sources: vec!["source_a".into(), "source_b".into()],
reductions: vec![],
};
assert!(strategy.is_stack());
assert_eq!(strategy.source_names().len(), 2);
}
#[test]
fn test_composition_plan_simple() {
let mut plan = CompositionPlan::new();
plan.unit_strategies.insert(
"value_a".into(),
CompositionStrategy::Direct {
source: "source".into(),
},
);
plan.join_units.push("value_a".into());
assert!(plan.is_simple());
assert!(!plan.requires_stacking());
assert!(!plan.requires_joining());
}
#[test]
fn test_composition_plan_with_join() {
let mut plan = CompositionPlan::new();
plan.unit_strategies.insert(
"value_a".into(),
CompositionStrategy::Direct {
source: "source_a".into(),
},
);
plan.unit_strategies.insert(
"value_b".into(),
CompositionStrategy::Direct {
source: "source_b".into(),
},
);
plan.join_units = vec!["value_a".into(), "value_b".into()];
assert!(!plan.is_simple());
assert!(plan.requires_joining());
assert!(!plan.requires_stacking());
}
#[test]
fn test_composition_plan_with_stack() {
let mut plan = CompositionPlan::new();
plan.unit_strategies.insert(
"value_a".into(),
CompositionStrategy::Stack {
sources: vec!["source_a".into(), "source_b".into()],
reductions: vec![],
},
);
plan.join_units.push("value_a".into());
assert!(plan.requires_stacking());
assert!(!plan.requires_joining());
}
#[test]
fn test_composition_plan_errors() {
let mut plan = CompositionPlan::new();
plan.unit_strategies.insert(
"value_a".into(),
CompositionStrategy::Incompatible {
unit: "value_a".into(),
reason: "No source provides this unit".into(),
},
);
assert!(plan.has_errors());
assert_eq!(plan.errors().len(), 1);
}
}