andygrove commented on code in PR #27097:
URL: https://github.com/apache/datafusion/pull/27097#discussion_r4473625865
##########
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:
Done, thanks.
##########
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:
I measured it: with the counter on `BufferIdMap`, a batch of four columns
took 7–9% longer, about 30 ns against 32 ns in `column_count/4`, `row_count/*`
and `array_layout/primitive`, in each of 8 alternating rounds. The other cases
were within 3.5%. So I kept both, and both scans now use the indexed loop, as
`decrement` already did: `.position()` in both took 4–7% longer on those
batches. On `count_uncount`, moving `increment` to the loop is a wash, 5%
faster with 4 columns and 5% slower with 16. The `BufferIdMap` doc no longer
mentions `position()`.
##########
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:
Done, thanks.
##########
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.
Review Comment:
Done, thanks. I've also fixed the PR description, which said that buffers
outside FFI always agree on capacity.
--
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]
