Eliminate Repartitioning for Small Datasets

[gene-bordegaray]

Describe the bug

Small datasets undergo unnecessary repartitioning.

For example, a dataset of 5 rows will be split across multiple certain queries, leading to overhead that exceeds the parallelism benefit.

To Reproduce

An example can be found in the aggregate_repartition.slt file. This test creates two tables of 5 rows each in Parquet and CSV format, and an aggregation on one of their columns. Despite the small size, the query plan repartitions the data:

Setup the data:

COPY (
  SELECT * FROM (VALUES 
    ('prod', 100, 'A'),
    ('dev', 200, 'B'),
    ('test', 150, 'A'),
    ('prod', 300, 'C'),
    ('dev', 250, 'B')
  ) AS t(env, value, category)
)
TO 'test_files/scratch/aggregate_repartition/dim.csv'
STORED AS CSV
OPTIONS ('format.has_header' 'true');

COPY (
  SELECT * FROM (VALUES 
    ('prod', 100, 'A'),
    ('dev', 200, 'B'),
    ('test', 150, 'A'),
    ('prod', 300, 'C'),
    ('dev', 250, 'B')
  ) AS t(env, value, category)
)
TO 'test_files/scratch/aggregate_repartition/dim.parquet'
STORED AS PARQUET;

CREATE EXTERNAL TABLE dim_csv
STORED AS CSV 
LOCATION 'test_files/scratch/aggregate_repartition/dim.csv'
OPTIONS ('format.has_header' 'true');

CREATE EXTERNAL TABLE dim_parquet
STORED AS PARQUET 
LOCATION 'test_files/scratch/aggregate_repartition/dim.parquet';

CSV Physical Plan:

EXPLAIN SELECT env, count(*) FROM dim_csv GROUP BY env;

physical_plan
01)ProjectionExec: expr=[env@0 as env, count(Int64(1))@1 as count(*)]
02)--AggregateExec: mode=FinalPartitioned, gby=[env@0 as env], aggr=[count(Int64(1))]
03)----CoalesceBatchesExec: target_batch_size=8192
04)------RepartitionExec: partitioning=Hash([env@0], 4), input_partitions=4
05)--------AggregateExec: mode=Partial, gby=[env@0 as env], aggr=[count(Int64(1))]
06)----------RepartitionExec: partitioning=RoundRobinBatch(4), input_partitions=1
07)------------DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/aggregate_repartition/dim.csv]]}, projection=[env], file_type=csv, has_header=true

Parquet Physical Plan:

EXPLAIN SELECT env, count(*) FROM dim_parquet GROUP BY env;

physical_plan
01)ProjectionExec: expr=[env@0 as env, count(Int64(1))@1 as count(*)]
02)--AggregateExec: mode=FinalPartitioned, gby=[env@0 as env], aggr=[count(Int64(1))]
03)----CoalesceBatchesExec: target_batch_size=8192
04)------RepartitionExec: partitioning=Hash([env@0], 4), input_partitions=1
05)--------AggregateExec: mode=Partial, gby=[env@0 as env], aggr=[count(Int64(1))]
06)----------DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/aggregate_repartition/dim.parquet]]}, projection=[env], file_type=parquet

Expected behavior

The data sets should not be repartitioned when this small in size, thus should yield the plans:

CSV Physical Plan:

EXPLAIN SELECT env, count(*) FROM dim_csv GROUP BY env;

physical_plan
01)ProjectionExec: expr=[env@0 as env, count(Int64(1))@1 as count(*)]
02)--AggregateExec: mode=FinalPartitioned, gby=[env@0 as env], aggr=[count(Int64(1))]
03)----DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/aggregate_repartition/dim.csv]]}, projection=[env], file_type=csv, has_header=true

Parquet Physical Plan:

EXPLAIN SELECT env, count(*) FROM dim_parquet GROUP BY env;

physical_plan
01)ProjectionExec: expr=[env@0 as env, count(Int64(1))@1 as count(*)]
02)--AggregateExec: mode=FinalPartitioned, gby=[env@0 as env], aggr=[count(Int64(1))]
03)----DataSourceExec: file_groups={1 group: [[WORKSPACE_ROOT/datafusion/sqllogictest/test_files/scratch/aggregate_repartition/dim.parquet]]}, projection=[env], file_type=parquet

Additional context

The problem roots from enforce_distribution.rs:get_repartition_requirement_status.
Here round-robin partitioning is added based on statistics on the file:

• we are optimistic with CSV files and return true since CSV return no statistics. We could look at adding a better way to estimate if a CSV file needs repartitioning at the file level
• We have exact statistics on Parquet files, thus we do not round robin repartition at the file level, but we still hash repartition later on in the query. We can look at how hash repartitioning is handles based on file statistics and improve this.

[gene-bordegaray]

cc: @NGA-TRAN

[NGA-TRAN]

• we are optimistic with CSV files and return true since CSV return no statistics. We could look at adding a better way to estimate if a CSV file needs repartitioning at the file level
• We have exact statistics on Parquet files, thus we do not round robin repartition at the file level, but we still hash repartition later on in the query. We can look at how hash repartitioning is handles based on file statistics and improve this.

I suggest we start with Parquet, since we already have precise statistics available. CSV is more complex and can be tackled separately once we’ve gathered the necessary stats.

[gene-bordegaray]

• we are optimistic with CSV files and return true since CSV return no statistics. We could look at adding a better way to estimate if a CSV file needs repartitioning at the file level
• We have exact statistics on Parquet files, thus we do not round robin repartition at the file level, but we still hash repartition later on in the query. We can look at how hash repartitioning is handles based on file statistics and improve this.

I suggest we start with Parquet, since we already have precise statistics available. CSV is more complex and can be tackled separately once we’ve gathered the necessary stats.

ok, and yes the parquet will be much less work. Should I create two separate issues for these?

[NGA-TRAN]

ok, and yes the parquet will be much less work. Should I create two separate issues for these?

No need. Once someone closes this issue, they will open a new one as needed

[LiaCastaneda]

I'm wondering where the fix to this should be, Rather than in the optimizer EnforceDistribution or maybe in the planner? I wonder why its not selecting AggregateMode::Single (not partition) directly. This is where the node is created in the physical planner.

[feniljain]

Issue sounds very similar to 18513, are they related by any chance?

[NGA-TRAN]

@feniljain

Issue sounds very similar to #18513, are they related by any chance?

It may be similar issue. I wonder after this fix (see PR), will it get better? See the summary of the fix here

@LiaCastaneda

I'm wondering where the fix to this should be, Rather than in the optimizer EnforceDistribution or maybe in the planner? I wonder why its not selecting AggregateMode::Single (not partition) directly. This is where the node is created in the physical planner.

I suspect the fix might be tricky—by the time we detect that the file is small, we may have already committed to repartitioning elsewhere. This ticket likely requires a deeper investigation to fully understand the planning flow and its implications.

[feniljain]

It may be similar issue. I wonder after #18341 (seehttps://github.com//pull/18521#issuecomment-3516585320), will it get better? Seehttps://github.com//issues/18341#issuecomment-3505130845

I tested on latest main. It is still happening, also I am mentioning the issue due to this statement in it:

I'm not 100% sure this is a bug, but it seems like repartitioning small files is potentially pretty wasteful, this also seems somewhat related to # 4967, which is why I'm filing it.

though title of issue does say something else: Output ordering with a filter over FileSource seems to cause some unexpected repartitioning behavior

[gene-bordegaray]

Great catch @LiaCastaneda. I read read into some of the code and think a great place to look into is the condition which partitions on aggregates. This condition is in physical_planner.rs on line 801:

let can_repartition = !groups.is_empty() 
                    && session_state.config().target_partitions() > 1
                    && session_state.config().repartition_aggregations();

here it could be useful to have another condition checking if we actually benefit from repartitioning. This new condition could check statistics of the data source and make a decision based on this. Say in my example we were to see we have one input partition, then we check is this partition actually large enough based on statistic to make the AggregatePartitioned.

I believe catching this at the physical planner level would be best though.

[LiaCastaneda]

yeah this looks like a tricky problem 😿 , TLDR the issue is that here the Aggregate node has FinalPartitioned mode which enforces HashPartitioned mode which always enforces RepartitionExec right?

So the two options:

• If you do it on the optimizer rule you would already have commited on using FinalPartitioned on AggregateExec and unless you traverse the plan tree you wont be able to change it, you can skip repartitionning still, but it might break correctness assumptions elsewhere as Nga says. You will need to check its backwards compatible.
• Doing it in the physical planner implies you will be able to decide the mode upfront before committing the plan and since the physical plan is created bottom-up child stats are available at decision time so its possible. The downside of this I guess is that you will need to fetch the stats in the planner which adds complexity & i'm not sure stats are used during physcial planner as of now.