Skip to main content

akar_processor/physical/
missing_ops.rs

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
6/// Accumulate — materializes all input into a single contiguous chunk in memory.
7///
8/// This is required for operations that need random access to all rows,
9/// such as hash join build side, correlated subqueries, etc.
10pub 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
41/// Union — concatenates results from two child pipelines.
42///
43/// This operator is not currently used in the pipeline (Union is handled inline
44/// in QueryProcessor::execute_internal). Kept for API compatibility with C++.
45pub 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        // Not used - Union is handled inline in execute_internal
56        Ok(Vec::new())
57    }
58}
59
60/// ResultCollector — collects all input DataChunks into a single result set.
61///
62/// This is the final operator in the query pipeline. It consolidates
63/// all chunks into a single `Vec<DataChunk>` ready for client return.
64/// If there are multiple input chunks, they are merged into one.
65pub 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
99/// DummySink — consumes and discards all input.
100///
101/// Used as a pipeline terminal when results are not needed (e.g.,
102/// DDL statements, EXPLAIN without ANALYZE).
103pub 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
115/// DummySimpleSink — like DummySink but passes through a single empty chunk.
116///
117/// Some pipeline consumers expect at least one DataChunk to be returned.
118pub 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
130/// Profile — wraps an operator with timing instrumentation.
131///
132/// Measures wall-clock execution time of the wrapped operator.
133/// The elapsed time is stored in `elapsed` field for later inspection
134/// via EXPLAIN ANALYZE or profiling output.
135pub 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
168/// Partitioner — splits input DataChunks into fixed-size morsels for parallelism.
169///
170/// Each morsel (batch of `morsel_size` rows) can be processed independently
171/// by downstream operators. The mapper uses this to distribute work across
172/// threads via Rayon.
173///
174/// A morsel size of 0 disables splitting (pass-through).
175pub 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    /// Split a single DataChunk into morsels of `morsel_size` rows.
185    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        // 25 rows with morsel_size 10 → 3 morsels (10 + 10 + 5)
282        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        // Verify data integrity: first morsel
287        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        // Second morsel
290        assert_eq!(result[1].get_value(0, 0), Some(akar_common::types::Value::Int64(10)));
291        // Third morsel
292        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}