Some group by query is 6~7x slower than DuckDB

[wegamekinglc]

Describe the bug

Hi team, I have encountered a performance issue when I run same query on a big table with datafusion comparing with DuckDB.

I will try to simplify my case and replicate the issue in my following codes.

To Reproduce

import timeit
import numpy as np
import pyarrow as pa
import datafusion
from datafusion import SessionContext
import duckdb

print(duckdb.__version__)
print(datafusion.__version__)

# prepare data

batches = 100000

names = list("abcdefghijklmnopqrstuvwxyz")
names = [n + m for n in names for m in names]

names_array = pa.concat_arrays([pa.array(names)] * batches)
values_array = pa.concat_arrays([pa.array(np.random.randint(1, 100, len(names))) for _ in range(batches)])

pa_table = pa.Table.from_arrays([names_array, values_array], names=["name", "value"])

# prepare query
sql = "select name, sum(value) as value FROM pa_table group by name;"
n_round = 10

# duckb
elapsed = timeit.timeit('duckdb.sql(sql).to_arrow_table()', number=n_round , globals=globals())
duckdb_per_round = elapsed  / n_round

# datafusion
ctx = SessionContext()
_ = ctx.from_arrow(pa_table, "pa_table")
elapsed  = timeit.timeit('ctx.sql(sql).to_arrow_table()', number=n_round , globals=globals())
datafusion_per_round= elapsed  / n_round

# result
print(f"{'duckdb':<12}: {duckdb_per_round * 1000:.2f}ms")
print(f"{'datafusion':<12}: {datafusion_per_round * 1000:.2f}ms")

the output will look like:

1.3.1
47.0.0
duckdb      : 152.15ms
datafusion  : 1002.04ms

Expected behavior

No response

Additional context

No response

[jatin510]

take

[MrPowers]

I reran the groupby h2o benchmarks on the 1e8 dataset stored in Parquet to compare DataFusion v47.0.0 and DuckDB v1.3.1, here are the results:

DataFusion was faster for some queries and DuckDB was faster for others, but I didn't notice a large difference on any of the queries.

Can you rerun these benchmarks with the data stored in Parquet?

[alamb]

python3 -m venv test_venv
source test_venv/bin/activate
pip install duckdb datafusion numpy
python3 repro.py

I see the output

duckdb      : 246.51ms
datafusion  : 931.72ms

One thing I noticed is that only 1 CPU core is used

I could not seem to get a useful profile (the stack traces I had didn't show a stack frame)

Profiling command

samply record  python3 repro.py

Here is the query

select name, sum(value) as value FROM pa_table group by name;

Here is the explain plan

+---------------+------------------------------------------------------------------------------------------+
| plan_type     | plan                                                                                     |
+---------------+------------------------------------------------------------------------------------------+
| logical_plan  | Projection: pa_table.name, sum(pa_table.value) AS value                                  |
|               |   Aggregate: groupBy=[[pa_table.name]], aggr=[[sum(pa_table.value)]]                     |
|               |     TableScan: pa_table projection=[name, value]                                         |
| physical_plan | ProjectionExec: expr=[name@0 as name, sum(pa_table.value)@1 as value]                    |
|               |   AggregateExec: mode=FinalPartitioned, gby=[name@0 as name], aggr=[sum(pa_table.value)] |
|               |     CoalesceBatchesExec: target_batch_size=8192                                          |
|               |       RepartitionExec: partitioning=Hash([name@0], 16), input_partitions=16              |
|               |         AggregateExec: mode=Partial, gby=[name@0 as name], aggr=[sum(pa_table.value)]    |
|               |           RepartitionExec: partitioning=RoundRobinBatch(16), input_partitions=1          |
|               |             DataSourceExec: partitions=1, partition_sizes=[1]                            |
|               |                                                                                          |
+---------------+------------------------------------------------------------------------------------------+

Here is the explain(analyze=True):

+-------------------+------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+
| plan_type         | plan                                                                                                                                                                                                                                                 |
+-------------------+------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+
| Plan with Metrics | ProjectionExec: expr=[name@0 as name, sum(pa_table.value)@1 as value], metrics=[output_rows=676, elapsed_compute=5.041µs]                                                                                                                            |
|                   |   AggregateExec: mode=FinalPartitioned, gby=[name@0 as name], aggr=[sum(pa_table.value)], metrics=[output_rows=676, elapsed_compute=460.916µs, spill_count=0, spilled_bytes=0, spilled_rows=0, peak_mem_used=159424]                                 |
|                   |     CoalesceBatchesExec: target_batch_size=8192, metrics=[output_rows=676, elapsed_compute=29.718µs]                                                                                                                                                 |
|                   |       RepartitionExec: partitioning=Hash([name@0], 16), input_partitions=16, metrics=[fetch_time=1.057575333s, repartition_time=41.56µs, send_time=25.826µs]                                                                                         |
|                   |         AggregateExec: mode=Partial, gby=[name@0 as name], aggr=[sum(pa_table.value)], metrics=[output_rows=676, elapsed_compute=1.056854918s, spill_count=0, spilled_bytes=0, spilled_rows=0, skipped_aggregation_rows=0, peak_mem_used=1614708316] |
|                   |           RepartitionExec: partitioning=RoundRobinBatch(16), input_partitions=1, metrics=[fetch_time=6.542µs, repartition_time=1ns, send_time=6.181µs]                                                                                               |
|                   |             DataSourceExec: partitions=1, partition_sizes=[1], metrics=[]                                                                                                                                                                            |
|                   |                                                                                                                                                                                                                                                      |
+-------------------+------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------+
datafusion  : 910.99ms

Note that the first AggregateExec takes most of the time: elapsed_compute=1.056854918s

Conclusion

This plan shows the proper repartitioning (by 16) so I would expect it to keep all 16 cores busy.

However, only a single core was busy, so I conclude the bottleneck is in the input DataSource -- it seems to me that it can't feed the data fast enough

[alamb]

What I would like to do next is to get some real profiling -- probably by building datafusion-python from source and seeing what is going on

[Dandandan]

From the code it seems to me pa.concat_arrays is converting the input arrays into a single array.

DataFusion won't currently split the batch into smaller batches as of now, so the aggregation for a single array/batch will use a single thread only (RepartitionExec will send the input batch to a single output stream).

[zhuqi-lucas]

apache/datafusion#15532 (comment)

It may gain some performance if we can split big batch to small, and fit the L1 cache, from the above comments.

[alamb]

I tried changing the code to use ChunkedArray rather than a single array

-names_array = pa.concat_array([pa.array(names)] * batches)
+names_array = pa.chunked_array([pa.array(names)] * batches)

So the table now looks like

names_array = pa.chunked_array([pa.array(names)] * batches)
values_array = pa.chunked_array([pa.array(np.random.randint(1, 100, len(names))) for _ in range(batches)])

pa_table = pa.Table.from_arrays([names_array, values_array], names=["name", "value"])

And then actually I see the reverse performance (duckdb is slower 🤯 ):

(venv) andrewlamb@Andrews-MacBook-Pro-3:~/Downloads$ python repro.py
1.3.2
47.0.0
duckdb      : 981.68ms
datafusion  : 292.68ms

[alamb]

So one thing we could potentially do here is make RepartitionExec smarter for large batches, and if the input is really large split it into smaller batch sizes or something 🤔

[zhuqi-lucas]

So one thing we could potentially do here is make RepartitionExec smarter for large batches, and if the input is really large split it into smaller batch sizes or something 🤔

Looks like the right direction.

[wegamekinglc]

yes, i notice that when I run the scripts DuckDB can fully utilize my cpus while datafusion can't.

[ozankabak]

Isn't the easiest solution is to update DataSourceExec to adhere to target batch size?

[alamb]

Isn't the easiest solution is to update DataSourceExec to adhere to target batch size?

Maybe -- I will file a ticket to explain the problem with a DataFusion only reproducer