[GitHub] [arrow-rs] viirya commented on a diff in pull request #1648: Pretty Print `UnionArray`s

[GitBox <gi...@apache.org>]

viirya commented on code in PR #1648:
URL: https://github.com/apache/arrow-rs/pull/1648#discussion_r865555032


##########
arrow/src/util/pretty.rs:
##########
@@ -647,6 +648,146 @@ mod tests {
         Ok(())
     }
 
+    #[test]
+    fn test_pretty_format_dense_union() -> Result<()> {
+        let mut builder = UnionBuilder::new_dense(4);
+        builder.append::<Int32Type>("a", 1).unwrap();
+        builder.append::<Float64Type>("b", 3.2234).unwrap();
+        builder.append_null::<Float64Type>("b").unwrap();
+        builder.append_null::<Int32Type>("a").unwrap();
+        let union = builder.build().unwrap();
+
+        let schema = Schema::new(vec![Field::new(
+            "Teamsters",
+            DataType::Union(
+                vec![
+                    Field::new("a", DataType::Int32, false),
+                    Field::new("b", DataType::Float64, false),
+                ],
+                UnionMode::Dense,
+            ),
+            false,
+        )]);
+
+        let batch =
+            RecordBatch::try_new(Arc::new(schema), vec![Arc::new(union)]).unwrap();
+        let table = pretty_format_batches(&[batch])?.to_string();
+        let actual: Vec<&str> = table.lines().collect();
+        let expected = vec![
+            "+------------+",
+            "| Teamsters  |",
+            "+------------+",
+            "| {a=1}      |",
+            "| {b=3.2234} |",
+            "| {b=}       |",
+            "| {a=}       |",
+            "+------------+",
+        ];
+
+        assert_eq!(expected, actual);
+        Ok(())
+    }
+
+    #[test]
+    fn test_pretty_format_sparse_union() -> Result<()> {
+        let mut builder = UnionBuilder::new_sparse(4);
+        builder.append::<Int32Type>("a", 1).unwrap();
+        builder.append::<Float64Type>("b", 3.2234).unwrap();
+        builder.append_null::<Float64Type>("b").unwrap();
+        builder.append_null::<Int32Type>("a").unwrap();
+        let union = builder.build().unwrap();
+
+        let schema = Schema::new(vec![Field::new(
+            "Teamsters",
+            DataType::Union(
+                vec![
+                    Field::new("a", DataType::Int32, false),
+                    Field::new("b", DataType::Float64, false),
+                ],
+                UnionMode::Sparse,
+            ),
+            false,
+        )]);
+
+        let batch =
+            RecordBatch::try_new(Arc::new(schema), vec![Arc::new(union)]).unwrap();
+        let table = pretty_format_batches(&[batch])?.to_string();
+        let actual: Vec<&str> = table.lines().collect();
+        let expected = vec![
+            "+------------+",
+            "| Teamsters  |",
+            "+------------+",
+            "| {a=1}      |",
+            "| {b=3.2234} |",
+            "| {b=}       |",
+            "| {a=}       |",
+            "+------------+",
+        ];
+
+        assert_eq!(expected, actual);
+        Ok(())
+    }
+
+    #[test]
+    fn test_pretty_format_nested_union() -> Result<()> {
+        //Inner UnionArray
+        let mut builder = UnionBuilder::new_dense(4);
+        builder.append::<Int32Type>("b", 1).unwrap();
+        builder.append::<Float64Type>("c", 3.2234).unwrap();
+        builder.append_null::<Float64Type>("c").unwrap();
+        builder.append_null::<Int32Type>("b").unwrap();

Review Comment:
   Can you let outer show one null value from inner UnionArray too?



##########
arrow/src/util/display.rs:
##########
@@ -395,13 +396,49 @@ pub fn array_value_to_string(column: &array::ArrayRef, row: usize) -> Result<Str
 
             Ok(s)
         }
+        DataType::Union(field_vec, mode) => union_to_string(column, row, field_vec, mode),
         _ => Err(ArrowError::InvalidArgumentError(format!(
             "Pretty printing not implemented for {:?} type",
             column.data_type()
         ))),
     }
 }
 
+/// Converts the value of the union array at `row` to a String
+fn union_to_string(
+    column: &array::ArrayRef,
+    row: usize,
+    fields: &[Field],
+    mode: &UnionMode,
+) -> Result<String> {
+    let list = column
+        .as_any()
+        .downcast_ref::<array::UnionArray>()
+        .ok_or_else(|| {
+            ArrowError::InvalidArgumentError(
+                "Repl error: could not convert union column to union array.".to_string(),
+            )
+        })?;
+    let type_id = list.type_id(row);
+    let name = fields
+        .get(type_id as usize)
+        .ok_or_else(|| {
+            ArrowError::InvalidArgumentError(
+                "Repl error: could not get field for the corresponding type in union array."
+                .to_string(),

Review Comment:
   Maybe have `type_id` in error string?



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: github-unsubscribe@arrow.apache.org

For queries about this service, please contact Infrastructure at:
users@infra.apache.org