Extra repartitions in physical plan and strange optimizer behavior overall

[ravlio]

Describe the bug

Hi everyone!

While developing my project on top of DataFusion I faced with unexpected physical plan after the optimizer was applied. I have a custom PartitionedAggregateExec node that has required_input_distribution set to Distribution::HashPartitioned on columns project_id and user_id This doesn't help, optimizer does RepartitionExec: partitioning=RoundRobinBatch. Moreover, I tried to add the Repartition node to the logical plan but it was erased by the optimizer.

Additionally, projection doesn't push down and sort disappears.

To Reproduce

My logical plan:

TableScan: ?table? projection=[project_id, user_id, created_at, event_id, event, str_0]
  PartitionedAggregate: ...
    Filter: project_id = Int64(1) AND created_at >= TimestampNanosecond(1705419428144118000, None) AND created_at <= TimestampNanosecond(1706283428144118000, None) AND event = UInt16(13)
      Sort: project_id ASC NULLS LAST, user_id ASC NULLS LAST
        Repartition: Hash(project_id, user_id) partition_count=12
          Projection: project_id, user_id, created_at, event
            TableScan: ?table? projection=[project_id, user_id, created_at, event_id, event, str_0]]

physical plan:

PartitionedAggregateExec, metrics=[]
  CoalesceBatchesExec: target_batch_size=8192, metrics=[]
    RepartitionExec: partitioning=Hash([project_id@0, user_id@1], 12), input_partitions=12, metrics=[]
      RepartitionExec: partitioning=RoundRobinBatch(12), input_partitions=1, metrics=[]
        CoalescePartitionsExec, metrics=[]
          ProjectionExec: expr=[project_id@0 as project_id, user_id@1 as user_id, created_at@2 as created_at, event@4 as event], metrics=[]
            CoalesceBatchesExec: target_batch_size=8192, metrics=[]
              FilterExec: project_id@0 = 1 AND created_at@2 >= 1705419428144118000 AND created_at@2 <= 1706283428144118000 AND event@4 = 13, metrics=[]
                RepartitionExec: partitioning=RoundRobinBatch(12), input_partitions=1, metrics=[]
                  ParquetExec: file_groups={1 group: [[file.parquet]]}, projection=[project_id, user_id, created_at, event_id, event, str_0], predicate=project_id@0 = 1 AND created_at@2 >= 1705419428144118000 AND created_at@2 <= 1706283428144118000 AND event@4 = 13, pruning_predicate=project_id_min@0 <= 1 AND 1 <= project_id_max@1 AND created_at_max@2 >= 1705419428144118000 AND created_at_min@3 <= 1706283428144118000 AND event_min@4 <= 13 AND 13 <= event_max@5, metrics=[num_predicate_creation_errors=0]

A lot of repartitions here, no projection on ParquetExec (see extra str_0 in ParquetExec), and the sort has disappeared.

Expected behavior

I expected something like that. Sorry, it might contain mistakes since I wrote it by hand:

PartitionedAggregateExec, metrics=[]
  CoalesceBatchesExec: target_batch_size=8192, metrics=[]
    SortExec:
      RepartitionExec: partitioning=Hash([project_id@0, user_id@1], 12), input_partitions=12, metrics=[]
        FilterExec: project_id@0 = 1 AND created_at@2 >= 1705419428144118000 AND created_at@2 <= 1706283428144118000 AND event@4 = 13, metrics=[]
          ParquetExec: file_groups={1 group: [[file.parquet]]}, projection=[project_id, user_id, created_at, event_id, event], predicate=project_id@0 = 1 AND created_at@2 >= 1705419428144118000 AND created_at@2 <= 1706283428144118000 AND event@4 = 13, pruning_predicate=project_id_min@0 <= 1 AND 1 <= project_id_max@1 AND created_at_max@2 >= 1705419428144118000 AND created_at_min@3 <= 1706283428144118000 AND event_min@4 <= 13 AND 13 <= event_max@5, metrics=[num_predicate_creation_errors=0]

E.g. read certain cols and filter from parquet -> partition by hash -> sort -> my node

Additional context

BTW, can I do explain verbose, from the Rust code?

Thanks!

[alamb]

The original plan you show has a TableScan at the top -- is this a projection? Or is it a view definition somehow?

My reading of the plan

TableScan: ?table? projection=[project_id, user_id, created_at, event_id, event, str_0] <---- this says it needs all columns
  PartitionedAggregate: ...
    Filter: project_id = Int64(1) AND created_at >= TimestampNanosecond(1705419428144118000, None) AND created_at <= TimestampNanosecond(1706283428144118000, None) AND event = UInt16(13)
      Sort: project_id ASC NULLS LAST, user_id ASC NULLS LAST
        Repartition: Hash(project_id, user_id) partition_count=12
          Projection: project_id, user_id, created_at, event
            TableScan: ?table? projection=[project_id, user_id, created_at, event_id, event, str_0]]

As for the repartitioning i think that is happening to try and run the filter in parallel It is somewhat messy but not obviously wrong to me

The optimizer is going to try and satisfy the requirements stated by the topmost ExecutionPlan in your tree. Thus I think figuring out what the top most plan node is and what it is requesting of its input is the best way to understand what DataFusion is doing in this case

[ravlio]

The original plan you show has a TableScan at the top -- is this a projection? Or is it a view definition somehow?

No. str_0 is an extra column that I wanted to get rid of. My actual logical plan is:

logical plan: TableScan: ?table? projection=[project_id, user_id, created_at, event_id, event, str_0, str_1, str_2, str_3, str_4, str_5, str_6, str_7, str_8, str_9, str_10, str_11, str_12, str_13, str_14, str_15, str_16, str_17, str_18, str_19, str_20, str_21, str_22, str_23, str_24, ts_0, str_25, str_26, str_27, str_28, d_0, d_1, d_2, d_3, i8_0, i8_1, d_4]
Pivot
  Sort: date_trunc(Utf8("day"), created_at) AS created_at ASC NULLS LAST, str_20 ASC NULLS LAST
    Unpivot
      PartitionedAggregate: , agg: Count { filter: None, groups: Some([(Alias(Alias { expr: ScalarFunction(ScalarFunction { func_def: BuiltIn(DateTrunc), args: [Literal(Utf8("day")), Column(Column { relation: None, name: "created_at" })] }), relation: None, name: "created_at" }), SortField { data_type: Timestamp(Nanosecond, None) }), (Column(Column { relation: None, name: "str_20" }), SortField { data_type: Utf8 })]), predicate: Column { relation: None, name: "event" }, partition_col: Column { relation: None, name: "user_id" }, distinct: false } as "0_0"
        Filter: project_id = Int64(1) AND created_at >= TimestampNanosecond(1705582297292025000, None) AND created_at <= TimestampNanosecond(1706446297292025000, None) AND event = UInt16(13)
          Sort: project_id ASC NULLS LAST, user_id ASC NULLS LAST
            Projection: project_id, user_id, created_at, event, str_20 <--- explicit projection
              TableScan: ?table? projection=[project_id, user_id, created_at, event_id, event, str_0, str_1, str_2, str_3, str_4, str_5, str_6, str_7, str_8, str_9, str_10, str_11, str_12, str_13, str_14, str_15, str_16, str_17, str_18, str_19, str_20, str_21, str_22, str_23, str_24, ts_0, str_25, str_26, str_27, str_28, d_0, d_1, d_2, d_3, i8_0, i8_1, d_4]

You may see a lot of extra columns that I can't get rid of. Even with explicit projection that doesn't push down. I wanted only project_id, user_id, created_at, event, str_20 but I got everything.

And PartitionedAggregate has required_input_distribution = Hash(project_id,user_id) but it is not included to the plan...

[alamb]

I am not sure what is going on here.

Could it be that the Pivot/Unpivot operations somehow are configured to preserve all the columns (and thus the optimizer wouldn't be able to remove their use). This would be if the output of Pivot::expressions() refered to all the columns. Maybe we should extend their display to show what expressions are being used

One thing that might help debug here, depending on how you are invoking the optimzer, is to look at how DataFusion uses this

https://github.com/apache/arrow-datafusion/blob/1097dc0345079e81ba217ef23fd408a5c13c3d6c/datafusion/core/src/physical_optimizer/output_requirements.rs#L91

When it is optimizing internally, it uses OutputRequirementExec to describe what properties are needed at the output

[mustafasrepo]

As far as I can tell, physical plan produced is correct in terms of partitioning requirements.

PartitionedAggregateExec, metrics=[]
  CoalesceBatchesExec: target_batch_size=8192, metrics=[]
    RepartitionExec: partitioning=Hash([project_id@0, user_id@1], 12), input_partitions=12, metrics=[]
      RepartitionExec: partitioning=RoundRobinBatch(12), input_partitions=1, metrics=[]
        CoalescePartitionsExec, metrics=[]
          ProjectionExec: expr=[project_id@0 as project_id, user_id@1 as user_id, created_at@2 as created_at, event@4 as event], metrics=[]
            CoalesceBatchesExec: target_batch_size=8192, metrics=[]
              FilterExec: project_id@0 = 1 AND created_at@2 >= 1705419428144118000 AND created_at@2 <= 1706283428144118000 AND event@4 = 13, metrics=[]
                RepartitionExec: partitioning=RoundRobinBatch(12), input_partitions=1, metrics=[]
                  ParquetExec: file_groups={1 group: [[file.parquet]]}, projection=[project_id, user_id, created_at, event_id, event, str_0], predicate=project_id@0 = 1 AND created_at@2 >= 1705419428144118000 AND created_at@2 <= 1706283428144118000 AND event@4 = 13, pruning_predicate=project_id_min@0 <= 1 AND 1 <= project_id_max@1 AND created_at_max@2 >= 1705419428144118000 AND created_at_min@3 <= 1706283428144118000 AND event_min@4 <= 13 AND 13 <= event_max@5, metrics=[num_predicate_creation_errors=0]

where
RepartitionExec: partitioning=Hash([project_id@0, user_id@1], 12), input_partitions=12, metrics=[]
is added to satisfy partitioning requirement of the PartitionedAggregateExec. However, the plan is weird though. We should be able to generate a better plan, where intermediate CoalescePartitionsExec removed.

Additionally we have an API required_input_ordering, this API should be implemented if an executor depends on any ordering. If PartitionedAggregateExec depends on any ordering (such as the order of its partition expressions), you have to specify this dependency with this API. This might be the root cause why Sort is removed from the PhysicalPlan.

[ravlio]

@alamb Pivot and Unpivot do not affect projection. I disabled them and got the same result. I also added an expression to PartitionedAggregateNode

    fn expressions(&self) -> Vec<Expr> {
        vec![Expr::Column(Column {
            relation: None,
            name: "project_id".to_string(),
        })]
    }

but it didn't help for some reason. Parquet keeps reading all the columns. I tried to remove everything except PartitionedAggregateNode to debug it but the result is the same. I'm trying to satisfy the optimizer requirements, but I did not find any other way to tell to optimizer about the fields that the node needs except fn expressions() -> Vec<Expr>

@mustafasrepo Thanks! Sort works now!