Adding Pretty Print Support For Fixed Size List (#958)

* Inferring 2. as Float64 for issue #929 * Adding pretty print support for fixed size list array * fixing linting errors * adding null row to test
apache · Nov 22, 2021 · b039cf4 · b039cf4
1 parent 9703f98
commit b039cf4
Show file tree

Hide file tree

Showing 2 changed files with 58 additions and 1 deletion.
diff --git 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::<array::FixedSizeListArray>()
+            .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::<Result<Vec<String>>>()?;
+        Ok(format!("[{}]", string_values.join(", ")))
+    }};
+}
+
 #[inline(always)]
 pub fn make_string_from_decimal(column: &Arc<dyn Array>, row: usize) -> Result<String> {
     let array = column
@@ -308,6 +324,7 @@ pub fn array_value_to_string(column: &array::ArrayRef, row: usize) -> Result<Str
                 column.data_type()
             ))),
         },
+        DataType::FixedSizeList(_, _) => 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
@@ -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`