1use crate::physical::types::{OperatorResult, PhysicalOperatorExec};
2use crate::processor::chunk_helpers::{extract_all_rows_from_chunks, rows_to_columns};
3use akar_common::types::Value;
4use akar_common::vector::DataChunk;
5
6pub struct PhysicalAccumulate;
11
12impl PhysicalOperatorExec for PhysicalAccumulate {
13 fn operator_type(&self) -> &str {
14 "accumulate"
15 }
16
17 fn execute(&self, input: Vec<DataChunk>) -> OperatorResult {
18 if input.is_empty() {
19 return Ok(input);
20 }
21
22 let all_rows = extract_all_rows_from_chunks(&input);
23 if all_rows.is_empty() {
24 return Ok(vec![DataChunk::new(vec![], vec![])]);
25 }
26
27 let (fields, field_types) = rows_to_columns(&all_rows);
28 let size = all_rows.len();
29 let field_names = input[0].field_names.clone();
30
31 Ok(vec![DataChunk {
32 fields,
33 field_types,
34 size,
35 field_names,
36 sel_vector: None,
37 }])
38 }
39}
40
41pub struct PhysicalUnion {
46 pub all: bool,
47}
48
49impl PhysicalOperatorExec for PhysicalUnion {
50 fn operator_type(&self) -> &str {
51 "union"
52 }
53
54 fn execute(&self, _input: Vec<DataChunk>) -> OperatorResult {
55 Ok(Vec::new())
57 }
58}
59
60pub struct ResultCollector;
66
67impl PhysicalOperatorExec for ResultCollector {
68 fn operator_type(&self) -> &str {
69 "result_collector"
70 }
71
72 fn execute(&self, input: Vec<DataChunk>) -> OperatorResult {
73 if input.is_empty() {
74 return Ok(input);
75 }
76 if input.len() == 1 {
77 return Ok(input);
78 }
79
80 let all_rows = extract_all_rows_from_chunks(&input);
81 if all_rows.is_empty() {
82 return Ok(vec![DataChunk::new(vec![], vec![])]);
83 }
84
85 let (fields, field_types) = rows_to_columns(&all_rows);
86 let size = all_rows.len();
87 let field_names = input[0].field_names.clone();
88
89 Ok(vec![DataChunk {
90 fields,
91 field_types,
92 size,
93 field_names,
94 sel_vector: None,
95 }])
96 }
97}
98
99pub struct DummySink;
104
105impl PhysicalOperatorExec for DummySink {
106 fn operator_type(&self) -> &str {
107 "dummy_sink"
108 }
109
110 fn execute(&self, _input: Vec<DataChunk>) -> OperatorResult {
111 Ok(Vec::new())
112 }
113}
114
115pub struct DummySimpleSink;
119
120impl PhysicalOperatorExec for DummySimpleSink {
121 fn operator_type(&self) -> &str {
122 "dummy_simple_sink"
123 }
124
125 fn execute(&self, _input: Vec<DataChunk>) -> OperatorResult {
126 Ok(vec![DataChunk::new(Vec::new(), Vec::new())])
127 }
128}
129
130pub struct Profile {
136 pub inner: Box<dyn PhysicalOperatorExec + Send + Sync>,
137 pub elapsed_nanos: std::sync::atomic::AtomicU64,
138}
139
140impl Profile {
141 pub fn new(inner: Box<dyn PhysicalOperatorExec + Send + Sync>) -> Self {
142 Self {
143 inner,
144 elapsed_nanos: std::sync::atomic::AtomicU64::new(0),
145 }
146 }
147
148 pub fn elapsed_ns(&self) -> u64 {
149 self.elapsed_nanos.load(std::sync::atomic::Ordering::Relaxed)
150 }
151}
152
153impl PhysicalOperatorExec for Profile {
154 fn operator_type(&self) -> &str {
155 "profile"
156 }
157
158 fn execute(&self, input: Vec<DataChunk>) -> OperatorResult {
159 let start = std::time::Instant::now();
160 let result = self.inner.execute(input);
161 let elapsed = start.elapsed();
162 self.elapsed_nanos
163 .fetch_add(elapsed.as_nanos() as u64, std::sync::atomic::Ordering::Relaxed);
164 result
165 }
166}
167
168pub struct Partitioner {
176 pub morsel_size: usize,
177}
178
179impl Partitioner {
180 pub fn new(morsel_size: usize) -> Self {
181 Self { morsel_size }
182 }
183
184 fn split_chunk(&self, chunk: &DataChunk) -> Vec<DataChunk> {
186 if self.morsel_size == 0 || chunk.size <= self.morsel_size {
187 return vec![chunk.clone()];
188 }
189 let num_morsels = chunk.size.div_ceil(self.morsel_size);
190 let mut morsels = Vec::with_capacity(num_morsels);
191 for start in (0..chunk.size).step_by(self.morsel_size) {
192 let end = (start + self.morsel_size).min(chunk.size);
193 let morsel_fields: Vec<arrow::array::ArrayRef> = chunk
194 .fields
195 .iter()
196 .enumerate()
197 .map(|(col_idx, _fv)| {
198 let phys_type = chunk.field_types[col_idx];
199 let mut morsel = akar_common::vector::ValueVector::new(phys_type, end - start);
200 for i in start..end {
201 let val = chunk.get_value(col_idx, i).unwrap_or(Value::Null);
202 let _ = morsel.set_value(i - start, &val);
203 }
204 akar_common::arrow_vector::ArrowVector::from_legacy(&morsel).array
205 })
206 .collect();
207 morsels.push(DataChunk {
208 fields: morsel_fields,
209 field_types: chunk.field_types.clone(),
210 size: end - start,
211 field_names: chunk.field_names.clone(),
212 sel_vector: None,
213 });
214 }
215 morsels
216 }
217}
218
219impl PhysicalOperatorExec for Partitioner {
220 fn operator_type(&self) -> &str {
221 "partitioner"
222 }
223
224 fn execute(&self, input: Vec<DataChunk>) -> OperatorResult {
225 if self.morsel_size == 0 {
226 return Ok(input);
227 }
228 let morsels: Vec<DataChunk> = input.iter().flat_map(|chunk| self.split_chunk(chunk)).collect();
229 Ok(morsels)
230 }
231}
232
233#[cfg(test)]
234mod tests {
235 use super::*;
236
237 #[test]
238 fn test_partitioner_pass_through_zero_size() {
239 let p = Partitioner::new(0);
240 let mut fv = akar_common::vector::ValueVector::new(akar_common::types::PhysicalTypeID::Int64, 10);
241 fv.resize(10);
242 let chunk = {
243 let arrow_fields = vec![akar_common::arrow_vector::ArrowVector::from_legacy(&fv).array];
244 let arrow_field_types = vec![fv.physical_type()];
245 DataChunk::new(arrow_fields, arrow_field_types).with_names(vec!["val".into()])
246 };
247 let result = p.execute(vec![chunk.clone()]).unwrap();
248 assert_eq!(result.len(), 1);
249 assert_eq!(result[0].size, 10);
250 }
251
252 #[test]
253 fn test_partitioner_no_split_small_chunk() {
254 let p = Partitioner::new(100);
255 let mut fv = akar_common::vector::ValueVector::new(akar_common::types::PhysicalTypeID::Int64, 5);
256 fv.resize(5);
257 let chunk = {
258 let arrow_fields = vec![akar_common::arrow_vector::ArrowVector::from_legacy(&fv).array];
259 let arrow_field_types = vec![fv.physical_type()];
260 DataChunk::new(arrow_fields, arrow_field_types).with_names(vec!["val".into()])
261 };
262 let result = p.execute(vec![chunk.clone()]).unwrap();
263 assert_eq!(result.len(), 1);
264 assert_eq!(result[0].size, 5);
265 }
266
267 #[test]
268 fn test_partitioner_splits_large_chunk() {
269 let p = Partitioner::new(10);
270 let mut fv = akar_common::vector::ValueVector::new(akar_common::types::PhysicalTypeID::Int64, 25);
271 fv.resize(25);
272 for i in 0..25 {
273 fv.set_i64(i, i as i64);
274 }
275 let chunk = {
276 let arrow_fields = vec![akar_common::arrow_vector::ArrowVector::from_legacy(&fv).array];
277 let arrow_field_types = vec![fv.physical_type()];
278 DataChunk::new(arrow_fields, arrow_field_types).with_names(vec!["val".into()])
279 };
280 let result = p.execute(vec![chunk]).unwrap();
281 assert_eq!(result.len(), 3);
283 assert_eq!(result[0].size, 10);
284 assert_eq!(result[1].size, 10);
285 assert_eq!(result[2].size, 5);
286 assert_eq!(result[0].get_value(0, 0), Some(akar_common::types::Value::Int64(0)));
288 assert_eq!(result[0].get_value(0, 9), Some(akar_common::types::Value::Int64(9)));
289 assert_eq!(result[1].get_value(0, 0), Some(akar_common::types::Value::Int64(10)));
291 assert_eq!(result[2].get_value(0, 0), Some(akar_common::types::Value::Int64(20)));
293 assert_eq!(result[2].get_value(0, 4), Some(akar_common::types::Value::Int64(24)));
294 }
295
296 #[test]
297 fn test_partitioner_5_rows_split_into_3_2() {
298 let p = Partitioner::new(3);
299 let mut fv = akar_common::vector::ValueVector::new(akar_common::types::PhysicalTypeID::Int64, 5);
300 fv.resize(5);
301 for i in 0..5 {
302 fv.set_i64(i, i as i64);
303 }
304 let chunk = {
305 let arrow_fields = vec![akar_common::arrow_vector::ArrowVector::from_legacy(&fv).array];
306 let arrow_field_types = vec![fv.physical_type()];
307 DataChunk::new(arrow_fields, arrow_field_types).with_names(vec!["val".into()])
308 };
309 let result = p.execute(vec![chunk]).unwrap();
310 assert_eq!(result.len(), 2);
311 assert_eq!(result[0].size, 3);
312 assert_eq!(result[1].size, 2);
313 assert_eq!(result[0].get_value(0, 2), Some(akar_common::types::Value::Int64(2)));
314 assert_eq!(result[1].get_value(0, 0), Some(akar_common::types::Value::Int64(3)));
315 }
316
317 #[test]
318 fn test_partitioner_multiple_chunks() {
319 let p = Partitioner::new(3);
320 let mut fv1 = akar_common::vector::ValueVector::new(akar_common::types::PhysicalTypeID::Int64, 5);
321 fv1.resize(5);
322 for i in 0..5 {
323 fv1.set_i64(i, i as i64);
324 }
325 let mut fv2 = akar_common::vector::ValueVector::new(akar_common::types::PhysicalTypeID::Int64, 5);
326 fv2.resize(5);
327 for i in 0..5 {
328 fv2.set_i64(i, (i + 5) as i64);
329 }
330
331 let chunks = vec![
332 {
333 let arrow_fields = vec![akar_common::arrow_vector::ArrowVector::from_legacy(&fv1).array];
334 let arrow_field_types = vec![fv1.physical_type()];
335 DataChunk::new(arrow_fields, arrow_field_types).with_names(vec!["val".into()])
336 },
337 {
338 let arrow_fields = vec![akar_common::arrow_vector::ArrowVector::from_legacy(&fv2).array];
339 let arrow_field_types = vec![fv2.physical_type()];
340 DataChunk::new(arrow_fields, arrow_field_types).with_names(vec!["val".into()])
341 },
342 ];
343 let result = p.execute(chunks).unwrap();
344 assert_eq!(result.len(), 4, "expected 4 morsels from 2 chunks of 5 rows each");
345 assert_eq!(result[2].get_value(0, 0), Some(akar_common::types::Value::Int64(5)));
346 assert_eq!(result[2].get_value(0, 2), Some(akar_common::types::Value::Int64(7)));
347 assert_eq!(result[3].get_value(0, 0), Some(akar_common::types::Value::Int64(8)));
348 }
349}