Skip to content

Commit f2ac2c6

Browse files
authored
Remove custom doubling strategy + add examples to VecAllocEx (#9058)
* Add examples to VecAllocExt and remove custom doubling strategy * Add comment and add another example of using accounted push
1 parent 47809d9 commit f2ac2c6

File tree

1 file changed

+56
-10
lines changed
  • datafusion/execution/src/memory_pool

1 file changed

+56
-10
lines changed

datafusion/execution/src/memory_pool/proxy.rs

Lines changed: 56 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -28,6 +28,35 @@ pub trait VecAllocExt {
2828
/// `accounting` by any newly allocated bytes.
2929
///
3030
/// Note that allocation counts capacity, not size
31+
///
32+
/// # Example:
33+
/// ```
34+
/// # use datafusion_execution::memory_pool::proxy::VecAllocExt;
35+
/// // use allocated to incrementally track how much memory is allocated in the vec
36+
/// let mut allocated = 0;
37+
/// let mut vec = Vec::new();
38+
/// // Push data into the vec and the accounting will be updated to reflect
39+
/// // memory allocation
40+
/// vec.push_accounted(1, &mut allocated);
41+
/// assert_eq!(allocated, 16); // space for 4 u32s
42+
/// vec.push_accounted(1, &mut allocated);
43+
/// assert_eq!(allocated, 16); // no new allocation needed
44+
///
45+
/// // push more data into the vec
46+
/// for _ in 0..10 { vec.push_accounted(1, &mut allocated); }
47+
/// assert_eq!(allocated, 64); // underlying vec has space for 10 u32s
48+
/// assert_eq!(vec.allocated_size(), 64);
49+
/// ```
50+
/// # Example with other allocations:
51+
/// ```
52+
/// # use datafusion_execution::memory_pool::proxy::VecAllocExt;
53+
/// // You can use the same allocated size to track memory allocated by
54+
/// // another source. For example
55+
/// let mut allocated = 27;
56+
/// let mut vec = Vec::new();
57+
/// vec.push_accounted(1, &mut allocated); // allocates 16 bytes for vec
58+
/// assert_eq!(allocated, 43); // 16 bytes for vec, 27 bytes for other
59+
/// ```
3160
fn push_accounted(&mut self, x: Self::T, accounting: &mut usize);
3261

3362
/// Return the amount of memory allocated by this Vec to store elements
@@ -36,24 +65,41 @@ pub trait VecAllocExt {
3665
/// Note this calculation is not recursive, and does not include any heap
3766
/// allocations contained within the Vec's elements. Does not include the
3867
/// size of `self`
68+
///
69+
/// # Example:
70+
/// ```
71+
/// # use datafusion_execution::memory_pool::proxy::VecAllocExt;
72+
/// let mut vec = Vec::new();
73+
/// // Push data into the vec and the accounting will be updated to reflect
74+
/// // memory allocation
75+
/// vec.push(1);
76+
/// assert_eq!(vec.allocated_size(), 16); // space for 4 u32s
77+
/// vec.push(1);
78+
/// assert_eq!(vec.allocated_size(), 16); // no new allocation needed
79+
///
80+
/// // push more data into the vec
81+
/// for _ in 0..10 { vec.push(1); }
82+
/// assert_eq!(vec.allocated_size(), 64); // space for 64 now
83+
/// ```
3984
fn allocated_size(&self) -> usize;
4085
}
4186

4287
impl<T> VecAllocExt for Vec<T> {
4388
type T = T;
4489

4590
fn push_accounted(&mut self, x: Self::T, accounting: &mut usize) {
46-
if self.capacity() == self.len() {
47-
// allocate more
48-
49-
// growth factor: 2, but at least 2 elements
50-
let bump_elements = (self.capacity() * 2).max(2);
51-
let bump_size = std::mem::size_of::<u32>() * bump_elements;
52-
self.reserve(bump_elements);
91+
let prev_capacty = self.capacity();
92+
self.push(x);
93+
let new_capacity = self.capacity();
94+
if new_capacity > prev_capacty {
95+
// capacity changed, so we allocated more
96+
let bump_size = (new_capacity - prev_capacty) * std::mem::size_of::<T>();
97+
// Note multiplication should never overflow because `push` would
98+
// have panic'd first, but the checked_add could potentially
99+
// overflow since accounting could be tracking additional values, and
100+
// could be greater than what is stored in the Vec
53101
*accounting = (*accounting).checked_add(bump_size).expect("overflow");
54102
}
55-
56-
self.push(x);
57103
}
58104
fn allocated_size(&self) -> usize {
59105
std::mem::size_of::<T>() * self.capacity()
@@ -69,7 +115,7 @@ pub trait RawTableAllocExt {
69115
/// `accounting` by any newly allocated bytes.
70116
///
71117
/// Returns the bucket where the element was inserted.
72-
/// Note that allocation counts capacity, not size.
118+
/// Note that allocation counts capacity, not size.
73119
///
74120
/// # Example:
75121
/// ```

0 commit comments

Comments
 (0)