adriangb commented on code in PR #24857:
URL: https://github.com/apache/datafusion/pull/24857#discussion_r3945459389
##########
datafusion/physical-expr-common/src/binary_map.rs:
##########
@@ -663,6 +749,155 @@ mod tests {
use arrow::array::{BinaryArray, LargeBinaryArray, StringArray};
use std::collections::HashMap;
+ /// A lower bound on the bytes a hashbrown table holding `entries` entries
+ /// of type `T` must allocate: one entry slot and one control byte each.
+ /// Derived independently of the production accounting so it can bracket
it.
+ fn min_table_bytes<T>(entries: usize) -> usize {
+ entries * (size_of::<T>() + 1)
+ }
+
+ #[test]
+ fn map_new_does_not_allocate() {
+ let map = ArrowBytesMap::<i32, ()>::new(OutputType::Utf8);
+
+ assert_eq!(map.map.capacity(), 0);
+ assert_eq!(map.map.allocation_size(), 0);
+ assert_eq!(map.buffer.capacity(), 0);
+ // Only the single leading zero offset is allocated.
+ assert!(map.size() < 128, "expected {} to be tiny", map.size());
+ }
+
+ #[test]
+ fn map_with_capacity_reports_the_real_hash_table_allocation() {
+ let map = ArrowBytesMap::<i32, ()>::with_capacity(
+ OutputType::Utf8,
+ INITIAL_MAP_CAPACITY,
+ );
+
+ assert!(map.map.capacity() >= INITIAL_MAP_CAPACITY);
+ assert_eq!(map.buffer.capacity(), INITIAL_BUFFER_CAPACITY);
+
+ // Before this accounting was corrected the map reported its hash table
+ // as costing zero bytes until the table grew past its pre-allocation.
+ let table_bytes = map.map.allocation_size();
+ let lower_bound = min_table_bytes::<Entry<i32,
()>>(map.map.capacity());
+ assert!(
+ table_bytes >= lower_bound,
+ "expected {table_bytes} to be at least {lower_bound}"
+ );
+ assert!(
+ table_bytes <= 2 * lower_bound + 64,
+ "expected {table_bytes} to be within a small factor of
{lower_bound}"
+ );
+ assert!(map.size() >= table_bytes + INITIAL_BUFFER_CAPACITY);
+ }
+
+ #[test]
+ fn take_preserves_the_capacity_the_map_was_built_with() {
+ let mut preallocated = ArrowBytesMap::<i32, ()>::with_capacity(
+ OutputType::Utf8,
+ INITIAL_MAP_CAPACITY,
+ );
+ let capacity = preallocated.map.capacity();
+ preallocated.take();
+ assert_eq!(preallocated.map.capacity(), capacity);
+ assert_eq!(preallocated.buffer.capacity(), INITIAL_BUFFER_CAPACITY);
+
+ let mut lazy = ArrowBytesMap::<i32, ()>::new(OutputType::Utf8);
+ lazy.take();
+ assert_eq!(lazy.map.capacity(), 0);
+ assert_eq!(lazy.buffer.capacity(), 0);
+ }
+
+ #[test]
+ fn clear_and_release_frees_the_preallocation_that_take_keeps() {
+ let mut map = ArrowBytesMap::<i32, ()>::with_capacity(
+ OutputType::Utf8,
+ INITIAL_MAP_CAPACITY,
+ );
+ let values: ArrayRef = Arc::new(StringArray::from_iter_values(
+ (0..1_000).map(|i| format!("distinct value number {i}")),
+ ));
+ map.insert_if_new(&values, |_| (), |_| ());
+
+ let populated_size = map.size();
+ assert!(populated_size > INITIAL_BUFFER_CAPACITY);
+
+ // `take` deliberately keeps the map warm, so it does not release the
+ // configured capacities.
+ map.take();
+ let taken_size = map.size();
+ assert!(
+ taken_size > INITIAL_BUFFER_CAPACITY,
+ "expected take to retain the warm up allocations, got {taken_size}"
+ );
+
+ map.clear_and_release();
+ let released_size = map.size();
+ assert_eq!(map.map.allocation_size(), 0);
+ assert_eq!(map.buffer.capacity(), 0);
+ assert!(
+ released_size < 128,
+ "expected the released map to report approximately zero bytes, got
{released_size}"
+ );
+
+ // The configured capacities survive, so the map warms back up when it
+ // is emitted from again.
+ map.take();
+ assert!(map.map.capacity() >= INITIAL_MAP_CAPACITY);
+ assert_eq!(map.buffer.capacity(), INITIAL_BUFFER_CAPACITY);
+ }
+
+ #[test]
+ fn lazy_and_pre_allocated_buffers_grow_on_the_same_ladder() {
Review Comment:
Thanks! Addressed in 7a1a4df.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
