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 b8ca86a8e3 Support MapBuilder in make_builder (#5210)
b8ca86a8e3 is described below
commit b8ca86a8e36ad34e58e1151b9c1cc1d1d094b1d2
Author: Liang-Chi Hsieh <[email protected]>
AuthorDate: Fri Dec 15 09:48:34 2023 -0800
Support MapBuilder in make_builder (#5210)
---
arrow-array/src/builder/map_builder.rs | 45 ++++++++++++++++++++++++++++++-
arrow-array/src/builder/struct_builder.rs | 34 ++++++++++++++++-------
2 files changed, 68 insertions(+), 11 deletions(-)
diff --git a/arrow-array/src/builder/map_builder.rs
b/arrow-array/src/builder/map_builder.rs
index 3a5244ed81..5785bb5faf 100644
--- a/arrow-array/src/builder/map_builder.rs
+++ b/arrow-array/src/builder/map_builder.rs
@@ -234,7 +234,8 @@ impl<K: ArrayBuilder, V: ArrayBuilder> ArrayBuilder for
MapBuilder<K, V> {
#[cfg(test)]
mod tests {
- use crate::builder::{Int32Builder, StringBuilder};
+ use crate::builder::{make_builder, Int32Builder, StringBuilder};
+ use crate::{Int32Array, StringArray};
use super::*;
@@ -248,4 +249,46 @@ mod tests {
builder.finish();
}
+
+ #[test]
+ fn test_boxed_map_builder() {
+ let keys_builder = make_builder(&DataType::Utf8, 5);
+ let values_builder = make_builder(&DataType::Int32, 5);
+
+ let mut builder = MapBuilder::new(None, keys_builder, values_builder);
+ builder
+ .keys()
+ .as_any_mut()
+ .downcast_mut::<StringBuilder>()
+ .expect("should be an StringBuilder")
+ .append_value("1");
+ builder
+ .values()
+ .as_any_mut()
+ .downcast_mut::<Int32Builder>()
+ .expect("should be an Int32Builder")
+ .append_value(42);
+ builder.append(true).unwrap();
+
+ let map_array = builder.finish();
+
+ assert_eq!(
+ map_array
+ .keys()
+ .as_any()
+ .downcast_ref::<StringArray>()
+ .expect("should be an StringArray")
+ .value(0),
+ "1"
+ );
+ assert_eq!(
+ map_array
+ .values()
+ .as_any()
+ .downcast_ref::<Int32Array>()
+ .expect("should be an Int32Array")
+ .value(0),
+ 42
+ );
+ }
}
diff --git a/arrow-array/src/builder/struct_builder.rs
b/arrow-array/src/builder/struct_builder.rs
index 960949a2f0..196ae8092b 100644
--- a/arrow-array/src/builder/struct_builder.rs
+++ b/arrow-array/src/builder/struct_builder.rs
@@ -177,6 +177,24 @@ pub fn make_builder(datatype: &DataType, capacity: usize)
-> Box<dyn ArrayBuilde
let builder = make_builder(field.data_type(), capacity);
Box::new(LargeListBuilder::with_capacity(builder, capacity))
}
+ DataType::Map(field, _) => match field.data_type() {
+ DataType::Struct(fields) => {
+ let map_field_names = MapFieldNames {
+ key: fields[0].name().clone(),
+ value: fields[1].name().clone(),
+ entry: field.name().clone(),
+ };
+ let key_builder = make_builder(fields[0].data_type(),
capacity);
+ let value_builder = make_builder(fields[1].data_type(),
capacity);
+ Box::new(MapBuilder::with_capacity(
+ Some(map_field_names),
+ key_builder,
+ value_builder,
+ capacity,
+ ))
+ }
+ t => panic!("The field of Map data type {t:?} should has a child
Struct field"),
+ },
DataType::Struct(fields) =>
Box::new(StructBuilder::from_fields(fields.clone(), capacity)),
t => panic!("Data type {t:?} is not currently supported"),
}
@@ -514,19 +532,15 @@ mod tests {
}
#[test]
- #[should_panic(
- expected = "Data type Map(Field { name: \"entries\", data_type:
Struct([Field { name: \"keys\", data_type: Int32, nullable: false, dict_id: 0,
dict_is_ordered: false, metadata: {} }, Field { name: \"values\", data_type:
UInt32, nullable: false, dict_id: 0, dict_is_ordered: false, metadata: {} }]),
nullable: false, dict_id: 0, dict_is_ordered: false, metadata: {} }, false) is
not currently supported"
- )]
+ #[should_panic(expected = "Data type Dictionary(Int32, Utf8) is not
currently supported")]
fn test_struct_array_builder_from_schema_unsupported_type() {
- let keys = Arc::new(Field::new("keys", DataType::Int32, false));
- let values = Arc::new(Field::new("values", DataType::UInt32, false));
- let struct_type = DataType::Struct(Fields::from(vec![keys, values]));
- let map_data_type =
- DataType::Map(Arc::new(Field::new("entries", struct_type, false)),
false);
-
let fields = vec![
Field::new("f1", DataType::Int16, false),
- Field::new("f2", map_data_type, false),
+ Field::new(
+ "f2",
+ DataType::Dictionary(Box::new(DataType::Int32),
Box::new(DataType::Utf8)),
+ false,
+ ),
];
let _ = StructBuilder::from_fields(fields, 5);
