alamb commented on code in PR #24319:
URL: https://github.com/apache/datafusion/pull/24319#discussion_r3791617447
##########
datafusion/common/src/utils/memory.rs:
##########
@@ -164,49 +177,229 @@ impl RecordBatchMemoryCounter {
/// Count `batch`, returning the memory used by its buffers that have not
/// been counted before.
pub fn count_batch(&mut self, batch: &RecordBatch) -> usize {
- let mut total_size = 0;
+ let previous_memory_usage = self.memory_usage;
for array in batch.columns() {
- let array_data = array.to_data();
- count_array_data_memory_size(
- &array_data,
- &mut self.counted_buffers,
- &mut total_size,
- );
+ self.count_array_memory_size(array.as_ref());
}
- self.memory_usage += total_size;
- total_size
+ self.memory_usage - previous_memory_usage
}
/// Total memory of the unique buffers of all batches counted so far.
pub fn memory_usage(&self) -> usize {
self.memory_usage
}
-}
-/// Count the memory usage of `array_data` and its children recursively.
-fn count_array_data_memory_size(
- array_data: &ArrayData,
- counted_buffers: &mut HashSet<NonZero<usize>>,
- total_size: &mut usize,
-) {
- // Count memory usage for `array_data`
- for buffer in array_data.buffers() {
- if counted_buffers.insert(buffer.data_ptr().addr()) {
- *total_size += buffer.capacity();
- } // Otherwise the buffer's memory is already counted
+ fn count_buffer_memory_size(&mut self, buffer: &Buffer) {
+ if self.counted_buffers.insert(buffer.data_ptr().addr()) {
+ self.memory_usage += buffer.capacity();
+ }
}
- if let Some(null_buffer) = array_data.nulls()
- &&
counted_buffers.insert(null_buffer.inner().inner().data_ptr().addr())
- {
- *total_size += null_buffer.inner().inner().capacity();
+ /// Count the memory usage of `array` and its children recursively.
+ fn count_array_memory_size(&mut self, array: &dyn Array) {
+ if let Some(nulls) = array.nulls() {
+ self.count_buffer_memory_size(nulls.buffer());
+ }
+
+ downcast_primitive_array! {
+ array => self.count_buffer_memory_size(array.values().inner()),
+ DataType::Null => {}
+ DataType::Boolean => {
+
self.count_buffer_memory_size(array.as_boolean().values().inner());
+ }
+ DataType::Binary => {
+ self.count_byte_array_memory_size(array.as_binary::<i32>());
+ }
+ DataType::LargeBinary => {
+ self.count_byte_array_memory_size(array.as_binary::<i64>());
+ }
+ DataType::Utf8 => {
+ self.count_byte_array_memory_size(array.as_string::<i32>());
+ }
+ DataType::LargeUtf8 => {
+ self.count_byte_array_memory_size(array.as_string::<i64>());
+ }
+ DataType::BinaryView => {
+ self.count_byte_view_array_memory_size(array.as_binary_view());
+ }
+ DataType::Utf8View => {
+ self.count_byte_view_array_memory_size(array.as_string_view());
+ }
+ DataType::FixedSizeBinary(_) => {
+
self.count_buffer_memory_size(array.as_fixed_size_binary().values());
+ }
+ DataType::List(_) => {
+ self.count_list_array_memory_size(array.as_list::<i32>());
+ }
+ DataType::LargeList(_) => {
+ self.count_list_array_memory_size(array.as_list::<i64>());
+ }
+ DataType::ListView(_) => {
+
self.count_list_view_array_memory_size(array.as_list_view::<i32>());
+ }
+ DataType::LargeListView(_) => {
+
self.count_list_view_array_memory_size(array.as_list_view::<i64>());
+ }
+ DataType::FixedSizeList(_, _) => {
+ self.count_array_memory_size(
+ array.as_fixed_size_list().values().as_ref(),
+ );
+ }
+ DataType::Struct(_) => {
+ for child in array.as_struct().columns() {
+ self.count_array_memory_size(child.as_ref());
+ }
+ }
+ DataType::Union(_, _) => {
+ let array = array.as_union();
+ self.count_buffer_memory_size(array.type_ids().inner());
+ if let Some(offsets) = array.offsets() {
+ self.count_buffer_memory_size(offsets.inner());
+ }
+ for (type_id, _) in array.fields().iter() {
+
self.count_array_memory_size(array.child(type_id).as_ref());
+ }
+ }
+ DataType::Dictionary(_, _) => {
+ let array = array.as_any_dictionary();
+ self.count_array_memory_size(array.keys());
+ self.count_array_memory_size(array.values().as_ref());
+ }
+ DataType::Map(_, _) => {
+ let array = array.as_map();
+ self.count_buffer_memory_size(array.offsets().inner().inner());
+ self.count_array_memory_size(array.entries());
+ }
+ DataType::RunEndEncoded(run_ends, _) => match run_ends.data_type()
{
+ DataType::Int16 => {
+
self.count_run_array_memory_size::<arrow::datatypes::Int16Type>(
+ array,
+ );
+ }
+ DataType::Int32 => {
+
self.count_run_array_memory_size::<arrow::datatypes::Int32Type>(
+ array,
+ );
+ }
+ DataType::Int64 => {
+
self.count_run_array_memory_size::<arrow::datatypes::Int64Type>(
+ array,
+ );
+ }
+ // Arrow only permits Int16, Int32, and Int64 run-end indexes.
A
+ // custom Array implementation may still expose malformed data;
+ // retain correct accounting for it without panicking.
Review Comment:
👍
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
