akar_processor/physical/order_aggregate/
topk.rs

1//! Auto-extracted from physical_operator.rs
2use crate::physical::common::{store_value_in_vector, value_cmp};
3use crate::physical::types::{OperatorResult, PhysicalOperatorExec};
4use akar_common::types::Value;
5use akar_common::vector::{DataChunk, ValueVector};
6use std::collections::BinaryHeap;
7
8// ==================== TopK ====================
9
10/// Fused ORDER BY + LIMIT using a BinaryHeap (O(n log k) vs O(n log n)).
11///
12/// Maintains a max-heap of size (limit + offset). Pops the worst entry
13/// when capacity is exceeded. Uses `DirectedSortKey` to encode sort
14/// direction into the comparison, so the BinaryHeap's natural max-heap
15/// behavior correctly retains the best entries.
16pub struct PhysicalTopK {
17    pub sort_keys: Vec<(u32, bool)>,
18    pub limit: u64,
19    pub offset: u64,
20}
21
22/// Wrapper for a sort-key value that embeds sort direction.
23#[derive(Debug, Clone)]
24enum DirectedSortKey {
25    Asc(Value),
26    Desc(Value),
27}
28
29impl Eq for DirectedSortKey {}
30impl PartialEq for DirectedSortKey {
31    fn eq(&self, other: &Self) -> bool {
32        self.cmp(other) == std::cmp::Ordering::Equal
33    }
34}
35impl PartialOrd for DirectedSortKey {
36    fn partial_cmp(&self, other: &Self) -> Option<std::cmp::Ordering> {
37        Some(self.cmp(other))
38    }
39}
40impl Ord for DirectedSortKey {
41    fn cmp(&self, other: &Self) -> std::cmp::Ordering {
42        match (self, other) {
43            (DirectedSortKey::Asc(a), DirectedSortKey::Asc(b)) => value_cmp(a, b),
44            (DirectedSortKey::Desc(a), DirectedSortKey::Desc(b)) => value_cmp(b, a),
45            _ => std::cmp::Ordering::Equal,
46        }
47    }
48}
49
50#[derive(Debug, Clone)]
51struct TopKHeapEntry {
52    sort_key: Vec<DirectedSortKey>,
53    row_idx: usize,
54}
55
56impl Eq for TopKHeapEntry {}
57impl PartialEq for TopKHeapEntry {
58    fn eq(&self, other: &Self) -> bool {
59        self.sort_key == other.sort_key
60    }
61}
62impl PartialOrd for TopKHeapEntry {
63    fn partial_cmp(&self, other: &Self) -> Option<std::cmp::Ordering> {
64        Some(self.cmp(other))
65    }
66}
67impl Ord for TopKHeapEntry {
68    fn cmp(&self, other: &Self) -> std::cmp::Ordering {
69        for (a, b) in self.sort_key.iter().zip(other.sort_key.iter()) {
70            let cmp = a.cmp(b);
71            if cmp != std::cmp::Ordering::Equal {
72                return cmp;
73            }
74        }
75        std::cmp::Ordering::Equal
76    }
77}
78
79impl PhysicalOperatorExec for PhysicalTopK {
80    fn operator_type(&self) -> &str {
81        "top_k"
82    }
83
84    fn execute(&self, input: Vec<DataChunk>) -> OperatorResult {
85        let capacity = (self.limit + self.offset) as usize;
86        if capacity == 0 || input.is_empty() {
87            return Ok(Vec::new());
88        }
89
90        let total_rows: usize = input.iter().map(|c| c.size).sum();
91        if total_rows == 0 {
92            return Ok(Vec::new());
93        }
94
95        let num_fields = input[0].num_fields();
96
97        // Collect all values for random access
98        let mut all_values: Vec<Vec<(Value, bool)>> = (0..num_fields).map(|_| Vec::with_capacity(total_rows)).collect();
99        for chunk in &input {
100            for row in 0..chunk.size {
101                for col in 0..num_fields {
102                    if let Some(field) = chunk.fields.get(col) {
103                        let val = chunk.get_value(col, row).unwrap_or(Value::Null);
104                        let is_null = field.is_null(row);
105                        all_values[col].push((val, is_null));
106                    }
107                }
108            }
109        }
110
111        // BinaryHeap (max-heap): worst entry at top, popped when > capacity.
112        let mut heap: BinaryHeap<TopKHeapEntry> = BinaryHeap::with_capacity(capacity.min(total_rows) + 1);
113
114        for row_idx in 0..total_rows {
115            let sort_key: Vec<DirectedSortKey> = self
116                .sort_keys
117                .iter()
118                .map(|&(col, asc)| {
119                    let val = if col as usize >= num_fields {
120                        Value::Null
121                    } else {
122                        all_values[col as usize][row_idx].0.clone()
123                    };
124                    if asc {
125                        DirectedSortKey::Asc(val)
126                    } else {
127                        DirectedSortKey::Desc(val)
128                    }
129                })
130                .collect();
131
132            heap.push(TopKHeapEntry { sort_key, row_idx });
133            if heap.len() > capacity {
134                heap.pop();
135            }
136        }
137
138        // into_sorted_vec returns ascending DirectedSortKey order = best-first
139        let sorted: Vec<TopKHeapEntry> = heap.into_sorted_vec();
140
141        // Apply offset + limit
142        let start = (self.offset as usize).min(sorted.len());
143        let end = (start + self.limit as usize).min(sorted.len());
144        let entries = &sorted[start..end];
145
146        if entries.is_empty() {
147            return Ok(Vec::new());
148        }
149
150        // Capture field_names from input chunks for output propagation
151        let field_names = input[0].field_names.clone();
152
153        // Build output chunks (up to 100 rows each)
154        let chunk_size = 100usize;
155        let mut output = Vec::new();
156        for chunk_start in (0..entries.len()).step_by(chunk_size) {
157            let chunk_end = (chunk_start + chunk_size).min(entries.len());
158            let size = chunk_end - chunk_start;
159            let mut fields = Vec::new();
160            for col in 0..num_fields {
161                let first_row = entries[chunk_start].row_idx;
162                let first_val = &all_values[col][first_row].0;
163                let phys_type = first_val.physical_type();
164                let mut v = ValueVector::new(phys_type, size);
165                v.resize(size);
166                for (out_idx, entry) in entries[chunk_start..chunk_end].iter().enumerate() {
167                    let (ref val, is_null) = all_values[col][entry.row_idx];
168                    if is_null || matches!(val, Value::Null) {
169                        v.set_null(out_idx, true);
170                    } else {
171                        store_value_in_vector(&mut v, out_idx, val)?;
172                    }
173                }
174                fields.push(v);
175            }
176            let arrow_fields = fields
177                .iter()
178                .map(|v| akar_common::arrow_vector::ArrowVector::from_legacy(v).array)
179                .collect::<Vec<_>>();
180            let arrow_field_types = fields.iter().map(|v| v.physical_type()).collect::<Vec<_>>();
181            output.push(DataChunk::new(arrow_fields, arrow_field_types).with_names(field_names.clone()));
182        }
183
184        Ok(output)
185    }
186}
akar_processor/physical/order_aggregate/topk.rs

akar_processor/physical/order_aggregate/
topk.rs
