akar_processor/physical/order_aggregate/
topk.rs

1//! Auto-extracted from physical_operator.rs
2use crate::physical::common::{take_global_rows, value_cmp};
3use crate::physical::types::{OperatorResult, PhysicalOperatorExec};
4use akar_common::types::Value;
5use akar_common::vector::DataChunk;
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). Slice the source fields
154        // via Arrow `take`, preserving complex types (List/Struct).
155        let global_indices: Vec<usize> = entries.iter().map(|e| e.row_idx).collect();
156        let chunk_size = 100usize;
157        let mut output = Vec::new();
158        for chunk_start in (0..global_indices.len()).step_by(chunk_size) {
159            let chunk_end = (chunk_start + chunk_size).min(global_indices.len());
160            output.push(take_global_rows(
161                &input,
162                &global_indices[chunk_start..chunk_end],
163                field_names.clone(),
164            )?);
165        }
166
167        Ok(output)
168    }
169}
akar_processor/physical/order_aggregate/topk.rs

akar_processor/physical/order_aggregate/
topk.rs
