GenericDistinctBuffer::size() underreports retained HashSet memory #26937

Describe the bug

GenericDistinctBuffer::size() can significantly underestimate retained memory because it calculates the HashSet allocation estimate using self.values.len() rather than self.values.capacity().

A HashSet may retain substantially more allocated capacity than its current number of distinct elements, especially when processing inputs containing many duplicates or merging overlapping distinct states.

The current implementation is:

pub fn size(&self) -> usize {
    let num_elements = self.values.len();
    let fixed_size = size_of_val(self) + size_of_val(&self.values);

    estimate_memory_size::<T::Native>(
        num_elements,
        fixed_size,
    ).unwrap()
}

This implementation has two accounting issues:

1. Retained capacity is undercounted. The estimate is based on the number of stored elements, even when the HashSet retains a much larger allocation. 
2. The HashSet structure is counted twice. size_of_val(self) already includes the embedded self.values field, so adding size_of_val(&self.values) introduces duplicate fixed-size accounting.

As a result, the reported memory usage can differ substantially from the allocation retained by the buffer.

To Reproduce

We reproduced the memory-accounting issue in two DISTINCT accumulator implementations using independently tracked heap allocations.

Both experiments processed 8,192 input values containing 128 distinct values.

The measurements were obtained on Rocky Linux 9.8 (x86_64), using Rust 1.99.1 and a System allocator tracking requested live heap bytes.

Generic Float64 directly reproduces the GenericDistinctBuffer::size() defect described in this issue. The corrected implementation uses hashbrown's allocation_size() to account for retained buckets and control bytes.

Primitive Int64 demonstrates the same capacity-accounting problem in a separate implementation backed by std::collections::HashSet. Its correction uses retained capacity and the existing memory estimator. On the tested platform, this leaves a 16-byte underestimate per allocated table.

Both implementations substantially underestimated retained memory before the correction.

We also verified the effect on memory-limit enforcement:

• With the original accounting, duplicate-heavy DISTINCT queries can retain substantially more memory than their MemoryReservation reports.
• With corrected accounting, restrictive memory budgets appropriately trigger ResourcesExhausted.
• Increasing the budget allows the queries to complete with the expected SQL results.
• Successful and failed query execution paths release their memory reservations correctly.

These results demonstrate improved memory-accounting accuracy and memory-budget enforcement. Exact allocation sizes may vary across platforms and allocator implementations.

Expected behavior

GenericDistinctBuffer::size() should provide a reasonable estimate of its retained memory allocation, including the capacity of its internal HashSet.

The estimate should:

• Account for retained HashSet capacity, including unused allocated buckets.
• Avoid double-counting the embedded HashSet structure.
• Correctly handle empty and non-empty collections.
• Remain consistent with the existing memory-accounting contract.

The estimate need not be byte-exact across different allocators, but it should not substantially underestimate retained capacity.

Additional context

Impact

Inaccurate accumulator memory estimates can cause DataFusion's MemoryPool to under-account retained memory.

This weakens memory-budget enforcement and may allow queries to consume substantially more memory than their reservations indicate.

The proposed correction improves accounting accuracy. It does not directly reduce actual allocations or change the SQL aggregation algorithm.

Related work

• Part of feat: analyze memory usage for size functions #23393 — The broader audit of inaccurate size() implementations in DataFusion.
• Related to Find a solution to prevent processes being OOM killed due to inaccurate memory accounting #25650 — The broader investigation into differences between MemoryPool reservations and actual allocations.
• Related to Fix aggregate accumulator capacity accounting #24099 — Previously merged capacity-accounting fixes for other aggregate accumulators.
• Related to fix: stop pre-allocating and undercounting a hash table per COUNT(DISTINCT) group #24857 — Previously merged changes addressing HashTable allocation and accounting for grouped string COUNT(DISTINCT) aggregates.

This issue concerns the remaining capacity-accounting defect in GenericDistinctBuffer and is independent of those earlier fixes.

The problem was discovered while investigating the multiple-distinct aggregation performance issue #20942.

Willingness to contribute

I plan to contribute a fix for this bug independently