1use crate::aes::Aesthetic;
2use crate::data::{DataFrame, Value};
3use crate::scale::ScaleSet;
4
5use super::Stat;
6
7pub struct StatCount;
9
10impl Stat for StatCount {
11 fn compute_group(&self, data: &DataFrame, _scales: &ScaleSet) -> DataFrame {
12 let x_col = match data.column("x") {
13 Some(c) => c,
14 None => return DataFrame::new(),
15 };
16
17 let mut counts: Vec<(String, usize)> = Vec::new();
19 for v in x_col {
20 let key = v.to_group_key();
21 if let Some(entry) = counts.iter_mut().find(|(k, _)| k == &key) {
22 entry.1 += 1;
23 } else {
24 counts.push((key, 1));
25 }
26 }
27
28 let mut result = DataFrame::new();
29 let x_values: Vec<Value> = counts.iter().map(|(k, _)| Value::Str(k.clone())).collect();
30
31 let first_x = x_col.first();
33 let x_values: Vec<Value> = if matches!(first_x, Some(Value::Float(_) | Value::Integer(_))) {
34 counts
35 .iter()
36 .map(|(k, _)| {
37 k.parse::<f64>()
38 .map(Value::Float)
39 .unwrap_or_else(|_| Value::Str(k.clone()))
40 })
41 .collect()
42 } else {
43 x_values
44 };
45
46 let y_values: Vec<Value> = counts
47 .iter()
48 .map(|(_, c)| Value::Float(*c as f64))
49 .collect();
50
51 result.add_column("x".to_string(), x_values);
52 result.add_column("y".to_string(), y_values);
53
54 if data.has_column("fill") {
56 if let Some(fill_col) = data.column("fill") {
57 if let Some(first) = fill_col.first() {
58 result.add_column("fill".to_string(), vec![first.clone(); counts.len()]);
59 }
60 }
61 }
62 if data.has_column("color") {
63 if let Some(color_col) = data.column("color") {
64 if let Some(first) = color_col.first() {
65 result.add_column("color".to_string(), vec![first.clone(); counts.len()]);
66 }
67 }
68 }
69
70 result
71 }
72
73 fn required_aes(&self) -> Vec<Aesthetic> {
74 vec![Aesthetic::X]
75 }
76
77 fn name(&self) -> &str {
78 "count"
79 }
80}