Skip to main content

akar_processor/physical/order_aggregate/
orderby.rs

1//! Auto-extracted from physical_operator.rs
2use crate::physical::common::{store_value_in_vector, value_cmp};
3use crate::physical::order_aggregate::BlockMergeSorter;
4use crate::physical::types::{OperatorResult, PhysicalOperatorExec};
5use akar_common::types::Value;
6use akar_common::vector::{DataChunk, ValueVector};
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    fn physical_type(&self, col: usize, global_row: usize) -> akar_common::types::PhysicalTypeID {
62        let (ci, local) = self.resolve(global_row);
63        self.chunks[ci]
64            .get_value(col, local)
65            .map(|v| v.physical_type())
66            .unwrap_or(akar_common::types::PhysicalTypeID::Int64)
67    }
68}
69
70// ==================== OrderBy ====================
71
72pub struct PhysicalOrderBy {
73    pub sort_keys: Vec<(u32, bool)>,
74}
75
76impl PhysicalOperatorExec for PhysicalOrderBy {
77    fn operator_type(&self) -> &str {
78        "order_by"
79    }
80
81    fn execute(&self, input: Vec<DataChunk>) -> OperatorResult {
82        if input.is_empty() {
83            return Ok(Vec::new());
84        }
85
86        let accessor = ChunkAccessor::new(&input);
87        let total_rows = accessor.total_rows();
88        if total_rows == 0 {
89            return Ok(input);
90        }
91
92        let num_fields = accessor.num_fields;
93        let field_names = input[0].field_names.clone();
94
95        // Use BlockMergeSorter for large data, simple sort for small
96        let block_size = 10000usize;
97        let indices = if total_rows > block_size && !self.sort_keys.is_empty() {
98            let sorter = BlockMergeSorter::new(block_size, self.sort_keys.clone());
99            // Multi-block path still needs collected key values for k-way merge
100            let mut all_values: Vec<Vec<(Value, bool)>> =
101                (0..num_fields).map(|_| Vec::with_capacity(total_rows)).collect();
102            for global_row in 0..total_rows {
103                for col in 0..num_fields {
104                    let val = accessor.get_value(col, global_row);
105                    let is_null = accessor.is_null(col, global_row);
106                    all_values[col].push((val, is_null));
107                }
108            }
109            sorter.sort(&all_values, num_fields)
110        } else {
111            let mut indices: Vec<usize> = (0..total_rows).collect();
112            if !self.sort_keys.is_empty() {
113                indices.sort_by(|a, b| {
114                    for &(col, ascending) in &self.sort_keys {
115                        let col = col as usize;
116                        if col >= num_fields {
117                            continue;
118                        }
119                        let va = accessor.get_value(col, *a);
120                        let vb = accessor.get_value(col, *b);
121                        let cmp = value_cmp(&va, &vb);
122                        if cmp != std::cmp::Ordering::Equal {
123                            return if ascending { cmp } else { cmp.reverse() };
124                        }
125                    }
126                    std::cmp::Ordering::Equal
127                });
128            }
129            indices
130        };
131
132        // Build sorted output chunks (up to 100 rows per chunk)
133        let chunk_size = 100usize;
134        let mut output = Vec::new();
135        for chunk_start in (0..total_rows).step_by(chunk_size) {
136            let chunk_end = (chunk_start + chunk_size).min(total_rows);
137            let size = chunk_end - chunk_start;
138            let mut fields = Vec::new();
139            for col in 0..num_fields {
140                let phys_type = accessor.physical_type(col, indices[chunk_start]);
141                let mut v = ValueVector::new(phys_type, size);
142                v.resize(size);
143                for (out_idx, &src_idx) in indices[chunk_start..chunk_end].iter().enumerate() {
144                    if accessor.is_null(col, src_idx) {
145                        v.set_null(out_idx, true);
146                    } else {
147                        let val = accessor.get_value(col, src_idx);
148                        store_value_in_vector(&mut v, out_idx, &val)?;
149                    }
150                }
151                fields.push(v);
152            }
153            let arrow_fields = fields
154                .iter()
155                .map(|v| akar_common::arrow_vector::ArrowVector::from_legacy(v).array)
156                .collect::<Vec<_>>();
157            let arrow_field_types = fields.iter().map(|v| v.physical_type()).collect::<Vec<_>>();
158            output.push(DataChunk::new(arrow_fields, arrow_field_types).with_names(field_names.clone()));
159        }
160        Ok(output)
161    }
162}