Skip to content

chore: add mimalloc tests to assert size functions#23399

Open
comphead wants to merge 3 commits into
apache:mainfrom
comphead:mem_size
Open

chore: add mimalloc tests to assert size functions#23399
comphead wants to merge 3 commits into
apache:mainfrom
comphead:mem_size

Conversation

@comphead

@comphead comphead commented Jul 8, 2026

Copy link
Copy Markdown
Contributor

Rationale for this change

Related to #23393

Rationale for this change

MemoryPool reservations trust fn size() on every accumulator/group-values/hash-map — but those bodies are hand-written arithmetic over the struct's fields and rot silently whenever the backing types
change (a Vec becomes a HashMap, a scratch buffer is added, a ScalarValue variant grows). Nothing catches the drift today. DataFusion needs an independent, deterministic ground truth to pin these
formulas down; a mimalloc-wrapping tracking GlobalAlloc gives it (exact layout.size() sum, no size-class rounding or thread-cache lag).

What changes are included in this PR?

  • Fix SlidingDistinctCountAccumulator::size() — was returning a flat size_of_val(self) (64 B) regardless of contents. Now mirrors DistinctCountAccumulator::full_size: self + HashMap bucket capacity
  • per-key inner heap + DataType inner heap.
  • test_utils::memory_profiler — reusable module:
    • mimalloc_tracking::TrackingAlloc — a GlobalAlloc over mimalloc::MiMalloc tracking live layout.size() bytes.
    • size_matrix — a datatype × size grid: distinct_key_cases() (5 flat leaves) + nested_key_cases() (5 List<T> combinations) × SIZES (small → xxlarge). Each Case declares its own
      max_deviation_pct budget.
  • Regression tests in count.rs (flat + nested), gated on the new mimalloc feature (off by default; forwards from test-utilsdatafusion-functions-aggregatedatafusion).
  • CImimalloc added to --features in .github/workflows/extended.yml.

Memory usage before / after

SlidingDistinctCountAccumulator::size() reported vs. real heap bytes (mimalloc tracking allocator, Int64 keys):

distinct keys before (reported) after (reported) actual heap before undercount after deviation
64 64 B 8 128 B 9 760 B −99.3% −16.7%
1 024 64 B 129 088 B 157 600 B −99.96% −18.1%
16 384 64 B 2 064 448 B 2 523 040 B −99.997% −18.2%
262 144 64 B 33 030 208 B 40 370 080 B −99.9998% −18.2%
1 048 576 64 B 132 120 640 B 161 480 608 B −99.99996% −18.2%

Full matrix across Int64 / Decimal128(38, 10) / Utf8 / Binary / Timestamp(ns, UTC) stays within a 30% budget; nested List<T> cases (all 5 leaves) fence at 4 000% — a documented ceiling on the
known upstream ScalarValue::List::size() overcount, kept as a regression fence.

Are these changes tested?

Two new #[test]s under the mimalloc feature: sliding_distinct_count_size_matches_allocator_flat and ..._nested. Each iterates the shared matrix and asserts against the per-case max_deviation_pct.
Existing sliding_distinct_count_accumulator_{basic,retract,merge_states} tests continue to pass.

Are there any user-facing changes?

No API changes. SlidingDistinctCountAccumulator reservations are now accurate — queries that previously slipped past MemoryPool limits may now spill or fail earlier (intended).

@github-actions github-actions Bot added development-process Related to development process of DataFusion core Core DataFusion crate functions Changes to functions implementation labels Jul 8, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

core Core DataFusion crate development-process Related to development process of DataFusion functions Changes to functions implementation

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant