Add ExplainMemory Trait for Memory Usage Reporting and Enhance Memory Debugging Tools

[kosiew]

Which issue does this PR close?

• Closes Add a fine grain memory usage tracking / break down of memory usage *within* each operator #16904

Rationale for this change

Datafusion’s memory management intentionally does not track every allocation. The MemoryPool docs explain that only “large” consumers (like grouping or joining) register allocations, while transient buffers are not tracked. This keeps memory accounting simpler and avoids overhead:

/// Rather than tracking all allocations, DataFusion takes a pragmatic approach:
/// Intermediate memory used as data streams through the system is not accounted
/// (it assumed to be "small") but the large consumers of memory must register
/// and constrain their use. This design trades off the additional code
/// complexity of memory tracking with limiting resource usage.

Therefore, DataFusion already has abstractions for tracking memory per consumer.

This PR leverages existing MemoryReservation sizes and report_top to implement debugging helpers that show how much memory each accumulator or structure uses by introducing a standardized way to inspect and debug memory usage across various components of the execution engine. The ExplainMemory trait provides a human-readable summary of memory usage, which is helpful for debugging and monitoring purposes. This also lays the groundwork for improving visibility into memory consumption patterns in complex queries.

What changes are included in this PR?

• Introduced a new module memory_report implementing the ExplainMemory trait.
• Blanket implementation of ExplainMemory for any Accumulator using its size() method.
• Specific implementation of ExplainMemory for MemoryReservation.
• A helper function report_top_consumers that attempts to downcast a memory pool and extract consumer memory usage stats.
• Unit tests verifying the behavior of the ExplainMemory trait.
• Implemented ExplainMemory for GroupedHashAggregateStream to provide detailed insight into memory usage of group aggregates.
• Enabled lz4 and zstd features in datafusion/physical-plan to allow test coverage of compression-related functionality.

Are these changes tested?

Yes:

• Added unit tests for both MemoryReservation and Accumulator implementations of ExplainMemory.
• Ensured the GroupedHashAggregateStream's memory explanation logic compiles and integrates cleanly.
• Existing and new tests validate the memory reporting logic.

Are there any user-facing changes?

Yes:

• Users can now call .explain_memory() on memory reservations and accumulators to obtain human-readable memory usage summaries.
• Enhanced memory reporting visibility for GroupedHashAggregateStream, aiding debugging and optimization efforts.

[rluvaton]

Left some thoughts, thank you for opening the PR!

[rluvaton]

Wouldn't this will cause rust to disallow you to implement it on other trait.

As I wouldn't be able to do:

trait MyTrait {}

impl<T: MyTrait> ExplainMemory for T {
  // ...
}

As rust will say to me that MyTrait can implement Accumulator

[kosiew]

Hey @rluvaton! 👋

You're totally right to be cautious — that blanket impl of ExplainMemory for T where T: Accumulator will indeed block any other downstream impls of ExplainMemory for overlapping traits like T: MyTrait, due to Rust's coherence and orphan rules.

💡 Workaround: You can safely achieve what you want using a newtype pattern. That way, you avoid trait conflicts while still providing a blanket impl.

Here’s a minimal example:

// Newtype wrapper
struct AccumExplainer<T: Accumulator>(T);

// Blanket impl only applies to the wrapper, not all T
impl<T: Accumulator> ExplainMemory for AccumExplainer<T> {
    fn explain(&self) -> String {
        format!("Memory usage: {}", self.0.accumulate())
    }
}

// Convenience (optional)
impl<T: Accumulator> From<T> for AccumExplainer<T> {
    fn from(inner: T) -> Self {
        Self(inner)
    }
}

You can now wrap any Accumulator type like this:

fn explain_it<T: Accumulator>(thing: T) {
    let wrapped = AccumExplainer::from(thing);
    println!("{}", wrapped.explain());
}

This gives you a clean, conflict-free path to implement ExplainMemory broadly without locking out other impls for T in the future.

[rluvaton]

Do you think we should have something similar to formatter? So we will have easier to read memory report?

[rluvaton]

I think it should include size as well so any trait that need to report the size will also can report the breakdown of the size.

[kosiew]

closed by accident.

[kosiew]

Reopening..