EnforceSorting resorts the inout of UnionExec unnecessarily

Describe the bug

Given the following input plan (I see this by enabling trace logging via RUST_LOG=trace:

SortExec: [tag@2 ASC NULLS LAST]
  ProjectionExec: expr=[bar@0 as bar, foo@1 as foo, tag@2 as tag, time@3 as time]
    DeduplicateExec: [tag@2 ASC,time@3 ASC]
      SortPreservingMergeExec: [tag@2 ASC,time@3 ASC]
        UnionExec
          ParquetExec: limit=None, partitions={1 group: [[d.parquet]]}, output_ordering=[tag@2 ASC, time@3 ASC], projection=[bar, foo, tag, time]
          SortExec: [tag@2 ASC,time@3 ASC]
            RecordBatchesExec: batches_groups=1 batches=1

Here is the input to enforce sorting:

Optimized physical plan by EnforceDistribution:
SortExec: [tag@2 ASC NULLS LAST]
  CoalescePartitionsExec
    ProjectionExec: expr=[bar@0 as bar, foo@1 as foo, tag@2 as tag, time@3 as time]
      RepartitionExec: partitioning=RoundRobinBatch(4)
        DeduplicateExec: [tag@2 ASC,time@3 ASC]
          SortPreservingMergeExec: [tag@2 ASC,time@3 ASC]
            UnionExec                                 <-- ** Note that the ParquetExec is already sorted correctly!
              ParquetExec: limit=None, partitions={1 group: [[d.parquet]]}, output_ordering=[tag@2 ASC, time@3 ASC], projection=[bar, foo, tag, time]
              SortExec: [tag@2 ASC,time@3 ASC]
                RecordBatchesExec: batches_groups=1 batches=1

And here is the output from EnforceSorting, where it has moved the SortExec up to the top of the union:

Optimized physical plan by EnforceSorting:
SortExec: [tag@2 ASC NULLS LAST]
  CoalescePartitionsExec
    ProjectionExec: expr=[bar@0 as bar, foo@1 as foo, tag@2 as tag, time@3 as time]
      RepartitionExec: partitioning=RoundRobinBatch(4)
        DeduplicateExec: [tag@2 ASC,time@3 ASC]
          SortPreservingMergeExec: [tag@2 ASC,time@3 ASC]
            SortExec: [tag@2 ASC,time@3 ASC]        <-- ** SortExec is moved to the output of Union, *resorting* the parquet file
              UnionExec
                ParquetExec: limit=None, partitions={1 group: [[1/1/1/1/57d6a92a-314a-4a32-a633-33bc3e1fe7a3.parquet]]}, output_ordering=[tag@2 ASC, time@3 ASC], projection=[bar, foo, tag, time]
                RecordBatchesExec: batches_groups=1 batches=1

To Reproduce
I have a reproducer from IOx -- see https://github.com/influxdata/influxdb_iox/pull/6528#discussion_r1070632410

Expected behavior
I expect the SortExec to be left where it is (at the input to the

Additional context
I found this in the context of upgrading DataFusion in IOx: https://github.com/influxdata/influxdb_iox/pull/6528

@mustafasrepo

cc @mustafasrepo and @mingmwang who I think have expertise in this area

I think the reason is that UnionExec doesn't implement fn maintains_input_order(&self) -> bool. But sometimes it maintains_input_order. If we implement fn maintains_input_order(&self) -> bool as below (Taken from self.output_ordering implementation). I think the problem will be solved.

fn maintains_input_order(&self) -> bool {
    let first_input_ordering = self.inputs[0].output_ordering();
    // If the Union is not partition aware and all the input ordering spec strictly equal with the first_input_ordering
    // `UnionExec` maintains input order
    //
    // It might be too strict here in the case that the input ordering are compatible but not exactly the same.
    // For example one input ordering has the ordering spec SortExpr('a','b','c') and the other has the ordering
    // spec SortExpr('a'), It is safe to derive the out ordering with the spec SortExpr('a').
    !self.partition_aware
        && first_input_ordering.is_some()
        && self
        .inputs
        .iter()
        .map(|plan| plan.output_ordering())
        .all(|ordering| {
            ordering.is_some()
                && sort_expr_list_eq_strict_order(
                ordering.unwrap(),
                first_input_ordering.unwrap(),
            )
        })
}

@mustafasrepo

Thank you @mustafasrepo -- I will attempt to write a reproducer for this and try our your proposed fix

proposed PR: #4946

Nice !!

Yes I am quite pleased with how sophisticated the sorting based optimizations are becoming

Navigation Menu

Repository navigation

`EnforceSorting` resorts the inout of UnionExec unnecessarily #4943

Metadata

Assignees

Labels

Type

Projects

Milestone

Relationships

Development

EnforceSorting resorts the inout of UnionExec unnecessarily #4943

Description

Activity

alamb commented on Jan 17, 2023

mustafasrepo commented on Jan 17, 2023

alamb commented on Jan 17, 2023

alamb commented on Jan 17, 2023

mingmwang commented on Jan 18, 2023

alamb commented on Jan 18, 2023

Metadata

Metadata

Assignees

Labels

Type

Projects

Milestone

Relationships

Development

Issue actions

`EnforceSorting` resorts the inout of UnionExec unnecessarily #4943