From e9be49d962560ce5b87544a2933d8b207322cf60 Mon Sep 17 00:00:00 2001 From: Matthew Turner Date: Wed, 1 Dec 2021 09:33:14 -0500 Subject: [PATCH] Change `pretty_format_batches` to return `Result` (#975) * Create write_batches function * Update pretty_format_batches and pretty_format_columns to return Display * fix import warning * Fix compile error Co-authored-by: Andrew Lamb --- arrow/src/util/pretty.rs | 79 +++++++++++++++++++++++++++++++++------- 1 file changed, 66 insertions(+), 13 deletions(-) diff --git a/arrow/src/util/pretty.rs b/arrow/src/util/pretty.rs index 8e9cc5feafa2..ac7caa8ac39a 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 std::fmt::Display; use comfy_table::{Cell, Table}; @@ -27,13 +28,16 @@ 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 { - Ok(create_table(results)?.to_string()) +pub fn pretty_format_batches(results: &[RecordBatch]) -> Result { + create_table(results) } ///! 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()) +pub fn pretty_format_columns( + col_name: &str, + results: &[ArrayRef], +) -> Result { + create_column(col_name, results) } ///! Prints a visual representation of record batches to stdout @@ -115,6 +119,7 @@ mod tests { use super::*; use crate::array::{DecimalBuilder, FixedSizeListBuilder, Int32Array}; + use std::fmt::Write; use std::sync::Arc; #[test] @@ -144,7 +149,7 @@ mod tests { ], )?; - let table = pretty_format_batches(&[batch])?; + let table = pretty_format_batches(&[batch])?.to_string(); let expected = vec![ "+---+-----+", @@ -176,7 +181,7 @@ mod tests { Arc::new(array::StringArray::from(vec![Some("e"), None, Some("g")])), ]; - let table = pretty_format_columns("a", &columns)?; + let table = pretty_format_columns("a", &columns)?.to_string(); let expected = vec![ "+---+", "| a |", "+---+", "| a |", "| b |", "| |", "| d |", "| e |", @@ -208,7 +213,7 @@ mod tests { // define data (null) let batch = RecordBatch::try_new(schema, arrays).unwrap(); - let table = pretty_format_batches(&[batch]).unwrap(); + let table = pretty_format_batches(&[batch]).unwrap().to_string(); let expected = vec![ "+---+---+---+", @@ -244,7 +249,7 @@ mod tests { let batch = RecordBatch::try_new(schema, vec![array])?; - let table = pretty_format_batches(&[batch])?; + let table = pretty_format_batches(&[batch])?.to_string(); let expected = vec![ "+-------+", @@ -285,7 +290,7 @@ mod tests { let array = Arc::new(builder.finish()); let batch = RecordBatch::try_new(schema, vec![array])?; - let table = pretty_format_batches(&[batch])?; + let table = pretty_format_batches(&[batch])?.to_string(); let expected = vec![ "+-----------+", "| d1 |", @@ -320,7 +325,9 @@ mod tests { )])); let batch = RecordBatch::try_new(schema, vec![Arc::new(array)]).unwrap(); - let table = pretty_format_batches(&[batch]).expect("formatting batches"); + let table = pretty_format_batches(&[batch]) + .expect("formatting batches") + .to_string(); let expected = $EXPECTED_RESULT; let actual: Vec<&str> = table.lines().collect(); @@ -494,7 +501,7 @@ mod tests { let batch = RecordBatch::try_new(schema, vec![dm])?; - let table = pretty_format_batches(&[batch])?; + let table = pretty_format_batches(&[batch])?.to_string(); let expected = vec![ "+-------+", @@ -535,7 +542,7 @@ mod tests { let batch = RecordBatch::try_new(schema, vec![dm])?; - let table = pretty_format_batches(&[batch])?; + let table = pretty_format_batches(&[batch])?.to_string(); let expected = vec![ "+------+", "| f |", "+------+", "| 101 |", "| |", "| 200 |", "| 3040 |", "+------+", @@ -589,7 +596,7 @@ mod tests { RecordBatch::try_new(Arc::new(schema), vec![Arc::new(c1), Arc::new(c2)]) .unwrap(); - let table = pretty_format_batches(&[batch])?; + let table = pretty_format_batches(&[batch])?.to_string(); let expected = vec![ r#"+-------------------------------------+----+"#, r#"| c1 | c2 |"#, @@ -605,4 +612,50 @@ mod tests { Ok(()) } + + #[test] + fn test_writing_formatted_batches() -> Result<()> { + // define a schema. + let schema = Arc::new(Schema::new(vec![ + Field::new("a", DataType::Utf8, true), + Field::new("b", DataType::Int32, true), + ])); + + // define data. + let batch = RecordBatch::try_new( + schema, + vec![ + Arc::new(array::StringArray::from(vec![ + Some("a"), + Some("b"), + None, + Some("d"), + ])), + Arc::new(array::Int32Array::from(vec![ + Some(1), + None, + Some(10), + Some(100), + ])), + ], + )?; + + let mut buf = String::new(); + write!(&mut buf, "{}", pretty_format_batches(&[batch])?.to_string()).unwrap(); + + let s = vec![ + "+---+-----+", + "| a | b |", + "+---+-----+", + "| a | 1 |", + "| b | |", + "| | 10 |", + "| d | 100 |", + "+---+-----+", + ]; + let expected = String::from(s.join("\n")); + assert_eq!(expected, buf); + + Ok(()) + } }