Skip to main content

ggplot_rs/stat/
count.rs

1use crate::aes::Aesthetic;
2use crate::data::{DataFrame, Value};
3use crate::scale::ScaleSet;
4
5use super::Stat;
6
7/// Counts occurrences of each unique x value.
8pub 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        // Count unique x values
18        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        // Try to preserve original value types
32        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.clone());
53        // Expose the count under its ggplot stat name for after_stat expressions.
54        result.add_column("count".to_string(), y_values);
55
56        // Carry over group columns
57        if data.has_column("fill") {
58            if let Some(fill_col) = data.column("fill") {
59                if let Some(first) = fill_col.first() {
60                    result.add_column("fill".to_string(), vec![first.clone(); counts.len()]);
61                }
62            }
63        }
64        if data.has_column("color") {
65            if let Some(color_col) = data.column("color") {
66                if let Some(first) = color_col.first() {
67                    result.add_column("color".to_string(), vec![first.clone(); counts.len()]);
68                }
69            }
70        }
71
72        result
73    }
74
75    fn required_aes(&self) -> Vec<Aesthetic> {
76        vec![Aesthetic::X]
77    }
78
79    fn name(&self) -> &str {
80        "count"
81    }
82}