groupby/
lib.rs

1// Copyright (C) The Strand-Braid Authors
2// SPDX-License-Identifier: MIT OR Apache-2.0
3
4use std::collections::{BTreeMap, VecDeque};
5use strand_withkey::WithKey;
6
7// TODO: better error handling. Do not wrap Result types, but handle Results automatically.
8
9/// A type which reads ahead by `bufsize` elements and sorts within.
10pub struct BufferedSortIter<K, I, T, E>
11where
12    I: Iterator<Item = std::result::Result<T, E>>,
13    T: WithKey<K>,
14{
15    single_iter: I,
16    sorted_buf: BTreeMap<K, VecDeque<T>>,
17    done_reading: bool,
18    highest_key: Option<K>,
19}
20
21impl<K: std::cmp::Ord + Clone, I, T, E: std::fmt::Debug> BufferedSortIter<K, I, T, E>
22where
23    I: Iterator<Item = std::result::Result<T, E>>,
24    T: WithKey<K>,
25{
26    pub fn new(single_iter: I, bufsize: usize) -> Result<Self, E> {
27        let sorted_buf = BTreeMap::new();
28        let mut result = Self {
29            single_iter,
30            sorted_buf,
31            done_reading: false,
32            highest_key: None,
33        };
34
35        let mut count = 0;
36        while count < bufsize {
37            count += 1;
38            if !result.read_next()? {
39                break;
40            }
41        }
42        Ok(result)
43    }
44
45    pub fn get_ref(&self) -> &I {
46        &self.single_iter
47    }
48
49    pub fn into_inner(self) -> I {
50        self.single_iter
51    }
52
53    // read a single item from the iterator. returns true if should read again.
54    #[inline]
55    fn read_next(&mut self) -> Result<bool, E> {
56        if self.done_reading {
57            return Ok(false);
58        }
59        match self.single_iter.next() {
60            None => {
61                self.done_reading = true;
62                Ok(false)
63            }
64            Some(result_el) => {
65                let el = result_el?;
66                let key = el.key();
67                let rows_entry = &mut self.sorted_buf.entry(key).or_default();
68                rows_entry.push_back(el);
69                Ok(true)
70            }
71        }
72    }
73
74    // get the first item, return whether the key must be removed and the item
75    #[inline]
76    fn empty_first(&mut self) -> Option<std::result::Result<T, E>> {
77        let mut remove_key = None;
78        let result: T = {
79            let mut first = self.sorted_buf.iter_mut().next();
80            match first {
81                None => return None,
82                Some((this_key, ref mut this_el_vec)) => {
83                    if let Some(ref hk) = self.highest_key {
84                        assert!(this_key >= hk, "failed to sort data (bufsize too small?)");
85                    }
86                    self.highest_key = Some(this_key.clone());
87
88                    let first_el = this_el_vec.pop_front();
89                    if this_el_vec.is_empty() {
90                        remove_key = Some(this_key.clone());
91                    }
92                    first_el.unwrap()
93                }
94            }
95        };
96
97        if let Some(rk) = remove_key {
98            self.sorted_buf.remove(&rk);
99        }
100
101        Some(Ok(result))
102    }
103}
104
105impl<
106    K: std::cmp::Ord + std::fmt::Debug + std::cmp::PartialEq + std::cmp::PartialOrd + Clone,
107    I,
108    T,
109    E: std::fmt::Debug,
110> Iterator for BufferedSortIter<K, I, T, E>
111where
112    I: Iterator<Item = std::result::Result<T, E>>,
113    T: WithKey<K>,
114{
115    type Item = std::result::Result<T, E>;
116    fn next(&mut self) -> std::option::Option<<Self as Iterator>::Item> {
117        match self.read_next() {
118            Ok(_) => {}
119            Err(e) => {
120                // The error may be desynchronized somewhat from element in iterator
121                // from which it came, but at least it gets
122                // propagated up without panic.
123                return Some(Err(e));
124            }
125        };
126        self.empty_first()
127    }
128}
129
130/// A type takes a sorted iterator and returns grouped results.
131///
132/// Panics if the input iterator is not sorted.
133pub struct AscendingGroupIter<K, I, T, E>
134where
135    I: Iterator<Item = std::result::Result<T, E>>,
136    T: WithKey<K>,
137{
138    single_iter: I,
139    peek: Option<std::result::Result<T, E>>,
140    key_type: std::marker::PhantomData<K>,
141}
142
143impl<K, I, T, E> AscendingGroupIter<K, I, T, E>
144where
145    I: Iterator<Item = std::result::Result<T, E>>,
146    T: WithKey<K>,
147{
148    pub fn new(mut single_iter: I) -> Self {
149        let peek = single_iter.next();
150        Self {
151            single_iter,
152            peek,
153            key_type: std::marker::PhantomData,
154        }
155    }
156
157    pub fn get_ref(&self) -> &I {
158        &self.single_iter
159    }
160
161    pub fn into_inner(self) -> I {
162        self.single_iter
163    }
164}
165
166#[derive(Debug, PartialEq, Eq)]
167pub struct GroupedRows<K, T: WithKey<K>> {
168    pub group_key: K,
169    pub rows: Vec<T>,
170}
171
172impl<K: std::fmt::Debug + std::cmp::PartialEq + std::cmp::PartialOrd, I, T: WithKey<K>, E> Iterator
173    for AscendingGroupIter<K, I, T, E>
174where
175    I: Iterator<Item = std::result::Result<T, E>>,
176    T: WithKey<K>,
177{
178    type Item = std::result::Result<GroupedRows<K, T>, E>;
179    fn next(&mut self) -> std::option::Option<<Self as Iterator>::Item> {
180        match self.peek.take() {
181            None => None, // no more rows, return None
182            Some(row_result) => {
183                match row_result {
184                    Ok(next_seed) => {
185                        // first row of new group_key
186                        let mut item = GroupedRows {
187                            group_key: next_seed.key(),
188                            rows: vec![next_seed],
189                        };
190                        loop {
191                            // more rows of new group_key until seed for next
192                            match self.single_iter.next() {
193                                None => break, // no more rows in file, finish this group_key
194                                Some(result_val) => {
195                                    match result_val {
196                                        Ok(val) => {
197                                            // got another row
198                                            if val.key() == item.group_key {
199                                                item.rows.push(val);
200                                            } else {
201                                                if val.key().partial_cmp(&item.group_key)
202                                                    != Some(std::cmp::Ordering::Greater)
203                                                {
204                                                    panic!(
205                                                        "key is not monotonically ascending ({:?} < {:?})",
206                                                        val.key(),
207                                                        item.group_key
208                                                    );
209                                                }
210                                                self.peek = Some(Ok(val));
211                                                break;
212                                            }
213                                        }
214                                        Err(e) => return Some(Err(e)),
215                                    }
216                                }
217                            }
218                        }
219                        Some(Ok(item))
220                    }
221                    Err(e) => Some(Err(e)),
222                }
223            }
224        }
225    }
226}
227
228#[cfg(test)]
229mod tests {
230    use super::*;
231
232    #[derive(PartialEq)]
233    struct Foo {
234        x: u8,
235    }
236    impl WithKey<u8> for Foo {
237        fn key(&self) -> u8 {
238            self.x
239        }
240    }
241
242    #[test]
243    fn groupby_empty() {
244        let foos: Vec<Foo> = vec![];
245        let foos_iter = foos.into_iter().map(|x| {
246            let r: Result<Foo, u8> = Ok(x);
247            r
248        });
249        let mut group_iter = AscendingGroupIter::new(foos_iter);
250        assert!(group_iter.next().is_none());
251    }
252
253    #[test]
254    fn groupby_monotonic() {
255        let foos = vec![
256            Foo { x: 1 },
257            Foo { x: 1 },
258            Foo { x: 2 },
259            Foo { x: 2 },
260            Foo { x: 3 },
261        ];
262        let foos_iter = foos.into_iter().map(|x| {
263            let r: Result<Foo, u8> = Ok(x);
264            r
265        });
266        let mut group_iter = AscendingGroupIter::new(foos_iter);
267        assert!(
268            group_iter.next()
269                == Some(Ok(GroupedRows {
270                    group_key: 1,
271                    rows: vec![Foo { x: 1 }, Foo { x: 1 }]
272                }))
273        );
274        assert!(
275            group_iter.next()
276                == Some(Ok(GroupedRows {
277                    group_key: 2,
278                    rows: vec![Foo { x: 2 }, Foo { x: 2 }]
279                }))
280        );
281        assert!(
282            group_iter.next()
283                == Some(Ok(GroupedRows {
284                    group_key: 3,
285                    rows: vec![Foo { x: 3 }]
286                }))
287        );
288        assert!(group_iter.next().is_none());
289    }
290
291    #[test]
292    #[should_panic]
293    fn groupby_nonmonotonic() {
294        let foos = vec![Foo { x: 1 }, Foo { x: 3 }, Foo { x: 2 }];
295        let foos_iter = foos.into_iter().map(|x| {
296            let r: Result<Foo, u8> = Ok(x);
297            r
298        });
299        let mut group_iter = AscendingGroupIter::new(foos_iter);
300        assert!(
301            group_iter.next()
302                == Some(Ok(GroupedRows {
303                    group_key: 1,
304                    rows: vec![Foo { x: 1 }, Foo { x: 1 }]
305                }))
306        );
307        assert!(
308            group_iter.next()
309                == Some(Ok(GroupedRows {
310                    group_key: 2,
311                    rows: vec![Foo { x: 2 }, Foo { x: 2 }]
312                }))
313        );
314        assert!(
315            group_iter.next()
316                == Some(Ok(GroupedRows {
317                    group_key: 3,
318                    rows: vec![Foo { x: 3 }]
319                }))
320        );
321        assert!(group_iter.next().is_none());
322    }
323
324    #[test]
325    fn buffered_sort_empty() {
326        let foos = vec![];
327        let foos_iter = foos.into_iter().map(|x| {
328            let r: Result<Foo, u8> = Ok(x);
329            r
330        });
331        let mut sorted_iter = BufferedSortIter::new(foos_iter, 100).unwrap();
332        assert!(sorted_iter.next().is_none());
333    }
334
335    #[test]
336    fn buffered_sort_results() {
337        let foos = vec![
338            Foo { x: 1 },
339            Foo { x: 3 },
340            Foo { x: 1 },
341            Foo { x: 4 },
342            Foo { x: 3 },
343            Foo { x: 3 },
344            Foo { x: 3 },
345            Foo { x: 3 },
346            Foo { x: 1 },
347        ];
348        let foos_iter = foos.into_iter().map(|x| {
349            let r: Result<Foo, u8> = Ok(x);
350            r
351        });
352        let mut sorted_iter = BufferedSortIter::new(foos_iter, 100).unwrap();
353        assert!(sorted_iter.next() == Some(Ok(Foo { x: 1 })));
354        assert!(sorted_iter.next() == Some(Ok(Foo { x: 1 })));
355        assert!(sorted_iter.next() == Some(Ok(Foo { x: 1 })));
356        assert!(sorted_iter.next() == Some(Ok(Foo { x: 3 })));
357        assert!(sorted_iter.next() == Some(Ok(Foo { x: 3 })));
358        assert!(sorted_iter.next() == Some(Ok(Foo { x: 3 })));
359        assert!(sorted_iter.next() == Some(Ok(Foo { x: 3 })));
360        assert!(sorted_iter.next() == Some(Ok(Foo { x: 3 })));
361        assert!(sorted_iter.next() == Some(Ok(Foo { x: 4 })));
362        assert!(sorted_iter.next().is_none());
363    }
364
365    #[test]
366    #[should_panic]
367    fn buffered_sort_results_too_spread() {
368        // the bufsize is only 2 here but the disordered items extend further apart, so this should fail.
369        let foos = vec![
370            Foo { x: 1 },
371            Foo { x: 3 },
372            Foo { x: 1 },
373            Foo { x: 4 },
374            Foo { x: 3 },
375            Foo { x: 3 },
376            Foo { x: 3 },
377            Foo { x: 3 },
378            Foo { x: 1 },
379        ];
380        let foos_iter = foos.into_iter().map(|x| {
381            let r: Result<Foo, u8> = Ok(x);
382            r
383        });
384        let mut sorted_iter = BufferedSortIter::new(foos_iter, 2).unwrap();
385        assert!(sorted_iter.next() == Some(Ok(Foo { x: 1 })));
386        assert!(sorted_iter.next() == Some(Ok(Foo { x: 1 })));
387        assert!(sorted_iter.next() == Some(Ok(Foo { x: 1 })));
388        assert!(sorted_iter.next() == Some(Ok(Foo { x: 3 })));
389        assert!(sorted_iter.next() == Some(Ok(Foo { x: 3 })));
390        assert!(sorted_iter.next() == Some(Ok(Foo { x: 3 })));
391        assert!(sorted_iter.next() == Some(Ok(Foo { x: 3 })));
392        assert!(sorted_iter.next() == Some(Ok(Foo { x: 3 })));
393        assert!(sorted_iter.next() == Some(Ok(Foo { x: 4 })));
394        assert!(sorted_iter.next().is_none());
395    }
396
397    #[test]
398    fn buffered_sort_results_partial() {
399        // here bufsize is 4, which is enough
400        let foos = vec![
401            Foo { x: 1 },
402            Foo { x: 3 },
403            Foo { x: 4 },
404            Foo { x: 1 },
405            Foo { x: 3 },
406            Foo { x: 3 },
407            Foo { x: 3 },
408            Foo { x: 3 },
409        ];
410        let foos_iter = foos.into_iter().map(|x| {
411            let r: Result<Foo, u8> = Ok(x);
412            r
413        });
414        let mut sorted_iter = BufferedSortIter::new(foos_iter, 4).unwrap();
415        assert!(sorted_iter.next() == Some(Ok(Foo { x: 1 })));
416        assert!(sorted_iter.next() == Some(Ok(Foo { x: 1 })));
417        assert!(sorted_iter.next() == Some(Ok(Foo { x: 3 })));
418        assert!(sorted_iter.next() == Some(Ok(Foo { x: 3 })));
419        assert!(sorted_iter.next() == Some(Ok(Foo { x: 3 })));
420        assert!(sorted_iter.next() == Some(Ok(Foo { x: 3 })));
421        assert!(sorted_iter.next() == Some(Ok(Foo { x: 3 })));
422        assert!(sorted_iter.next() == Some(Ok(Foo { x: 4 })));
423        assert!(sorted_iter.next().is_none());
424    }
425}
groupby/lib.rs

groupby/
lib.rs
