@@ -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
4287impl < 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