@@ -27,7 +27,7 @@ use datafusion_common::Result;
2727use datafusion_physical_expr:: aggregate:: utils:: adjust_output_array;
2828use half:: f16;
2929use std:: cmp:: Ordering ;
30- use std:: fmt:: { Debug , Formatter } ;
30+ use std:: fmt:: { Debug , Display , Formatter } ;
3131use 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> {
361385impl < 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#"
467491val=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#"
487511val=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#"
505529val=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#"
515539val=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#"
536560val=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#"
547571val=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#"
568592val=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#"
589613val=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#"
612636val=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#"
621645val=2 idx=0, bucket=1
622646└── val=1 idx=1, bucket=2
0 commit comments