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

[valkum]

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.

[alamb]

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

[valkum]

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.

[alamb]

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 🤔

[valkum]

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.

[Omega359]

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

[alamb]

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

[duongcongtoai]

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

[duongcongtoai]

        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] ....

[valkum]

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.

[duongcongtoai]

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

[duongcongtoai]

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