Slow aggregrate query with array_agg, Polars is 4 times faster for equal query

Describe the bug

I encountered a slow query, which I already mentioned in Discord but got no answer. I went ahead and created a repro case with data I can share that showcases this slow query.

Plans:

DataFrame()
+---------------+---------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+
| plan_type     | plan                                                                                                                                                                                                                                              |
+---------------+---------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+
| logical_plan  | Projection: df.name, df.group, CASE WHEN array_element(markets, Int64(1)) IS NOT NULL THEN markets ELSE List() END AS CASE WHEN markets[Int64(1)] IS NOT NULL THEN markets ELSE NULL END                                                          |
|               |   Aggregate: groupBy=[[df.name, df.group]], aggr=[[array_agg(market) AS markets]]                                                                                                                                                                 |
|               |     Projection: df.name, df.group, CASE WHEN df.market IS NOT NULL AND df.price IS NOT NULL THEN named_struct(Utf8("market"), CAST(df.market AS Dictionary(UInt16, Utf8)), Utf8("price"), df.price) ELSE Struct({market:,price:}) END AS market   |
|               |       TableScan: df projection=[name, group, market, price]                                                                                                                                                                                       |
| physical_plan | ProjectionExec: expr=[name@0 as name, group@1 as group, CASE WHEN array_element(markets@2, 1) IS NOT NULL THEN markets@2 END as CASE WHEN markets[Int64(1)] IS NOT NULL THEN markets ELSE NULL END]                                               |
|               |   AggregateExec: mode=FinalPartitioned, gby=[name@0 as name, group@1 as group], aggr=[markets], ordering_mode=Sorted                                                                                                                              |
|               |     SortExec: expr=[name@0 ASC NULLS LAST, group@1 ASC NULLS LAST], preserve_partitioning=[true]                                                                                                                                                  |
|               |       CoalesceBatchesExec: target_batch_size=8192                                                                                                                                                                                                 |
|               |         RepartitionExec: partitioning=Hash([name@0, group@1], 10), input_partitions=10                                                                                                                                                            |
|               |           AggregateExec: mode=Partial, gby=[name@0 as name, group@1 as group], aggr=[markets], ordering_mode=Sorted                                                                                                                               |
|               |             ProjectionExec: expr=[name@0 as name, group@1 as group, CASE WHEN market@2 IS NOT NULL AND price@3 IS NOT NULL THEN named_struct(market, CAST(market@2 AS Dictionary(UInt16, Utf8)), price, price@3) END as market]                   |
|               |               DataSourceExec: file_groups={10 groups: [, ...]}, projection=[name, group, market, price], output_ordering=[name@0 ASC NULLS LAST, group@1 ASC NULLS LAST], file_type=parquet                                                       |
|               |                                                                                                                                                                                                                                                   |
+---------------+---------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+

Hyperfine reports:

Benchmark 1: uv run test_polars.py sample-1m.parquet
  Time (mean ± σ):     300.2 ms ± 104.0 ms    [User: 538.6 ms, System: 140.1 ms]
  Range (min … max):   249.2 ms … 585.3 ms    10 runs
 
  Warning: The first benchmarking run for this command was significantly slower than the rest (585.3 ms). This could be caused by (filesystem) caches that were not filled until after the first run. You should consider using the '--warmup' option to fill those caches before the actual benchmark. Alternatively, use the '--prepare' option to clear the caches before each timing run.
 
Benchmark 2: uv run test_datafusion.py sample-1m.parquet
  Time (mean ± σ):      1.320 s ±  0.062 s    [User: 4.919 s, System: 0.258 s]
  Range (min … max):    1.264 s …  1.444 s    10 runs
 
Summary
  uv run test_polars.py sample-1m.parquet ran
    4.40 ± 1.54 times faster than uv run test_datafusion.py sample-1m.parquet

To Reproduce

See https://github.com/valkum/polars-datafusion-comparison/tree/main for full repro case.

Expected behavior

No response

Additional context

I am not sure if a similar file and task is part of the benchmarks. As the 1m sample file is 14MB in size, thus it makes no sense to include such a file in the git repo. A smaller sample file (e.g. 10k) only results a 1.49x difference by Polars. Feel free to store a generated file to include this in the benchmarks.

@valkum

Thank you @valkum

I tried the reproducer locally and I do also see the difference reported. I rewrote the query into the equivalent SQL as that was easier to profile for me:

Here is the input:

repo.zip

COPY (
  WITH df as (
    SELECT
      name, group,
      CASE WHEN market IS NOT NULL AND price IS NOT NULL THEN named_struct('market', market, 'price', price) ELSE NULL END as market
    FROM 'sample-1m.parquet'
  ),
  df2 as (
    SELECT
      name,
      group,
      array_agg(market) as markets
    FROM df
    GROUP BY
      name,
      group
  )
  SELECT
    name,
    group,
    CASE WHEN markets[0] IS NOT NULL THEN markets ELSE NULL END
  FROM df2
) TO 'output-datafusion.parquet';

When I did some profiling with samply

samply record -- ~/Software/datafusion-cli/datafusion-cli-49.0.0 -f report.sql

Half the time is spent in the array_agg implementation

So I think a lot of the difference would be fixed with a better array_agg implementation

here is one idea of how to do it:

Implement GroupsAccumulator for array_agg aggregation function #10145

Thanks a lot for your investigation. I am OOO until Wednesday. I have to check if I have enough time to work on an updated attempt. But I certainly want to give it a try.

A GroupsAccumulator will be non trivial for ArrayAgg, I recommend you start with a simple type like Int64 first, and then we can make it generic for all primitives and then other types.

This query accumulates arrays of structs 🤔

Yes. We have a one to many relationship of some data and it doesn't make sense to store this normalized in a different file. My understanding of the parquet format, or rather dremel encoding, is that this should work rather good.
This is an intermediate step on our way to rework and modernize our data pipelines. So in the end this might not even be used very long, but we have this intermediate input file right now, and I thought it would be a good test to see if datafusion is fit for our task.

There was a previous attempt at a GroupsAccumulator for ArrayAgg (#11096) but the performance of that PR was mixed.

I don't know enough about how polars is implemented, but they may also use the ListView type (similar to StringView, but for lists) which would make copying / aggregating lists potentially much faster

I am looking for some topic about performance optimizing, and would like to take this ticket

        for row_idx in 0..num_rows {
            // Create the empty accumulator for converting
            let mut converted_accumulator = (self.factory)()?;

            // Convert row to states
            let values_to_accumulate =
                slice_and_maybe_filter(values, opt_filter, &[row_idx, row_idx + 1])?;
            converted_accumulator.update_batch(&values_to_accumulate)?;
            let states = converted_accumulator.state()?;

            // Resize results to have enough columns according to the converted states
            results.resize_with(states.len(), || Vec::with_capacity(num_rows));

            // Add the states to results
            for (idx, state_val) in states.into_iter().enumerate() {
                results[idx].push(state_val);
            }
        }

convert_to_state is invoked when skip partial aggregation is enabled at runtime (which make partial aggregation does nothing), but this "do-nothing" cost is a bit too much. I think as you mention we can try leverage the zero-copy functionality of ListArrays, which reference the original input values and adding some offsets of ones. For example given the input

[1,2,3,4,5,....]

Then convert_to_state is expected to return [1],[2],[3],[4],[5] ....

I haven't had time to start on this properly. If you have the time, feel free to look into this. I am looking forward to your results.

I did a tweak inside that function and it looks good now, still performance is not that impressive, i'll take a look at other functions like evaluate and update_batch

 hyperfine "~/proj/rust/build/release-nonlto/datafusion-cli-opt -f report.sql" "~/proj/rust/build/release-nonlto/datafusion-cli-49.0.0 -f report.sql" "~/proj/rust/build/release-nonlto/datafusion-cli-50.0.0 -f report.sql"
Benchmark 1: ~/proj/rust/build/release-nonlto/datafusion-cli-opt -f report.sql
  Time (mean ± σ):      2.994 s ±  0.053 s    [User: 23.054 s, System: 2.006 s]
  Range (min … max):    2.909 s …  3.056 s    10 runs
 
Benchmark 2: ~/proj/rust/build/release-nonlto/datafusion-cli-49.0.0 -f report.sql
  Time (mean ± σ):      6.318 s ±  0.104 s    [User: 29.943 s, System: 2.245 s]
  Range (min … max):    6.183 s …  6.529 s    10 runs
 
Benchmark 3: ~/proj/rust/build/release-nonlto/datafusion-cli-50.0.0 -f report.sql
  Time (mean ± σ):      4.456 s ±  0.062 s    [User: 29.947 s, System: 2.012 s]
  Range (min … max):    4.394 s …  4.562 s    10 runs
 
Summary
  ~/proj/rust/build/release-nonlto/datafusion-cli-opt -f report.sql ran
    1.49 ± 0.03 times faster than ~/proj/rust/build/release-nonlto/datafusion-cli-50.0.0 -f report.sql
    2.11 ± 0.05 times faster than ~/proj/rust/build/release-nonlto/datafusion-cli-49.0.0 -f report.sql

I also notice that AccumulatorState::size takes up to 30% of CPU

fn evaluate(&mut self, emit_to: EmitTo) -> Result<ArrayRef> {
    let vec_size_pre = self.states.allocated_size();

    let states = emit_to.take_needed(&mut self.states);

    let results: Vec<ScalarValue> = states
        .into_iter()
        .map(|mut state| {
            self.free_allocation(state.size()); ----> 12%
            state.accumulator.evaluate() ->
        })
        .collect::<Result<_>>()?;


fn invoke_per_accumulator
            let values_to_accumulate = slice_and_maybe_filter(
                &values,
                opt_filter.as_ref().map(|f| f.as_boolean()),
                offsets,
            )?;
            f(state.accumulator.as_mut(), &values_to_accumulate)?;

            // clear out the state so they are empty for next
            // iteration
            state.indices.clear();
            sizes_post += state.size();  ----> 6.2%
        }


fn state()
        for mut state in states {
            self.free_allocation(state.size()); ----> 11%
            let accumulator_state = state.accumulator.state()?; 
            results.resize_with(accumulator_state.len(), Vec::new);
            for (idx, state_val) in accumulator_state.into_iter().enumerate() {
                results[idx].push(state_val);
            }
        }

Yes i'm working on it

Yes i'm working on it

Awesome!

Indeed polar use binary view to build the agg result

array agg for the specific query mentioned in this issue does not work so well with existing nullstate method which works only with primitive types, as well as this attemp, the struct can be any combination of subfield (not to mention nested field), in this case it is a dictionary field and a utf8 field. Slicing row of a struct arrays can be expensive, i wonder how polar store the data of this column during their execution, or do they consider the column as a struct type at all

Slicing row of a struct arrays can be expensive, i wonder how polar store the data of this column during their execution, or do they consider the column as a struct type at all

What I would suggest we do initially in DataFusion is:

Retain references to all input ArrayRefs and the corresponding group_indexes on aggregation phase
On output phase, form a new underlying values array by copying the relevant input rows in order (using the interleave kernel)

That approach requires no slicing and will work for any arbitrary array type

This design will require 2x the input size at peak (for the input and the inprogress output) but I think is probably significantly better than using the Accumulators (which also holds each input array entirely, but also has a bunch of individual allocations per row to store the ScalarValues

Here is a ascii art idea:

GroupsAccumulator State

   ┌────────┐     ┌─────┐                             
   │   A    │     │  0  │                             
   ├────────┤     ├─────┤                             
   │   B    │     │  1  │                             
   ├────────┤     ├─────┤                             
   │   C    │     │  0  │                             
   ├────────┤     ├─────┤                             
   │   D    │     │  2  │                             
   ├────────┤     ├─────┤               State         
   │   E    │     │  3  │                             
   └────────┘     └─────┘       Accumulator keeps all 
                                 input `ArrayRef` and 
                                  the corresponding   
   ┌────────┐     ┌─────┐       group_indexes for each
   │   V    │     │  0  │                row          
   ├────────┤     ├─────┤                             
   │   W    │     │  0  │                             
   ├────────┤     ├─────┤                             
   │   X    │     │  4  │                             
   ├────────┤     ├─────┤                             
   │   Y    │     │  2  │                             
   ├────────┤     ├─────┤                             
   │   Z    │     │  4  │                             
   └────────┘     └─────┘                             
                                                      
                                                      
                                                      
     ArrayRefs     group                              
                  indices

GroupsAccumulator output (emit):

                                                              
                                           Output ListArray   
                                                              
   ┌────────┐     ┌─────┐                                     
   │   A    │     │  0  │───                                  
   ├────────┤     ├─────┤   │      ┌ ─ ─ ▶ [A, C, V, W]       
   │   C    │     │  0  │                                     
   ├────────┤     ├─────┤   │─ ─ ─ ┘                          
   │   V    │     │  0  │                                     
   ├────────┤     ├─────┤   │                                 
   │   W    │     │  0  │───      ┌ ─ ─ ─ ▶ [B]               
   ├────────┤     ├─────┤                                     
   │   B    │     │  1  │ ─ ─ ─ ─ ┘                           
   ├────────┤     ├─────┤                                     
   │   D    │     │  2  │───      ┌ ─ ─ ▶  [D, Y]             
   ├────────┤     ├─────┤   │─ ─ ─                            
   │   Y    │     │  2  │───                                  
   ├────────┤     ├─────┤          ─ ─ ─▶  [E, X]             
   │   E    │     │  3  │───      │                           
   ├────────┤     ├─────┤   │─ ─ ─                            
   │   X    │     │  3  │───                                  
   ├────────┤     ├─────┤         ─ ─ ─ ─ ▶ [Z]               
   │   Z    │     │  4  │─ ─ ─ ─ ┘                            
   └────────┘     └─────┘                                     
                                Output: use the               
 Output Values    group       interleave kernel to            
     array       indexes       reorder the input              
                              arrays so all values            
                              from the same group             
                              index are contiguous            
   New ArrayRef    group          and form the                
                  indices        corresponding                
                                ListArray output

@alamb

Thank you @alamb for the detailed explanation. I'll try this out maybe tmr

Benchmark 1: uv run polar.py sample-1m.parquet
  Time (mean ± σ):     258.6 ms ±  32.2 ms    [User: 514.8 ms, System: 218.1 ms]
  Range (min … max):   238.5 ms … 348.3 ms    10 runs
 
  Warning: The first benchmarking run for this command was significantly slower than the rest (348.3 ms). This could be caused by (filesystem) caches that were not filled until after the first run. You should consider using the '--warmup' option to fill those caches before the actual benchmark. Alternatively, use the '--prepare' option to clear the caches before each timing run.
 
Benchmark 2: uv run df.py sample-1m.parquet
  Time (mean ± σ):     345.3 ms ±   8.4 ms    [User: 2194.9 ms, System: 241.7 ms]
  Range (min … max):   331.7 ms … 360.7 ms    10 runs
 
Summary
  uv run polar.py sample-1m.parquet ran
    1.34 ± 0.17 times faster than uv run df.py sample-1m.parquet

There was significant improvement after using interleave, but there is still some overflowing error writing to Parquet (so i changed the script to evaluating in-memory), i'm fixing that and push for review soon

thread 'tokio-runtime-worker' (2611103) panicked at /home/toai/.cargo/registry/src/index.crates.io-1949cf8c6b5b557f/arrow-data-56.1.0/src/transform/mod.rs:676:31:
MutableArrayData::new is infallible: DictionaryKeyOverflowError
stack backtrace:
   0: __rustc::rust_begin_unwind
   1: core::panicking::panic_fmt
   2: core::result::unwrap_failed
   3: arrow_data::transform::MutableArrayData::with_capacities
   4: <alloc::vec::Vec<T> as alloc::vec::spec_from_iter_nested::SpecFromIterNested<T,I>>::from_iter
   5: arrow_data::transform::MutableArrayData::with_capacities
   6: arrow_data::transform::MutableArrayData::with_capacities
   7: arrow_select::interleave::interleave_fallback
   8: <core::iter::adapters::GenericShunt<I,R> as core::iter::traits::iterator::Iterator>::next
   9: datafusion_physical_plan::sorts::builder::BatchBuilder::build_record_batch
  10: <datafusion_physical_plan::sorts::merge::SortPreservingMergeStream<C> as futures_core::stream::Stream>::poll_next
  11: datafusion_common_runtime::trace_utils::trace_future::{{closure}}
  12: <futures_util::future::future::Map<Fut,F> as core::future::future::Future>::poll

looks related: #17445

toai@salamancabrothehood:~/proj/rust/playpy$ bash test.sh sample-1m.parquet 
Benchmark 1: uv run polar.py sample-1m.parquet
  Time (mean ± σ):     867.0 ms ±  28.5 ms    [User: 2599.8 ms, System: 800.0 ms]
  Range (min … max):   835.4 ms … 916.0 ms    10 runs
 
Benchmark 2: uv run df.py sample-1m.parquet
  Time (mean ± σ):      1.445 s ±  0.020 s    [User: 4.682 s, System: 0.537 s]
  Range (min … max):    1.400 s …  1.472 s    10 runs
 
Summary
  uv run polar.py sample-1m.parquet ran
    1.67 ± 0.06 times faster than uv run df.py sample-1m.parquet

I applied a patch in arrow-rs to fix dictionary overflow error and here is the performance improvement

ArrayAgg ran within the first 600-700ms, the rest is the time for writing Parquet file, which is another point for improvement.

With this result, i think this PR is good to push for review: #17915

perf: Optimize `array_agg()` using `GroupsAccumulator` (#20504)

perf: Optimize `array_agg()` using `GroupsAccumulator` (apache#20504)

Navigation Menu

Repository navigation

Slow aggregrate query with `array_agg`, Polars is 4 times faster for equal query #17446

Describe the bug

To Reproduce

Expected behavior

Additional context

4 remaining items

Metadata

Assignees

Labels

Type

Projects

Milestone

Relationships

Development

Slow aggregrate query with array_agg, Polars is 4 times faster for equal query #17446

Description

Describe the bug

To Reproduce

Expected behavior

Additional context

Activity

alamb commented on Sep 5, 2025

valkum commented on Sep 5, 2025

alamb commented on Sep 6, 2025

valkum commented on Sep 6, 2025

Omega359 commented on Sep 12, 2025

alamb commented on Sep 13, 2025

duongcongtoai commented on Sep 20, 2025

duongcongtoai commented on Sep 22, 2025

valkum commented on Sep 23, 2025

duongcongtoai commented on Sep 26, 2025

duongcongtoai commented on Sep 27, 2025

4 remaining items

duongcongtoai commented on Sep 29, 2025

vegarsti commented on Sep 29, 2025

duongcongtoai commented on Sep 30, 2025

duongcongtoai commented on Sep 30, 2025

alamb commented on Oct 1, 2025

duongcongtoai commented on Oct 2, 2025

duongcongtoai commented on Oct 3, 2025

duongcongtoai commented on Oct 3, 2025

duongcongtoai commented on Nov 7, 2025

Metadata

Metadata

Assignees

Labels

Type

Projects

Milestone

Relationships

Development

Issue actions

Slow aggregrate query with `array_agg`, Polars is 4 times faster for equal query #17446