diff --git a/arrow-row/src/lib.rs b/arrow-row/src/lib.rs index 72a295627ed2..aa6543485fe3 100644 --- a/arrow-row/src/lib.rs +++ b/arrow-row/src/lib.rs @@ -1131,8 +1131,8 @@ impl Rows { pub fn size(&self) -> usize { // Size of fields is accounted for as part of RowConverter std::mem::size_of::() - + self.buffer.len() - + self.offsets.len() * std::mem::size_of::() + + self.buffer.capacity() + + self.offsets.capacity() * std::mem::size_of::() } /// Create a [BinaryArray] from the [Rows] data without reallocating the @@ -4050,4 +4050,47 @@ mod tests { // "a" < "z" assert!(rows.row(3) < rows.row(1)); } + + #[test] + fn rows_size_should_count_for_capacity() { + let row_converter = RowConverter::new(vec![SortField::new(DataType::UInt8)]).unwrap(); + + let empty_rows_size_with_preallocate_rows_and_data = { + let rows = row_converter.empty_rows(1000, 1000); + + rows.size() + }; + let empty_rows_size_with_preallocate_rows = { + let rows = row_converter.empty_rows(1000, 0); + + rows.size() + }; + let empty_rows_size_with_preallocate_data = { + let rows = row_converter.empty_rows(0, 1000); + + rows.size() + }; + let empty_rows_size_without_preallocate = { + let rows = row_converter.empty_rows(0, 0); + + rows.size() + }; + + assert!( + empty_rows_size_with_preallocate_rows_and_data > empty_rows_size_with_preallocate_rows, + "{empty_rows_size_with_preallocate_rows_and_data} should be larger than {empty_rows_size_with_preallocate_rows}" + ); + assert!( + empty_rows_size_with_preallocate_rows_and_data > empty_rows_size_with_preallocate_data, + "{empty_rows_size_with_preallocate_rows_and_data} should be larger than {empty_rows_size_with_preallocate_data}" + ); + assert!( + empty_rows_size_with_preallocate_rows > empty_rows_size_without_preallocate, + "{empty_rows_size_with_preallocate_rows} should be larger than {empty_rows_size_without_preallocate}" + ); + assert!( + empty_rows_size_with_preallocate_data > empty_rows_size_without_preallocate, + "{empty_rows_size_with_preallocate_data} should be larger than {empty_rows_size_without_preallocate}" + ); + } }