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);
53
54        // Carry over group columns
55        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}