diff --git a/arrow/src/util/display.rs b/arrow/src/util/display.rs index bb75a3a66081..a65cb5faba3b 100644 --- a/arrow/src/util/display.rs +++ b/arrow/src/util/display.rs @@ -205,6 +205,34 @@ pub fn make_string_from_decimal(column: &Arc, row: usize) -> Result, + row: usize, +) -> Result<()> { + target.push('"'); + target.push_str(name); + target.push_str("\": "); + + if field_col.is_null(row) { + target.push_str("null"); + } else { + match field_col.data_type() { + DataType::Utf8 | DataType::LargeUtf8 => { + target.push('"'); + target.push_str(array_value_to_string(field_col, row)?.as_str()); + target.push('"'); + } + _ => { + target.push_str(array_value_to_string(field_col, row)?.as_str()); + } + } + } + + Ok(()) +} + /// Get the value at the given row in an array as a String. /// /// Note this function is quite inefficient and is unlikely to be @@ -280,6 +308,31 @@ pub fn array_value_to_string(column: &array::ArrayRef, row: usize) -> Result { + let st = column + .as_any() + .downcast_ref::() + .ok_or_else(|| { + ArrowError::InvalidArgumentError( + "Repl error: could not convert struct column to struct array." + .to_string(), + ) + })?; + + let mut s = String::new(); + s.push('{'); + let mut kv_iter = st.columns().into_iter().zip(st.column_names().into_iter()); + if let Some((col, name)) = kv_iter.next() { + append_struct_field_string(&mut s, name, col, row)?; + } + for (col, name) in kv_iter { + s.push_str(", "); + append_struct_field_string(&mut s, name, col, row)?; + } + s.push('}'); + + Ok(s) + } _ => Err(ArrowError::InvalidArgumentError(format!( "Pretty printing not implemented for {:?} type", column.data_type() diff --git a/arrow/src/util/pretty.rs b/arrow/src/util/pretty.rs index d307406289e2..f9599448c799 100644 --- a/arrow/src/util/pretty.rs +++ b/arrow/src/util/pretty.rs @@ -106,9 +106,9 @@ mod tests { use crate::{ array::{ self, new_null_array, Array, Date32Array, Date64Array, PrimitiveBuilder, - StringBuilder, StringDictionaryBuilder, Time32MillisecondArray, - Time32SecondArray, Time64MicrosecondArray, Time64NanosecondArray, - TimestampMicrosecondArray, TimestampMillisecondArray, + StringArray, StringBuilder, StringDictionaryBuilder, StructArray, + Time32MillisecondArray, Time32SecondArray, Time64MicrosecondArray, + Time64NanosecondArray, TimestampMicrosecondArray, TimestampMillisecondArray, TimestampNanosecondArray, TimestampSecondArray, }, datatypes::{DataType, Field, Int32Type, Schema}, @@ -507,4 +507,63 @@ mod tests { Ok(()) } + + #[test] + fn test_pretty_format_struct() -> Result<()> { + let schema = Schema::new(vec![ + Field::new( + "c1", + DataType::Struct(vec![ + Field::new("c11", DataType::Int32, false), + Field::new( + "c12", + DataType::Struct(vec![Field::new("c121", DataType::Utf8, false)]), + false, + ), + ]), + false, + ), + Field::new("c2", DataType::Utf8, false), + ]); + + let c1 = StructArray::from(vec![ + ( + Field::new("c11", DataType::Int32, false), + Arc::new(Int32Array::from(vec![Some(1), None, Some(5)])) as ArrayRef, + ), + ( + Field::new( + "c12", + DataType::Struct(vec![Field::new("c121", DataType::Utf8, false)]), + false, + ), + Arc::new(StructArray::from(vec![( + Field::new("c121", DataType::Utf8, false), + Arc::new(StringArray::from(vec![Some("e"), Some("f"), Some("g")])) + as ArrayRef, + )])) as ArrayRef, + ), + ]); + let c2 = StringArray::from(vec![Some("a"), Some("b"), Some("c")]); + + let batch = + RecordBatch::try_new(Arc::new(schema), vec![Arc::new(c1), Arc::new(c2)]) + .unwrap(); + + let table = pretty_format_batches(&[batch])?; + let expected = vec![ + r#"+-------------------------------------+----+"#, + r#"| c1 | c2 |"#, + r#"+-------------------------------------+----+"#, + r#"| {"c11": 1, "c12": {"c121": "e"}} | a |"#, + r#"| {"c11": null, "c12": {"c121": "f"}} | b |"#, + r#"| {"c11": 5, "c12": {"c121": "g"}} | c |"#, + r#"+-------------------------------------+----+"#, + ]; + + let actual: Vec<&str> = table.lines().collect(); + assert_eq!(expected, actual, "Actual result:\n{}", table); + + Ok(()) + } }