fix: Incorrect memory accounting in `array_agg` function #16519

sfluor · 2025-06-23T16:15:33Z

Which issue does this PR close?

See this issue: #16517

Closes Incorrect memory accounting in array_agg function #16517.

What changes are included in this PR?

Fixes the over-accounting in array_agg functions

Are these changes tested?

Added a test that shows the problem.

Are there any user-facing changes?

Cloning the data might lead to slightly higher "real" memory usage.

See this issue: apache#16517

gabotechs · 2025-06-25T06:57:43Z

datafusion/functions-aggregate/src/array_agg.rs

+                    // The ArrayRef might be holding a reference to its original input buffer, so
+                    // storing it here directly copied/compacted avoids over accounting memory
+                    // not used here.
+                    self.values
+                        .push(make_array(copy_array_data(&values.to_data())));
                }


🤔 I'm not sure if this will solve the issue. Keep in mind that the merge_batch method argument receives the states of other accumulators, which already hold "compacted" data, so I'd expect this compaction here to be unnecessary.

#16519 (comment)

I wonder if we should add a special case to copy_array_data to avoid copying the data when it already is only a single row / has no offset 🤔

Right now it seems to copy the data unconditionally which is a non trivial overhead on each row 🤔

datafusion/datafusion/common/src/scalar/mod.rs

Lines 3564 to 3567 in a87d6f2

pub fn copy_array_data(src_data: &ArrayData) -> ArrayData {

let mut copy = MutableArrayData::new(vec![&src_data], true, src_data.len());

copy.extend(0, 0, src_data.len());

copy.freeze()

Perhaps we can do that as a follow of PR

gabotechs · 2025-06-25T07:01:05Z

datafusion/functions-aggregate/src/array_agg.rs

+        acc1.merge_batch(&[Arc::new(a1.slice(0, 1))])?;
+        acc2.merge_batch(&[Arc::new(a2.slice(0, 1))])?;
+
+        acc1 = merge(acc1, acc2)?;


The merge_batch functions do not receive arbitrary data, it receives the results of calling state() in other accumulators. A fairer test would be to do something like:

Suggested change

acc1.merge_batch(&[Arc::new(a1.slice(0, 1))])?;

acc2.merge_batch(&[Arc::new(a2.slice(0, 1))])?;

acc1 = merge(acc1, acc2)?;

acc1.update_batch(&[Arc::new(a1.slice(0, 1))])?;

acc2.update_batch(&[Arc::new(a2.slice(0, 1))])?;

acc1 = merge(acc1, acc2)?;

With this, you would notice that the test result is the same regardless of the changes in this PR

#16519 (comment)

gabotechs · 2025-06-25T14:23:55Z

EDIT: I'm seeing that there's cases where the merge_batch method is used for something other than merging states:

datafusion/datafusion/physical-plan/src/aggregates/no_grouping.rs

Lines 232 to 239 in 9278233

    
           let res = match mode { 
        
               AggregateMode::Partial 
        
               | AggregateMode::Single 
        
               | AggregateMode::SinglePartitioned => accum.update_batch(&values), 
        
               AggregateMode::Final | AggregateMode::FinalPartitioned => { 
        
                   accum.merge_batch(&values) 
        
               } 
        
           };

Which means that probably we do want compaction to happen also in the merge_batch function. I think this is good to go then 👍

gabotechs

This is good for a review cc @alamb

alamb · 2025-06-25T16:16:27Z

Thanks for the code and review @gabotechs and @sfluor

alamb · 2025-06-25T16:17:17Z

There appears to be an array_agg benchmark -- I will run that on this PR to see what it shows

alamb · 2025-06-25T19:32:20Z

🤖 ./gh_compare_branch_bench.sh Benchmark Script Running
Linux aal-dev 6.11.0-1015-gcp #15~24.04.1-Ubuntu SMP Thu Apr 24 20:41:05 UTC 2025 x86_64 x86_64 x86_64 GNU/Linux
Comparing sami/fix-overaccounting-of-memory-in-array-agg (1da3e04) to b6c8cc5 diff
BENCH_NAME=array_agg
BENCH_COMMAND=cargo bench --bench array_agg
BENCH_FILTER=
BENCH_BRANCH_NAME=sami_fix-overaccounting-of-memory-in-array-agg
Results will be posted here when complete

alamb · 2025-06-25T19:40:57Z

🤖: Benchmark completed

Details

group                                                                              main                                   sami_fix-overaccounting-of-memory-in-array-agg
-----                                                                              ----                                   ----------------------------------------------
array_agg i64 merge_batch 30% nulls, 0% of nulls point to a zero length array      1.00    547.4±7.12µs        ? ?/sec    9.08      5.0±0.03ms        ? ?/sec
array_agg i64 merge_batch 30% nulls, 100% of nulls point to a zero length array    1.00      6.1±0.03µs        ? ?/sec    7.26     44.5±0.85µs        ? ?/sec
array_agg i64 merge_batch 30% nulls, 50% of nulls point to a zero length array     1.00    547.7±1.36µs        ? ?/sec    9.28      5.1±0.02ms        ? ?/sec
array_agg i64 merge_batch 30% nulls, 90% of nulls point to a zero length array     1.00    548.2±0.87µs        ? ?/sec    9.38      5.1±0.02ms        ? ?/sec
array_agg i64 merge_batch 30% nulls, 99% of nulls point to a zero length array     1.00    561.4±4.24µs        ? ?/sec    9.12      5.1±0.06ms        ? ?/sec
array_agg i64 merge_batch 70% nulls, 0% of nulls point to a zero length array      1.00    243.2±1.64µs        ? ?/sec    8.13  1977.2±10.35µs        ? ?/sec
array_agg i64 merge_batch 70% nulls, 100% of nulls point to a zero length array    1.00      5.9±0.02µs        ? ?/sec    3.96     23.2±0.43µs        ? ?/sec
array_agg i64 merge_batch 70% nulls, 50% of nulls point to a zero length array     1.00    242.1±0.26µs        ? ?/sec    8.39      2.0±0.02ms        ? ?/sec
array_agg i64 merge_batch 70% nulls, 90% of nulls point to a zero length array     1.00    243.5±0.44µs        ? ?/sec    8.08  1968.8±15.78µs        ? ?/sec
array_agg i64 merge_batch 70% nulls, 99% of nulls point to a zero length array     1.00    243.0±0.62µs        ? ?/sec    8.26      2.0±0.01ms        ? ?/sec
array_agg i64 merge_batch all nulls, 100% of nulls point to a zero length array    1.00     86.9±0.11ns        ? ?/sec    1.01     87.6±0.14ns        ? ?/sec
array_agg i64 merge_batch all nulls, 90% of nulls point to a zero length array     1.00     86.9±0.09ns        ? ?/sec    1.01     87.6±0.15ns        ? ?/sec
array_agg i64 merge_batch no nulls                                                 1.00    100.8±0.12ns        ? ?/sec    544.14    54.8±3.04µs        ? ?/sec

Bnchmarks show significant regression

alamb · 2025-06-26T17:58:53Z

🤔 the benchmarks show a significant regression in performance (10x in some cases)

I think we need to resolve that prior to merging this in

We have some documentation on how to profile here: https://datafusion.apache.org/library-user-guide/profiling.html

The benchmarks can be run locally like this

cargo bench --bench array_agg

I think it might be worth explroing: #16519 (comment)

github-actions bot added the functions Changes to functions implementation label Jun 23, 2025

fix: Incorrect memory accounting in array_agg function

1da3e04

See this issue: apache#16517

sfluor force-pushed the sami/fix-overaccounting-of-memory-in-array-agg branch from d4a69ce to 1da3e04 Compare June 23, 2025 16:17

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

fix: Incorrect memory accounting in `array_agg` function #16519

fix: Incorrect memory accounting in `array_agg` function #16519

Uh oh!

sfluor commented Jun 23, 2025 •

edited

Loading

Uh oh!

gabotechs Jun 25, 2025

Uh oh!

gabotechs Jun 25, 2025

Uh oh!

alamb Jun 25, 2025

Uh oh!

gabotechs Jun 25, 2025 •

edited

Loading

Uh oh!

gabotechs Jun 25, 2025

Uh oh!

gabotechs commented Jun 25, 2025

Uh oh!

gabotechs left a comment

Uh oh!

alamb commented Jun 25, 2025

Uh oh!

alamb commented Jun 25, 2025

Uh oh!

This comment was marked as outdated.

This comment was marked as outdated.

This comment was marked as outdated.

alamb commented Jun 25, 2025

Uh oh!

alamb commented Jun 25, 2025

Uh oh!

alamb commented Jun 26, 2025 •

edited

Loading

Uh oh!

Uh oh!

	pub fn copy_array_data(src_data: &ArrayData) -> ArrayData {
	let mut copy = MutableArrayData::new(vec![&src_data], true, src_data.len());
	copy.extend(0, 0, src_data.len());
	copy.freeze()

fix: Incorrect memory accounting in array_agg function #16519

Are you sure you want to change the base?

fix: Incorrect memory accounting in array_agg function #16519

Uh oh!

Conversation

sfluor commented Jun 23, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Which issue does this PR close?

What changes are included in this PR?

Are these changes tested?

Are there any user-facing changes?

Uh oh!

gabotechs Jun 25, 2025

Choose a reason for hiding this comment

Uh oh!

gabotechs Jun 25, 2025

Choose a reason for hiding this comment

Uh oh!

alamb Jun 25, 2025

Choose a reason for hiding this comment

Uh oh!

gabotechs Jun 25, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

gabotechs Jun 25, 2025

Choose a reason for hiding this comment

Uh oh!

gabotechs commented Jun 25, 2025

Uh oh!

gabotechs left a comment

Choose a reason for hiding this comment

Uh oh!

alamb commented Jun 25, 2025

Uh oh!

alamb commented Jun 25, 2025

Uh oh!

This comment was marked as outdated.

This comment was marked as outdated.

This comment was marked as outdated.

alamb commented Jun 25, 2025

Uh oh!

alamb commented Jun 25, 2025

Uh oh!

alamb commented Jun 26, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Uh oh!

fix: Incorrect memory accounting in `array_agg` function #16519

fix: Incorrect memory accounting in `array_agg` function #16519

sfluor commented Jun 23, 2025 •

edited

Loading

gabotechs Jun 25, 2025 •

edited

Loading

alamb commented Jun 26, 2025 •

edited

Loading