avantgardnerio commented on code in PR #2845:
URL:
https://github.com/apache/datafusion-ballista/pull/2845#discussion_r4542929887
##########
ballista/core/src/execution_plans/buffer.rs:
##########
@@ -315,11 +316,15 @@ fn build_dam_stream(
// `get_array_memory_size()` ever changed between calls.
let mut buffered: VecDeque<(usize, RecordBatch)> = VecDeque::new();
let mut dam_broken = false;
+ // Counts each buffer once across the buffered batches, so zero-copy
+ // slices of one large array (as aggregates emit) are not charged
+ // once per slice.
+ let mut counter = RecordBatchMemoryCounter::new();
// Fill phase: read and buffer until pool pressure or end-of-stream.
while let Some(batch) = input.next().await {
let batch = batch?;
- let size = batch.get_array_memory_size();
+ let size = counter.count_batch(&batch);
Review Comment:
Brent's Claude here.
Nice catch on the over-reservation. I think the counter swings it the other
way during the drain though. `RecordBatchMemoryCounter` charges a shared buffer
to the first slice only. When that slice is popped, `res.shrink(size)` releases
the whole buffer while the other nine slices still hold it, so the pool reads
close to zero for memory that is still live. The new test checks for exactly
that state after the first pull.
DataFusion's docs point at `RecordBatchMemoryTracker` for this case, since
it refcounts each buffer and `uncount_batch` frees it when the last slice
leaves. It lands in DataFusion 56, which is only days away. Would it work to
rebase this onto #2718 and drain with
`res.shrink(tracker.uncount_batch(&batch))`? The test could then expect
`whole_size` to stay reserved until the last slice is gone.
--
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]
