trueno_db/
topk.rs

1//! Top-K selection algorithms
2//!
3//! **Problem**: `ORDER BY ... LIMIT K` is O(N log N). Top-K selection is O(N).
4//!
5//! **Solution**: Min-heap based Top-K selection algorithm
6//!
7//! **Performance Impact** (1M files):
8//! - Full sort: 2.3 seconds
9//! - Top-K selection: 0.08 seconds
10//! - **Speedup**: 28.75x
11//!
12//! Toyota Way Principles:
13//! - **Kaizen**: Algorithmic improvement (O(N log N) → O(N))
14//! - **Muda elimination**: Avoid unnecessary full sort
15//! - **Genchi Genbutsu**: Actual performance measurements guide optimization
16//!
17//! References:
18//! - ../paiml-mcp-agent-toolkit/docs/specifications/trueno-db-integration-review-response.md Issue #2
19
20use crate::Error;
21use arrow::array::{
22    Array, ArrayRef, Float32Array, Float64Array, Int32Array, Int64Array, StringArray,
23};
24use arrow::compute::SortOptions;
25use arrow::record_batch::RecordBatch;
26use std::cmp::Ordering;
27use std::collections::BinaryHeap;
28use std::sync::Arc;
29
30/// Sort order for Top-K selection
31#[derive(Debug, Clone, Copy, PartialEq, Eq)]
32pub enum SortOrder {
33    /// Ascending order (smallest K values)
34    Ascending,
35    /// Descending order (largest K values)
36    Descending,
37}
38
39impl From<SortOrder> for SortOptions {
40    fn from(order: SortOrder) -> Self {
41        Self { descending: matches!(order, SortOrder::Descending), nulls_first: false }
42    }
43}
44
45/// Trait for Top-K selection on record batches
46pub trait TopKSelection {
47    /// Select top K rows by a specific column
48    ///
49    /// # Arguments
50    /// * `column_index` - Index of the column to sort by
51    /// * `k` - Number of rows to select
52    /// * `order` - Sort order (Ascending or Descending)
53    ///
54    /// # Returns
55    /// A new `RecordBatch` containing the top K rows
56    ///
57    /// # Errors
58    /// Returns error if:
59    /// - Column index is out of bounds
60    /// - Column data type is not sortable
61    /// - K is zero
62    ///
63    /// # Examples
64    ///
65    /// ```rust
66    /// use trueno_db::topk::{TopKSelection, SortOrder};
67    /// use arrow::array::{Float64Array, RecordBatch};
68    /// use arrow::datatypes::{DataType, Field, Schema};
69    /// use std::sync::Arc;
70    ///
71    /// # fn main() -> Result<(), Box<dyn std::error::Error>> {
72    /// let schema = Arc::new(Schema::new(vec![
73    ///     Field::new("score", DataType::Float64, false),
74    /// ]));
75    /// let batch = RecordBatch::try_new(
76    ///     schema,
77    ///     vec![Arc::new(Float64Array::from(vec![1.0, 5.0, 3.0, 9.0, 2.0]))],
78    /// )?;
79    ///
80    /// // Get top 3 highest scores
81    /// let top3 = batch.top_k(0, 3, SortOrder::Descending)?;
82    /// assert_eq!(top3.num_rows(), 3);
83    /// # Ok(())
84    /// # }
85    /// ```
86    fn top_k(&self, column_index: usize, k: usize, order: SortOrder) -> crate::Result<RecordBatch>;
87}
88
89impl TopKSelection for RecordBatch {
90    fn top_k(&self, column_index: usize, k: usize, order: SortOrder) -> crate::Result<RecordBatch> {
91        // Validate inputs
92        if k == 0 {
93            return Err(Error::InvalidInput("k must be greater than 0".to_string()));
94        }
95
96        if column_index >= self.num_columns() {
97            return Err(Error::InvalidInput(format!(
98                "Column index {} out of bounds (batch has {} columns)",
99                column_index,
100                self.num_columns()
101            )));
102        }
103
104        // If k >= num_rows, just sort and return all rows
105        if k >= self.num_rows() {
106            return sort_all_rows(self, column_index, order);
107        }
108
109        // Use heap-based Top-K selection
110        let column = self.column(column_index);
111        let indices = select_top_k_indices(column, k, order)?;
112
113        // Build result batch from selected indices
114        build_batch_from_indices(self, &indices)
115    }
116}
117
118/// Select top K indices using min-heap algorithm
119///
120/// Time complexity: O(N log K) where N = number of rows, K = selection size
121/// Space complexity: O(K) for the heap
122fn select_top_k_indices(
123    column: &ArrayRef,
124    k: usize,
125    order: SortOrder,
126) -> crate::Result<Vec<usize>> {
127    match column.data_type() {
128        arrow::datatypes::DataType::Int32 => {
129            let array = column.as_any().downcast_ref::<Int32Array>().ok_or_else(|| {
130                Error::Other("Failed to downcast Int32 column to Int32Array".to_string())
131            })?;
132            select_top_k_typed(array.len(), k, order, |i| array.is_null(i), |i| array.value(i))
133        }
134        arrow::datatypes::DataType::Int64 => {
135            let array = column.as_any().downcast_ref::<Int64Array>().ok_or_else(|| {
136                Error::Other("Failed to downcast Int64 column to Int64Array".to_string())
137            })?;
138            select_top_k_typed(array.len(), k, order, |i| array.is_null(i), |i| array.value(i))
139        }
140        arrow::datatypes::DataType::Float32 => {
141            let array = column.as_any().downcast_ref::<Float32Array>().ok_or_else(|| {
142                Error::Other("Failed to downcast Float32 column to Float32Array".to_string())
143            })?;
144            select_top_k_typed(array.len(), k, order, |i| array.is_null(i), |i| array.value(i))
145        }
146        arrow::datatypes::DataType::Float64 => {
147            let array = column.as_any().downcast_ref::<Float64Array>().ok_or_else(|| {
148                Error::Other("Failed to downcast Float64 column to Float64Array".to_string())
149            })?;
150            select_top_k_typed(array.len(), k, order, |i| array.is_null(i), |i| array.value(i))
151        }
152        dt => Err(Error::InvalidInput(format!("Top-K not supported for data type: {dt:?}"))),
153    }
154}
155
156// Heap item for descending order (min-heap: keep smallest at top, so we can find largest K)
157#[derive(Debug)]
158struct MinHeapItem<V> {
159    value: V,
160    index: usize,
161}
162
163impl<V: PartialOrd> PartialEq for MinHeapItem<V> {
164    fn eq(&self, other: &Self) -> bool {
165        self.value.partial_cmp(&other.value) == Some(Ordering::Equal)
166    }
167}
168
169impl<V: PartialOrd> Eq for MinHeapItem<V> {}
170
171impl<V: PartialOrd> Ord for MinHeapItem<V> {
172    fn cmp(&self, other: &Self) -> Ordering {
173        // Reverse comparison for min-heap (smallest at top)
174        other.value.partial_cmp(&self.value).unwrap_or(Ordering::Equal)
175    }
176}
177
178impl<V: PartialOrd> PartialOrd for MinHeapItem<V> {
179    fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
180        Some(self.cmp(other))
181    }
182}
183
184// Heap item for ascending order (max-heap: keep largest at top, so we can find smallest K)
185#[derive(Debug)]
186struct MaxHeapItem<V> {
187    value: V,
188    index: usize,
189}
190
191impl<V: PartialOrd> PartialEq for MaxHeapItem<V> {
192    fn eq(&self, other: &Self) -> bool {
193        self.value.partial_cmp(&other.value) == Some(Ordering::Equal)
194    }
195}
196
197impl<V: PartialOrd> Eq for MaxHeapItem<V> {}
198
199impl<V: PartialOrd> Ord for MaxHeapItem<V> {
200    fn cmp(&self, other: &Self) -> Ordering {
201        // Normal comparison for max-heap (largest at top)
202        self.value.partial_cmp(&other.value).unwrap_or(Ordering::Equal)
203    }
204}
205
206impl<V: PartialOrd> PartialOrd for MaxHeapItem<V> {
207    fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
208        Some(self.cmp(other))
209    }
210}
211
212/// Collect top-K indices from a min-heap (descending order: find largest K values)
213fn collect_top_k_descending<V: PartialOrd>(
214    len: usize,
215    k: usize,
216    is_null: impl Fn(usize) -> bool,
217    get_value: impl Fn(usize) -> V,
218) -> Vec<usize> {
219    let mut heap: BinaryHeap<MinHeapItem<V>> = BinaryHeap::with_capacity(k);
220
221    for index in 0..len {
222        if !is_null(index) {
223            let value = get_value(index);
224            if heap.len() < k {
225                heap.push(MinHeapItem { value, index });
226            } else if let Some(top) = heap.peek() {
227                if value.partial_cmp(&top.value) == Some(Ordering::Greater) {
228                    heap.pop();
229                    heap.push(MinHeapItem { value, index });
230                }
231            }
232        }
233    }
234
235    let mut result: Vec<_> = heap.into_vec();
236    result.sort_by(|a, b| b.value.partial_cmp(&a.value).unwrap_or(Ordering::Equal));
237    result.into_iter().map(|item| item.index).collect()
238}
239
240/// Collect top-K indices from a max-heap (ascending order: find smallest K values)
241fn collect_top_k_ascending<V: PartialOrd>(
242    len: usize,
243    k: usize,
244    is_null: impl Fn(usize) -> bool,
245    get_value: impl Fn(usize) -> V,
246) -> Vec<usize> {
247    let mut heap: BinaryHeap<MaxHeapItem<V>> = BinaryHeap::with_capacity(k);
248
249    for index in 0..len {
250        if !is_null(index) {
251            let value = get_value(index);
252            if heap.len() < k {
253                heap.push(MaxHeapItem { value, index });
254            } else if let Some(top) = heap.peek() {
255                if value.partial_cmp(&top.value) == Some(Ordering::Less) {
256                    heap.pop();
257                    heap.push(MaxHeapItem { value, index });
258                }
259            }
260        }
261    }
262
263    let mut result: Vec<_> = heap.into_vec();
264    result.sort_by(|a, b| a.value.partial_cmp(&b.value).unwrap_or(Ordering::Equal));
265    result.into_iter().map(|item| item.index).collect()
266}
267
268/// Generic top-K selection for any Arrow array with `PartialOrd` values
269#[allow(clippy::unnecessary_wraps)]
270fn select_top_k_typed<V: PartialOrd>(
271    len: usize,
272    k: usize,
273    order: SortOrder,
274    is_null: impl Fn(usize) -> bool,
275    get_value: impl Fn(usize) -> V,
276) -> crate::Result<Vec<usize>> {
277    let indices = match order {
278        SortOrder::Descending => collect_top_k_descending(len, k, is_null, get_value),
279        SortOrder::Ascending => collect_top_k_ascending(len, k, is_null, get_value),
280    };
281    Ok(indices)
282}
283
284/// Build a new record batch from selected row indices
285fn build_batch_from_indices(batch: &RecordBatch, indices: &[usize]) -> crate::Result<RecordBatch> {
286    use arrow::datatypes::DataType;
287
288    let mut new_columns: Vec<ArrayRef> = Vec::with_capacity(batch.num_columns());
289
290    for col_idx in 0..batch.num_columns() {
291        let column = batch.column(col_idx);
292
293        let new_array: ArrayRef = match column.data_type() {
294            DataType::Int32 => {
295                let array = column.as_any().downcast_ref::<Int32Array>().ok_or_else(|| {
296                    Error::Other("Failed to downcast Int32 column to Int32Array".to_string())
297                })?;
298                let values: Vec<i32> = indices.iter().map(|&idx| array.value(idx)).collect();
299                Arc::new(Int32Array::from(values))
300            }
301            DataType::Int64 => {
302                let array = column.as_any().downcast_ref::<Int64Array>().ok_or_else(|| {
303                    Error::Other("Failed to downcast Int64 column to Int64Array".to_string())
304                })?;
305                let values: Vec<i64> = indices.iter().map(|&idx| array.value(idx)).collect();
306                Arc::new(Int64Array::from(values))
307            }
308            DataType::Float32 => {
309                let array = column.as_any().downcast_ref::<Float32Array>().ok_or_else(|| {
310                    Error::Other("Failed to downcast Float32 column to Float32Array".to_string())
311                })?;
312                let values: Vec<f32> = indices.iter().map(|&idx| array.value(idx)).collect();
313                Arc::new(Float32Array::from(values))
314            }
315            DataType::Float64 => {
316                let array = column.as_any().downcast_ref::<Float64Array>().ok_or_else(|| {
317                    Error::Other("Failed to downcast Float64 column to Float64Array".to_string())
318                })?;
319                let values: Vec<f64> = indices.iter().map(|&idx| array.value(idx)).collect();
320                Arc::new(Float64Array::from(values))
321            }
322            DataType::Utf8 => {
323                let array = column.as_any().downcast_ref::<StringArray>().ok_or_else(|| {
324                    Error::Other("Failed to downcast Utf8 column to StringArray".to_string())
325                })?;
326                let values: Vec<&str> = indices.iter().map(|&idx| array.value(idx)).collect();
327                Arc::new(StringArray::from(values))
328            }
329            dt => {
330                return Err(Error::InvalidInput(format!(
331                    "Top-K not implemented for column data type: {dt:?}"
332                )));
333            }
334        };
335
336        new_columns.push(new_array);
337    }
338
339    RecordBatch::try_new(batch.schema(), new_columns)
340        .map_err(|e| Error::StorageError(format!("Failed to create result batch: {e}")))
341}
342
343/// Fallback: sort all rows when k >= `num_rows`
344fn sort_all_rows(
345    batch: &RecordBatch,
346    column_index: usize,
347    order: SortOrder,
348) -> crate::Result<RecordBatch> {
349    use arrow::compute::sort_to_indices;
350
351    let sort_options = SortOptions::from(order);
352    let indices = sort_to_indices(batch.column(column_index).as_ref(), Some(sort_options), None)
353        .map_err(|e| Error::StorageError(format!("Failed to sort: {e}")))?;
354
355    // Convert indices to usize vec
356    let indices_array =
357        indices.as_any().downcast_ref::<arrow::array::UInt32Array>().ok_or_else(|| {
358            Error::Other(
359                "Failed to downcast sort indices to UInt32Array (expected from sort_to_indices)"
360                    .to_string(),
361            )
362        })?;
363    let indices_vec: Vec<usize> =
364        (0..indices_array.len()).map(|i| indices_array.value(i) as usize).collect();
365
366    build_batch_from_indices(batch, &indices_vec)
367}
368
369#[cfg(test)]
370#[allow(
371    clippy::cast_possible_truncation,
372    clippy::cast_possible_wrap,
373    clippy::cast_precision_loss,
374    clippy::float_cmp,
375    clippy::redundant_closure
376)]
377mod tests {
378    use super::*;
379    use arrow::datatypes::{DataType, Field, Schema};
380    use std::sync::Arc;
381
382    fn create_test_batch(values: Vec<f64>) -> RecordBatch {
383        let schema = Arc::new(Schema::new(vec![
384            Field::new("id", DataType::Int32, false),
385            Field::new("score", DataType::Float64, false),
386        ]));
387
388        let ids: Vec<i32> = (0..values.len() as i32).collect();
389
390        RecordBatch::try_new(
391            schema,
392            vec![Arc::new(Int32Array::from(ids)), Arc::new(Float64Array::from(values))],
393        )
394        .unwrap()
395    }
396
397    #[test]
398    fn test_top_k_descending_basic() {
399        // Test: Get top 3 highest scores
400        let batch = create_test_batch(vec![1.0, 5.0, 3.0, 9.0, 2.0]);
401        let result = batch.top_k(1, 3, SortOrder::Descending).unwrap();
402
403        assert_eq!(result.num_rows(), 3);
404
405        let scores = result.column(1).as_any().downcast_ref::<Float64Array>().unwrap();
406        assert_eq!(scores.value(0), 9.0);
407        assert_eq!(scores.value(1), 5.0);
408        assert_eq!(scores.value(2), 3.0);
409    }
410
411    #[test]
412    fn test_top_k_ascending_basic() {
413        // Test: Get top 3 lowest scores
414        let batch = create_test_batch(vec![1.0, 5.0, 3.0, 9.0, 2.0]);
415        let result = batch.top_k(1, 3, SortOrder::Ascending).unwrap();
416
417        assert_eq!(result.num_rows(), 3);
418
419        let scores = result.column(1).as_any().downcast_ref::<Float64Array>().unwrap();
420        assert_eq!(scores.value(0), 1.0);
421        assert_eq!(scores.value(1), 2.0);
422        assert_eq!(scores.value(2), 3.0);
423    }
424
425    #[test]
426    fn test_top_k_k_equals_length() {
427        // Edge case: k equals number of rows (should return sorted batch)
428        let batch = create_test_batch(vec![3.0, 1.0, 2.0]);
429        let result = batch.top_k(1, 3, SortOrder::Descending).unwrap();
430
431        assert_eq!(result.num_rows(), 3);
432
433        let scores = result.column(1).as_any().downcast_ref::<Float64Array>().unwrap();
434        assert_eq!(scores.value(0), 3.0);
435        assert_eq!(scores.value(1), 2.0);
436        assert_eq!(scores.value(2), 1.0);
437    }
438
439    #[test]
440    fn test_top_k_k_greater_than_length() {
441        // Edge case: k > number of rows (should return all rows sorted)
442        let batch = create_test_batch(vec![3.0, 1.0, 2.0]);
443        let result = batch.top_k(1, 10, SortOrder::Descending).unwrap();
444
445        assert_eq!(result.num_rows(), 3);
446
447        let scores = result.column(1).as_any().downcast_ref::<Float64Array>().unwrap();
448        assert_eq!(scores.value(0), 3.0);
449        assert_eq!(scores.value(1), 2.0);
450        assert_eq!(scores.value(2), 1.0);
451    }
452
453    #[test]
454    fn test_top_k_k_zero_fails() {
455        // Error case: k = 0 should fail
456        let batch = create_test_batch(vec![1.0, 2.0, 3.0]);
457        let result = batch.top_k(1, 0, SortOrder::Descending);
458
459        assert!(result.is_err());
460        assert!(result.unwrap_err().to_string().contains("must be greater than 0"));
461    }
462
463    #[test]
464    fn test_top_k_invalid_column_index() {
465        // Error case: invalid column index
466        let batch = create_test_batch(vec![1.0, 2.0, 3.0]);
467        let result = batch.top_k(99, 2, SortOrder::Descending);
468
469        assert!(result.is_err());
470        assert!(result.unwrap_err().to_string().contains("out of bounds"));
471    }
472
473    #[test]
474    fn test_top_k_preserves_row_integrity() {
475        // Test: Ensure all columns stay aligned (row integrity)
476        let batch = create_test_batch(vec![1.0, 5.0, 3.0]);
477        let result = batch.top_k(1, 2, SortOrder::Descending).unwrap();
478
479        let ids = result.column(0).as_any().downcast_ref::<Int32Array>().unwrap();
480        let scores = result.column(1).as_any().downcast_ref::<Float64Array>().unwrap();
481
482        // Top 2: scores 5.0 (id=1) and 3.0 (id=2)
483        assert_eq!(scores.value(0), 5.0);
484        assert_eq!(ids.value(0), 1);
485
486        assert_eq!(scores.value(1), 3.0);
487        assert_eq!(ids.value(1), 2);
488    }
489
490    #[test]
491    fn test_top_k_large_dataset() {
492        // Performance test: 1M rows (should be O(N) vs O(N log N))
493        let values: Vec<f64> = (0..1_000_000).map(|i| f64::from(i)).collect();
494        let batch = create_test_batch(values);
495
496        let result = batch.top_k(1, 10, SortOrder::Descending).unwrap();
497
498        assert_eq!(result.num_rows(), 10);
499
500        let scores = result.column(1).as_any().downcast_ref::<Float64Array>().unwrap();
501        // Top 10 should be 999999, 999998, ..., 999990
502        for i in 0..10 {
503            assert_eq!(scores.value(i), 999_999.0 - i as f64);
504        }
505
506        // Performance assertion removed — flaky on shared CI runners.
507        // Target for release builds: <80ms for 1M rows.
508        // Verify with `cargo bench` not wall-clock assertions in tests.
509        // GH-739: 633ms on loaded runner triggered false failure.
510    }
511
512    // Property-based tests
513    #[cfg(test)]
514    mod property_tests {
515        use super::*;
516        use proptest::prelude::*;
517
518        proptest! {
519            /// Property: Top-K always returns exactly K rows (or fewer if input is smaller)
520            #[test]
521            fn prop_top_k_returns_k_rows(
522                values in prop::collection::vec(0.0f64..1000.0, 10..1000),
523                k in 1usize..100
524            ) {
525                let batch = create_test_batch(values.clone());
526                let result = batch.top_k(1, k, SortOrder::Descending).unwrap();
527
528                let expected_rows = k.min(values.len());
529                prop_assert_eq!(result.num_rows(), expected_rows);
530            }
531
532            /// Property: Top-K descending returns values in descending order
533            #[test]
534            fn prop_top_k_descending_is_sorted(
535                values in prop::collection::vec(0.0f64..1000.0, 10..1000),
536                k in 1usize..100
537            ) {
538                let batch = create_test_batch(values);
539                let result = batch.top_k(1, k, SortOrder::Descending).unwrap();
540
541                let scores = result.column(1).as_any().downcast_ref::<Float64Array>().unwrap();
542
543                // Check descending order
544                for i in 0..scores.len().saturating_sub(1) {
545                    prop_assert!(
546                        scores.value(i) >= scores.value(i + 1),
547                        "Not in descending order: {} < {}",
548                        scores.value(i),
549                        scores.value(i + 1)
550                    );
551                }
552            }
553
554            /// Property: Top-K ascending returns values in ascending order
555            #[test]
556            fn prop_top_k_ascending_is_sorted(
557                values in prop::collection::vec(0.0f64..1000.0, 10..1000),
558                k in 1usize..100
559            ) {
560                let batch = create_test_batch(values);
561                let result = batch.top_k(1, k, SortOrder::Ascending).unwrap();
562
563                let scores = result.column(1).as_any().downcast_ref::<Float64Array>().unwrap();
564
565                // Check ascending order
566                for i in 0..scores.len().saturating_sub(1) {
567                    prop_assert!(
568                        scores.value(i) <= scores.value(i + 1),
569                        "Not in ascending order: {} > {}",
570                        scores.value(i),
571                        scores.value(i + 1)
572                    );
573                }
574            }
575        }
576    }
577
578    // Additional tests for all data types
579    #[test]
580    fn test_top_k_int32() {
581        use arrow::array::Int32Array;
582        use arrow::datatypes::{DataType, Field, Schema};
583        use std::sync::Arc;
584
585        let schema = Schema::new(vec![Field::new("value", DataType::Int32, false)]);
586        let values = Int32Array::from(vec![5, 2, 8, 1, 9, 3]);
587        let batch = RecordBatch::try_new(Arc::new(schema), vec![Arc::new(values)]).unwrap();
588
589        let result = batch.top_k(0, 3, SortOrder::Descending).unwrap();
590        assert_eq!(result.num_rows(), 3);
591
592        let col = result.column(0).as_any().downcast_ref::<Int32Array>().unwrap();
593        assert_eq!(col.value(0), 9);
594        assert_eq!(col.value(1), 8);
595        assert_eq!(col.value(2), 5);
596    }
597
598    #[test]
599    fn test_top_k_int32_ascending() {
600        use arrow::array::Int32Array;
601        use arrow::datatypes::{DataType, Field, Schema};
602        use std::sync::Arc;
603
604        let schema = Schema::new(vec![Field::new("value", DataType::Int32, false)]);
605        let values = Int32Array::from(vec![5, 2, 8, 1, 9, 3]);
606        let batch = RecordBatch::try_new(Arc::new(schema), vec![Arc::new(values)]).unwrap();
607
608        let result = batch.top_k(0, 3, SortOrder::Ascending).unwrap();
609        assert_eq!(result.num_rows(), 3);
610
611        let col = result.column(0).as_any().downcast_ref::<Int32Array>().unwrap();
612        assert_eq!(col.value(0), 1);
613        assert_eq!(col.value(1), 2);
614        assert_eq!(col.value(2), 3);
615    }
616
617    #[test]
618    fn test_top_k_int64() {
619        use arrow::array::Int64Array;
620        use arrow::datatypes::{DataType, Field, Schema};
621        use std::sync::Arc;
622
623        let schema = Schema::new(vec![Field::new("value", DataType::Int64, false)]);
624        let values = Int64Array::from(vec![100i64, 200, 50, 300, 150]);
625        let batch = RecordBatch::try_new(Arc::new(schema), vec![Arc::new(values)]).unwrap();
626
627        let result = batch.top_k(0, 2, SortOrder::Ascending).unwrap();
628        assert_eq!(result.num_rows(), 2);
629
630        let col = result.column(0).as_any().downcast_ref::<Int64Array>().unwrap();
631        assert_eq!(col.value(0), 50);
632        assert_eq!(col.value(1), 100);
633    }
634
635    #[test]
636    fn test_top_k_int64_descending() {
637        use arrow::array::Int64Array;
638        use arrow::datatypes::{DataType, Field, Schema};
639        use std::sync::Arc;
640
641        let schema = Schema::new(vec![Field::new("value", DataType::Int64, false)]);
642        let values = Int64Array::from(vec![100i64, 200, 50, 300, 150]);
643        let batch = RecordBatch::try_new(Arc::new(schema), vec![Arc::new(values)]).unwrap();
644
645        let result = batch.top_k(0, 2, SortOrder::Descending).unwrap();
646        assert_eq!(result.num_rows(), 2);
647
648        let col = result.column(0).as_any().downcast_ref::<Int64Array>().unwrap();
649        assert_eq!(col.value(0), 300);
650        assert_eq!(col.value(1), 200);
651    }
652
653    #[test]
654    fn test_top_k_float32() {
655        use arrow::array::Float32Array;
656        use arrow::datatypes::{DataType, Field, Schema};
657        use std::sync::Arc;
658
659        let schema = Schema::new(vec![Field::new("value", DataType::Float32, false)]);
660        let values = Float32Array::from(vec![1.5f32, 2.7, 0.3, 4.2, 3.1]);
661        let batch = RecordBatch::try_new(Arc::new(schema), vec![Arc::new(values)]).unwrap();
662
663        let result = batch.top_k(0, 3, SortOrder::Descending).unwrap();
664        assert_eq!(result.num_rows(), 3);
665
666        let col = result.column(0).as_any().downcast_ref::<Float32Array>().unwrap();
667        assert!((col.value(0) - 4.2).abs() < 0.001);
668        assert!((col.value(1) - 3.1).abs() < 0.001);
669        assert!((col.value(2) - 2.7).abs() < 0.001);
670    }
671
672    #[test]
673    fn test_top_k_float32_ascending() {
674        use arrow::array::Float32Array;
675        use arrow::datatypes::{DataType, Field, Schema};
676        use std::sync::Arc;
677
678        let schema = Schema::new(vec![Field::new("value", DataType::Float32, false)]);
679        let values = Float32Array::from(vec![1.5f32, 2.7, 0.3, 4.2, 3.1]);
680        let batch = RecordBatch::try_new(Arc::new(schema), vec![Arc::new(values)]).unwrap();
681
682        let result = batch.top_k(0, 3, SortOrder::Ascending).unwrap();
683        assert_eq!(result.num_rows(), 3);
684
685        let col = result.column(0).as_any().downcast_ref::<Float32Array>().unwrap();
686        assert!((col.value(0) - 0.3).abs() < 0.001);
687        assert!((col.value(1) - 1.5).abs() < 0.001);
688        assert!((col.value(2) - 2.7).abs() < 0.001);
689    }
690
691    #[test]
692    fn test_top_k_unsupported_type() {
693        use arrow::array::StringArray;
694        use arrow::datatypes::{DataType, Field, Schema};
695        use std::sync::Arc;
696
697        let schema = Schema::new(vec![Field::new("value", DataType::Utf8, false)]);
698        let values = StringArray::from(vec!["a", "b", "c"]);
699        let batch = RecordBatch::try_new(Arc::new(schema), vec![Arc::new(values)]).unwrap();
700
701        let result = batch.top_k(0, 2, SortOrder::Descending);
702        assert!(result.is_err());
703        assert!(result.unwrap_err().to_string().contains("Top-K not supported for data type"));
704    }
705
706    // ========================================================================
707    // Heap Item Trait Tests (for coverage of MinHeapItem/MaxHeapItem)
708    // ========================================================================
709
710    #[test]
711    fn test_min_heap_item_eq() {
712        let item1 = MinHeapItem { value: 42i32, index: 0 };
713        let item2 = MinHeapItem { value: 42i32, index: 1 };
714        let item3 = MinHeapItem { value: 43i32, index: 2 };
715
716        assert_eq!(item1, item2);
717        assert_ne!(item1, item3);
718    }
719
720    #[test]
721    fn test_min_heap_item_ord() {
722        let item1 = MinHeapItem { value: 10i32, index: 0 };
723        let item2 = MinHeapItem { value: 20i32, index: 1 };
724        let item3 = MinHeapItem { value: 30i32, index: 2 };
725
726        // Min-heap: reverse ordering (smaller values at top)
727        assert!(item3 < item2); // 30 < 20 in min-heap ordering
728        assert!(item2 < item1); // 20 < 10 in min-heap ordering
729    }
730
731    #[test]
732    fn test_min_heap_item_partial_ord() {
733        let item1 = MinHeapItem { value: 5i32, index: 0 };
734        let item2 = MinHeapItem { value: 10i32, index: 1 };
735
736        assert!(item1.partial_cmp(&item2) == Some(Ordering::Greater));
737    }
738
739    #[test]
740    fn test_max_heap_item_eq() {
741        let item1 = MaxHeapItem { value: 42i32, index: 0 };
742        let item2 = MaxHeapItem { value: 42i32, index: 1 };
743        let item3 = MaxHeapItem { value: 43i32, index: 2 };
744
745        assert_eq!(item1, item2);
746        assert_ne!(item1, item3);
747    }
748
749    #[test]
750    fn test_max_heap_item_ord() {
751        let item1 = MaxHeapItem { value: 10i32, index: 0 };
752        let item2 = MaxHeapItem { value: 20i32, index: 1 };
753        let item3 = MaxHeapItem { value: 30i32, index: 2 };
754
755        // Max-heap: normal ordering (larger values at top)
756        assert!(item3 > item2);
757        assert!(item2 > item1);
758    }
759
760    #[test]
761    fn test_max_heap_item_partial_ord() {
762        let item1 = MaxHeapItem { value: 5i32, index: 0 };
763        let item2 = MaxHeapItem { value: 10i32, index: 1 };
764
765        assert!(item1.partial_cmp(&item2) == Some(Ordering::Less));
766    }
767
768    #[test]
769    fn test_heap_item_with_floats() {
770        let item1 = MinHeapItem { value: 1.5f64, index: 0 };
771        let item2 = MinHeapItem { value: 2.5f64, index: 1 };
772
773        assert_ne!(item1, item2);
774        assert!(item2 < item1); // Min-heap: reverse ordering
775    }
776
777    #[test]
778    fn test_heap_item_eq_method_with_floats() {
779        let item1 = MaxHeapItem { value: 3.25f64, index: 0 };
780        let item2 = MaxHeapItem { value: 3.25f64, index: 1 };
781        let item3 = MaxHeapItem { value: 2.75f64, index: 2 };
782
783        assert!(item1.eq(&item2));
784        assert!(!item1.eq(&item3));
785    }
786}
trueno_db/topk.rs

trueno_db/
topk.rs
