diff --git a/arrow/src/util/display.rs b/arrow/src/util/display.rs index fbc0bd579e0..09872a79661 100644 --- a/arrow/src/util/display.rs +++ b/arrow/src/util/display.rs @@ -194,6 +194,22 @@ macro_rules! make_string_from_list { }}; } +macro_rules! make_string_from_fixed_size_list { + ($column: ident, $row: ident) => {{ + let list = $column + .as_any() + .downcast_ref::() + .ok_or(ArrowError::InvalidArgumentError(format!( + "Repl error: could not convert list column to list array." + )))? + .value($row); + let string_values = (0..list.len()) + .map(|i| array_value_to_string(&list.clone(), i)) + .collect::>>()?; + Ok(format!("[{}]", string_values.join(", "))) + }}; +} + #[inline(always)] pub fn make_string_from_decimal(column: &Arc, row: usize) -> Result { let array = column @@ -308,6 +324,7 @@ pub fn array_value_to_string(column: &array::ArrayRef, row: usize) -> Result make_string_from_fixed_size_list!(column, row), DataType::Struct(_) => { let st = column .as_any() diff --git a/arrow/src/util/pretty.rs b/arrow/src/util/pretty.rs index 28bb0165ddd..8e9cc5feafa 100644 --- a/arrow/src/util/pretty.rs +++ b/arrow/src/util/pretty.rs @@ -114,7 +114,7 @@ mod tests { }; use super::*; - use crate::array::{DecimalBuilder, Int32Array}; + use crate::array::{DecimalBuilder, FixedSizeListBuilder, Int32Array}; use std::sync::Arc; #[test] @@ -263,6 +263,46 @@ mod tests { Ok(()) } + #[test] + fn test_pretty_format_fixed_size_list() -> Result<()> { + // define a schema. + let field_type = DataType::FixedSizeList( + Box::new(Field::new("item", DataType::Int32, true)), + 3, + ); + let schema = Arc::new(Schema::new(vec![Field::new("d1", field_type, true)])); + + let keys_builder = Int32Array::builder(3); + let mut builder = FixedSizeListBuilder::new(keys_builder, 3); + + builder.values().append_slice(&[1, 2, 3]).unwrap(); + builder.append(true).unwrap(); + builder.values().append_slice(&[4, 5, 6]).unwrap(); + builder.append(false).unwrap(); + builder.values().append_slice(&[7, 8, 9]).unwrap(); + builder.append(true).unwrap(); + + let array = Arc::new(builder.finish()); + + let batch = RecordBatch::try_new(schema, vec![array])?; + let table = pretty_format_batches(&[batch])?; + let expected = vec![ + "+-----------+", + "| d1 |", + "+-----------+", + "| [1, 2, 3] |", + "| |", + "| [7, 8, 9] |", + "+-----------+", + ]; + + let actual: Vec<&str> = table.lines().collect(); + + assert_eq!(expected, actual, "Actual result:\n{}", table); + + Ok(()) + } + /// Generate an array with type $ARRAYTYPE with a numeric value of /// $VALUE, and compare $EXPECTED_RESULT to the output of /// formatting that array with `pretty_format_batches`