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 (first-seen order; O(1) per row, borrowed keys).
18        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        // Try to preserve original value types
31        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        // Expose the count under its ggplot stat name for after_stat expressions.
53        result.add_column("count".to_string(), y_values);
54
55        // Carry over group columns
56        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}