diff --git a/rust/arrow/src/util/pretty.rs b/rust/arrow/src/util/pretty.rs index 089aa2f0251..dc564ddb4aa 100644 --- a/rust/arrow/src/util/pretty.rs +++ b/rust/arrow/src/util/pretty.rs @@ -18,7 +18,7 @@ //! Utilities for printing record batches use crate::array; -use crate::array::PrimitiveArrayOps; +use crate::array::{Array, PrimitiveArrayOps}; use crate::datatypes::{DataType, TimeUnit}; use crate::record_batch::RecordBatch; @@ -71,24 +71,22 @@ fn create_table(results: &[RecordBatch]) -> Result { macro_rules! make_string { ($array_type:ty, $column: ident, $row: ident) => {{ - Ok($column - .as_any() - .downcast_ref::<$array_type>() - .unwrap() - .value($row) - .to_string()) + let array = $column.as_any().downcast_ref::<$array_type>().unwrap(); + + let s = if array.is_null($row) { + "".to_string() + } else { + array.value($row).to_string() + }; + + Ok(s) }}; } /// Get the value at the given row in an array as a string fn array_value_to_string(column: array::ArrayRef, row: usize) -> Result { match column.data_type() { - DataType::Utf8 => Ok(column - .as_any() - .downcast_ref::() - .unwrap() - .value(row) - .to_string()), + DataType::Utf8 => make_string!(array::StringArray, column, row), DataType::Boolean => make_string!(array::BooleanArray, column, row), DataType::Int16 => make_string!(array::Int16Array, column, row), DataType::Int32 => make_string!(array::Int32Array, column, row), @@ -143,16 +141,26 @@ mod tests { fn test_pretty_format_batches() -> Result<()> { // define a schema. let schema = Arc::new(Schema::new(vec![ - Field::new("a", DataType::Utf8, false), - Field::new("b", DataType::Int32, false), + 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!["a", "b", "c", "d"])), - Arc::new(array::Int32Array::from(vec![1, 10, 10, 100])), + 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), + ])), ], )?; @@ -163,15 +171,15 @@ mod tests { "| a | b |", "+---+-----+", "| a | 1 |", - "| b | 10 |", - "| c | 10 |", + "| b | |", + "| | 10 |", "| d | 100 |", "+---+-----+", ]; let actual: Vec<&str> = table.lines().collect(); - assert_eq!(expected, actual); + assert_eq!(expected, actual, "Actual result:\n{}", table); Ok(()) }