Coverage Report

Created: 2024-10-13 08:39

/Users/andrewlamb/Software/datafusion/datafusion/physical-plan/src/aggregates/topk/heap.rs
Line
Count
Source (jump to first uncovered line)
1
// Licensed to the Apache Software Foundation (ASF) under one
2
// or more contributor license agreements.  See the NOTICE file
3
// distributed with this work for additional information
4
// regarding copyright ownership.  The ASF licenses this file
5
// to you under the Apache License, Version 2.0 (the
6
// "License"); you may not use this file except in compliance
7
// with the License.  You may obtain a copy of the License at
8
//
9
//   http://www.apache.org/licenses/LICENSE-2.0
10
//
11
// Unless required by applicable law or agreed to in writing,
12
// software distributed under the License is distributed on an
13
// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
14
// KIND, either express or implied.  See the License for the
15
// specific language governing permissions and limitations
16
// under the License.
17
18
//! A custom binary heap implementation for performant top K aggregation
19
20
use arrow::datatypes::i256;
21
use arrow_array::cast::AsArray;
22
use arrow_array::{downcast_primitive, ArrayRef, ArrowPrimitiveType, PrimitiveArray};
23
use arrow_buffer::{IntervalDayTime, IntervalMonthDayNano};
24
use arrow_schema::DataType;
25
use datafusion_common::DataFusionError;
26
use datafusion_common::Result;
27
use datafusion_physical_expr::aggregate::utils::adjust_output_array;
28
use half::f16;
29
use std::cmp::Ordering;
30
use std::fmt::{Debug, Display, Formatter};
31
use std::sync::Arc;
32
33
/// A custom version of `Ord` that only exists to we can implement it for the Values in our heap
34
pub trait Comparable {
35
    fn comp(&self, other: &Self) -> Ordering;
36
}
37
38
impl Comparable for Option<String> {
39
0
    fn comp(&self, other: &Self) -> Ordering {
40
0
        self.cmp(other)
41
0
    }
42
}
43
44
/// A "type alias" for Values which are stored in our heap
45
pub trait ValueType: Comparable + Clone + Debug {}
46
47
impl<T> ValueType for T where T: Comparable + Clone + Debug {}
48
49
/// An entry in our heap, which contains both the value and a index into an external HashTable
50
struct HeapItem<VAL: ValueType> {
51
    val: VAL,
52
    map_idx: usize,
53
}
54
55
/// A custom heap implementation that allows several things that couldn't be achieved with
56
/// `collections::BinaryHeap`:
57
/// 1. It allows values to be updated at arbitrary positions (when group values change)
58
/// 2. It can be either a min or max heap
59
/// 3. It can use our `HeapItem` type & `Comparable` trait
60
/// 4. It is specialized to grow to a certain limit, then always replace without grow & shrink
61
struct TopKHeap<VAL: ValueType> {
62
    desc: bool,
63
    len: usize,
64
    capacity: usize,
65
    heap: Vec<Option<HeapItem<VAL>>>,
66
}
67
68
/// An interface to hide the generic type signature of TopKHeap behind arrow arrays
69
pub trait ArrowHeap {
70
    fn set_batch(&mut self, vals: ArrayRef);
71
    fn is_worse(&self, idx: usize) -> bool;
72
    fn worst_map_idx(&self) -> usize;
73
    fn renumber(&mut self, heap_to_map: &[(usize, usize)]);
74
    fn insert(&mut self, row_idx: usize, map_idx: usize, map: &mut Vec<(usize, usize)>);
75
    fn replace_if_better(
76
        &mut self,
77
        heap_idx: usize,
78
        row_idx: usize,
79
        map: &mut Vec<(usize, usize)>,
80
    );
81
    fn drain(&mut self) -> (ArrayRef, Vec<usize>);
82
}
83
84
/// An implementation of `ArrowHeap` that deals with primitive values
85
pub struct PrimitiveHeap<VAL: ArrowPrimitiveType>
86
where
87
    <VAL as ArrowPrimitiveType>::Native: Comparable,
88
{
89
    batch: ArrayRef,
90
    heap: TopKHeap<VAL::Native>,
91
    desc: bool,
92
    data_type: DataType,
93
}
94
95
impl<VAL: ArrowPrimitiveType> PrimitiveHeap<VAL>
96
where
97
    <VAL as ArrowPrimitiveType>::Native: Comparable,
98
{
99
10
    pub fn new(limit: usize, desc: bool, data_type: DataType) -> Self {
100
10
        let owned: ArrayRef = Arc::new(PrimitiveArray::<VAL>::builder(0).finish());
101
10
        Self {
102
10
            batch: owned,
103
10
            heap: TopKHeap::new(limit, desc),
104
10
            desc,
105
10
            data_type,
106
10
        }
107
10
    }
108
}
109
110
impl<VAL: ArrowPrimitiveType> ArrowHeap for PrimitiveHeap<VAL>
111
where
112
    <VAL as ArrowPrimitiveType>::Native: Comparable,
113
{
114
10
    fn set_batch(&mut self, vals: ArrayRef) {
115
10
        self.batch = vals;
116
10
    }
117
118
20
    fn is_worse(&self, row_idx: usize) -> bool {
119
20
        if !self.heap.is_full() {
120
15
            return false;
121
5
        }
122
5
        let vals = self.batch.as_primitive::<VAL>();
123
5
        let new_val = vals.value(row_idx);
124
5
        let worst_val = self.heap.worst_val().expect("Missing root");
125
5
        (!self.desc && 
new_val > *worst_val2
) || (
self.desc4
&&
new_val < *worst_val3
)
126
20
    }
127
128
18
    fn worst_map_idx(&self) -> usize {
129
18
        self.heap.worst_map_idx()
130
18
    }
131
132
13
    fn renumber(&mut self, heap_to_map: &[(usize, usize)]) {
133
13
        self.heap.renumber(heap_to_map);
134
13
    }
135
136
13
    fn insert(&mut self, row_idx: usize, map_idx: usize, map: &mut Vec<(usize, usize)>) {
137
13
        let vals = self.batch.as_primitive::<VAL>();
138
13
        let new_val = vals.value(row_idx);
139
13
        self.heap.append_or_replace(new_val, map_idx, map);
140
13
    }
141
142
5
    fn replace_if_better(
143
5
        &mut self,
144
5
        heap_idx: usize,
145
5
        row_idx: usize,
146
5
        map: &mut Vec<(usize, usize)>,
147
5
    ) {
148
5
        let vals = self.batch.as_primitive::<VAL>();
149
5
        let new_val = vals.value(row_idx);
150
5
        self.heap.replace_if_better(heap_idx, new_val, map);
151
5
    }
152
153
10
    fn drain(&mut self) -> (ArrayRef, Vec<usize>) {
154
10
        let (vals, map_idxs) = self.heap.drain();
155
10
        let vals = Arc::new(PrimitiveArray::<VAL>::from_iter_values(vals));
156
10
        let vals = adjust_output_array(&self.data_type, vals).expect("Type is incorrect");
157
10
        (vals, map_idxs)
158
10
    }
159
}
160
161
impl<VAL: ValueType> TopKHeap<VAL> {
162
17
    pub fn new(limit: usize, desc: bool) -> Self {
163
17
        Self {
164
17
            desc,
165
17
            capacity: limit,
166
17
            len: 0,
167
89
            heap: (0..=limit).map(|_| None).collect::<Vec<_>>(),
168
17
        }
169
17
    }
170
171
6
    pub fn worst_val(&self) -> Option<&VAL> {
172
6
        let root = self.heap.first()
?0
;
173
6
        let hi = match root {
174
0
            None => return None,
175
6
            Some(hi) => hi,
176
6
        };
177
6
        Some(&hi.val)
178
6
    }
179
180
19
    pub fn worst_map_idx(&self) -> usize {
181
19
        self.heap[0].as_ref().map(|hi| 
hi.map_idx9
).unwrap_or(0)
182
19
    }
183
184
50
    pub fn is_full(&self) -> bool {
185
50
        self.len >= self.capacity
186
50
    }
187
188
38
    pub fn len(&self) -> usize {
189
38
        self.len
190
38
    }
191
192
30
    pub fn append_or_replace(
193
30
        &mut self,
194
30
        new_val: VAL,
195
30
        map_idx: usize,
196
30
        map: &mut Vec<(usize, usize)>,
197
30
    ) {
198
30
        if self.is_full() {
199
3
            self.replace_root(new_val, map_idx, map);
200
27
        } else {
201
27
            self.append(new_val, map_idx, map);
202
27
        }
203
30
    }
204
205
27
    fn append(&mut self, new_val: VAL, map_idx: usize, mapper: &mut Vec<(usize, usize)>) {
206
27
        let hi = HeapItem::new(new_val, map_idx);
207
27
        self.heap[self.len] = Some(hi);
208
27
        self.heapify_up(self.len, mapper);
209
27
        self.len += 1;
210
27
    }
211
212
24
    fn pop(&mut self, map: &mut Vec<(usize, usize)>) -> Option<HeapItem<VAL>> {
213
24
        if self.len() == 0 {
214
11
            return None;
215
13
        }
216
13
        if self.len() == 1 {
217
11
            self.len = 0;
218
11
            return self.heap[0].take();
219
2
        }
220
2
        self.swap(0, self.len - 1, map);
221
2
        let former_root = self.heap[self.len - 1].take();
222
2
        self.len -= 1;
223
2
        self.heapify_down(0, map);
224
2
        former_root
225
24
    }
226
227
11
    pub fn drain(&mut self) -> (Vec<VAL>, Vec<usize>) {
228
11
        let mut map = Vec::with_capacity(self.len);
229
11
        let mut vals = Vec::with_capacity(self.len);
230
11
        let mut map_idxs = Vec::with_capacity(self.len);
231
24
        while let Some(
worst_hi13
) = self.pop(&mut map) {
232
13
            vals.push(worst_hi.val);
233
13
            map_idxs.push(worst_hi.map_idx);
234
13
        }
235
11
        vals.reverse();
236
11
        map_idxs.reverse();
237
11
        (vals, map_idxs)
238
11
    }
239
240
3
    fn replace_root(
241
3
        &mut self,
242
3
        new_val: VAL,
243
3
        map_idx: usize,
244
3
        mapper: &mut Vec<(usize, usize)>,
245
3
    ) {
246
3
        let hi = self.heap[0].as_mut().expect("No root");
247
3
        hi.val = new_val;
248
3
        hi.map_idx = map_idx;
249
3
        self.heapify_down(0, mapper);
250
3
    }
251
252
6
    pub fn replace_if_better(
253
6
        &mut self,
254
6
        heap_idx: usize,
255
6
        new_val: VAL,
256
6
        mapper: &mut Vec<(usize, usize)>,
257
6
    ) {
258
6
        let existing = self.heap[heap_idx].as_mut().expect("Missing heap item");
259
6
        if (!self.desc && 
new_val.comp(&existing.val) != Ordering::Less3
)
260
5
            || (self.desc && 
new_val.comp(&existing.val) != Ordering::Greater3
)
261
        {
262
2
            return;
263
4
        }
264
4
        existing.val = new_val;
265
4
        self.heapify_down(heap_idx, mapper);
266
6
    }
267
268
14
    pub fn renumber(&mut self, heap_to_map: &[(usize, usize)]) {
269
14
        for (
heap_idx, map_idx2
) in heap_to_map.iter() {
270
2
            if let Some(Some(hi)) = self.heap.get_mut(*heap_idx) {
271
2
                hi.map_idx = *map_idx;
272
2
            }
0
273
        }
274
14
    }
275
276
27
    fn heapify_up(&mut self, mut idx: usize, mapper: &mut Vec<(usize, usize)>) {
277
27
        let desc = self.desc;
278
37
        while idx != 0 {
279
11
            let parent_idx = (idx - 1) / 2;
280
11
            let node = self.heap[idx].as_ref().expect("No heap item");
281
11
            let parent = self.heap[parent_idx].as_ref().expect("No heap item");
282
11
            if (!desc && 
node.val.comp(&parent.val) != Ordering::Greater10
)
283
11
                || (desc && 
node.val.comp(&parent.val) != Ordering::Less1
)
284
            {
285
1
                return;
286
10
            }
287
10
            self.swap(idx, parent_idx, mapper);
288
10
            idx = parent_idx;
289
        }
290
27
    }
291
292
14
    fn swap(&mut self, a_idx: usize, b_idx: usize, mapper: &mut Vec<(usize, usize)>) {
293
14
        let a_hi = self.heap[a_idx].take().expect("Missing heap entry");
294
14
        let b_hi = self.heap[b_idx].take().expect("Missing heap entry");
295
14
296
14
        mapper.push((a_hi.map_idx, b_idx));
297
14
        mapper.push((b_hi.map_idx, a_idx));
298
14
299
14
        self.heap[a_idx] = Some(b_hi);
300
14
        self.heap[b_idx] = Some(a_hi);
301
14
    }
302
303
11
    fn heapify_down(&mut self, node_idx: usize, mapper: &mut Vec<(usize, usize)>) {
304
11
        let left_child = node_idx * 2 + 1;
305
11
        let desc = self.desc;
306
11
        let entry = self.heap.get(node_idx).expect("Missing node!");
307
11
        let entry = entry.as_ref().expect("Missing node!");
308
11
        let mut best_idx = node_idx;
309
11
        let mut best_val = &entry.val;
310
22
        for child_idx in 
left_child..=left_child + 111
{
311
22
            if let Some(Some(
child3
)) = self.heap.get(child_idx) {
312
3
                if (!desc && child.val.comp(best_val) == Ordering::Greater)
313
0
                    || (desc && child.val.comp(best_val) == Ordering::Less)
314
3
                {
315
3
                    best_val = &child.val;
316
3
                    best_idx = child_idx;
317
3
                }
0
318
19
            }
319
        }
320
11
        if best_val.comp(&entry.val) != Ordering::Equal {
321
2
            self.swap(best_idx, node_idx, mapper);
322
2
            self.heapify_down(best_idx, mapper);
323
9
        }
324
11
    }
325
326
25
    fn _tree_print(
327
25
        &self,
328
25
        idx: usize,
329
25
        prefix: String,
330
25
        is_tail: bool,
331
25
        output: &mut String,
332
25
    ) {
333
25
        if let Some(Some(hi)) = self.heap.get(idx) {
334
25
            let connector = if idx != 0 {
335
15
                if is_tail {
336
11
                    "└── "
337
                } else {
338
4
                    "├── "
339
                }
340
            } else {
341
10
                ""
342
            };
343
25
            output.push_str(&format!(
344
25
                "{}{}val={:?} idx={}, bucket={}\n",
345
25
                prefix, connector, hi.val, idx, hi.map_idx
346
25
            ));
347
25
            let new_prefix = if is_tail { 
""21
} else {
"│ "4
};
348
25
            let child_prefix = format!("{}{}", prefix, new_prefix);
349
25
350
25
            let left_idx = idx * 2 + 1;
351
25
            let right_idx = idx * 2 + 2;
352
25
353
25
            let left_exists = left_idx < self.len;
354
25
            let right_exists = right_idx < self.len;
355
25
356
25
            if left_exists {
357
11
                self._tree_print(left_idx, child_prefix.clone(), !right_exists, output);
358
14
            }
359
25
            if right_exists {
360
4
                self._tree_print(right_idx, child_prefix, true, output);
361
21
            }
362
0
        }
363
25
    }
364
}
365
366
impl<VAL: ValueType> Display for TopKHeap<VAL> {
367
10
    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
368
10
        let mut output = String::new();
369
10
        if self.heap.first().is_some() {
370
10
            self._tree_print(0, String::new(), true, &mut output);
371
10
        }
0
372
10
        write!(f, "{}", output)
373
10
    }
374
}
375
376
impl<VAL: ValueType> HeapItem<VAL> {
377
27
    pub fn new(val: VAL, buk_idx: usize) -> Self {
378
27
        Self {
379
27
            val,
380
27
            map_idx: buk_idx,
381
27
        }
382
27
    }
383
}
384
385
impl<VAL: ValueType> Debug for HeapItem<VAL> {
386
0
    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
387
0
        f.write_str("bucket=")?;
388
0
        Debug::fmt(&self.map_idx, f)?;
389
0
        f.write_str(" val=")?;
390
0
        Debug::fmt(&self.val, f)?;
391
0
        f.write_str("\n")?;
392
0
        Ok(())
393
0
    }
394
}
395
396
impl<VAL: ValueType> Eq for HeapItem<VAL> {}
397
398
impl<VAL: ValueType> PartialEq<Self> for HeapItem<VAL> {
399
0
    fn eq(&self, other: &Self) -> bool {
400
0
        self.cmp(other) == Ordering::Equal
401
0
    }
402
}
403
404
impl<VAL: ValueType> PartialOrd<Self> for HeapItem<VAL> {
405
0
    fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
406
0
        Some(self.cmp(other))
407
0
    }
408
}
409
410
impl<VAL: ValueType> Ord for HeapItem<VAL> {
411
0
    fn cmp(&self, other: &Self) -> Ordering {
412
0
        let res = self.val.comp(&other.val);
413
0
        if res != Ordering::Equal {
414
0
            return res;
415
0
        }
416
0
        self.map_idx.cmp(&other.map_idx)
417
0
    }
418
}
419
420
macro_rules! compare_float {
421
    ($($t:ty),+) => {
422
        $(impl Comparable for Option<$t> {
423
0
            fn comp(&self, other: &Self) -> Ordering {
424
0
                match (self, other) {
425
0
                    (Some(me), Some(other)) => me.total_cmp(other),
426
0
                    (Some(_), None) => Ordering::Greater,
427
0
                    (None, Some(_)) => Ordering::Less,
428
0
                    (None, None) => Ordering::Equal,
429
                }
430
0
            }
431
        })+
432
433
        $(impl Comparable for $t {
434
0
            fn comp(&self, other: &Self) -> Ordering {
435
0
                self.total_cmp(other)
436
0
            }
437
        })+
438
    };
439
}
440
441
macro_rules! compare_integer {
442
    ($($t:ty),+) => {
443
        $(impl Comparable for Option<$t> {
444
0
            fn comp(&self, other: &Self) -> Ordering {
445
0
                self.cmp(other)
446
0
            }
447
        })+
448
449
        $(impl Comparable for $t {
450
31
            fn comp(&self, other: &Self) -> Ordering {
451
31
                self.cmp(other)
452
31
            }
453
        })+
454
    };
455
}
456
457
compare_integer!(i8, i16, i32, i64, i128, i256);
458
compare_integer!(u8, u16, u32, u64);
459
compare_integer!(IntervalDayTime, IntervalMonthDayNano);
460
compare_float!(f16, f32, f64);
461
462
10
pub fn new_heap(
463
10
    limit: usize,
464
10
    desc: bool,
465
10
    vt: DataType,
466
10
) -> Result<Box<dyn ArrowHeap + Send>> {
467
    macro_rules! downcast_helper {
468
        ($vt:ty, $d:ident) => {
469
            return Ok(Box::new(PrimitiveHeap::<$vt>::new(limit, desc, vt)))
470
        };
471
    }
472
473
0
    downcast_primitive! {
474
0
        vt => (downcast_helper, vt),
475
0
        _ => {}
476
0
    }
477
0
478
0
    Err(DataFusionError::Execution(format!(
479
0
        "Can't group type: {vt:?}"
480
0
    )))
481
10
}
482
483
#[cfg(test)]
484
mod tests {
485
    use super::*;
486
487
    #[test]
488
1
    fn should_append() -> Result<()> {
489
1
        let mut map = vec![];
490
1
        let mut heap = TopKHeap::new(10, false);
491
1
        heap.append_or_replace(1, 1, &mut map);
492
1
493
1
        let actual = heap.to_string();
494
1
        let expected = r#"
495
1
val=1 idx=0, bucket=1
496
1
        "#;
497
1
        assert_eq!(actual.trim(), expected.trim());
498
499
1
        Ok(())
500
1
    }
501
502
    #[test]
503
1
    fn should_heapify_up() -> Result<()> {
504
1
        let mut map = vec![];
505
1
        let mut heap = TopKHeap::new(10, false);
506
1
507
1
        heap.append_or_replace(1, 1, &mut map);
508
1
        assert_eq!(map, vec![]);
509
510
1
        heap.append_or_replace(2, 2, &mut map);
511
1
        assert_eq!(map, vec![(2, 0), (1, 1)]);
512
513
1
        let actual = heap.to_string();
514
1
        let expected = r#"
515
1
val=2 idx=0, bucket=2
516
1
└── val=1 idx=1, bucket=1
517
1
        "#;
518
1
        assert_eq!(actual.trim(), expected.trim());
519
520
1
        Ok(())
521
1
    }
522
523
    #[test]
524
1
    fn should_heapify_down() -> Result<()> {
525
1
        let mut map = vec![];
526
1
        let mut heap = TopKHeap::new(3, false);
527
1
528
1
        heap.append_or_replace(1, 1, &mut map);
529
1
        heap.append_or_replace(2, 2, &mut map);
530
1
        heap.append_or_replace(3, 3, &mut map);
531
1
        let actual = heap.to_string();
532
1
        let expected = r#"
533
1
val=3 idx=0, bucket=3
534
1
├── val=1 idx=1, bucket=1
535
1
└── val=2 idx=2, bucket=2
536
1
        "#;
537
1
        assert_eq!(actual.trim(), expected.trim());
538
539
1
        let mut map = vec![];
540
1
        heap.append_or_replace(0, 0, &mut map);
541
1
        let actual = heap.to_string();
542
1
        let expected = r#"
543
1
val=2 idx=0, bucket=2
544
1
├── val=1 idx=1, bucket=1
545
1
└── val=0 idx=2, bucket=0
546
1
        "#;
547
1
        assert_eq!(actual.trim(), expected.trim());
548
1
        assert_eq!(map, vec![(2, 0), (0, 2)]);
549
550
1
        Ok(())
551
1
    }
552
553
    #[test]
554
1
    fn should_replace() -> Result<()> {
555
1
        let mut map = vec![];
556
1
        let mut heap = TopKHeap::new(4, false);
557
1
558
1
        heap.append_or_replace(1, 1, &mut map);
559
1
        heap.append_or_replace(2, 2, &mut map);
560
1
        heap.append_or_replace(3, 3, &mut map);
561
1
        heap.append_or_replace(4, 4, &mut map);
562
1
        let actual = heap.to_string();
563
1
        let expected = r#"
564
1
val=4 idx=0, bucket=4
565
1
├── val=3 idx=1, bucket=3
566
1
│   └── val=1 idx=3, bucket=1
567
1
└── val=2 idx=2, bucket=2
568
1
        "#;
569
1
        assert_eq!(actual.trim(), expected.trim());
570
571
1
        let mut map = vec![];
572
1
        heap.replace_if_better(1, 0, &mut map);
573
1
        let actual = heap.to_string();
574
1
        let expected = r#"
575
1
val=4 idx=0, bucket=4
576
1
├── val=1 idx=1, bucket=1
577
1
│   └── val=0 idx=3, bucket=3
578
1
└── val=2 idx=2, bucket=2
579
1
        "#;
580
1
        assert_eq!(actual.trim(), expected.trim());
581
1
        assert_eq!(map, vec![(1, 1), (3, 3)]);
582
583
1
        Ok(())
584
1
    }
585
586
    #[test]
587
1
    fn should_find_worst() -> Result<()> {
588
1
        let mut map = vec![];
589
1
        let mut heap = TopKHeap::new(10, false);
590
1
591
1
        heap.append_or_replace(1, 1, &mut map);
592
1
        heap.append_or_replace(2, 2, &mut map);
593
1
594
1
        let actual = heap.to_string();
595
1
        let expected = r#"
596
1
val=2 idx=0, bucket=2
597
1
└── val=1 idx=1, bucket=1
598
1
        "#;
599
1
        assert_eq!(actual.trim(), expected.trim());
600
601
1
        assert_eq!(heap.worst_val(), Some(&2));
602
1
        assert_eq!(heap.worst_map_idx(), 2);
603
604
1
        Ok(())
605
1
    }
606
607
    #[test]
608
1
    fn should_drain() -> Result<()> {
609
1
        let mut map = vec![];
610
1
        let mut heap = TopKHeap::new(10, false);
611
1
612
1
        heap.append_or_replace(1, 1, &mut map);
613
1
        heap.append_or_replace(2, 2, &mut map);
614
1
615
1
        let actual = heap.to_string();
616
1
        let expected = r#"
617
1
val=2 idx=0, bucket=2
618
1
└── val=1 idx=1, bucket=1
619
1
        "#;
620
1
        assert_eq!(actual.trim(), expected.trim());
621
622
1
        let (vals, map_idxs) = heap.drain();
623
1
        assert_eq!(vals, vec![1, 2]);
624
1
        assert_eq!(map_idxs, vec![1, 2]);
625
1
        assert_eq!(heap.len(), 0);
626
627
1
        Ok(())
628
1
    }
629
630
    #[test]
631
1
    fn should_renumber() -> Result<()> {
632
1
        let mut map = vec![];
633
1
        let mut heap = TopKHeap::new(10, false);
634
1
635
1
        heap.append_or_replace(1, 1, &mut map);
636
1
        heap.append_or_replace(2, 2, &mut map);
637
1
638
1
        let actual = heap.to_string();
639
1
        let expected = r#"
640
1
val=2 idx=0, bucket=2
641
1
└── val=1 idx=1, bucket=1
642
1
        "#;
643
1
        assert_eq!(actual.trim(), expected.trim());
644
645
1
        let numbers = vec![(0, 1), (1, 2)];
646
1
        heap.renumber(numbers.as_slice());
647
1
        let actual = heap.to_string();
648
1
        let expected = r#"
649
1
val=2 idx=0, bucket=1
650
1
└── val=1 idx=1, bucket=2
651
1
        "#;
652
1
        assert_eq!(actual.trim(), expected.trim());
653
654
1
        Ok(())
655
1
    }
656
}