perf(functions-aggregate): coalesce small ordered ARRAY_AGG batches

[TinyMurky]

Which issue does this PR close?

• Closes Ordered ARRAY_AGG accumulator retains one Arrow array per update_batch call, inflating per-row memory for grouped aggregation #25199

Rationale for this change

#24392 changed OrderSensitiveArrayAggAccumulator to retain payloads as Arrow arrays instead of converting each value to a ScalarValue. This significantly reduced memory usage for reasonably sized input batches.

However, the accumulator currently stores one ArrayRef per update_batch() call. Each array has a fixed cost from its ArrayData, buffer allocation, and Arc, regardless of how many rows it contains.

Ordered ARRAY_AGG does not support a native GroupsAccumulator because groups_accumulator_supported requires order_bys to be empty. Grouped execution therefore falls back to GroupsAccumulatorAdapter, which calls update_batch() once per group per input batch.

With a high-cardinality GROUP BY, these calls often contain only one or two rows. Without this change, each call retains a separate ArrayRef, so the accumulator pays the fixed per-array allocation cost for many small Arrow arrays, increasing the per-row memory footprint.

What changes are included in this PR?

• Coalesce consecutive small ordered ARRAY_AGG payload batches up to 64 rows.
• Add tests covering:
  • coalescing exactly up to the threshold;
  • creating a new batch after the threshold is reached;
  • entry indices across coalesced and newly created batches;
  • ordering across multiple coalesced batches;
  • coalescing payloads received through partial-state merge_batch().
• Add ordered ARRAY_AGG benchmarks covering:
  • 1, 8, 64, and 2,048 rows per update_batch();
  • random input;
  • preordered input.

Retained memory

Retained memory was measured using Accumulator::size() after inserting 2,048 Int64 payloads with an Int64 ordering key and before calling evaluate(). The same measurement code was used for the baseline after #24392 and for this change.

Rows/update	Batches	Total retained (with this change)	Bytes/row (with this change)	Total retained (baseline #24392	Bytes/row (baseline #24392)
1	32	104,701 B	51.12 B/row	445,181 B	217.37 B/row
8	32	90,365 B	44.12 B/row	115,453 B	56.37 B/row
64	32	88,573 B	43.25 B/row	88,573 B	43.25 B/row
2048	1	84,901 B	41.46 B/row	84,901 B	41.46 B/row

The worst-case one-row update footprint decreases from 217.37 B/row to 51.12 B/row. Inputs already at or above the 64-row coalescing threshold retain the existing memory footprint.

Are these changes tested?

Added unit tests covering:

• coalescing small batches exactly up to the threshold;
• creating a new batch when the threshold would be exceeded;
• preserving the correct batch_idx and row_idx for entries;
• sorting values across multiple coalesced batches;
• coalescing small partial-state payloads passed through merge_batch().

Added Criterion benchmarks for ordered ARRAY_AGG using random and preordered input with 1, 8, 64, and 2,048 rows per update_batch().

The following commands have been executed and passed:

• cargo test -p datafusion-functions-aggregate --lib array_agg::tests
• cargo bench -p datafusion-functions-aggregate --bench array_agg --no-run
• cargo bench -p datafusion-functions-aggregate --bench array_agg --ordered_array_agg
• cargo test --profile=ci --test sqllogictests
• cargo test -p datafusion
• cargo test -p datafusion-cli

Are there any user-facing changes?

No

[codecov-commenter]

Codecov Report

❌ Patch coverage is 90.54054% with 7 lines in your changes missing coverage. Please review.
✅ Project coverage is 82.38%. Comparing base (d477a2b) to head (7cdd7b5).
⚠️ Report is 1 commits behind head on main.

Files with missing lines	Patch %	Lines
datafusion/functions-aggregate/src/array_agg.rs	90.54%	3 Missing and 4 partials ⚠️

Additional details and impacted files

@@           Coverage Diff            @@
##             main   #25497    +/-   ##
========================================
  Coverage   82.37%   82.38%            
========================================
  Files        1138     1138            
  Lines      433505   433803   +298     
  Branches   433505   433803   +298     
========================================
+ Hits       357102   357385   +283     
- Misses      54850    54856     +6     
- Partials    21553    21562     +9     

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.
• 📦 JS Bundle Analysis: Save yourself from yourself by tracking and limiting bundle sizes in JS merges.

[TinyMurky]

I used this function to test the retained memories

    /// To calculate how maay bytes per row do `OrderSensitiveArrayAggAccumulator` need in different rows per update_batch
    #[test]
    #[ignore = "manual retained-memory measurement"]
    fn report_ordered_array_agg_retained_memory() -> Result<()> {
        use arrow::array::Int64Array;

        const TOTAL_ROWS: usize = 2048;

        for rows_per_update in [1, 8, 64, TOTAL_ROWS] {
            let mut accumulator = ordered_accumulator(
                DataType::Int64,
                DataType::Int64,
                SortOptions::new(false, false),
                false, // input is not declared preordered
                false, // not reversed
            )?;

            for offset in (0..TOTAL_ROWS).step_by(rows_per_update) {
                let len = rows_per_update.min(TOTAL_ROWS - offset);

                let values = (offset..offset + len)
                    .map(|value| value as i64)
                    .collect::<Vec<_>>();

                let payload = Arc::new(Int64Array::from(values)) as ArrayRef;

                accumulator.update_batch(&[Arc::clone(&payload), payload])?;
            }

            let total_bytes = accumulator.size();
            let bytes_per_row = total_bytes as f64 / TOTAL_ROWS as f64;

            eprintln!(
                "rows/update={rows_per_update:>4}, \
               batches={:>4}, \
               total={total_bytes:>8} B, \
               bytes/row={bytes_per_row:.2}",
                accumulator.batches.len(),
            );
        }

        Ok(())
    }

[jayzhan211]

Thanks @TinyMurky , here is a suggestion

[jayzhan211]

Row-only threshold → every 1-row update re-copies the whole tail batch (≤63 rows) regardless of width. Measured, 2048 × 1-row update_batch, Utf8 payload, update only, main → PR: 16 KB strings 1.25 → 17.0 ms (13.7×) for 1.2% less size(); 256 KB strings 12.3 → 288 ms (23×) for 0.07%. Per-array overhead saved is only ~170 B, so coalescing should stop once the tail is no longer tiny.

With the cap below: 16 KB 1.69 ms, 256 KB 18.9 ms, Int64 size() unchanged (104,661 vs 445,141 on main).

+/// Skip coalescing once the tail batch is large enough that the fixed
+/// per-array overhead is negligible relative to the payload.
+const ORDERED_ARRAY_AGG_COALESCE_BYTES: usize = 4096;
+
         let (batch_idx, row_offset) = match self.batches.last() {
             Some(last_batch)
-                if last_batch.len() + row_count <= ORDERED_ARRAY_AGG_COALESCE_ROWS =>
+                if last_batch.len() + row_count <= ORDERED_ARRAY_AGG_COALESCE_ROWS
+                    && last_batch.get_array_memory_size()
+                        + values.get_array_memory_size()
+                        <= ORDERED_ARRAY_AGG_COALESCE_BYTES =>
             {

Please add a wide-payload case (e.g. 4 KB Utf8, 1 row per update_batch) to the new bench. Separately, the PR's own bench shows +29–42% at 1 row/update and +18–35% at 8 rows/update for Int64 vs main (concat fixed cost, not memcpy) — worth stating in the description as the CPU-for-memory trade