harmos_storage/query/transform/
sort.rs

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

harmos_storage/query/transform/
sort.rs
