ariel-miculas commented on PR #23777: URL: https://github.com/apache/datafusion/pull/23777#issuecomment-5588460282
I see two paths forward: * measure the output bytes using the original large RecordBatch created at emit time, that yields all of the subsequent small RecordBatches as slices, POC here: https://github.com/apache/datafusion/pull/25086 * use the approach in this PR to opt-in to using the RecordBatchMemoryCounter, but this needs fixing the issue with reused allocations for different Buffers For the second approach, we could store a non-owning handle to the allocation (an Arc Weak pointer) to determine whether the allocation is still alive or not: ``` fn count_batch(&mut self, batch: &Batch) -> usize { let before = self.memory_usage; for buffer in batch.columns() { let addr = buffer.data_addr(); let is_live_share = self .seen .get(&addr) .is_some_and(|weak| weak.upgrade().is_some()); if is_live_share { continue; } // Either never seen, or seen but since freed and the address // recycled. Both mean this is new memory. self.seen.insert(addr, buffer.weak()); self.memory_usage += buffer.capacity(); } self.memory_usage - before } ``` where buffer.weak creates a weak pointer to the Arc: ``` pub fn weak(&self) -> Weak<Bytes> { Arc::downgrade(&self.data) } ``` The downside is that we need to patch arrow-rs to have access to the underlying Arc, so we can create the weak reference. @2010YOUY01 what do you think is the better approach between these two? -- 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]
