@@ -31,27 +31,51 @@ use super::VIEW_SIZE_BYTES;
3131/// variable-width values, and the elements of its nested values. Buffers that
3232/// every row shares, such as validity bitmaps and dictionary values, are not
3333/// counted.
34- pub ( crate ) struct SpilledRowSizes ( ColumnSize ) ;
34+ pub ( crate ) struct SpilledRowSizes {
35+ /// Bytes that every row takes in the fixed-width columns
36+ fixed : usize ,
37+ /// The other columns
38+ variable : Vec < ColumnSize > ,
39+ /// See [`Self::max_row`]
40+ max_row : usize ,
41+ }
3542
3643impl SpilledRowSizes {
3744 pub ( crate ) fn new ( batch : & RecordBatch ) -> Self {
38- Self ( ColumnSize :: Struct (
39- batch
40- . columns ( )
41- . iter ( )
42- . map ( |column| ColumnSize :: new ( column. as_ref ( ) ) )
43- . collect ( ) ,
44- ) )
45+ let mut fixed = 0 ;
46+ let mut variable = vec ! [ ] ;
47+ for column in batch. columns ( ) {
48+ let column = ColumnSize :: new ( column. as_ref ( ) ) ;
49+ match column. fixed ( ) {
50+ Some ( width) => fixed += width,
51+ None => variable. push ( column) ,
52+ }
53+ }
54+ let max_row = variable. iter ( ) . fold ( fixed, |bytes, column| {
55+ bytes. saturating_add ( column. max_row ( ) )
56+ } ) ;
57+ Self {
58+ fixed,
59+ variable,
60+ max_row,
61+ }
4562 }
4663
4764 /// Returns the size in bytes of row `row`.
4865 pub ( crate ) fn row ( & self , row : usize ) -> usize {
49- self . 0 . rows ( row, row + 1 )
66+ self . variable . iter ( ) . fold ( self . fixed , |bytes, column| {
67+ bytes + column. rows ( row, row + 1 )
68+ } )
5069 }
5170
5271 /// Returns the size in bytes of every row, if they all have the same size.
5372 pub ( crate ) fn fixed ( & self ) -> Option < usize > {
54- self . 0 . fixed ( )
73+ self . variable . is_empty ( ) . then_some ( self . fixed )
74+ }
75+
76+ /// Returns an upper bound on the size in bytes of any row.
77+ pub ( crate ) fn max_row ( & self ) -> usize {
78+ self . max_row
5579 }
5680}
5781
@@ -129,6 +153,30 @@ impl ColumnSize {
129153 }
130154 }
131155
156+ /// Returns an upper bound on the size in bytes of a row.
157+ fn max_row ( & self ) -> usize {
158+ match self {
159+ Self :: Fixed ( width) => * width,
160+ Self :: Bytes ( offsets) => offsets. width ( ) + offsets. max_len ( ) ,
161+ Self :: Views ( views) => {
162+ let max_len = views
163+ . iter ( )
164+ . map ( |& view| view as u32 )
165+ . filter ( |& len| len > MAX_INLINE_VIEW_LEN )
166+ . max ( )
167+ . unwrap_or_default ( ) ;
168+ VIEW_SIZE_BYTES + max_len as usize
169+ }
170+ Self :: List ( offsets, values) => offsets
171+ . width ( )
172+ . saturating_add ( offsets. max_len ( ) . saturating_mul ( values. max_row ( ) ) ) ,
173+ Self :: FixedSizeList ( size, values) => size. saturating_mul ( values. max_row ( ) ) ,
174+ Self :: Struct ( fields) => fields
175+ . iter ( )
176+ . fold ( 0 , |bytes, field| bytes. saturating_add ( field. max_row ( ) ) ) ,
177+ }
178+ }
179+
132180 /// Returns the size in bytes of rows `start..end`.
133181 fn rows ( & self , start : usize , end : usize ) -> usize {
134182 let num_rows = end - start;
@@ -181,6 +229,21 @@ impl Offsets {
181229 Self :: Large ( offsets) => offsets[ index] as usize ,
182230 }
183231 }
232+
233+ /// Returns the largest distance between consecutive offsets.
234+ fn max_len ( & self ) -> usize {
235+ match self {
236+ Self :: Small ( offsets) => offsets
237+ . windows ( 2 )
238+ . map ( |pair| ( pair[ 1 ] - pair[ 0 ] ) as usize )
239+ . max ( ) ,
240+ Self :: Large ( offsets) => offsets
241+ . windows ( 2 )
242+ . map ( |pair| ( pair[ 1 ] - pair[ 0 ] ) as usize )
243+ . max ( ) ,
244+ }
245+ . unwrap_or_default ( )
246+ }
184247}
185248
186249impl From < & OffsetBuffer < i32 > > for Offsets {
@@ -247,6 +310,36 @@ mod tests {
247310 ) ;
248311 }
249312
313+ #[ test]
314+ fn largest_row_bounds_every_row ( ) {
315+ let ints: ArrayRef = Arc :: new ( Int64Array :: from ( vec ! [ 1 , 2 , 3 ] ) ) ;
316+ let strings: ArrayRef = Arc :: new ( StringArray :: from ( vec ! [ "a" , "bbbb" , "" ] ) ) ;
317+ // 8 bytes of Int64, plus a 4 byte offset and the longest string
318+ assert_eq ! (
319+ SpilledRowSizes :: new( & batch( vec![ ints, strings] ) ) . max_row( ) ,
320+ 16
321+ ) ;
322+
323+ let long = "a string longer than twelve" ;
324+ let views: ArrayRef = Arc :: new ( StringViewArray :: from ( vec ! [ "short" , long] ) ) ;
325+ assert_eq ! (
326+ SpilledRowSizes :: new( & batch( vec![ views] ) ) . max_row( ) ,
327+ 16 + long. len( )
328+ ) ;
329+
330+ // An offset plus the most elements in a row, each of the largest size
331+ let list: ArrayRef =
332+ Arc :: new ( ListArray :: from_iter_primitive :: < Int64Type , _ , _ > ( vec ! [
333+ Some ( vec![ Some ( 1 ) , Some ( 2 ) , Some ( 3 ) ] ) ,
334+ Some ( vec![ ] ) ,
335+ Some ( vec![ Some ( 4 ) ] ) ,
336+ ] ) ) ;
337+ assert_eq ! (
338+ SpilledRowSizes :: new( & batch( vec![ list] ) ) . max_row( ) ,
339+ 4 + 3 * 8
340+ ) ;
341+ }
342+
250343 #[ test]
251344 fn views_count_their_bytes_unless_inlined ( ) {
252345 let long = "a string longer than twelve" ;
0 commit comments