Nook001 opened a new pull request, #26941: URL: https://github.com/apache/datafusion/pull/26941
## Which issue does this PR close? - Closes #26937. - Part of #23393; related to #25650. - Discovered while investigating #20942. This PR addresses memory accounting, not the multi-DISTINCT performance work. ## Rationale for this change Duplicate-heavy `DISTINCT` aggregation can retain substantially more `HashSet` memory than the accumulator reports to DataFusion's `MemoryPool`. A single batch of 8,192 `Float64` values containing 128 distinct values retains 147,472 bytes of `HashSet` heap allocation, while the existing `GenericDistinctBuffer::size()` accounts for only 2,344 bytes of heap. This can allow queries to exceed their intended memory reservation without receiving `ResourcesExhausted`. Conversely, embedded `HashSet` structures are counted twice in some accumulator sizes, resulting in over-accounting of small or empty states. ## What changes are included in this PR? - Use `hashbrown`'s existing `allocation_size()` in `GenericDistinctBuffer::size()` to include retained buckets, control bytes, and unused capacity while counting `Self` once. - Account for embedded child structures exactly once in floating-point `DISTINCT` `COUNT`/`SUM`/`AVG`, decimal `AVG`, and `DISTINCT` variance accumulators. - Include the separately boxed dictionary `COUNT` wrapper's own fixed size. - Use `std::collections::HashSet::capacity()` rather than `len()` in primitive `DISTINCT COUNT`, including correct handling of unallocated empty tables. Aggregation algorithms, intermediate states, SQL results, and public APIs remain unchanged. ## What is the testing strategy for this PR? ### Regression tests The PR includes: - Unit tests for unallocated and reserved-empty tables, retained capacity after insertion or clearing, nested accumulator ownership, and boxed dictionary accounting. - SQLLogicTest coverage for integer and floating-point `DISTINCT` aggregation, including memory-budget exhaustion, nullable inputs, empty inputs, and successful execution with sufficient memory. ### Independent allocation measurements For 8,192 inputs containing 128 distinct values: | Accumulator | Actual retained heap | Previous estimate | Fixed estimate | | :--- | :--- | :--- | :--- | | Generic `Float64` | 147,472 B | 2,344 B | 147,472 B | | Primitive `Int64` | 147,472 B | 2,344 B | 147,456 B | These estimates exclude the accumulator's 64-byte inline structure to compare heap allocations on the same basis. #### Rocky Linux 9.8 — AMD EPYC 9654P, kernel 5.14, Rust 1.99.1 - 128 before/after allocator observations. - 88 SQL comparisons covering results, typed memory-limit errors, boundary budgets, and reservation cleanup. - Ablation checks confirming that the new regression tests detect the corresponding defects. #### CachyOS — AMD Ryzen 5 5600, kernel 7.3.0-rc6, Rust 1.99.1 - Five targeted unit tests passed. - The `aggregate.slt` run completed without reported failures. - An independent `System` allocator probe covered 41 scenarios, including empty tables, duplicate-heavy batches, reserved capacity, clearing, and chunked input. - 8,200 repeated observations across both versions and accumulator types confirmed consistent results. - Generic `Float64` matched requested retained heap allocation in every tested scenario. - Allocated Primitive `Int64` tables consistently underestimated retained heap by 16 bytes. - Representative measurements matched Rocky Linux byte-for-byte. The allocation measurements track requested live heap bytes rather than RSS, allocator metadata, or fragmentation. The standard-library `HashSet` estimator remains approximate. On the tested `x86_64` toolchains, nonempty tables retain a 16-byte estimation residual. ### Repository checks On Rocky Linux: - `cargo fmt --all`: passed. - All-target/all-feature Clippy with `-D warnings`: passed. - All 21 `dev/rust_lint.sh` checks: passed. - All 545 SQLLogicTest files: passed. The extended workspace test run is not fully green: - 13,344 passed. - 1 failed. - 8 ignored. The only failure, `fuzz_cases::sort_query_fuzz::sort_query_fuzzer_runner`, also reproduces deterministically on unmodified upstream. Its reproducer uses a UTF8/dictionary `ORDER BY` query without aggregation and an unbounded memory pool. This unrelated failure remains unresolved. ### Additional validation Native Windows/MSVC validation is in progress. Results will be added before marking this PR ready for review. ## Are there any user-facing changes? Queries whose retained `DISTINCT` state was previously undercharged may now correctly return `ResourcesExhausted` under restrictive memory budgets. Correcting duplicate fixed-size accounting also avoids unnecessary rejections of some empty or small states. This change improves retained-state accounting. It does not reduce physical allocations or guarantee a process RSS hard limit. -- 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]
