alamb commented on code in PR #24319:
URL: https://github.com/apache/datafusion/pull/24319#discussion_r3784785704
##########
datafusion/common/src/utils/memory.rs:
##########
@@ -185,31 +194,225 @@ impl RecordBatchMemoryCounter {
}
}
-/// 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>>,
+/// Tracks a small number of buffers inline, avoiding a heap allocation for
+/// typical batches, and promotes to a hash set when more buffers are seen.
+#[derive(Debug)]
+struct BufferIdSet {
+ inline: [Option<NonZero<usize>>; INLINE_BUFFER_IDS],
+ len: usize,
+ overflow: Option<HashSet<NonZero<usize>>>,
+}
+
+impl Default for BufferIdSet {
+ fn default() -> Self {
+ Self {
+ inline: [None; INLINE_BUFFER_IDS],
+ len: 0,
+ overflow: None,
+ }
+ }
+}
+
+impl BufferIdSet {
+ fn insert(&mut self, buffer_id: NonZero<usize>) -> bool {
+ if let Some(overflow) = &mut self.overflow {
+ return overflow.insert(buffer_id);
+ }
+
+ if self.inline[..self.len].contains(&Some(buffer_id)) {
+ return false;
+ }
+
+ if self.len < INLINE_BUFFER_IDS {
+ self.inline[self.len] = Some(buffer_id);
+ self.len += 1;
+ return true;
+ }
+
+ let mut overflow = HashSet::with_capacity(INLINE_BUFFER_IDS + 1);
+ overflow.extend(self.inline.iter().flatten().copied());
+ let inserted = overflow.insert(buffer_id);
+ self.overflow = Some(overflow);
+ inserted
+ }
+}
+
+fn count_buffer_memory_size(
+ buffer: &Buffer,
+ counted_buffers: &mut BufferIdSet,
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
+ if counted_buffers.insert(buffer.data_ptr().addr()) {
+ *total_size += 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(
+ array: &dyn Array,
+ counted_buffers: &mut BufferIdSet,
+ total_size: &mut usize,
+) {
+ if let Some(nulls) = array.nulls() {
+ count_buffer_memory_size(nulls.buffer(), counted_buffers, total_size);
}
- // Count all children `ArrayData` recursively
- for child in array_data.child_data() {
- count_array_data_memory_size(child, counted_buffers, total_size);
+ downcast_primitive_array! {
+ array => count_buffer_memory_size(
+ array.values().inner(),
+ counted_buffers,
+ total_size,
+ ),
+ DataType::Null => {}
+ DataType::Boolean => count_buffer_memory_size(
+ array.as_boolean().values().inner(),
+ counted_buffers,
+ total_size,
+ ),
+ DataType::Binary => count_byte_array_memory_size(
+ array.as_binary::<i32>(),
+ counted_buffers,
+ total_size,
+ ),
+ DataType::LargeBinary => count_byte_array_memory_size(
+ array.as_binary::<i64>(),
+ counted_buffers,
+ total_size,
+ ),
+ DataType::Utf8 => count_byte_array_memory_size(
+ array.as_string::<i32>(),
+ counted_buffers,
+ total_size,
+ ),
+ DataType::LargeUtf8 => count_byte_array_memory_size(
+ array.as_string::<i64>(),
+ counted_buffers,
+ total_size,
+ ),
+ DataType::BinaryView => {
+ let array = array.as_binary_view();
+ count_buffer_memory_size(array.views().inner(), counted_buffers,
total_size);
+ for buffer in array.data_buffers() {
+ count_buffer_memory_size(buffer, counted_buffers, total_size);
+ }
+ }
+ DataType::Utf8View => {
+ let array = array.as_string_view();
+ count_buffer_memory_size(array.views().inner(), counted_buffers,
total_size);
+ for buffer in array.data_buffers() {
+ count_buffer_memory_size(buffer, counted_buffers, total_size);
+ }
+ }
+ DataType::FixedSizeBinary(_) => count_buffer_memory_size(
+ array.as_fixed_size_binary().values(),
+ counted_buffers,
+ total_size,
+ ),
+ DataType::List(_) => count_list_array_memory_size(
+ array.as_list::<i32>(),
+ counted_buffers,
+ total_size,
+ ),
+ DataType::LargeList(_) => count_list_array_memory_size(
+ array.as_list::<i64>(),
+ counted_buffers,
+ total_size,
+ ),
+ DataType::ListView(_) => {
+ let array = array.as_list_view::<i32>();
+ count_buffer_memory_size(array.offsets().inner(), counted_buffers,
total_size);
+ count_buffer_memory_size(array.sizes().inner(), counted_buffers,
total_size);
+ count_array_memory_size(array.values().as_ref(), counted_buffers,
total_size);
+ }
+ DataType::LargeListView(_) => {
+ let array = array.as_list_view::<i64>();
+ count_buffer_memory_size(array.offsets().inner(), counted_buffers,
total_size);
+ count_buffer_memory_size(array.sizes().inner(), counted_buffers,
total_size);
+ count_array_memory_size(array.values().as_ref(), counted_buffers,
total_size);
+ }
+ DataType::FixedSizeList(_, _) => count_array_memory_size(
+ array.as_fixed_size_list().values().as_ref(),
+ counted_buffers,
+ total_size,
+ ),
+ DataType::Struct(_) => {
+ for child in array.as_struct().columns() {
+ count_array_memory_size(child.as_ref(), counted_buffers,
total_size);
+ }
+ }
+ DataType::Union(_, _) => {
+ let array = array.as_union();
+ count_buffer_memory_size(array.type_ids().inner(),
counted_buffers, total_size);
+ if let Some(offsets) = array.offsets() {
+ count_buffer_memory_size(offsets.inner(), counted_buffers,
total_size);
+ }
+ for (type_id, _) in array.fields().iter() {
+ count_array_memory_size(
+ array.child(type_id).as_ref(),
+ counted_buffers,
+ total_size,
+ );
+ }
+ }
+ DataType::Dictionary(_, _) => {
+ let array = array.as_any_dictionary();
+ count_array_memory_size(array.keys(), counted_buffers, total_size);
+ count_array_memory_size(array.values().as_ref(), counted_buffers,
total_size);
+ }
+ DataType::Map(_, _) => {
+ let array = array.as_map();
+ count_buffer_memory_size(
+ array.offsets().inner().inner(),
+ counted_buffers,
+ total_size,
+ );
+ count_array_memory_size(array.entries(), counted_buffers,
total_size);
+ }
+ DataType::RunEndEncoded(_, _) => downcast_run_array! {
+ array => {
+ count_buffer_memory_size(
+ array.run_ends().inner().inner(),
+ counted_buffers,
+ total_size,
+ );
+ count_array_memory_size(
+ array.values().as_ref(),
+ counted_buffers,
+ total_size,
+ );
+ },
+ _ => unreachable!(),
Review Comment:
why is this unreachable? An unreachable results in a panic at runtime
I suggest:
1. explicitly list out all array types that are not uspported (rather than
`_` catch all) so it is explict what is not being handled
2. (probably) implement support for the relevant types (or leave a TODO
comment referencing a ticket to support them)
##########
datafusion/common/src/utils/memory.rs:
##########
@@ -185,31 +194,225 @@ impl RecordBatchMemoryCounter {
}
}
-/// 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>>,
+/// Tracks a small number of buffers inline, avoiding a heap allocation for
+/// typical batches, and promotes to a hash set when more buffers are seen.
+#[derive(Debug)]
+struct BufferIdSet {
+ inline: [Option<NonZero<usize>>; INLINE_BUFFER_IDS],
+ len: usize,
+ overflow: Option<HashSet<NonZero<usize>>>,
+}
+
+impl Default for BufferIdSet {
+ fn default() -> Self {
+ Self {
+ inline: [None; INLINE_BUFFER_IDS],
+ len: 0,
+ overflow: None,
+ }
+ }
+}
+
+impl BufferIdSet {
+ fn insert(&mut self, buffer_id: NonZero<usize>) -> bool {
+ if let Some(overflow) = &mut self.overflow {
+ return overflow.insert(buffer_id);
+ }
+
+ if self.inline[..self.len].contains(&Some(buffer_id)) {
+ return false;
+ }
+
+ if self.len < INLINE_BUFFER_IDS {
+ self.inline[self.len] = Some(buffer_id);
+ self.len += 1;
+ return true;
+ }
+
+ let mut overflow = HashSet::with_capacity(INLINE_BUFFER_IDS + 1);
+ overflow.extend(self.inline.iter().flatten().copied());
+ let inserted = overflow.insert(buffer_id);
+ self.overflow = Some(overflow);
+ inserted
+ }
+}
+
+fn count_buffer_memory_size(
Review Comment:
This looks like it could be a method on RecordBatchMemoryCounter as it is
taking two mut fields from it and updating them (I think it is logically part
of that logic). Not required, just a suggestion
##########
datafusion/common/benches/record_batch_memory.rs:
##########
@@ -0,0 +1,146 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements. See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership. The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied. See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+use std::hint::black_box;
+use std::sync::Arc;
+
+use arrow::array::{ArrayRef, Int64Array, ListArray, StructArray};
+use arrow::datatypes::{DataType, Field, Int64Type, Schema};
+use arrow::record_batch::RecordBatch;
+use criterion::{BenchmarkId, Criterion, criterion_group, criterion_main};
+use datafusion_common::utils::memory::get_record_batch_memory_size;
+
+fn make_batch(columns: Vec<ArrayRef>) -> RecordBatch {
+ let fields = columns
+ .iter()
+ .enumerate()
+ .map(|(index, column)| {
+ Field::new(format!("col_{index}"), column.data_type().clone(),
false)
+ })
+ .collect::<Vec<_>>();
+
+ RecordBatch::try_new(Arc::new(Schema::new(fields)), columns).unwrap()
+}
+
+fn make_primitive_batch(num_rows: usize, num_columns: usize) -> RecordBatch {
+ let columns = (0..num_columns)
+ .map(|index| {
+ Arc::new(Int64Array::from_iter_values(
+ (0..num_rows).map(|value| value as i64 + index as i64),
+ )) as ArrayRef
+ })
+ .collect::<Vec<_>>();
+
+ make_batch(columns)
+}
+
+fn make_list_batch(num_rows: usize, num_columns: usize) -> RecordBatch {
+ let columns = (0..num_columns)
+ .map(|column| {
+ Arc::new(ListArray::from_iter_primitive::<Int64Type, _, _>(
+ (0..num_rows).map(|row| {
+ let value = row as i64 + column as i64;
+ Some(vec![Some(value), Some(value + 1)])
+ }),
+ )) as ArrayRef
+ })
+ .collect::<Vec<_>>();
+
+ make_batch(columns)
+}
+
+fn make_struct_batch(num_rows: usize, num_columns: usize) -> RecordBatch {
+ let columns = (0..num_columns)
+ .map(|column| {
+ let left = Arc::new(Int64Array::from_iter_values(
+ (0..num_rows).map(|row| row as i64 + column as i64),
+ )) as ArrayRef;
+ let right = Arc::new(Int64Array::from_iter_values(
+ (0..num_rows).map(|row| row as i64 - column as i64),
+ )) as ArrayRef;
+
+ Arc::new(StructArray::from(vec![
+ (Arc::new(Field::new("left", DataType::Int64, false)), left),
+ (Arc::new(Field::new("right", DataType::Int64, false)), right),
+ ])) as ArrayRef
+ })
+ .collect::<Vec<_>>();
+
+ make_batch(columns)
+}
+
+fn benchmark_column_count(c: &mut Criterion) {
+ let mut group = c.benchmark_group("record_batch_memory_size/column_count");
+
+ for num_columns in [1, 4, 16, 64] {
+ let batch = make_primitive_batch(8192, num_columns);
+ group.bench_with_input(
+ BenchmarkId::from_parameter(num_columns),
+ &batch,
+ |bencher, batch| {
+ bencher.iter(||
get_record_batch_memory_size(black_box(batch)));
+ },
+ );
+ }
+
+ group.finish();
+}
+
+fn benchmark_row_count(c: &mut Criterion) {
+ let mut group = c.benchmark_group("record_batch_memory_size/row_count");
+
+ for num_rows in [1, 128, 8192, 65_536] {
+ let batch = make_primitive_batch(num_rows, 4);
+ group.bench_with_input(
+ BenchmarkId::from_parameter(num_rows),
+ &batch,
+ |bencher, batch| {
+ bencher.iter(||
get_record_batch_memory_size(black_box(batch)));
+ },
+ );
+ }
+
+ group.finish();
+}
+
+fn benchmark_array_layout(c: &mut Criterion) {
+ let mut group = c.benchmark_group("record_batch_memory_size/array_layout");
+
+ for (name, batch) in [
+ ("primitive", make_primitive_batch(8192, 4)),
+ ("list", make_list_batch(8192, 4)),
+ ("struct", make_struct_batch(8192, 4)),
+ ] {
+ group.bench_with_input(
+ BenchmarkId::from_parameter(name),
+ &batch,
+ |bencher, batch| {
+ bencher.iter(||
get_record_batch_memory_size(black_box(batch)));
+ },
+ );
+ }
+
+ group.finish();
+}
+
+criterion_group!(
Review Comment:
If we want to keep this benchmark, can you please add some comments
explaining what it is testing (specifically the overhead of measuring memory
usage in record batches)?
##########
datafusion/common/src/utils/memory.rs:
##########
@@ -185,31 +194,225 @@ impl RecordBatchMemoryCounter {
}
}
-/// 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>>,
+/// Tracks a small number of buffers inline, avoiding a heap allocation for
Review Comment:
👍
##########
datafusion/common/src/utils/memory.rs:
##########
@@ -185,31 +194,225 @@ impl RecordBatchMemoryCounter {
}
}
-/// 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>>,
+/// Tracks a small number of buffers inline, avoiding a heap allocation for
+/// typical batches, and promotes to a hash set when more buffers are seen.
+#[derive(Debug)]
+struct BufferIdSet {
+ inline: [Option<NonZero<usize>>; INLINE_BUFFER_IDS],
+ len: usize,
+ overflow: Option<HashSet<NonZero<usize>>>,
+}
+
+impl Default for BufferIdSet {
+ fn default() -> Self {
+ Self {
+ inline: [None; INLINE_BUFFER_IDS],
+ len: 0,
+ overflow: None,
+ }
+ }
+}
+
+impl BufferIdSet {
+ fn insert(&mut self, buffer_id: NonZero<usize>) -> bool {
+ if let Some(overflow) = &mut self.overflow {
+ return overflow.insert(buffer_id);
+ }
+
+ if self.inline[..self.len].contains(&Some(buffer_id)) {
+ return false;
+ }
+
+ if self.len < INLINE_BUFFER_IDS {
+ self.inline[self.len] = Some(buffer_id);
+ self.len += 1;
+ return true;
+ }
+
+ let mut overflow = HashSet::with_capacity(INLINE_BUFFER_IDS + 1);
+ overflow.extend(self.inline.iter().flatten().copied());
+ let inserted = overflow.insert(buffer_id);
+ self.overflow = Some(overflow);
+ inserted
+ }
+}
+
+fn count_buffer_memory_size(
+ buffer: &Buffer,
+ counted_buffers: &mut BufferIdSet,
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
+ if counted_buffers.insert(buffer.data_ptr().addr()) {
+ *total_size += 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(
+ array: &dyn Array,
+ counted_buffers: &mut BufferIdSet,
+ total_size: &mut usize,
+) {
+ if let Some(nulls) = array.nulls() {
+ count_buffer_memory_size(nulls.buffer(), counted_buffers, total_size);
}
- // Count all children `ArrayData` recursively
- for child in array_data.child_data() {
- count_array_data_memory_size(child, counted_buffers, total_size);
+ downcast_primitive_array! {
+ array => count_buffer_memory_size(
+ array.values().inner(),
+ counted_buffers,
+ total_size,
+ ),
+ DataType::Null => {}
+ DataType::Boolean => count_buffer_memory_size(
+ array.as_boolean().values().inner(),
+ counted_buffers,
+ total_size,
+ ),
+ DataType::Binary => count_byte_array_memory_size(
+ array.as_binary::<i32>(),
+ counted_buffers,
+ total_size,
+ ),
+ DataType::LargeBinary => count_byte_array_memory_size(
+ array.as_binary::<i64>(),
+ counted_buffers,
+ total_size,
+ ),
+ DataType::Utf8 => count_byte_array_memory_size(
+ array.as_string::<i32>(),
+ counted_buffers,
+ total_size,
+ ),
+ DataType::LargeUtf8 => count_byte_array_memory_size(
+ array.as_string::<i64>(),
+ counted_buffers,
+ total_size,
+ ),
+ DataType::BinaryView => {
+ let array = array.as_binary_view();
+ count_buffer_memory_size(array.views().inner(), counted_buffers,
total_size);
+ for buffer in array.data_buffers() {
+ count_buffer_memory_size(buffer, counted_buffers, total_size);
+ }
+ }
+ DataType::Utf8View => {
+ let array = array.as_string_view();
+ count_buffer_memory_size(array.views().inner(), counted_buffers,
total_size);
+ for buffer in array.data_buffers() {
+ count_buffer_memory_size(buffer, counted_buffers, total_size);
+ }
+ }
+ DataType::FixedSizeBinary(_) => count_buffer_memory_size(
+ array.as_fixed_size_binary().values(),
+ counted_buffers,
+ total_size,
+ ),
+ DataType::List(_) => count_list_array_memory_size(
+ array.as_list::<i32>(),
+ counted_buffers,
+ total_size,
+ ),
+ DataType::LargeList(_) => count_list_array_memory_size(
+ array.as_list::<i64>(),
+ counted_buffers,
+ total_size,
+ ),
+ DataType::ListView(_) => {
+ let array = array.as_list_view::<i32>();
+ count_buffer_memory_size(array.offsets().inner(), counted_buffers,
total_size);
+ count_buffer_memory_size(array.sizes().inner(), counted_buffers,
total_size);
+ count_array_memory_size(array.values().as_ref(), counted_buffers,
total_size);
+ }
+ DataType::LargeListView(_) => {
+ let array = array.as_list_view::<i64>();
+ count_buffer_memory_size(array.offsets().inner(), counted_buffers,
total_size);
+ count_buffer_memory_size(array.sizes().inner(), counted_buffers,
total_size);
+ count_array_memory_size(array.values().as_ref(), counted_buffers,
total_size);
+ }
+ DataType::FixedSizeList(_, _) => count_array_memory_size(
+ array.as_fixed_size_list().values().as_ref(),
+ counted_buffers,
+ total_size,
+ ),
+ DataType::Struct(_) => {
+ for child in array.as_struct().columns() {
+ count_array_memory_size(child.as_ref(), counted_buffers,
total_size);
+ }
+ }
+ DataType::Union(_, _) => {
+ let array = array.as_union();
+ count_buffer_memory_size(array.type_ids().inner(),
counted_buffers, total_size);
+ if let Some(offsets) = array.offsets() {
+ count_buffer_memory_size(offsets.inner(), counted_buffers,
total_size);
+ }
+ for (type_id, _) in array.fields().iter() {
+ count_array_memory_size(
+ array.child(type_id).as_ref(),
+ counted_buffers,
+ total_size,
+ );
+ }
+ }
+ DataType::Dictionary(_, _) => {
+ let array = array.as_any_dictionary();
+ count_array_memory_size(array.keys(), counted_buffers, total_size);
+ count_array_memory_size(array.values().as_ref(), counted_buffers,
total_size);
+ }
+ DataType::Map(_, _) => {
+ let array = array.as_map();
+ count_buffer_memory_size(
+ array.offsets().inner().inner(),
+ counted_buffers,
+ total_size,
+ );
+ count_array_memory_size(array.entries(), counted_buffers,
total_size);
+ }
+ DataType::RunEndEncoded(_, _) => downcast_run_array! {
+ array => {
+ count_buffer_memory_size(
+ array.run_ends().inner().inner(),
+ counted_buffers,
+ total_size,
+ );
+ count_array_memory_size(
+ array.values().as_ref(),
+ counted_buffers,
+ total_size,
+ );
+ },
+ _ => unreachable!(),
+ }
+ _ => unreachable!("unsupported array type: {}", array.data_type()),
}
}
+fn count_byte_array_memory_size<T: ByteArrayType>(
+ array: &arrow::array::GenericByteArray<T>,
+ counted_buffers: &mut BufferIdSet,
Review Comment:
similarly I think this could be a method on RecordBatchMemoryCounter
##########
datafusion/common/src/utils/memory.rs:
##########
@@ -185,31 +194,225 @@ impl RecordBatchMemoryCounter {
}
}
-/// 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>>,
+/// Tracks a small number of buffers inline, avoiding a heap allocation for
+/// typical batches, and promotes to a hash set when more buffers are seen.
+#[derive(Debug)]
+struct BufferIdSet {
+ inline: [Option<NonZero<usize>>; INLINE_BUFFER_IDS],
+ len: usize,
+ overflow: Option<HashSet<NonZero<usize>>>,
+}
+
+impl Default for BufferIdSet {
+ fn default() -> Self {
+ Self {
+ inline: [None; INLINE_BUFFER_IDS],
+ len: 0,
+ overflow: None,
+ }
+ }
+}
+
+impl BufferIdSet {
+ fn insert(&mut self, buffer_id: NonZero<usize>) -> bool {
+ if let Some(overflow) = &mut self.overflow {
+ return overflow.insert(buffer_id);
+ }
+
+ if self.inline[..self.len].contains(&Some(buffer_id)) {
+ return false;
+ }
+
+ if self.len < INLINE_BUFFER_IDS {
+ self.inline[self.len] = Some(buffer_id);
+ self.len += 1;
+ return true;
+ }
+
+ let mut overflow = HashSet::with_capacity(INLINE_BUFFER_IDS + 1);
+ overflow.extend(self.inline.iter().flatten().copied());
+ let inserted = overflow.insert(buffer_id);
+ self.overflow = Some(overflow);
+ inserted
+ }
+}
+
+fn count_buffer_memory_size(
+ buffer: &Buffer,
+ counted_buffers: &mut BufferIdSet,
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
+ if counted_buffers.insert(buffer.data_ptr().addr()) {
+ *total_size += 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(
+ array: &dyn Array,
+ counted_buffers: &mut BufferIdSet,
+ total_size: &mut usize,
+) {
+ if let Some(nulls) = array.nulls() {
+ count_buffer_memory_size(nulls.buffer(), counted_buffers, total_size);
}
- // Count all children `ArrayData` recursively
- for child in array_data.child_data() {
- count_array_data_memory_size(child, counted_buffers, total_size);
+ downcast_primitive_array! {
+ array => count_buffer_memory_size(
+ array.values().inner(),
+ counted_buffers,
+ total_size,
+ ),
+ DataType::Null => {}
+ DataType::Boolean => count_buffer_memory_size(
+ array.as_boolean().values().inner(),
+ counted_buffers,
+ total_size,
+ ),
+ DataType::Binary => count_byte_array_memory_size(
+ array.as_binary::<i32>(),
+ counted_buffers,
+ total_size,
+ ),
+ DataType::LargeBinary => count_byte_array_memory_size(
+ array.as_binary::<i64>(),
+ counted_buffers,
+ total_size,
+ ),
+ DataType::Utf8 => count_byte_array_memory_size(
+ array.as_string::<i32>(),
+ counted_buffers,
+ total_size,
+ ),
+ DataType::LargeUtf8 => count_byte_array_memory_size(
+ array.as_string::<i64>(),
+ counted_buffers,
+ total_size,
+ ),
+ DataType::BinaryView => {
+ let array = array.as_binary_view();
+ count_buffer_memory_size(array.views().inner(), counted_buffers,
total_size);
+ for buffer in array.data_buffers() {
+ count_buffer_memory_size(buffer, counted_buffers, total_size);
+ }
+ }
+ DataType::Utf8View => {
+ let array = array.as_string_view();
+ count_buffer_memory_size(array.views().inner(), counted_buffers,
total_size);
+ for buffer in array.data_buffers() {
+ count_buffer_memory_size(buffer, counted_buffers, total_size);
+ }
+ }
+ DataType::FixedSizeBinary(_) => count_buffer_memory_size(
+ array.as_fixed_size_binary().values(),
+ counted_buffers,
+ total_size,
+ ),
+ DataType::List(_) => count_list_array_memory_size(
+ array.as_list::<i32>(),
+ counted_buffers,
+ total_size,
+ ),
+ DataType::LargeList(_) => count_list_array_memory_size(
+ array.as_list::<i64>(),
+ counted_buffers,
+ total_size,
+ ),
+ DataType::ListView(_) => {
+ let array = array.as_list_view::<i32>();
+ count_buffer_memory_size(array.offsets().inner(), counted_buffers,
total_size);
+ count_buffer_memory_size(array.sizes().inner(), counted_buffers,
total_size);
+ count_array_memory_size(array.values().as_ref(), counted_buffers,
total_size);
+ }
+ DataType::LargeListView(_) => {
+ let array = array.as_list_view::<i64>();
+ count_buffer_memory_size(array.offsets().inner(), counted_buffers,
total_size);
+ count_buffer_memory_size(array.sizes().inner(), counted_buffers,
total_size);
+ count_array_memory_size(array.values().as_ref(), counted_buffers,
total_size);
+ }
+ DataType::FixedSizeList(_, _) => count_array_memory_size(
+ array.as_fixed_size_list().values().as_ref(),
+ counted_buffers,
+ total_size,
+ ),
+ DataType::Struct(_) => {
+ for child in array.as_struct().columns() {
+ count_array_memory_size(child.as_ref(), counted_buffers,
total_size);
+ }
+ }
+ DataType::Union(_, _) => {
+ let array = array.as_union();
+ count_buffer_memory_size(array.type_ids().inner(),
counted_buffers, total_size);
+ if let Some(offsets) = array.offsets() {
+ count_buffer_memory_size(offsets.inner(), counted_buffers,
total_size);
+ }
+ for (type_id, _) in array.fields().iter() {
+ count_array_memory_size(
+ array.child(type_id).as_ref(),
+ counted_buffers,
+ total_size,
+ );
+ }
+ }
+ DataType::Dictionary(_, _) => {
+ let array = array.as_any_dictionary();
+ count_array_memory_size(array.keys(), counted_buffers, total_size);
+ count_array_memory_size(array.values().as_ref(), counted_buffers,
total_size);
+ }
+ DataType::Map(_, _) => {
+ let array = array.as_map();
+ count_buffer_memory_size(
+ array.offsets().inner().inner(),
+ counted_buffers,
+ total_size,
+ );
+ count_array_memory_size(array.entries(), counted_buffers,
total_size);
+ }
+ DataType::RunEndEncoded(_, _) => downcast_run_array! {
+ array => {
+ count_buffer_memory_size(
+ array.run_ends().inner().inner(),
+ counted_buffers,
+ total_size,
+ );
+ count_array_memory_size(
+ array.values().as_ref(),
+ counted_buffers,
+ total_size,
+ );
+ },
+ _ => unreachable!(),
+ }
+ _ => unreachable!("unsupported array type: {}", array.data_type()),
}
}
+fn count_byte_array_memory_size<T: ByteArrayType>(
+ array: &arrow::array::GenericByteArray<T>,
+ counted_buffers: &mut BufferIdSet,
+ total_size: &mut usize,
+) {
+ count_buffer_memory_size(
+ array.offsets().inner().inner(),
+ counted_buffers,
+ total_size,
+ );
+ count_buffer_memory_size(array.values(), counted_buffers, total_size);
+}
+
+fn count_list_array_memory_size<O: arrow::array::OffsetSizeTrait>(
Review Comment:
list view and byte view arrays also probably need special handling
--
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]
