2010YOUY01 commented on code in PR #22862:
URL: https://github.com/apache/datafusion/pull/22862#discussion_r3384735604
##########
datafusion/common/src/utils/memory.rs:
##########
@@ -133,12 +133,29 @@ pub fn estimate_memory_size<T>(num_elements: usize,
fixed_size: usize) -> Result
pub fn get_record_batch_memory_size(batch: &RecordBatch) -> usize {
// Store pointers to `Buffer`'s start memory address (instead of actual
// used data region's pointer represented by current `Array`)
- let mut counted_buffers: HashSet<NonNull<u8>> = HashSet::new();
+ let mut counted_buffers: HashSet<usize> = HashSet::new();
+ count_record_batch_memory_size(batch, &mut counted_buffers)
+}
+
+/// Calculate the memory used by `batch`, counting only `Buffer`s whose start
+/// addresses are not already present in `counted_buffers`, and record the
+/// batch's buffer addresses in the set.
+///
+/// Use this instead of [`get_record_batch_memory_size`] to account for the
total
+/// memory of a sequence of batches that may share underlying buffers, e.g.
+/// zero-copy slices of one larger batch. Calling
+/// [`get_record_batch_memory_size`] on each such slice counts the shared
+/// buffers once per slice, while sharing `counted_buffers` across the calls
+/// counts each buffer exactly once.
+pub fn count_record_batch_memory_size(
Review Comment:
We could make it a deeper module like
```rust
/// Tracks memory already accounted for across multiple `RecordBatch`es.
///
/// Some batches may share the same underlying Arrow buffers, for example
when
/// they are zero-copy slices of a larger batch. This counter remembers
buffer
/// start addresses so shared buffers are counted only once.
#[derive(Default)]
struct RecordBatchMemoryCounter {
accounted_buffers: HashSet<usize>,
memory_usage: usize,
}
impl RecordBatchMemoryCounter {
/// Accounts for buffers in `batch` that have not already been seen.
fn count_batch(&mut self, batch: &RecordBatch) {
self.memory_usage += count_record_batch_memory_size(
batch,
&mut self.accounted_buffers,
);
}
/// Returns the total memory accounted for so far.
fn memory_usage(&self) -> usize {
self.memory_usage
}
}
```
The benefit is that users such as `HashJoinExec` do not need to know about
the implementation details (e.g. the `HashSet` used to track buffer addresses).
It also makes the intent clearer for readers who are not already familiar with
this context.
Since this issue is not specific to hash joins, a simpler abstraction could
make it easier to reuse.
--
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]
