Skip to content

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

Open
TinyMurky wants to merge 1 commit into
apache:mainfrom
TinyMurky:ordered-array-agg-small-batch-coalescing
Open

TinyMurky wants to merge 1 commit into
apache:mainfrom
TinyMurky:ordered-array-agg-small-batch-coalescing

Conversation

@TinyMurky

Copy link
Copy Markdown
Contributor

Which issue does this PR close?

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

- Closes [apache#25199](apache#25199)

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.

  - 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 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  apache#24392 and for this change.

| Rows/update | Batches | Total retained (with this change) | Bytes/row
(with this change) | Total retained (baseline
[apache#24392](apache#24392)) | Bytes/row
(baseline [apache#24392](apache#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.

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
@codecov-commenter

Copy link
Copy Markdown

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

Copy link
Copy Markdown
Contributor Author

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(())
    }

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

functions Changes to functions implementation

Projects

None yet

Development

Successfully merging this pull request may close these issues.

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

2 participants