harmos_storage/query/transform/
sort.rs

1/* ***
2 *
3 * Harmos
4 *
5 * SPDX-License-Identifier: Apache-2.0
6 *
7 * ***
8 */
9
10// Storage Arrow Query Contract
11
12/* === Modules === */
13
14/* === Re-Exports === */
15
16pub use harmos_domain::model::storage::query::transform::sort::{Sort, SortKey};
17
18/* === Imports === */
19
20use crate::storage::query::transform::FieldPathExecution;
21// Standard Library Imports
22use std::cmp::Ordering;
23// Third Party Imports
24use arrow_array::RecordBatch;
25use arrow_schema::Schema;
26// Local Imports
27use super::{invalid, scalar_type, take};
28use crate::Result;
29use crate::storage::query::accumulate::sort_values;
30use harmos_domain::model::storage::query::scalar::ScalarValue;
31
32/* === Constants === */
33
34/* === Macros === */
35
36/* === Type Aliases === */
37
38/* === Structs === */
39
40/* === Enums === */
41
42/* === Traits === */
43
44pub trait SortExecution {
45    fn output_schema(&self, input: &Schema) -> Result<Schema>;
46
47    fn apply(&self, input: &RecordBatch, check: &mut impl FnMut() -> Result<()>) -> Result<RecordBatch>;
48}
49
50/* === Implementations === */
51
52impl SortExecution for Sort {
53    fn output_schema(&self, input: &Schema) -> Result<Schema> {
54        if self.keys.is_empty() || self.keys.len() > 64 {
55            return Err(invalid("sort requires 1..64 keys"));
56        }
57        for key in &self.keys {
58            if !scalar_type(key.field.field(input)?.data_type()) {
59                return Err(invalid("sort requires scalar fields"));
60            }
61        }
62        Ok(input.clone())
63    }
64
65    fn apply(&self, input: &RecordBatch, check: &mut impl FnMut() -> Result<()>) -> Result<RecordBatch> {
66        let mut indices = (0..input.num_rows() as u64).collect::<Vec<_>>();
67        sort_values(&mut indices, |left, right| {
68            check()?;
69            for key in &self.keys {
70                let left = key.field.value(input, left as usize)?;
71                let right = key.field.value(input, right as usize)?;
72                let absent = |value: &ScalarValue| {
73                    matches!(value, ScalarValue::Null | ScalarValue::Missing)
74                        || matches!(value, ScalarValue::Float(value) if !value.is_finite())
75                };
76                let null_order = match (absent(&left), absent(&right)) {
77                    (true, true) => Ordering::Equal,
78                    (true, false) => {
79                        if key.nulls_first {
80                            Ordering::Less
81                        } else {
82                            Ordering::Greater
83                        }
84                    }
85                    (false, true) => {
86                        if key.nulls_first {
87                            Ordering::Greater
88                        } else {
89                            Ordering::Less
90                        }
91                    }
92                    (false, false) => {
93                        let order = left.compare(&right)?;
94                        if key.descending { order.reverse() } else { order }
95                    }
96                };
97                if !null_order.is_eq() {
98                    return Ok(null_order);
99                }
100            }
101            // Original row order is a final tie key, independent of sort direction.
102            Ok(left.cmp(&right))
103        })?;
104        take(input, indices)
105    }
106}
107
108/* === Public Functions === */
109
110/* === Private Functions === */
111
112/* === Tests === */
113
114#[cfg(test)]
115mod tests {
116    use crate::query::transform::{
117        FieldPath, Transform, TransformExecution,
118        project::{Project, Projection},
119        sort::{Sort, SortKey},
120    };
121    use arrow_array::{Array, Int64Array, RecordBatch};
122    use arrow_schema::{DataType, Field, Schema};
123    use std::sync::Arc;
124
125    fn batch() -> RecordBatch {
126        let schema = Arc::new(Schema::new(vec![Field::new("key", DataType::Int64, false), Field::new("row", DataType::Int64, false)]));
127        RecordBatch::try_new(schema, vec![Arc::new(Int64Array::from(vec![2, 1, 1, 2])), Arc::new(Int64Array::from(vec![0, 1, 2, 3]))])
128            .unwrap()
129    }
130
131    #[test]
132    fn sort_keeps_ties() {
133        // Sort Keeps Stable Ties Before Projection Selects The Original Row Order.
134        let original = batch();
135        let sorted = Transform::Sort(Sort { keys: vec![SortKey { field: FieldPath::from("key"), descending: false, nulls_first: false }] })
136            .apply(&original, &mut || Ok(()))
137            .unwrap();
138        let projected = Transform::Project(Project { fields: vec![Projection { field: FieldPath::from("row"), name: "selected".into() }] })
139            .apply(&sorted, &mut || Ok(()))
140            .unwrap();
141        let values = projected.column(0).as_any().downcast_ref::<Int64Array>().unwrap();
142        assert_eq!(values.values().as_ref(), &[1, 2, 0, 3]);
143    }
144}
harmos_storage/query/transform/sort.rs

harmos_storage/query/transform/
sort.rs
