comphead commented on code in PR #27097:
URL: https://github.com/apache/datafusion/pull/27097#discussion_r4469987025
##########
datafusion/common/src/utils/memory.rs:
##########
@@ -237,8 +237,8 @@ impl BufferVisitor for RecordBatchMemoryCounter {
/// Operators that retain batches incrementally and drop them later (sort,
/// window, sort-merge join, TopK) count a batch when they retain it and
/// [`Self::uncount_batch`] it when they drop it. The tracker keeps a reference
-/// count per buffer: a buffer's capacity is added when its count goes from 0
-/// to 1 and subtracted when it returns to 0.
+/// count per buffer: its capacity is added when it is first counted and
+/// released when its count returns to 0.
Review Comment:
This is no longer exact. A later count can raise the charge for a buffer.
```suggestion
/// count per buffer: it is charged the largest capacity seen while its count
/// is positive, and released when its count returns to 0.
```
##########
datafusion/common/src/utils/memory.rs:
##########
@@ -604,93 +602,115 @@ 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.
+/// typical batches, and promotes to a hash map when more buffers are seen.
+///
+/// Keeps the largest capacity seen per address: FFI-imported slices of one
+/// allocation share an address but report the length they were imported with.
#[derive(Debug)]
struct BufferIdSet {
Review Comment:
With `capacities`, this is `BufferIdMap` without `counts`, and `insert`
repeats the scan, append and promotion logic from `increment`. Did you measure
the counter on `BufferIdMap`? That would remove this type and one of the two
promotion paths. The cost is initializing 64 more bytes per counter, so keeping
both is fine if `record_batch_memory` shows a difference.
If they stay separate, use one scan form in both. `insert` now uses the
indexed loop, `increment` still uses `.position()`, and the `BufferIdMap` doc
still refers to the `position()` scan.
##########
datafusion/common/src/utils/memory.rs:
##########
@@ -699,42 +719,43 @@ impl BufferIdMap {
.position(|id| *id == Some(buffer_id))
{
self.counts[i] += 1;
- return false; // Already tracked, count > 1 now.
+ return grow_counted_capacity(&mut self.capacities[i], capacity);
}
// New buffer: try inline first.
if self.len < INLINE_BUFFER_IDS {
self.ids[self.len] = Some(buffer_id);
self.counts[self.len] = 1;
+ self.capacities[self.len] = capacity;
self.len += 1;
- return true;
+ return capacity;
}
// Promote to overflow.
let mut overflow = HashMap::with_capacity(INLINE_BUFFER_IDS + 1);
for i in 0..self.len {
if let Some(id) = self.ids[i] {
- overflow.insert(id, self.counts[i]);
+ overflow.insert(id, (self.counts[i], self.capacities[i]));
}
}
- overflow.insert(buffer_id, 1);
+ overflow.insert(buffer_id, (1, capacity));
self.overflow = Some(overflow);
- true
+ capacity
}
- /// Decrement the reference count for `buffer_id`. Returns `true` when the
- /// count reached zero (i.e., the buffer is no longer tracked). Returns
- /// `false` if the buffer was not tracked or its count is still positive.
- fn decrement(&mut self, buffer_id: NonZero<usize>) -> bool {
+ /// Decrement the reference count for `buffer_id`, returning its counted
+ /// capacity when the count reaches zero and 0 otherwise.
+ fn decrement(&mut self, buffer_id: NonZero<usize>) -> usize {
if let Some(overflow) = &mut self.overflow {
- if let Some(count) = overflow.get_mut(&buffer_id) {
+ if let Some((count, counted)) = overflow.get_mut(&buffer_id) {
*count -= 1;
if *count == 0 {
+ let counted = *counted;
overflow.remove(&buffer_id);
- return true;
+ return counted;
}
}
- return false;
+ return 0;
Review Comment:
This path now reads the counted capacity, so it can use one hash lookup
instead of `get_mut` followed by `remove` (needs `use crate::hash_map::Entry`):
```suggestion
let Entry::Occupied(mut entry) = overflow.entry(buffer_id) else {
return 0;
};
let (count, _) = entry.get_mut();
*count -= 1;
return if *count == 0 { entry.remove().1 } else { 0 };
```
--
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]
