support struct array in pretty display (#579) (#609)

* support struct array in pretty display

* print null value as null

to be consistent with postgresql

* add test for struct array pretty print

* Revert "print null value as null"

This reverts commit 3d620f6.

* handle null value as sepcial case in struct display

Co-authored-by: QP Hou <[email protected]>

arrow/src/util/display.rs

@@ -205,6 +205,34 @@ pub fn make_string_from_decimal(column: &Arc<dyn Array>, row: usize) -> Result<S
     Ok(formatted_decimal)
 }
 
+fn append_struct_field_string(
+    target: &mut String,
+    name: &str,
+    field_col: &Arc<dyn Array>,
+    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<Str
                 column.data_type()
             ))),
         },
+        DataType::Struct(_) => {
+            let st = column
+                .as_any()
+                .downcast_ref::<array::StructArray>()
+                .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()

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(())
+    }
 }