@@ -28,6 +28,35 @@ pub trait VecAllocExt {
28
28
/// `accounting` by any newly allocated bytes.
29
29
///
30
30
/// 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
+ /// ```
31
60
fn push_accounted ( & mut self , x : Self :: T , accounting : & mut usize ) ;
32
61
33
62
/// Return the amount of memory allocated by this Vec to store elements
@@ -36,24 +65,41 @@ pub trait VecAllocExt {
36
65
/// Note this calculation is not recursive, and does not include any heap
37
66
/// allocations contained within the Vec's elements. Does not include the
38
67
/// 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
+ /// ```
39
84
fn allocated_size ( & self ) -> usize ;
40
85
}
41
86
42
87
impl < T > VecAllocExt for Vec < T > {
43
88
type T = T ;
44
89
45
90
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
53
101
* accounting = ( * accounting) . checked_add ( bump_size) . expect ( "overflow" ) ;
54
102
}
55
-
56
- self . push ( x) ;
57
103
}
58
104
fn allocated_size ( & self ) -> usize {
59
105
std:: mem:: size_of :: < T > ( ) * self . capacity ( )
@@ -69,7 +115,7 @@ pub trait RawTableAllocExt {
69
115
/// `accounting` by any newly allocated bytes.
70
116
///
71
117
/// Returns the bucket where the element was inserted.
72
- /// Note that allocation counts capacity, not size.
118
+ /// Note that allocation counts capacity, not size.
73
119
///
74
120
/// # Example:
75
121
/// ```
0 commit comments