diff --git a/arrow-array/src/builder/primitive_dictionary_builder.rs b/arrow-array/src/builder/primitive_dictionary_builder.rs index 282f0ae9d5b..f4a6662462e 100644 --- a/arrow-array/src/builder/primitive_dictionary_builder.rs +++ b/arrow-array/src/builder/primitive_dictionary_builder.rs @@ -126,10 +126,11 @@ where keys_builder.is_empty() && values_builder.is_empty(), "keys and values builders must be empty" ); + let values_capacity = values_builder.capacity(); Self { keys_builder, values_builder, - map: HashMap::new(), + map: HashMap::with_capacity(values_capacity), } } @@ -633,4 +634,19 @@ mod tests { assert_eq!(values, [None, None]); } + + #[test] + fn creating_dictionary_from_builders_should_use_values_capacity_for_the_map() { + let builder = PrimitiveDictionaryBuilder::::new_from_empty_builders( + PrimitiveBuilder::with_capacity(1).with_data_type(DataType::Int32), + PrimitiveBuilder::with_capacity(2).with_data_type(DataType::Timestamp(arrow_schema::TimeUnit::Microsecond, Some("+08:00".into()))), + ); + + assert!( + builder.map.capacity() >= builder.values_builder.capacity(), + "map capacity {} should be at least the values capacity {}", + builder.map.capacity(), + builder.values_builder.capacity() + ) + } }