Repository navigation
fix: subtract emitted memory size in MinMaxStructState::emit_to(EmitTo::First) - #27115
dragonknightk wants to merge 1 commit into
Conversation
github-actions
Bot
added
sqllogictest
…o::First) ## Which issue does this PR close? - Closes issue 26891. ## Rationale for this change Grouped `max`/`min` over structs with ordered aggregation can fail with Resources exhausted even when the data is small. When groups are emitted, the memory accounting subtracts the row count instead of the byte size, so emitted groups keep being counted. Specifically, `MinMaxStructState` tracks the memory of stored values in `total_data_bytes` via `get_array_memory_size()`, but `EmitTo::First(n)` subtracted the `len()` of the emitted values, which is only 1 per group. Ordered aggregation calls `First(n)` after every batch, so the reported size grows with the number of groups seen. ## What changes are included in this PR? Subtract `get_array_memory_size()` of the emitted values instead. ## What is the testing strategy for this PR? Added a sqllogictest in `ordered_aggregate_spill.slt` that fails with Resources exhausted before the fix and returns the correct result after it, plus a unit test. ## Are there any user-facing changes? No API changes. ## AI Disclosure Human-led contribution. Grok Bot, working for dragonknightk, assisted with this change. All submitted content was reviewed by a human before submission. (๑•̀ㅂ•́)و✧
8a3b6f9 to
5937ed3
Compare
|
Hi @parthchandra! This is my first contribution to DataFusion, so please bear with me if I've missed anything (´｡• ᵕ •｡`) This PR fixes #26891, which you opened: when groups are emitted, the accumulator now subtracts the actual bytes of the emitted values instead of the row count. I noticed you ran into the same issue in apache/datafusion-comet#7191, and the fix here follows the same idea as your approach there. Since I'm a first-time contributor, the upstream CI needs a committer's approval before it runs. The Rust, Dev and Dependencies checks have all passed on my fork. When you have a moment, could you click "Approve and run workflows"? Thanks a lot! |
parthchandra
left a comment
parthchandra
left a comment
There was a problem hiding this comment.
Thanks for the fix @dragonknightk. Looks good, pending ci.
@comphead could you also take a look?
| let first_data_capacity: usize = first_min_maxes | ||
| .iter() | ||
| .map(|opt| opt.as_ref().map(|s| s.len()).unwrap_or(0)) | ||
| .map(|opt| { |
There was a problem hiding this comment.
Looks good. This matches how set_value adds memory (line 197 and 203 use get_array_memory_size()), so now add and subtract use the same unit. The old s.len() was always 1 per group, which is why memory kept growing. Nice fix.
|
|
||
| let mut accumulator = MinMaxStructAccumulator::new_min(data_type.clone()); | ||
| accumulator | ||
| .update_batch(&values, &[0, 1, 2, 3], None, 4) |
There was a problem hiding this comment.
Good test. The assert_eq!(accumulator.size(), 0) after emitting everything is the key check, it would fail before the fix.
Which issue does this PR close?
Rationale for this change
Grouped
max/minover structs with ordered aggregation can fail with Resources exhausted even when the data is small. When groups are emitted, the memory accounting subtracts the row count instead of the byte size, so emitted groups keep being counted.Specifically,
MinMaxStructStatetracks the memory of stored values intotal_data_bytesviaget_array_memory_size(), butEmitTo::First(n)subtracted thelen()of the emitted values, which is only 1 per group. Ordered aggregation callsFirst(n)after every batch, so the reported size grows with the number of groups seen.What changes are included in this PR?
Subtract
get_array_memory_size()of the emitted values instead.What is the testing strategy for this PR?
Added a sqllogictest in
ordered_aggregate_spill.sltthat fails with Resources exhausted before the fix and returns the correct result after it, plus a unit test.Are there any user-facing changes?
No API changes.
AI Disclosure
Human-led contribution. Grok Bot, working for dragonknightk, assisted with this change. All submitted content was reviewed by a human before submission. (๑•̀ㅂ•́)و✧