volas_time/
agg.rs

1//! Typed aggregators for cumulation (replaces stock-pandas's Python callables).
2
3use std::collections::HashMap;
4
5use volas_core::{Column, CombineOp, Result, VolasError};
6
7/// A column aggregator over a period's rows.
8#[derive(Clone, Copy, Debug, PartialEq, Eq)]
9pub enum Agg {
10    /// First value.
11    First,
12    /// Maximum (numeric / datetime).
13    Max,
14    /// Minimum (numeric / datetime).
15    Min,
16    /// Last value.
17    Last,
18    /// Sum (numeric).
19    Sum,
20}
21
22impl Agg {
23    /// Parse an aggregator name (`first` / `max` / `min` / `last` / `sum`).
24    pub fn from_name(s: &str) -> Result<Agg> {
25        Ok(match s {
26            "first" => Agg::First,
27            "max" => Agg::Max,
28            "min" => Agg::Min,
29            "last" => Agg::Last,
30            "sum" => Agg::Sum,
31            _ => return Err(VolasError::Value(format!("unknown aggregator '{s}'"))),
32        })
33    }
34
35    fn name(&self) -> &'static str {
36        match self {
37            Agg::First => "first",
38            Agg::Max => "max",
39            Agg::Min => "min",
40            Agg::Last => "last",
41            Agg::Sum => "sum",
42        }
43    }
44
45    /// The in-place single-cell combine op equivalent to this aggregator — the
46    /// incremental dual of [`Self::reduce`], used by the live tf-fold to update the
47    /// forming bar without re-reducing the whole period.
48    pub fn as_combine_op(&self) -> CombineOp {
49        match self {
50            Agg::First => CombineOp::Keep,
51            Agg::Last => CombineOp::Replace,
52            Agg::Max => CombineOp::Max,
53            Agg::Min => CombineOp::Min,
54            Agg::Sum => CombineOp::Sum,
55        }
56    }
57
58    /// Reduce the values at `kept` positions of `col` to a 1-element column of
59    /// the appropriate dtype.
60    pub fn reduce(&self, col: &Column, kept: &[usize]) -> Result<Column> {
61        match self {
62            Agg::First => Ok(col.take(&[kept[0]])),
63            Agg::Last => Ok(col.take(&[*kept.last().unwrap()])),
64            Agg::Max | Agg::Min | Agg::Sum => self.reduce_numeric(col, kept),
65        }
66    }
67
68    fn reduce_numeric(&self, col: &Column, kept: &[usize]) -> Result<Column> {
69        match col {
70            Column::F64(v) => {
71                let it = kept.iter().map(|&i| v[i]);
72                let val = match self {
73                    Agg::Max => it.fold(f64::NEG_INFINITY, f64::max),
74                    Agg::Min => it.fold(f64::INFINITY, f64::min),
75                    Agg::Sum => it.sum(),
76                    _ => unreachable!(), // LCOV_EXCL_LINE
77                };
78                Ok(Column::f64(vec![val]))
79            }
80            Column::I64(v, _) | Column::Datetime(v) => {
81                let it = kept.iter().map(|&i| v[i]);
82                let val = match self {
83                    Agg::Max => it.max().unwrap(),
84                    Agg::Min => it.min().unwrap(),
85                    Agg::Sum => it.sum(),
86                    _ => unreachable!(), // LCOV_EXCL_LINE
87                };
88                Ok(if matches!(col, Column::Datetime(_)) {
89                    Column::datetime(vec![val])
90                } else {
91                    Column::i64(vec![val])
92                })
93            }
94            other => Err(VolasError::DType(format!(
95                "cannot {} a {} column",
96                self.name(),
97                other.dtype()
98            ))),
99        }
100    }
101}
102
103/// Per-column aggregator assignment, with a default for unlisted columns.
104#[derive(Clone, Debug)]
105pub struct AggSpec {
106    default: Agg,
107    by_name: HashMap<String, Agg>,
108}
109
110impl AggSpec {
111    /// The OHLCV defaults: open=first, high=max, low=min, close=last, volume=sum;
112    /// every other column falls back to `last`.
113    pub fn ohlcv() -> Self {
114        let mut by_name = HashMap::new();
115        for (name, agg) in [
116            ("open", Agg::First),
117            ("high", Agg::Max),
118            ("low", Agg::Min),
119            ("close", Agg::Last),
120            ("volume", Agg::Sum),
121        ] {
122            by_name.insert(name.to_string(), agg);
123        }
124        AggSpec {
125            default: Agg::Last,
126            by_name,
127        }
128    }
129
130    /// Override (or add) a column's aggregator.
131    pub fn set(&mut self, name: impl Into<String>, agg: Agg) {
132        self.by_name.insert(name.into(), agg);
133    }
134
135    /// The aggregator for a column (the default if unset).
136    pub fn agg_for(&self, name: &str) -> Agg {
137        self.by_name.get(name).copied().unwrap_or(self.default)
138    }
139}
140
141impl Default for AggSpec {
142    fn default() -> Self {
143        AggSpec::ohlcv()
144    }
145}
146
147#[cfg(test)]
148mod tests {
149    use super::*;
150
151    #[test]
152    fn spec_defaults_and_overrides() {
153        let mut spec = AggSpec::ohlcv();
154        assert_eq!(spec.agg_for("open"), Agg::First);
155        assert_eq!(spec.agg_for("volume"), Agg::Sum);
156        assert_eq!(spec.agg_for("anything"), Agg::Last);
157        spec.set("volume", Agg::Last);
158        assert_eq!(spec.agg_for("volume"), Agg::Last);
159    }
160
161    #[test]
162    fn reduce_numeric_and_first_last() {
163        let col = Column::f64(vec![3.0, 1.0, 4.0, 1.0]);
164        let kept = [0, 1, 2, 3];
165        assert_eq!(
166            Agg::First.reduce(&col, &kept).unwrap(),
167            Column::f64(vec![3.0])
168        );
169        assert_eq!(
170            Agg::Last.reduce(&col, &kept).unwrap(),
171            Column::f64(vec![1.0])
172        );
173        assert_eq!(
174            Agg::Max.reduce(&col, &kept).unwrap(),
175            Column::f64(vec![4.0])
176        );
177        assert_eq!(
178            Agg::Min.reduce(&col, &kept).unwrap(),
179            Column::f64(vec![1.0])
180        );
181        assert_eq!(
182            Agg::Sum.reduce(&col, &kept).unwrap(),
183            Column::f64(vec![9.0])
184        );
185    }
186
187    #[test]
188    fn as_combine_op_maps_each_aggregator() {
189        assert_eq!(Agg::First.as_combine_op(), CombineOp::Keep);
190        assert_eq!(Agg::Last.as_combine_op(), CombineOp::Replace);
191        assert_eq!(Agg::Max.as_combine_op(), CombineOp::Max);
192        assert_eq!(Agg::Min.as_combine_op(), CombineOp::Min);
193        assert_eq!(Agg::Sum.as_combine_op(), CombineOp::Sum);
194    }
195
196    #[test]
197    fn sum_on_string_errors() {
198        let col = Column::str(vec!["a".into(), "b".into()]);
199        assert!(Agg::Sum.reduce(&col, &[0, 1]).is_err());
200        // first/last still work on strings
201        assert_eq!(
202            Agg::First.reduce(&col, &[0, 1]).unwrap(),
203            Column::str(vec!["a".into()])
204        );
205    }
206
207    #[test]
208    fn from_name_parses_all_and_rejects_unknown() {
209        assert_eq!(Agg::from_name("first").unwrap(), Agg::First);
210        assert_eq!(Agg::from_name("max").unwrap(), Agg::Max);
211        assert_eq!(Agg::from_name("min").unwrap(), Agg::Min);
212        assert_eq!(Agg::from_name("last").unwrap(), Agg::Last);
213        assert_eq!(Agg::from_name("sum").unwrap(), Agg::Sum);
214        assert!(Agg::from_name("median").is_err());
215    }
216
217    #[test]
218    fn reduce_over_i64_and_datetime_columns() {
219        let ints = Column::i64(vec![3, 1, 4, 1]);
220        assert_eq!(
221            Agg::Max.reduce(&ints, &[0, 1, 2, 3]).unwrap(),
222            Column::i64(vec![4])
223        );
224        assert_eq!(
225            Agg::Min.reduce(&ints, &[0, 1, 2, 3]).unwrap(),
226            Column::i64(vec![1])
227        );
228        assert_eq!(
229            Agg::Sum.reduce(&ints, &[0, 1, 2, 3]).unwrap(),
230            Column::i64(vec![9])
231        );
232
233        // Datetime max/min stay in the datetime dtype.
234        let ts = Column::datetime(vec![100, 500, 200]);
235        assert_eq!(
236            Agg::Max.reduce(&ts, &[0, 1, 2]).unwrap(),
237            Column::datetime(vec![500])
238        );
239        assert_eq!(
240            Agg::Min.reduce(&ts, &[0, 1, 2]).unwrap(),
241            Column::datetime(vec![100])
242        );
243
244        // Sum on a bool column is a dtype error (hits the fallback arm).
245        let flags = Column::bool(vec![true, false]);
246        assert!(Agg::Sum.reduce(&flags, &[0, 1]).is_err());
247    }
248
249    #[test]
250    fn default_spec_and_max_min_dtype_errors() {
251        assert_eq!(AggSpec::default().agg_for("open"), Agg::First);
252        // Max / Min on a non-numeric column hit the dtype-error path (exercises name()).
253        let flags = Column::bool(vec![true, false]);
254        assert!(Agg::Max.reduce(&flags, &[0, 1]).is_err());
255        assert!(Agg::Min.reduce(&flags, &[0, 1]).is_err());
256    }
257
258    #[test]
259    fn agg_name_first_last() {
260        // first / last names (the other variants are exercised via reduce errors).
261        assert_eq!(Agg::First.name(), "first");
262        assert_eq!(Agg::Last.name(), "last");
263    }
264}
volas_time/agg.rs

volas_time/
agg.rs
