Recurse into Dictionary value type in DataType::is_nested (#3083)

apache · Nov 10, 2022 · ed20bf1 · ed20bf1
1 parent ce5e26f
commit ed20bf1
Showing 1 changed file with 38 additions and 9 deletions.
diff --git a/arrow-schema/src/datatype.rs b/arrow-schema/src/datatype.rs
@@ -328,18 +328,20 @@ impl DataType {
         )
     }
 
-    /// Returns true if this type is nested (List, FixedSizeList, LargeList, Struct, Union, or Map)
+    /// Returns true if this type is nested (List, FixedSizeList, LargeList, Struct, Union,
+    /// or Map), or a dictionary of a nested type
     pub fn is_nested(t: &DataType) -> bool {
         use DataType::*;
-        matches!(
-            t,
+        match t {
+            Dictionary(_, v) => DataType::is_nested(v.as_ref()),
             List(_)
-                | FixedSizeList(_, _)
-                | LargeList(_)
-                | Struct(_)
-                | Union(_, _, _)
-                | Map(_, _)
-        )
+            | FixedSizeList(_, _)
+            | LargeList(_)
+            | Struct(_)
+            | Union(_, _, _)
+            | Map(_, _) => true,
+            _ => false,
+        }
     }
 
     /// Compares the datatype with another, ignoring nested field names
@@ -489,4 +491,31 @@ mod tests {
             ),
         ]);
     }
+
+    #[test]
+    fn test_nested() {
+        let list = DataType::List(Box::new(Field::new("foo", DataType::Utf8, true)));
+
+        assert!(!DataType::is_nested(&DataType::Boolean));
+        assert!(!DataType::is_nested(&DataType::Int32));
+        assert!(!DataType::is_nested(&DataType::Utf8));
+        assert!(DataType::is_nested(&list));
+
+        assert!(!DataType::is_nested(&DataType::Dictionary(
+            Box::new(DataType::Int32),
+            Box::new(DataType::Boolean)
+        )));
+        assert!(!DataType::is_nested(&DataType::Dictionary(
+            Box::new(DataType::Int32),
+            Box::new(DataType::Int64)
+        )));
+        assert!(!DataType::is_nested(&DataType::Dictionary(
+            Box::new(DataType::Int32),
+            Box::new(DataType::LargeUtf8)
+        )));
+        assert!(DataType::is_nested(&DataType::Dictionary(
+            Box::new(DataType::Int32),
+            Box::new(list)
+        )));
+    }
 }