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: indexmap::IndexMap<std::borrow::Cow<'_, str>, usize> =
19 indexmap::IndexMap::new();
20 for v in x_col {
21 *counts.entry(v.key_str()).or_insert(0) += 1;
22 }
23
24 let mut result = DataFrame::new();
25 let x_values: Vec<Value> = counts
26 .iter()
27 .map(|(k, _)| Value::Str(k.to_string()))
28 .collect();
29
30 let first_x = x_col.first();
32 let x_values: Vec<Value> = if matches!(first_x, Some(Value::Float(_) | Value::Integer(_))) {
33 counts
34 .iter()
35 .map(|(k, _)| {
36 k.parse::<f64>()
37 .map(Value::Float)
38 .unwrap_or_else(|_| Value::Str(k.to_string()))
39 })
40 .collect()
41 } else {
42 x_values
43 };
44
45 let y_values: Vec<Value> = counts
46 .iter()
47 .map(|(_, c)| Value::Float(*c as f64))
48 .collect();
49
50 result.add_column("x".to_string(), x_values);
51 result.add_column("y".to_string(), y_values.clone());
52 result.add_column("count".to_string(), y_values);
54
55 if data.has_column("fill") {
57 if let Some(fill_col) = data.column("fill") {
58 if let Some(first) = fill_col.first() {
59 result.add_column("fill".to_string(), vec![first.clone(); counts.len()]);
60 }
61 }
62 }
63 if data.has_column("color") {
64 if let Some(color_col) = data.column("color") {
65 if let Some(first) = color_col.first() {
66 result.add_column("color".to_string(), vec![first.clone(); counts.len()]);
67 }
68 }
69 }
70
71 result
72 }
73
74 fn required_aes(&self) -> Vec<Aesthetic> {
75 vec![Aesthetic::X]
76 }
77
78 fn name(&self) -> &str {
79 "count"
80 }
81}