use serde_json::json;
use crate::{
Cluster, ClusterId, Diagnostics, FaceError, Record, Strategy,
cluster_tree::{
AxisPlan, build_tree_impl, extend_id,
util::{resolve_numeric, score_range},
},
};
pub(super) fn cluster<D: Diagnostics + ?Sized>(
plan: &AxisPlan,
items: Vec<Record>,
parent: &ClusterId,
diag: &mut D,
) -> Result<Vec<Cluster>, FaceError> {
let Strategy::Quantiles { count } = &plan.axis.strategy else {
unreachable!("quantiles clusterer is only called for Strategy::Quantiles")
};
let count = (*count).max(1) as usize;
let mut resolved = resolve_numeric(items, &plan.axis.field, diag);
if resolved.is_empty() {
return Ok(Vec::new());
}
resolved.sort_by(|a, b| a.0.total_cmp(&b.0));
let n = resolved.len();
if n == 1 || count == 1 {
let lo = resolved[0].0;
let hi = resolved[n - 1].0;
let group: Vec<Record> = resolved.into_iter().map(|(_, r)| r).collect();
return Ok(vec![build_cluster(
plan,
parent,
format_label(lo, hi, &[lo, hi]),
lo,
hi,
group,
diag,
)?]);
}
let bucket_starts: Vec<usize> = (0..=count).map(|i| i * n / count).collect();
let mut buckets: Vec<Vec<(f64, Record)>> = (0..count).map(|_| Vec::new()).collect();
let mut bucket_idx = 0usize;
for (i, item) in resolved.into_iter().enumerate() {
while bucket_idx + 1 < count && i >= bucket_starts[bucket_idx + 1] {
bucket_idx += 1;
}
buckets[bucket_idx].push(item);
}
let mut merged: Vec<Vec<(f64, Record)>> = Vec::with_capacity(count);
for bucket in buckets {
if bucket.is_empty() {
continue;
}
let lo = bucket.first().expect("non-empty checked").0;
let hi = bucket.last().expect("non-empty checked").0;
if let Some(prev) = merged.last_mut() {
let prev_lo = prev.first().expect("non-empty by construction").0;
let prev_hi = prev.last().expect("non-empty by construction").0;
if prev_lo == prev_hi && lo == hi && prev_lo == lo {
prev.extend(bucket);
continue;
}
}
merged.push(bucket);
}
let mut out = Vec::with_capacity(merged.len());
for bucket in merged {
let lo = bucket.first().expect("non-empty by construction").0;
let hi = bucket.last().expect("non-empty by construction").0;
let label = format_label(lo, hi, &[lo, hi]);
let group: Vec<Record> = bucket.into_iter().map(|(_, r)| r).collect();
out.push(build_cluster(plan, parent, label, lo, hi, group, diag)?);
}
out.reverse();
Ok(out)
}
fn build_cluster<D: Diagnostics + ?Sized>(
plan: &AxisPlan,
parent: &ClusterId,
label: String,
lo: f64,
hi: f64,
group: Vec<Record>,
diag: &mut D,
) -> Result<Cluster, FaceError> {
let id = extend_id(parent, &plan.axis.field, &label);
let total = group.len() as u64;
let (score_min, score_max) = score_range(&group);
let children = if plan.within.is_empty() {
Vec::new()
} else {
let mut acc: Vec<Cluster> = Vec::new();
for child_plan in &plan.within {
let child_clusters = build_tree_impl(child_plan, group.clone(), &id, diag)?;
acc.extend(child_clusters);
}
acc
};
Ok(Cluster {
id,
label,
axis: plan.axis.field.clone(),
value: json!({"min": lo, "max": hi}),
total,
score_min,
score_max,
clusters: children,
})
}
fn format_label(lo: f64, hi: f64, bounds: &[f64]) -> String {
if bounds.iter().all(|b| b.fract() == 0.0) {
format!("{}\u{2013}{}", lo as i64, hi as i64)
} else {
format!("{lo:.2}\u{2013}{hi:.2}")
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::{Axis, NullDiagnostics, Strategy};
use serde_json::json;
fn axis(field: &str, count: u8) -> Axis {
Axis {
field: field.into(),
strategy: Strategy::Quantiles { count },
auto: false,
}
}
#[test]
fn equal_population_partitioning() {
let plan = AxisPlan::leaf(axis("v", 4));
let items: Vec<serde_json::Value> = (0..8).map(|i| json!({"v": i as f64})).collect();
let mut diag = NullDiagnostics;
let clusters = super::cluster(
&plan,
Record::from_items(items, None),
&ClusterId::default(),
&mut diag,
)
.unwrap();
assert_eq!(clusters.len(), 4);
for c in &clusters {
assert_eq!(c.total, 2, "expected equal-population buckets, got {c:?}");
}
assert_eq!(clusters[0].label, "6\u{2013}7");
assert_eq!(clusters[3].label, "0\u{2013}1");
}
#[test]
fn handles_duplicates() {
let plan = AxisPlan::leaf(axis("v", 4));
let items = vec![
json!({"v": 1.0}),
json!({"v": 1.0}),
json!({"v": 1.0}),
json!({"v": 1.0}),
json!({"v": 5.0}),
json!({"v": 5.0}),
json!({"v": 5.0}),
json!({"v": 5.0}),
];
let mut diag = NullDiagnostics;
let clusters = super::cluster(
&plan,
Record::from_items(items, None),
&ClusterId::default(),
&mut diag,
)
.unwrap();
assert_eq!(clusters.len(), 2);
let total: u64 = clusters.iter().map(|c| c.total).sum();
assert_eq!(total, 8);
assert_eq!(clusters[0].value, json!({"min": 5.0, "max": 5.0}));
assert_eq!(clusters[1].value, json!({"min": 1.0, "max": 1.0}));
}
#[test]
fn fewer_records_than_buckets_drops_empties() {
let plan = AxisPlan::leaf(axis("v", 5));
let items = vec![json!({"v": 0.0}), json!({"v": 1.0}), json!({"v": 2.0})];
let mut diag = NullDiagnostics;
let clusters = super::cluster(
&plan,
Record::from_items(items, None),
&ClusterId::default(),
&mut diag,
)
.unwrap();
let total: u64 = clusters.iter().map(|c| c.total).sum();
assert_eq!(total, 3);
assert!(
clusters.iter().all(|c| c.total >= 1),
"no empty buckets should be emitted"
);
}
#[test]
fn empty_records_returns_empty() {
let plan = AxisPlan::leaf(axis("v", 4));
let items: Vec<serde_json::Value> = vec![]; let mut diag = NullDiagnostics;
let clusters = super::cluster(
&plan,
Record::from_items(items, None),
&ClusterId::default(),
&mut diag,
)
.unwrap();
assert!(clusters.is_empty());
}
}