[arrow-rs] branch master updated: add map array to pretty print (#3339)

[tu...@apache.org]

This is an automated email from the ASF dual-hosted git repository.

tustvold pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/arrow-rs.git


The following commit(s) were added to refs/heads/master by this push:
     new a93859b07 add map array to pretty print (#3339)
a93859b07 is described below

commit a93859b07516b91511ffe3106a423b9af4b69f34
Author: askoa <11...@users.noreply.github.com>
AuthorDate: Tue Dec 13 16:07:20 2022 -0500

    add map array to pretty print (#3339)
    
    Co-authored-by: askoa <as...@local>
---
 arrow-cast/src/display.rs | 77 +++++++++++++++++++++++++++++++++++++++++++++++
 1 file changed, 77 insertions(+)

diff --git a/arrow-cast/src/display.rs b/arrow-cast/src/display.rs
index 287065eb6..10709994d 100644
--- a/arrow-cast/src/display.rs
+++ b/arrow-cast/src/display.rs
@@ -281,6 +281,29 @@ fn append_struct_field_string(
     Ok(())
 }
 
+fn append_map_field_string(
+    target: &mut String,
+    field_col: &Arc<dyn Array>,
+    row: usize,
+) -> Result<(), ArrowError> {
+    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
@@ -430,6 +453,38 @@ pub fn array_value_to_string(
 
             Ok(s)
         }
+        DataType::Map(_, _) => {
+            let map_array =
+                column.as_any().downcast_ref::<MapArray>().ok_or_else(|| {
+                    ArrowError::InvalidArgumentError(
+                        "Repl error: could not convert column to map array.".to_string(),
+                    )
+                })?;
+            let map_entry = map_array.value(row);
+            let st = map_entry
+                .as_any()
+                .downcast_ref::<StructArray>()
+                .ok_or_else(|| {
+                    ArrowError::InvalidArgumentError(
+                        "Repl error: could not convert map entry to struct array."
+                            .to_string(),
+                    )
+                })?;
+            let mut s = String::new();
+            s.push('{');
+            let entries_count = st.column(0).len();
+            for i in 0..entries_count {
+                if i > 0 {
+                    s.push_str(", ");
+                }
+                append_map_field_string(&mut s, st.column(0), i)?;
+                s.push_str(": ");
+                append_map_field_string(&mut s, st.column(1), i)?;
+            }
+            s.push('}');
+
+            Ok(s)
+        }
         DataType::Union(field_vec, type_ids, mode) => {
             union_to_string(column, row, field_vec, type_ids, mode)
         }
@@ -527,6 +582,28 @@ pub fn lexical_to_string<N: lexical_core::ToLexical>(n: N) -> String {
 mod tests {
     use super::*;
 
+    #[test]
+    fn test_map_arry_to_string() {
+        let keys = vec!["a", "b", "c", "d", "e", "f", "g", "h"];
+        let values_data = UInt32Array::from(vec![0u32, 10, 20, 30, 40, 50, 60, 70]);
+
+        // Construct a buffer for value offsets, for the nested array:
+        //  [[a, b, c], [d, e, f], [g, h]]
+        let entry_offsets = [0, 3, 6, 8];
+
+        let map_array = MapArray::new_from_strings(
+            keys.clone().into_iter(),
+            &values_data,
+            &entry_offsets,
+        )
+        .unwrap();
+        let param = Arc::new(map_array) as ArrayRef;
+        assert_eq!(
+            "{\"d\": 30, \"e\": 40, \"f\": 50}",
+            array_value_to_string(&param, 1).unwrap()
+        );
+    }
+
     #[test]
     fn test_array_value_to_string_duration() {
         let ns_array =