jayzhan211 commented on code in PR #26149:
URL: https://github.com/apache/datafusion/pull/26149#discussion_r4235831674
##########
datafusion/common/src/utils/memory.rs:
##########
@@ -464,46 +586,111 @@ fn array_children(array: &ArrayRef) -> Vec<&ArrayRef> {
}
}
-/// 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.
+/// Reference-counted buffer tracker with an inline fast path.
+///
+/// Tracks buffer addresses with their reference counts. The first
+/// [`INLINE_BUFFER_IDS`] distinct buffers are stored in a fixed-size inline
+/// array to avoid heap allocation for typical batches. When more buffers are
+/// encountered, the tracker promotes to a heap-allocated `HashMap`.
#[derive(Debug)]
-struct BufferIdSet {
- inline: [Option<NonZero<usize>>; INLINE_BUFFER_IDS],
+struct BufferIdMap {
+ /// Inline storage: `(address, count)` pairs.
+ inline: [(NonZero<usize>, u32); INLINE_BUFFER_IDS],
Review Comment:
Counting regresses against `main`, which #26140 asks to avoid.
`record_batch_memory`, interleaved, 2 runs each:
| case | main | PR |
|---|---|---|
| column_count/4 | 24.9 ns | 27.4 ns (+10%) |
| column_count/16 | 120.2 ns | 129.7 ns (+8%) |
| array_layout/list | 58.1 ns | 65.0 ns (+12%) |
| shared_slices/4 | 693 ns | 805 ns (+16%) |
| shared_slices/16 | 3.64 µs | 4.06 µs (+12%) |
The padded `[(NonZero<usize>, u32); 16]` doubles the inline array (128 → 256
B), and `get_record_batch_memory_size` builds a new counter for every batch.
Separate arrays plus `position()` bring the 16-buffer cases back to `main`, but
the 4-buffer cases stay about 10% slower, so more is needed:
```rs
struct BufferIdMap {
ids: [Option<NonZero<usize>>; INLINE_BUFFER_IDS],
counts: [u32; INLINE_BUFFER_IDS],
len: usize,
overflow: Option<hashbrown::HashMap<NonZero<usize>, u32>>,
}
```
Please also add `main` vs PR numbers to the description, which the issue
lists as an acceptance item
##########
datafusion/common/src/utils/memory.rs:
##########
@@ -410,6 +505,33 @@ fn count_unique_array_object_memory_size(
.sum::<usize>()
}
+/// Inverse of [`count_unique_array_object_memory_size`]: removes array
+/// identities from the tracked set and returns the released overhead.
+fn uncount_unique_array_object_memory_size(
Review Comment:
`counted_arrays` is still a set, so array overhead is released on the first
uncount even when another counted batch holds the same `ArrayRef`
(`batch.clone()`, dictionary `values()` shared across reader batches). Buffers
use refcounts; overhead should too.
Repro (fails on this PR: `memory_usage()` is 12 after the first uncount,
expected 108):
```rs
#[test]
fn overhead_released_while_still_held() {
let col: ArrayRef = Arc::new(Int32Array::from(vec![1, 2, 3]));
let schema = Arc::new(Schema::new(vec![Field::new("v", DataType::Int32,
false)]));
let b1 = RecordBatch::try_new(Arc::clone(&schema),
vec![Arc::clone(&col)]).unwrap();
let b2 = b1.clone();
let mut counter = RecordBatchMemoryCounter::new();
let full = counter.count_batch_with_array_overhead(&b1);
counter.count_batch_with_array_overhead(&b2);
assert_eq!(counter.uncount_batch_with_array_overhead(&b1), 0);
assert_eq!(counter.memory_usage(), full);
assert_eq!(counter.uncount_batch_with_array_overhead(&b2), full);
}
```
Fix: change `counted_arrays` to `HashMap<usize, u32>` (`crate::HashMap`) and
only recurse on 0→1 / 1→0, which keeps count and uncount symmetric:
```diff
- if !counted_arrays.insert(array_ptr) {
- return 0;
- }
+ let count = counted_arrays.entry(array_ptr).or_insert(0);
+ *count += 1;
+ if *count > 1 {
+ return 0;
+ }
```
```diff
- if !counted_arrays.remove(&array_ptr) {
- return 0;
- }
+ let Entry::Occupied(mut entry) = counted_arrays.entry(array_ptr) else {
+ return 0;
+ };
+ *entry.get_mut() -= 1;
+ if *entry.get() > 0 {
+ return 0;
+ }
+ entry.remove();
```
(`use hashbrown::hash_map::Entry;`.)
`test_uncount_batch_with_array_overhead_shared_across_is because `.slice()`
creates new `ArrayRef`s; please add the test above.
##########
datafusion/common/src/utils/memory.rs:
##########
@@ -1103,4 +1290,924 @@ mod record_batch_tests {
let size = get_record_batch_memory_size(&batch);
assert_eq!(size, 8208);
}
+
+ // ---- uncount tests ----
+
+ #[test]
+ fn test_uncount_batch_round_trip() {
Review Comment:
Could you trim the tests? They're ~900 of the added lines, which makes the
PR hard to review, and most per-type round trips repeat coverage the file
already has. `test_array_memory_size_matches_array_data_layouts` already loops
over every layout (view, run-end, union, map, dictionary, ...). Adding a round
trip to its helper covers all of them:
```diff
fn assert_array_memory_size_matches(array: &dyn Array) {
let mut counter = RecordBatchMemoryCounter::new();
- counter.visit_array_buffers(array, BufferOp::Count);
- assert_eq!(counter.memory_usage(), array_data_memory_size(array));
+ let counted = counter.count_array(array);
+ assert_eq!(counted, array_data_memory_size(array));
+ assert_eq!(counter.uncount_array(array), counted);
+ assert_eq!(counter.memory_usage(), 0);
}
```
With that, the standalone round trips can go:
`test_uncount_{batch,array}_round_trip`, `_nested_struct`, `_list_view_array`,
`_large_list_view_array`, `_null_array`, `_boolean_array`, `_binary_array`,
`_large_utf8_array`, `_fixed_size_binary`, `_fixed_size_list`,
`_array_with_nulls`, `_string_array_with_nulls`, `_union_array`,
`_run_end_encoded_array`, `_deeply_nested`, `_empty_array`, `_empty_batch`.
The shared-buffer tests (two slices, view / dictionary / list / map / union
/ run-end / string sharing) could become one table of `(first, second)` arrays
that share buffers. Each row would check: count both → uncount first releases
nothing → uncount second releases everything.
`test_uncount_binary_view_array_shared_data_buffers` duplicates the StringView
case. Please keep the randomized model test and the overflow-promotion,
never-counted, double-uncount and overhead tests.
--
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]
