ariel-miculas commented on PR #23606: URL: https://github.com/apache/datafusion/pull/23606#issuecomment-5184552103
I had to switch to PeakRecordingPool instead of TrackConsumersPool because the latter was reporting 0 bytes as the peak, the reason being that the memory consumer was being unregistered before we could get the actual value. This used to work previously, before the refactor to SortPreservingMergeStream in https://github.com/apache/datafusion/pull/23407 and https://github.com/apache/datafusion/pull/23976 Longer explanation: SortPreservingMergeStream's reservation is unregistered from the memory pool before the test reads .metrics(), because of how the stream is now implemented. Two commits on main rewrote SortPreservingMergeStream from a manual poll_next state machine into an async generator: - 096012e08 — "Simplifying SortPreservingMergeStream to use generators instead of state machine" (#23407) - 77b172e15 — "simplify SortPreservingMergeStream to be as textbook-like as possible" (#23702) Before: SortPreservingMergeStream implemented Stream directly (`fn poll_next(self: Pin<&mut Self>, ...)`). The struct — including its MemoryReservation — was boxed and returned as the SendableRecordBatchStream itself. It stayed alive for exactly as long as the caller held the stream, and was only dropped (triggering unregister()) when the caller dropped it — in the test, after tracking_pool.metrics() was read. Peak was correctly captured. After: the implementation moved to: ``` fn create_stream(mut self) -> impl Stream<Item = Result<RecordBatch>> { async_try_stream(|mut emitter| async move { /* ...self... */ }) } ``` self (and everything it owns, including the reservation) is now captured inside the generator's async block. Once that block runs to completion — which happens on the final poll_next() call that returns None, i.e. during the test's draining loop, not when the caller later drops the stream object — self is dropped right there. That drop cascades: MemoryReservation::drop → SharedRegistration::drop (once the last Arc clone goes away) → MemoryPool::unregister(). TrackConsumersPool::unregister removes the consumer from its tracked_consumers map entirely so by the time the test calls .metrics() after the loop, the consumer is already gone and its recorded peak is lost. -- 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]
