diff --git a/rust/arrow/src/util/pretty.rs b/rust/arrow/src/util/pretty.rs index d9177d907b7..7baf5590dbd 100644 --- a/rust/arrow/src/util/pretty.rs +++ b/rust/arrow/src/util/pretty.rs @@ -18,7 +18,7 @@ //! Utilities for printing record batches. Note this module is not //! available unless `feature = "prettyprint"` is enabled. -use crate::record_batch::RecordBatch; +use crate::{array::ArrayRef, record_batch::RecordBatch}; use prettytable::format; use prettytable::{Cell, Row, Table}; @@ -32,12 +32,23 @@ pub fn pretty_format_batches(results: &[RecordBatch]) -> Result { Ok(create_table(results)?.to_string()) } +///! Create a visual representation of columns +pub fn pretty_format_columns(col_name: &str, results: &[ArrayRef]) -> Result { + Ok(create_column(col_name, results)?.to_string()) +} + ///! Prints a visual representation of record batches to stdout pub fn print_batches(results: &[RecordBatch]) -> Result<()> { create_table(results)?.printstd(); Ok(()) } +///! Prints a visual representation of a list of column to stdout +pub fn print_columns(col_name: &str, results: &[ArrayRef]) -> Result<()> { + create_column(col_name, results)?.printstd(); + Ok(()) +} + ///! Convert a series of record batches into a table fn create_table(results: &[RecordBatch]) -> Result { let mut table = Table::new(); @@ -69,6 +80,28 @@ fn create_table(results: &[RecordBatch]) -> Result
{ Ok(table) } +fn create_column(field: &str, columns: &[ArrayRef]) -> Result
{ + let mut table = Table::new(); + table.set_format(*format::consts::FORMAT_NO_LINESEP_WITH_TITLE); + + if columns.is_empty() { + return Ok(table); + } + + let header = vec![Cell::new(field)]; + table.set_titles(Row::new(header)); + + for col in columns { + for row in 0..col.len() { + let mut cells = Vec::new(); + cells.push(Cell::new(&array_value_to_string(&col, row)?)); + table.add_row(Row::new(cells)); + } + } + + Ok(table) +} + #[cfg(test)] mod tests { use crate::{ @@ -132,6 +165,32 @@ mod tests { Ok(()) } + #[test] + fn test_pretty_format_columns() -> Result<()> { + let columns = vec![ + Arc::new(array::StringArray::from(vec![ + Some("a"), + Some("b"), + None, + Some("d"), + ])) as ArrayRef, + Arc::new(array::StringArray::from(vec![Some("e"), None, Some("g")])), + ]; + + let table = pretty_format_columns("a", &columns)?; + + let expected = vec![ + "+---+", "| a |", "+---+", "| a |", "| b |", "| |", "| d |", "| e |", + "| |", "| g |", "+---+", + ]; + + let actual: Vec<&str> = table.lines().collect(); + + assert_eq!(expected, actual, "Actual result:\n{}", table); + + Ok(()) + } + #[test] fn test_pretty_format_null() { let schema = Arc::new(Schema::new(vec![