txwyy123 opened a new pull request, #26149: URL: https://github.com/apache/datafusion/pull/26149
## Which issue does this PR close? Closes #26140. ## Rationale for this change `RecordBatchMemoryCounter` can only count (add) batches but never release them. Operators that retain batches incrementally and drop them later (sort, window, sort-merge join, TopK) can't use it -- they keep their own estimates instead, leading to over-counting (sort: 1488 MB reported for 346 MB actual) or no counting at all (window: 200 MB limit ignored, 1.27 GB RSS). ## What changes are included in this PR? - **`BufferIdSet` → `BufferIdMap`**: internal storage changed from an insert-only set to a reference-counted map (`(NonZero<usize>, u32)` tuples inline, `HashMap<NonZero<usize>, u32>` overflow). Inline fast path for ≤16 buffers preserved. - **New public API**: `uncount_batch()`, `uncount_array()`, `uncount_batch_with_array_overhead()` -- each returns the bytes released. - **Shared visitor**: Extracted `visit_array_buffers(array, op)` with `BufferOp::Count | Uncount` direction, replacing duplicated per-type buffer walks. Same type dispatch (primitives, boolean, binary, utf8, views, lists, list-views, fixed-size, struct, union, dictionary, map, run-end-encoded, fallback). - **52 tests** (22 pre-existing + 30 new) covering every acceptance criterion. ## Are these changes tested? Yes -- comprehensive test suite covering all acceptance criteria: | Criterion | Tests | |---|---| | Count/uncount round trip | `test_uncount_batch_round_trip`, `test_uncount_array_round_trip`, `test_uncount_batch_with_array_overhead_round_trip` | | Two-slice example from issue | `test_uncount_two_slices_example_from_issue` | | View arrays sharing data buffers | `test_uncount_view_array_slices_sharing_data_buffers`, `test_uncount_binary_view_array_shared_data_buffers` | | Dictionaries sharing values | `test_uncount_dictionaries_sharing_values` | | Nested types | struct, list (shared child), map (shared children), union, union slices, run-end-encoded, run-end slices, fixed-size-binary, fixed-size-list, deeply nested (List\<Struct\<Dict\>\>), list-view, large-list-view | | >16 buffers + removals | `test_uncount_with_overflow_promotion` | | Randomized sequence vs reference model | 200-op sequence across 10 array types (Int32, Int64, String, StringView, Float64, Dict, List, Boolean + slices) | | Edge cases | null bitmaps, empty arrays, empty batches, NullArray, boolean (bit-packed), double-uncount (no underflow), never-counted no-op, array-overhead shared across batches | ## Are there any user-facing changes? No. Purely additive API -- existing `count_*` methods and `get_record_batch_memory_size` return identical results. No caller changes needed. ## Benchmark `record_batch_memory` benchmark shows no regression for counting: ``` column_count/1 22 ns column_count/64 2.1 µs shared_slices/4 1.8 µs shared_slices/64 30.1 µs count_uncount/4 3.7 µs (new -- full count+uncount cycle, 32 slices × 4 columns) ``` -- 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]
