Skip to main content

akar_processor/physical/order_aggregate/
orderby.rs

1//! Auto-extracted from physical_operator.rs
2use crate::physical::common::{take_global_rows, value_cmp};
3use crate::physical::order_aggregate::BlockMergeSorter;
4use crate::physical::types::{OperatorResult, PhysicalOperatorExec};
5use akar_common::types::Value;
6use akar_common::vector::DataChunk;
7
8// ==================== ChunkAccessor ====================
9
10/// Provides random-access by global row index across multiple DataChunks.
11/// Eliminates the need to pre-collect all values into Vec<Vec<(Value, bool)>>.
12struct ChunkAccessor<'a> {
13    chunks: &'a [DataChunk],
14    offsets: Vec<usize>,
15    num_fields: usize,
16}
17
18impl<'a> ChunkAccessor<'a> {
19    fn new(chunks: &'a [DataChunk]) -> Self {
20        let mut offsets = Vec::with_capacity(chunks.len());
21        let mut cum = 0usize;
22        for c in chunks {
23            offsets.push(cum);
24            cum += c.size;
25        }
26        let num_fields = chunks.first().map(|c| c.num_fields()).unwrap_or(0);
27        Self {
28            chunks,
29            offsets,
30            num_fields,
31        }
32    }
33
34    fn total_rows(&self) -> usize {
35        self.offsets
36            .last()
37            .map(|&o| o + self.chunks.last().unwrap().size)
38            .unwrap_or(0)
39    }
40
41    fn resolve(&self, global_row: usize) -> (usize, usize) {
42        for (ci, chunk) in self.chunks.iter().enumerate() {
43            let offset = self.offsets[ci];
44            if global_row < offset + chunk.size {
45                return (ci, global_row - offset);
46            }
47        }
48        (self.chunks.len() - 1, 0)
49    }
50
51    fn get_value(&self, col: usize, global_row: usize) -> Value {
52        let (ci, local) = self.resolve(global_row);
53        self.chunks[ci].get_value(col, local).unwrap_or(Value::Null)
54    }
55
56    fn is_null(&self, col: usize, global_row: usize) -> bool {
57        let (ci, local) = self.resolve(global_row);
58        self.chunks[ci].is_null(col, local)
59    }
60}
61
62// ==================== OrderBy ====================
63
64pub struct PhysicalOrderBy {
65    pub sort_keys: Vec<(u32, bool)>,
66}
67
68impl PhysicalOperatorExec for PhysicalOrderBy {
69    fn operator_type(&self) -> &str {
70        "order_by"
71    }
72
73    fn execute(&self, input: Vec<DataChunk>) -> OperatorResult {
74        if input.is_empty() {
75            return Ok(Vec::new());
76        }
77
78        let accessor = ChunkAccessor::new(&input);
79        let total_rows = accessor.total_rows();
80        if total_rows == 0 {
81            return Ok(input);
82        }
83
84        let num_fields = accessor.num_fields;
85        let field_names = input[0].field_names.clone();
86
87        // Use BlockMergeSorter for large data, simple sort for small
88        let block_size = 10000usize;
89        let indices = if total_rows > block_size && !self.sort_keys.is_empty() {
90            // Collect key values for k-way merge. Only the sort-key columns are
91            // materialized (not the full O(rows × fields) matrix); indices are
92            // remapped to the compact column set (P79).
93            let mut sort_cols: Vec<u32> = Vec::new();
94            for &(col, _) in &self.sort_keys {
95                if (col as usize) < num_fields && !sort_cols.contains(&col) {
96                    sort_cols.push(col);
97                }
98            }
99            if sort_cols.is_empty() {
100                (0..total_rows).collect()
101            } else {
102                let mut key_values: Vec<Vec<(Value, bool)>> =
103                    sort_cols.iter().map(|_| Vec::with_capacity(total_rows)).collect();
104                for global_row in 0..total_rows {
105                    for (ci, &col) in sort_cols.iter().enumerate() {
106                        let val = accessor.get_value(col as usize, global_row);
107                        let is_null = accessor.is_null(col as usize, global_row);
108                        key_values[ci].push((val, is_null));
109                    }
110                }
111                let compact_keys: Vec<(u32, bool)> = self
112                    .sort_keys
113                    .iter()
114                    .filter(|&&(col, _)| (col as usize) < num_fields)
115                    .map(|&(col, asc)| {
116                        let idx = sort_cols.iter().position(|&c| c == col).unwrap() as u32;
117                        (idx, asc)
118                    })
119                    .collect();
120                let sorter = BlockMergeSorter::new(block_size, compact_keys);
121                sorter.sort(&key_values, sort_cols.len())
122            }
123        } else {
124            let mut indices: Vec<usize> = (0..total_rows).collect();
125            if !self.sort_keys.is_empty() {
126                indices.sort_by(|a, b| {
127                    for &(col, ascending) in &self.sort_keys {
128                        let col = col as usize;
129                        if col >= num_fields {
130                            continue;
131                        }
132                        let va = accessor.get_value(col, *a);
133                        let vb = accessor.get_value(col, *b);
134                        let cmp = value_cmp(&va, &vb);
135                        if cmp != std::cmp::Ordering::Equal {
136                            return if ascending { cmp } else { cmp.reverse() };
137                        }
138                    }
139                    std::cmp::Ordering::Equal
140                });
141            }
142            indices
143        };
144
145        // Build sorted output chunks (up to 100 rows per chunk). Slice the
146        // source fields via Arrow `take`, preserving complex types (List/Struct).
147        let chunk_size = 100usize;
148        let mut output = Vec::new();
149        for chunk_start in (0..total_rows).step_by(chunk_size) {
150            let chunk_end = (chunk_start + chunk_size).min(total_rows);
151            output.push(take_global_rows(
152                &input,
153                &indices[chunk_start..chunk_end],
154                field_names.clone(),
155            )?);
156        }
157        Ok(output)
158    }
159}