akar_processor/physical/order_aggregate/
orderby.rs1use 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
8struct 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
70pub 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 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 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 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}