1818//! Utilities for printing record batches. Note this module is not
1919//! available unless `feature = "prettyprint"` is enabled.
2020
21- use crate :: record_batch:: RecordBatch ;
21+ use crate :: { array :: ArrayRef , record_batch:: RecordBatch } ;
2222
2323use prettytable:: format;
2424use prettytable:: { Cell , Row , Table } ;
@@ -32,12 +32,23 @@ pub fn pretty_format_batches(results: &[RecordBatch]) -> Result<String> {
3232 Ok ( create_table ( results) ?. to_string ( ) )
3333}
3434
35+ ///! Create a visual representation of columns
36+ pub fn pretty_format_columns ( col_name : & str , results : & [ ArrayRef ] ) -> Result < String > {
37+ Ok ( create_column ( col_name, results) ?. to_string ( ) )
38+ }
39+
3540///! Prints a visual representation of record batches to stdout
3641pub fn print_batches ( results : & [ RecordBatch ] ) -> Result < ( ) > {
3742 create_table ( results) ?. printstd ( ) ;
3843 Ok ( ( ) )
3944}
4045
46+ ///! Prints a visual representation of a list of column to stdout
47+ pub fn print_columns ( col_name : & str , results : & [ ArrayRef ] ) -> Result < ( ) > {
48+ create_column ( col_name, results) ?. printstd ( ) ;
49+ Ok ( ( ) )
50+ }
51+
4152///! Convert a series of record batches into a table
4253fn create_table ( results : & [ RecordBatch ] ) -> Result < Table > {
4354 let mut table = Table :: new ( ) ;
@@ -69,6 +80,28 @@ fn create_table(results: &[RecordBatch]) -> Result<Table> {
6980 Ok ( table)
7081}
7182
83+ fn create_column ( field : & str , columns : & [ ArrayRef ] ) -> Result < Table > {
84+ let mut table = Table :: new ( ) ;
85+ table. set_format ( * format:: consts:: FORMAT_NO_LINESEP_WITH_TITLE ) ;
86+
87+ if columns. is_empty ( ) {
88+ return Ok ( table) ;
89+ }
90+
91+ let header = vec ! [ Cell :: new( field) ] ;
92+ table. set_titles ( Row :: new ( header) ) ;
93+
94+ for col in columns {
95+ for row in 0 ..col. len ( ) {
96+ let mut cells = Vec :: new ( ) ;
97+ cells. push ( Cell :: new ( & array_value_to_string ( & col, row) ?) ) ;
98+ table. add_row ( Row :: new ( cells) ) ;
99+ }
100+ }
101+
102+ Ok ( table)
103+ }
104+
72105#[ cfg( test) ]
73106mod tests {
74107 use crate :: {
@@ -132,6 +165,32 @@ mod tests {
132165 Ok ( ( ) )
133166 }
134167
168+ #[ test]
169+ fn test_pretty_format_columns ( ) -> Result < ( ) > {
170+ let columns = vec ! [
171+ Arc :: new( array:: StringArray :: from( vec![
172+ Some ( "a" ) ,
173+ Some ( "b" ) ,
174+ None ,
175+ Some ( "d" ) ,
176+ ] ) ) as ArrayRef ,
177+ Arc :: new( array:: StringArray :: from( vec![ Some ( "e" ) , None , Some ( "g" ) ] ) ) ,
178+ ] ;
179+
180+ let table = pretty_format_columns ( "a" , & columns) ?;
181+
182+ let expected = vec ! [
183+ "+---+" , "| a |" , "+---+" , "| a |" , "| b |" , "| |" , "| d |" , "| e |" ,
184+ "| |" , "| g |" , "+---+" ,
185+ ] ;
186+
187+ let actual: Vec < & str > = table. lines ( ) . collect ( ) ;
188+
189+ assert_eq ! ( expected, actual, "Actual result:\n {}" , table) ;
190+
191+ Ok ( ( ) )
192+ }
193+
135194 #[ test]
136195 fn test_pretty_format_null ( ) {
137196 let schema = Arc :: new ( Schema :: new ( vec ! [
0 commit comments