From 720bac884a91523265b3f90c6307a7225d73317e Mon Sep 17 00:00:00 2001 From: Raphael Taylor-Davies Date: Tue, 14 Feb 2023 15:10:41 +0000 Subject: [PATCH 1/3] Add format with options --- arrow/src/util/pretty.rs | 91 +++++++++++++++++++++++++++++++++------- 1 file changed, 76 insertions(+), 15 deletions(-) diff --git a/arrow/src/util/pretty.rs b/arrow/src/util/pretty.rs index 4defa71a779..a8317e1635b 100644 --- a/arrow/src/util/pretty.rs +++ b/arrow/src/util/pretty.rs @@ -19,6 +19,7 @@ //! available unless `feature = "prettyprint"` is enabled. use crate::{array::ArrayRef, record_batch::RecordBatch}; +use arrow_array::Array; use arrow_cast::display::{ArrayFormatter, FormatOptions}; use comfy_table::{Cell, Table}; use std::fmt::Display; @@ -27,33 +28,51 @@ use crate::error::Result; use super::display::array_value_to_string; -///! Create a visual representation of record batches +/// Create a visual representation of record batches pub fn pretty_format_batches(results: &[RecordBatch]) -> Result { - create_table(results) + let options = FormatOptions::default().with_display_error(true); + pretty_format_batches_with_options(results, &options) +} + +/// Create a visual representation of record batches +pub fn pretty_format_batches_with_options( + results: &[RecordBatch], + options: &FormatOptions, +) -> Result { + create_table(results, options) } -///! Create a visual representation of columns +/// Create a visual representation of columns pub fn pretty_format_columns( col_name: &str, results: &[ArrayRef], ) -> Result { - create_column(col_name, results) + let options = FormatOptions::default().with_display_error(true); + pretty_format_columns_with_options(col_name, results, &options) +} + +pub fn pretty_format_columns_with_options( + col_name: &str, + results: &[ArrayRef], + options: &FormatOptions, +) -> Result { + create_column(col_name, results, options) } -///! Prints a visual representation of record batches to stdout +/// Prints a visual representation of record batches to stdout pub fn print_batches(results: &[RecordBatch]) -> Result<()> { - println!("{}", create_table(results)?); + println!("{}", pretty_format_batches(results)?); Ok(()) } -///! Prints a visual representation of a list of column to stdout +/// Prints a visual representation of a list of column to stdout pub fn print_columns(col_name: &str, results: &[ArrayRef]) -> Result<()> { - println!("{}", create_column(col_name, results)?); + println!("{}", pretty_format_columns(col_name, results)?); Ok(()) } -///! Convert a series of record batches into a table -fn create_table(results: &[RecordBatch]) -> Result { +/// Convert a series of record batches into a table +fn create_table(results: &[RecordBatch], options: &FormatOptions) -> Result
{ let mut table = Table::new(); table.load_preset("||--+-++| ++++++"); @@ -69,13 +88,11 @@ fn create_table(results: &[RecordBatch]) -> Result
{ } table.set_header(header); - let options = FormatOptions::default().with_display_error(true); - for batch in results { let formatters = batch .columns() .iter() - .map(|c| ArrayFormatter::try_new(c.as_ref(), &options)) + .map(|c| ArrayFormatter::try_new(c.as_ref(), options)) .collect::>>()?; for row in 0..batch.num_rows() { @@ -90,7 +107,11 @@ fn create_table(results: &[RecordBatch]) -> Result
{ Ok(table) } -fn create_column(field: &str, columns: &[ArrayRef]) -> Result
{ +fn create_column( + field: &str, + columns: &[ArrayRef], + options: &FormatOptions, +) -> Result
{ let mut table = Table::new(); table.load_preset("||--+-++| ++++++"); @@ -102,8 +123,9 @@ fn create_column(field: &str, columns: &[ArrayRef]) -> Result
{ table.set_header(header); for col in columns { + let formatter = ArrayFormatter::try_new(col.as_ref(), options)?; for row in 0..col.len() { - let cells = vec![Cell::new(array_value_to_string(col, row)?)]; + let cells = vec![Cell::new(formatter.value(row))]; table.add_row(cells); } } @@ -1057,4 +1079,43 @@ mod tests { Ok(()) } + + #[test] + fn test_format_options() { + let options = FormatOptions::default().with_null("null"); + let array = Int32Array::from(vec![Some(1), Some(2), None, Some(3), Some(4)]); + let batch = + RecordBatch::try_from_iter([("my_column_name", Arc::new(array) as _)]) + .unwrap(); + + let column = pretty_format_columns_with_options( + "my_column_name", + &[batch.column(0).clone()], + &options, + ) + .unwrap() + .to_string(); + + let batch = pretty_format_batches_with_options(&[batch], &options) + .unwrap() + .to_string(); + + let expected = vec![ + "+----------------+", + "| my_column_name |", + "+----------------+", + "| 1 |", + "| 2 |", + "| null |", + "| 3 |", + "| 4 |", + "+----------------+", + ]; + + let actual: Vec<&str> = column.lines().collect(); + assert_eq!(expected, actual, "Actual result:\n{column}"); + + let actual: Vec<&str> = batch.lines().collect(); + assert_eq!(expected, actual, "Actual result:\n{batch}"); + } } From e2b3e435e20f158c7dca0aff4ba556b2f98b7646 Mon Sep 17 00:00:00 2001 From: Raphael Taylor-Davies Date: Tue, 14 Feb 2023 15:18:18 +0000 Subject: [PATCH 2/3] Clippy --- arrow/src/util/pretty.rs | 2 -- 1 file changed, 2 deletions(-) diff --git a/arrow/src/util/pretty.rs b/arrow/src/util/pretty.rs index a8317e1635b..73e9a89f3be 100644 --- a/arrow/src/util/pretty.rs +++ b/arrow/src/util/pretty.rs @@ -26,8 +26,6 @@ use std::fmt::Display; use crate::error::Result; -use super::display::array_value_to_string; - /// Create a visual representation of record batches pub fn pretty_format_batches(results: &[RecordBatch]) -> Result { let options = FormatOptions::default().with_display_error(true); From 0b7dfeb29925a8c42481ddca545fb5a1262ab925 Mon Sep 17 00:00:00 2001 From: Raphael Taylor-Davies Date: Tue, 14 Feb 2023 15:29:04 +0000 Subject: [PATCH 3/3] Fix tests --- arrow/src/util/pretty.rs | 1 + 1 file changed, 1 insertion(+) diff --git a/arrow/src/util/pretty.rs b/arrow/src/util/pretty.rs index 73e9a89f3be..21d03582685 100644 --- a/arrow/src/util/pretty.rs +++ b/arrow/src/util/pretty.rs @@ -153,6 +153,7 @@ mod tests { use arrow_array::builder::PrimitiveBuilder; use arrow_array::types::{ArrowTimestampType, TimestampSecondType}; + use arrow_cast::display::array_value_to_string; use half::f16; #[test]