versatile_data/
sort.rs

1use std::{cmp::Ordering, fmt::Debug, num::NonZeroU32};
2
3use idx_binary::{AvltrieeSearch, IdxBinary, IdxFileAvlTriee};
4
5use crate::{Data, FieldName, RowSet};
6
7pub trait CustomSort {
8    fn compare(&self, a: NonZeroU32, b: NonZeroU32) -> Ordering;
9    fn asc(&self) -> Vec<NonZeroU32>;
10    fn desc(&self) -> Vec<NonZeroU32>;
11}
12impl Debug for dyn CustomSort {
13    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
14        writeln!(f, "CustomSort")?;
15        Ok(())
16    }
17}
18
19pub struct NoCustomSort {}
20
21impl CustomSort for NoCustomSort {
22    fn compare(&self, _: NonZeroU32, _: NonZeroU32) -> Ordering {
23        unreachable!()
24    }
25
26    fn asc(&self) -> Vec<NonZeroU32> {
27        unreachable!()
28    }
29
30    fn desc(&self) -> Vec<NonZeroU32> {
31        unreachable!()
32    }
33}
34
35pub enum CustomOrderKey<C: CustomSort> {
36    Serial,
37    Row,
38    TermBegin,
39    TermEnd,
40    LastUpdated,
41    Field(FieldName),
42    Custom(C),
43}
44
45pub type OrderKey = CustomOrderKey<NoCustomSort>;
46
47pub enum Order<C: CustomSort> {
48    Asc(CustomOrderKey<C>),
49    Desc(CustomOrderKey<C>),
50}
51
52impl Data {
53    /// Sort search results.
54    pub fn sort<C: CustomSort>(&self, rows: &RowSet, orders: &[Order<C>]) -> Vec<NonZeroU32> {
55        let sub_orders = &orders[1..];
56        match &orders[0] {
57            Order::Asc(key) => self.sort_with_key(rows, key, sub_orders),
58            Order::Desc(key) => self.sort_with_key_desc(rows, key, sub_orders),
59        }
60    }
61
62    /// Compares two rows by one key. Rows without a value come last
63    /// in both ascending and descending order.
64    fn compare_by_key<C: CustomSort>(
65        &self,
66        key: &CustomOrderKey<C>,
67        a: NonZeroU32,
68        b: NonZeroU32,
69        desc: bool,
70    ) -> Ordering {
71        fn missing_last<T>(
72            a: Option<T>,
73            b: Option<T>,
74            desc: bool,
75            cmp: impl Fn(T, T) -> Ordering,
76        ) -> Ordering {
77            match (a, b) {
78                (Some(a), Some(b)) => {
79                    if desc {
80                        cmp(b, a)
81                    } else {
82                        cmp(a, b)
83                    }
84                }
85                (Some(_), None) => Ordering::Less,
86                (None, Some(_)) => Ordering::Greater,
87                (None, None) => Ordering::Equal,
88            }
89        }
90        let index = |f: &Option<idx_binary::IdxFile<u64>>| {
91            f.as_ref().map_or(Ordering::Equal, |f| {
92                missing_last(f.value(a), f.value(b), desc, Ord::cmp)
93            })
94        };
95        match key {
96            CustomOrderKey::Serial => {
97                missing_last(self.serial.value(a), self.serial.value(b), desc, Ord::cmp)
98            }
99            CustomOrderKey::Row => missing_last(Some(a), Some(b), desc, |a, b| a.cmp(&b)),
100            CustomOrderKey::TermBegin => index(&self.term_begin),
101            CustomOrderKey::TermEnd => index(&self.term_end),
102            CustomOrderKey::LastUpdated => index(&self.last_updated),
103            CustomOrderKey::Field(name) => self.fields.get(name).map_or(Ordering::Equal, |f| {
104                missing_last(f.value(a), f.value(b), desc, IdxBinary::cmp)
105            }),
106            CustomOrderKey::Custom(custom_order) => {
107                if desc {
108                    custom_order.compare(b, a)
109                } else {
110                    custom_order.compare(a, b)
111                }
112            }
113        }
114    }
115
116    fn subsort<C: CustomSort>(
117        &self,
118        mut tmp: Vec<NonZeroU32>,
119        sub_orders: &[Order<C>],
120    ) -> Vec<NonZeroU32> {
121        tmp.sort_by(|a, b| {
122            for order in sub_orders {
123                let ord = match order {
124                    Order::Asc(key) => self.compare_by_key(key, *a, *b, false),
125                    Order::Desc(key) => self.compare_by_key(key, *a, *b, true),
126                };
127                if ord != Ordering::Equal {
128                    return ord;
129                }
130            }
131            Ordering::Equal
132        });
133        tmp
134    }
135
136    /// Sorts rows in the order of the index. Rows that are not in the index
137    /// (rows without a value for the key) come last.
138    fn sort_with_triee_inner<T: PartialEq, I: ?Sized, C: CustomSort>(
139        &self,
140        rows: &RowSet,
141        triee: &IdxFileAvlTriee<T, I>,
142        iter: impl Iterator<Item = NonZeroU32>,
143        sub_orders: &[Order<C>],
144    ) -> Vec<NonZeroU32> {
145        let sorted_group = |tmp: Vec<NonZeroU32>| {
146            if tmp.len() <= 1 || sub_orders.is_empty() {
147                tmp
148            } else {
149                self.subsort(tmp, sub_orders)
150            }
151        };
152
153        let mut ret = Vec::with_capacity(rows.len());
154        let mut before: Option<&T> = None;
155        let mut tmp: Vec<NonZeroU32> = Vec::new();
156        for r in iter {
157            if rows.contains(&r) {
158                let value: &T = triee.node(r).unwrap();
159                if before.is_some_and(|before| before.ne(value)) {
160                    ret.extend(sorted_group(std::mem::take(&mut tmp)));
161                }
162                tmp.push(r);
163                before = Some(value);
164            }
165        }
166        ret.extend(sorted_group(tmp));
167
168        if ret.len() < rows.len() {
169            let in_index: RowSet = ret.iter().copied().collect();
170            let missing = rows
171                .iter()
172                .copied()
173                .filter(|r| !in_index.contains(r))
174                .collect();
175            ret.extend(sorted_group(missing));
176        }
177        ret
178    }
179
180    fn sort_with_triee<T: PartialEq, I: ?Sized, C: CustomSort>(
181        &self,
182        rows: &RowSet,
183        triee: &IdxFileAvlTriee<T, I>,
184        sub_orders: &[Order<C>],
185    ) -> Vec<NonZeroU32> {
186        self.sort_with_triee_inner(rows, triee, triee.iter(), sub_orders)
187    }
188
189    fn sort_with_triee_desc<T: PartialEq, I: ?Sized, C: CustomSort>(
190        &self,
191        rows: &RowSet,
192        index: &IdxFileAvlTriee<T, I>,
193        sub_orders: &[Order<C>],
194    ) -> Vec<NonZeroU32> {
195        self.sort_with_triee_inner(rows, index, index.desc_iter(), sub_orders)
196    }
197
198    fn sort_with_key<C: CustomSort>(
199        &self,
200        rows: &RowSet,
201        key: &CustomOrderKey<C>,
202        sub_orders: &[Order<C>],
203    ) -> Vec<NonZeroU32> {
204        match key {
205            CustomOrderKey::Serial => self.sort_with_triee::<u32, u32, C>(rows, &self.serial, &[]),
206            CustomOrderKey::Row => rows.iter().cloned().collect(),
207            CustomOrderKey::TermBegin => self.term_begin.as_ref().map_or_else(
208                || rows.iter().cloned().collect(),
209                |f| self.sort_with_triee(rows, f, sub_orders),
210            ),
211            CustomOrderKey::TermEnd => self.term_end.as_ref().map_or_else(
212                || rows.iter().cloned().collect(),
213                |f| self.sort_with_triee(rows, f, sub_orders),
214            ),
215            CustomOrderKey::LastUpdated => self.last_updated.as_ref().map_or_else(
216                || rows.iter().cloned().collect(),
217                |f| self.sort_with_triee(rows, f, sub_orders),
218            ),
219            CustomOrderKey::Field(name) => self.fields.get(name).map_or_else(
220                || rows.iter().cloned().collect(),
221                |f| self.sort_with_triee(rows, f.as_ref(), sub_orders),
222            ),
223            CustomOrderKey::Custom(custom_order) => custom_order.asc(),
224        }
225    }
226
227    fn sort_with_key_desc<C: CustomSort>(
228        &self,
229        rows: &RowSet,
230        key: &CustomOrderKey<C>,
231        sub_orders: &[Order<C>],
232    ) -> Vec<NonZeroU32> {
233        match key {
234            CustomOrderKey::Serial => {
235                self.sort_with_triee_desc::<u32, u32, C>(rows, &self.serial, &[])
236            }
237            CustomOrderKey::Row => rows.iter().rev().cloned().collect(),
238            CustomOrderKey::TermBegin => self.term_begin.as_ref().map_or_else(
239                || rows.iter().rev().cloned().collect(),
240                |f| self.sort_with_triee_desc(rows, f, sub_orders),
241            ),
242            CustomOrderKey::TermEnd => self.term_end.as_ref().map_or_else(
243                || rows.iter().rev().cloned().collect(),
244                |f| self.sort_with_triee_desc(rows, f, sub_orders),
245            ),
246            CustomOrderKey::LastUpdated => self.last_updated.as_ref().map_or_else(
247                || rows.iter().rev().cloned().collect(),
248                |f| self.sort_with_triee_desc(rows, f, sub_orders),
249            ),
250            CustomOrderKey::Field(name) => self.fields.get(name).map_or_else(
251                || rows.iter().rev().cloned().collect(),
252                |f| self.sort_with_triee_desc(rows, f.as_ref(), sub_orders),
253            ),
254            CustomOrderKey::Custom(custom_order) => custom_order.desc(),
255        }
256    }
257}
versatile_data/sort.rs

versatile_data/
sort.rs
