Skip to content

Commit a43cf79

Browse files
authored
remove termtree dependency (#11416)
* remove termtree dependency * impl Display for TopKHeap, replace uses of tree_print in tests * use to_string instead of format!
1 parent f5d88d1 commit a43cf79

File tree

2 files changed

+55
-32
lines changed

2 files changed

+55
-32
lines changed

datafusion/physical-plan/Cargo.toml

Lines changed: 0 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -66,7 +66,6 @@ tokio = { workspace = true }
6666
[dev-dependencies]
6767
rstest = { workspace = true }
6868
rstest_reuse = "0.7.0"
69-
termtree = "0.5.0"
7069
tokio = { workspace = true, features = [
7170
"rt-multi-thread",
7271
"fs",

datafusion/physical-plan/src/aggregates/topk/heap.rs

Lines changed: 55 additions & 31 deletions
Original file line numberDiff line numberDiff line change
@@ -27,7 +27,7 @@ use datafusion_common::Result;
2727
use datafusion_physical_expr::aggregate::utils::adjust_output_array;
2828
use half::f16;
2929
use std::cmp::Ordering;
30-
use std::fmt::{Debug, Formatter};
30+
use std::fmt::{Debug, Display, Formatter};
3131
use std::sync::Arc;
3232

3333
/// A custom version of `Ord` that only exists to we can implement it for the Values in our heap
@@ -323,29 +323,53 @@ impl<VAL: ValueType> TopKHeap<VAL> {
323323
}
324324
}
325325

326-
#[cfg(test)]
327-
fn _tree_print(&self, idx: usize) -> Option<termtree::Tree<String>> {
328-
let hi = self.heap.get(idx)?;
329-
match hi {
330-
None => None,
331-
Some(hi) => {
332-
let label =
333-
format!("val={:?} idx={}, bucket={}", hi.val, idx, hi.map_idx);
334-
let left = self._tree_print(idx * 2 + 1);
335-
let right = self._tree_print(idx * 2 + 2);
336-
let children = left.into_iter().chain(right);
337-
let me = termtree::Tree::new(label).with_leaves(children);
338-
Some(me)
326+
fn _tree_print(
327+
&self,
328+
idx: usize,
329+
prefix: String,
330+
is_tail: bool,
331+
output: &mut String,
332+
) {
333+
if let Some(Some(hi)) = self.heap.get(idx) {
334+
let connector = if idx != 0 {
335+
if is_tail {
336+
"└── "
337+
} else {
338+
"├── "
339+
}
340+
} else {
341+
""
342+
};
343+
output.push_str(&format!(
344+
"{}{}val={:?} idx={}, bucket={}\n",
345+
prefix, connector, hi.val, idx, hi.map_idx
346+
));
347+
let new_prefix = if is_tail { "" } else { "│ " };
348+
let child_prefix = format!("{}{}", prefix, new_prefix);
349+
350+
let left_idx = idx * 2 + 1;
351+
let right_idx = idx * 2 + 2;
352+
353+
let left_exists = left_idx < self.len;
354+
let right_exists = right_idx < self.len;
355+
356+
if left_exists {
357+
self._tree_print(left_idx, child_prefix.clone(), !right_exists, output);
358+
}
359+
if right_exists {
360+
self._tree_print(right_idx, child_prefix, true, output);
339361
}
340362
}
341363
}
364+
}
342365

343-
#[cfg(test)]
344-
fn tree_print(&self) -> String {
345-
match self._tree_print(0) {
346-
None => "".to_string(),
347-
Some(root) => format!("{}", root),
366+
impl<VAL: ValueType> Display for TopKHeap<VAL> {
367+
fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
368+
let mut output = String::new();
369+
if self.heap.first().is_some() {
370+
self._tree_print(0, String::new(), true, &mut output);
348371
}
372+
write!(f, "{}", output)
349373
}
350374
}
351375

@@ -361,9 +385,9 @@ impl<VAL: ValueType> HeapItem<VAL> {
361385
impl<VAL: ValueType> Debug for HeapItem<VAL> {
362386
fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
363387
f.write_str("bucket=")?;
364-
self.map_idx.fmt(f)?;
388+
Debug::fmt(&self.map_idx, f)?;
365389
f.write_str(" val=")?;
366-
self.val.fmt(f)?;
390+
Debug::fmt(&self.val, f)?;
367391
f.write_str("\n")?;
368392
Ok(())
369393
}
@@ -462,7 +486,7 @@ mod tests {
462486
let mut heap = TopKHeap::new(10, false);
463487
heap.append_or_replace(1, 1, &mut map);
464488

465-
let actual = heap.tree_print();
489+
let actual = heap.to_string();
466490
let expected = r#"
467491
val=1 idx=0, bucket=1
468492
"#;
@@ -482,7 +506,7 @@ val=1 idx=0, bucket=1
482506
heap.append_or_replace(2, 2, &mut map);
483507
assert_eq!(map, vec![(2, 0), (1, 1)]);
484508

485-
let actual = heap.tree_print();
509+
let actual = heap.to_string();
486510
let expected = r#"
487511
val=2 idx=0, bucket=2
488512
└── val=1 idx=1, bucket=1
@@ -500,7 +524,7 @@ val=2 idx=0, bucket=2
500524
heap.append_or_replace(1, 1, &mut map);
501525
heap.append_or_replace(2, 2, &mut map);
502526
heap.append_or_replace(3, 3, &mut map);
503-
let actual = heap.tree_print();
527+
let actual = heap.to_string();
504528
let expected = r#"
505529
val=3 idx=0, bucket=3
506530
├── val=1 idx=1, bucket=1
@@ -510,7 +534,7 @@ val=3 idx=0, bucket=3
510534

511535
let mut map = vec![];
512536
heap.append_or_replace(0, 0, &mut map);
513-
let actual = heap.tree_print();
537+
let actual = heap.to_string();
514538
let expected = r#"
515539
val=2 idx=0, bucket=2
516540
├── val=1 idx=1, bucket=1
@@ -531,7 +555,7 @@ val=2 idx=0, bucket=2
531555
heap.append_or_replace(2, 2, &mut map);
532556
heap.append_or_replace(3, 3, &mut map);
533557
heap.append_or_replace(4, 4, &mut map);
534-
let actual = heap.tree_print();
558+
let actual = heap.to_string();
535559
let expected = r#"
536560
val=4 idx=0, bucket=4
537561
├── val=3 idx=1, bucket=3
@@ -542,7 +566,7 @@ val=4 idx=0, bucket=4
542566

543567
let mut map = vec![];
544568
heap.replace_if_better(1, 0, &mut map);
545-
let actual = heap.tree_print();
569+
let actual = heap.to_string();
546570
let expected = r#"
547571
val=4 idx=0, bucket=4
548572
├── val=1 idx=1, bucket=1
@@ -563,7 +587,7 @@ val=4 idx=0, bucket=4
563587
heap.append_or_replace(1, 1, &mut map);
564588
heap.append_or_replace(2, 2, &mut map);
565589

566-
let actual = heap.tree_print();
590+
let actual = heap.to_string();
567591
let expected = r#"
568592
val=2 idx=0, bucket=2
569593
└── val=1 idx=1, bucket=1
@@ -584,7 +608,7 @@ val=2 idx=0, bucket=2
584608
heap.append_or_replace(1, 1, &mut map);
585609
heap.append_or_replace(2, 2, &mut map);
586610

587-
let actual = heap.tree_print();
611+
let actual = heap.to_string();
588612
let expected = r#"
589613
val=2 idx=0, bucket=2
590614
└── val=1 idx=1, bucket=1
@@ -607,7 +631,7 @@ val=2 idx=0, bucket=2
607631
heap.append_or_replace(1, 1, &mut map);
608632
heap.append_or_replace(2, 2, &mut map);
609633

610-
let actual = heap.tree_print();
634+
let actual = heap.to_string();
611635
let expected = r#"
612636
val=2 idx=0, bucket=2
613637
└── val=1 idx=1, bucket=1
@@ -616,7 +640,7 @@ val=2 idx=0, bucket=2
616640

617641
let numbers = vec![(0, 1), (1, 2)];
618642
heap.renumber(numbers.as_slice());
619-
let actual = heap.tree_print();
643+
let actual = heap.to_string();
620644
let expected = r#"
621645
val=2 idx=0, bucket=1
622646
└── val=1 idx=1, bucket=2

0 commit comments

Comments
 (0)