address comments and enhance tests

goldmedal · Jul 3, 2024 · 510d533 · 510d533
1 parent ef21385
commit 510d533
Show file tree

Hide file tree

Showing 2 changed files with 32 additions and 4 deletions.
diff --git a/datafusion/common/src/scalar/mod.rs b/datafusion/common/src/scalar/mod.rs
@@ -3560,6 +3560,9 @@ impl fmt::Debug for ScalarValue {
             ScalarValue::List(_) => write!(f, "List({self})"),
             ScalarValue::LargeList(_) => write!(f, "LargeList({self})"),
             ScalarValue::Struct(struct_arr) => {
+                // ScalarValue Struct should always have a single element
+                assert_eq!(struct_arr.len(), 1);
+
                 let columns = struct_arr.columns();
                 let fields = struct_arr.fields();
 
@@ -3579,8 +3582,6 @@ impl fmt::Debug for ScalarValue {
                 )
             }
             ScalarValue::Map(map_arr) => {
-                // ScalarValue Map should always have a single element
-                assert_eq!(map_arr.len(), 1);
                 write!(
                     f,
                     "Map([{}])",
@@ -6298,6 +6299,7 @@ mod tests {
             .unwrap();
 
         assert_eq!(s.to_string(), "{a:1,b:}");
+        assert_eq!(format!("{s:?}"), r#"Struct({a:1,b:})"#);
 
         let ScalarValue::Struct(arr) = s else {
             panic!("Expected struct");
@@ -6340,7 +6342,7 @@ mod tests {
     }
 
     #[test]
-    fn test_map_display() {
+    fn test_map_display_and_debug() {
         let string_builder = StringBuilder::new();
         let int_builder = Int32Builder::with_capacity(4);
         let mut builder = MapBuilder::new(None, string_builder, int_builder);
@@ -6359,6 +6361,10 @@ mod tests {
         let map_value = ScalarValue::Map(Arc::new(builder.finish()));
 
         assert_eq!(map_value.to_string(), "[{joe:1},{blogs:2,foo:4},{},NULL]");
+        assert_eq!(
+            format!("{map_value:?}"),
+            r#"Map([{"joe":"1"},{"blogs":"2","foo":"4"},{},NULL])"#
+        );
 
         let ScalarValue::Map(arr) = map_value else {
             panic!("Expected map");

diff --git a/datafusion/proto/tests/cases/roundtrip_logical_plan.rs b/datafusion/proto/tests/cases/roundtrip_logical_plan.rs
@@ -21,7 +21,9 @@ use std::fmt::{self, Debug, Formatter};
 use std::sync::Arc;
 use std::vec;
 
-use arrow::array::{ArrayRef, FixedSizeListArray};
+use arrow::array::{
+    ArrayRef, FixedSizeListArray, Int32Builder, MapArray, MapBuilder, StringBuilder,
+};
 use arrow::datatypes::{
     DataType, Field, Fields, Int32Type, IntervalDayTimeType, IntervalMonthDayNanoType,
     IntervalUnit, Schema, SchemaRef, TimeUnit, UnionFields, UnionMode,
@@ -1270,6 +1272,7 @@ fn round_trip_scalar_values() {
             true,
         ))
         .unwrap(),
+        ScalarValue::Map(Arc::new(create_map_array_test_case())),
         ScalarValue::FixedSizeBinary(b"bar".to_vec().len() as i32, Some(b"bar".to_vec())),
         ScalarValue::FixedSizeBinary(0, None),
         ScalarValue::FixedSizeBinary(5, None),
@@ -1292,6 +1295,25 @@ fn round_trip_scalar_values() {
     }
 }
 
+// create a map array [{joe:1}, {blogs:2, foo:4}, {}, null] for testing
+fn create_map_array_test_case() -> MapArray {
+    let string_builder = StringBuilder::new();
+    let int_builder = Int32Builder::with_capacity(4);
+    let mut builder = MapBuilder::new(None, string_builder, int_builder);
+    builder.keys().append_value("joe");
+    builder.values().append_value(1);
+    builder.append(true).unwrap();
+
+    builder.keys().append_value("blogs");
+    builder.values().append_value(2);
+    builder.keys().append_value("foo");
+    builder.values().append_value(4);
+    builder.append(true).unwrap();
+    builder.append(true).unwrap();
+    builder.append(false).unwrap();
+    builder.finish()
+}
+
 #[test]
 fn round_trip_scalar_types() {
     let should_pass: Vec<DataType> = vec![