Skip to content

Commit

Permalink
fix: use the values builder capacity for the hash map in `PrimitiveDi…
Browse files Browse the repository at this point in the history
…ctionaryBuilder::new_from_builders` (#7012)

* feat: allow setting custom value data type in `PrimitiveDictionaryBuilder`

Fixes #7011

* use the values capacity for the hash map

* update new_from_empty_builders and not new_from_builders
  • Loading branch information
rluvaton authored Jan 25, 2025
1 parent d6fa078 commit f59b94f
Showing 1 changed file with 17 additions and 1 deletion.
18 changes: 17 additions & 1 deletion arrow-array/src/builder/primitive_dictionary_builder.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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),
}
}

Expand Down Expand Up @@ -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::<Int32Type, crate::types::TimestampMicrosecondType>::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()
)
}
}

0 comments on commit f59b94f

Please sign in to comment.