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        // Stream rows into a bounded max-heap directly from the chunks instead
98        // of materializing all values into a full O(rows × fields) matrix. Only
99        // the sort-key values are needed for the heap (P79).
100        let mut heap: BinaryHeap<TopKHeapEntry> = BinaryHeap::with_capacity(capacity.min(total_rows) + 1);
101        let mut row_idx: usize = 0;
102        for chunk in &input {
103            for row in 0..chunk.size {
104                let sort_key: Vec<DirectedSortKey> = self
105                    .sort_keys
106                    .iter()
107                    .map(|&(col, asc)| {
108                        let val = if col as usize >= num_fields {
109                            Value::Null
110                        } else {
111                            chunk.get_value(col as usize, row).unwrap_or(Value::Null)
112                        };
113                        if asc {
114                            DirectedSortKey::Asc(val)
115                        } else {
116                            DirectedSortKey::Desc(val)
117                        }
118                    })
119                    .collect();
120                heap.push(TopKHeapEntry { sort_key, row_idx });
121                row_idx += 1;
122                if heap.len() > capacity {
123                    heap.pop();
124                }
125            }
126        }
127
128        // into_sorted_vec returns ascending DirectedSortKey order = best-first
129        let sorted: Vec<TopKHeapEntry> = heap.into_sorted_vec();
130
131        // Apply offset + limit
132        let start = (self.offset as usize).min(sorted.len());
133        let end = (start + self.limit as usize).min(sorted.len());
134        let entries = &sorted[start..end];
135
136        if entries.is_empty() {
137            return Ok(Vec::new());
138        }
139
140        // Capture field_names from input chunks for output propagation
141        let field_names = input[0].field_names.clone();
142
143        // Build output chunks (up to 100 rows each). Slice the source fields
144        // via Arrow `take`, preserving complex types (List/Struct).
145        let global_indices: Vec<usize> = entries.iter().map(|e| e.row_idx).collect();
146        let chunk_size = 100usize;
147        let mut output = Vec::new();
148        for chunk_start in (0..global_indices.len()).step_by(chunk_size) {
149            let chunk_end = (chunk_start + chunk_size).min(global_indices.len());
150            output.push(take_global_rows(
151                &input,
152                &global_indices[chunk_start..chunk_end],
153                field_names.clone(),
154            )?);
155        }
156
157        Ok(output)
158    }
159}
akar_processor/physical/order_aggregate/topk.rs

akar_processor/physical/order_aggregate/
topk.rs
