jayzhan211 commented on code in PR #26149:
URL: https://github.com/apache/datafusion/pull/26149#discussion_r4258326040
##########
datafusion/common/src/utils/memory.rs:
##########
@@ -222,175 +249,247 @@ impl RecordBatchMemoryCounter {
total_size
}
+ /// Release the buffers of `batch`, returning the bytes freed.
+ ///
+ /// Each buffer's reference count is decremented; the buffer's capacity is
+ /// returned (and subtracted from [`Self::memory_usage`]) only when its
+ /// count reaches zero. Buffers that were never counted are silently
+ /// ignored (their count is already zero).
+ pub fn uncount_batch(&mut self, batch: &RecordBatch) -> usize {
+ let previous_memory_usage = self.memory_usage;
+
+ for array in batch.columns() {
+ self.visit_array_buffers(array.as_ref(), BufferOp::Uncount);
+ }
+
+ previous_memory_usage - self.memory_usage
+ }
+
+ /// Release the buffers of `array`, returning the bytes freed.
+ ///
+ /// See [`Self::uncount_batch`] for semantics.
+ pub fn uncount_array(&mut self, array: &dyn Array) -> usize {
+ let previous_memory_usage = self.memory_usage;
+ self.visit_array_buffers(array, BufferOp::Uncount);
+ previous_memory_usage - self.memory_usage
+ }
+
+ /// Inverse of [`Self::count_batch_with_array_overhead`]: releases buffer
+ /// memory and array-object overhead for `batch`.
+ ///
+ /// Array-object overhead is released when the array identity (pointer) is
+ /// removed from the tracked set.
+ pub fn uncount_batch_with_array_overhead(&mut self, batch: &RecordBatch)
-> usize {
+ let mut total_released = self.uncount_batch(batch);
+ let mut array_overhead = 0;
+
+ for array in batch.columns() {
+ array_overhead +=
+ uncount_unique_array_object_memory_size(array, &mut
self.counted_arrays);
+ }
+
+ total_released += array_overhead;
+ self.memory_usage -= array_overhead;
+ total_released
+ }
+
/// Total memory of all counted allocations.
pub fn memory_usage(&self) -> usize {
self.memory_usage
}
- fn count_buffer_memory_size(&mut self, buffer: &Buffer) {
- if self.counted_buffers.insert(buffer.data_ptr().addr()) {
- self.memory_usage += buffer.capacity();
+ /// Apply `op` to `buffer`: increment or decrement its reference count and
+ /// adjust `memory_usage` accordingly.
+ fn apply_buffer_op(&mut self, buffer: &Buffer, op: BufferOp) {
Review Comment:
Splitting `ids`/`counts` alone doesn't fix the regression.
`record_batch_memory` at b4e0528, interleaved, 2 runs each: `shared_slices/4`
676-708 → 799-808 ns (+13-20%), `shared_slices/16` +10-14%,
`array_layout/{list,struct}` +8-12%, `column_count/16` +6-9%, `row_count/8192`
+5-8%. With the change below, tests pass, `/16` cases are at or below `main`,
and 4-buffer cases are within +4-7%:
```diff
+ #[inline(always)]
fn apply_buffer_op(&mut self, buffer: &Buffer, op: BufferOp) {
```
```diff
+ #[inline(always)]
fn increment(&mut self, buffer_id: NonZero<usize>) -> bool {
@@
- let target = Some(buffer_id);
- for i in 0..self.len {
- if self.ids[i] == target {
- self.counts[i] += 1;
- return false; // Already tracked, count > 1 now.
- }
- }
+ if let Some(i) = self.ids[..self.len]
+ .iter()
+ .position(|id| *id == Some(buffer_id))
+ {
+ self.counts[i] += 1;
+ return false;
+ }
```
--
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]
