Skip to content

Commit

Permalink
support struct array in pretty display (#579)
Browse files Browse the repository at this point in the history
* 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
  • Loading branch information
QP Hou committed Jul 23, 2021
1 parent 9f40f89 commit 87d2840
Show file tree
Hide file tree
Showing 2 changed files with 115 additions and 3 deletions.
53 changes: 53 additions & 0 deletions arrow/src/util/display.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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()
Expand Down
65 changes: 62 additions & 3 deletions arrow/src/util/pretty.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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},
Expand Down Expand Up @@ -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(())
}
}

0 comments on commit 87d2840

Please sign in to comment.