initial support for list of dictionaries

lancedb · Mar 8, 2023 · 35574ce · 35574ce
1 parent c2da942
commit 35574ce
Showing 1 changed file with 13 additions and 1 deletion.
diff --git a/rust/src/datatypes.rs b/rust/src/datatypes.rs
@@ -8,7 +8,7 @@ use std::fmt::{self};
 use arrow_array::types::{
     Int16Type, Int32Type, Int64Type, Int8Type, UInt16Type, UInt32Type, UInt64Type, UInt8Type,
 };
-use arrow_array::{cast::as_dictionary_array, ArrayRef, RecordBatch};
+use arrow_array::{cast::as_dictionary_array, ArrayRef, RecordBatch, ListArray, LargeListArray};
 use arrow_schema::{DataType, Field as ArrowField, Schema as ArrowSchema, TimeUnit};
 use async_recursion::async_recursion;
 
@@ -338,6 +338,18 @@ impl Field {
                     panic!("Unsupported dictionary key type: {}", key_type);
                 }
             },
+            DataType::List(values) => {
+                // should this be self.children[0]?
+                let mut f = Field::try_from(values.as_ref()).unwrap();
+                let list_arr = arr.as_any().downcast_ref::<ListArray>().unwrap();
+                f.set_dictionary(list_arr.values());
+            }
+            DataType::LargeList(values) => {
+                // should this be self.children[0]?
+                let mut f = Field::try_from(values.as_ref()).unwrap();
+                let list_arr = arr.as_any().downcast_ref::<LargeListArray>().unwrap();
+                f.set_dictionary(list_arr.values());
+            },
             _ => {
                 // Add nested struct support.
             }