Teach UnionExec to require its inputs sorted

[NGA-TRAN]

Is your feature request related to a problem or challenge?

In InfluxDB IOx, we want to improve this query plan:

Sort (col1, col2, col3, ... coln)
   Union
       Projection (col1, col2,  col3, ... coln)
       Projection (col1, col2,  col3, ... coln)
       Projection (col1, col2,  col3, ... coln)

Since col1 is the same for each projection, the sort can be pushed down below the Union and the new plan will be like this:

ProgressiveEval                            <-- new operator (available in InfluxDB IOx that output data in their order)
   Union
         Sort(col2, col3, ... coln)     <-- sort is pushed down but only from col2 to coln
              Projection (col1, col2,  col3, ... coln)
        Sort(col2, col3, ... coln)
             Projection (col1, col2,  col3, ... coln)
        .....
             .....
        Sort(col2, col3, ... coln)
             Projection (col1, col2,  col3, ... coln)

There are now many sorts but each only work on a subset of data in parallel. Also, the ProgressiveEval ensure the sort streams is ordered by col1 which is very cheap.

The above plan would work for us, however, we hit an issue in DataFusion that the Sorts under Union are always removed from the plan at the the enforce_sorting step https://github.com/apache/arrow-datafusion/blob/09f5a544d25f36ff1d65cc377123aee9b0e8f538/datafusion/core/src/physical_optimizer/enforce_sorting.rs#L361.

After some investigation, we found the reason the sorts under union are always removed because the Union does not implement required_input_ordering. It uses the default https://github.com/apache/arrow-datafusion/blob/179179c0b719a7f9e33d138ab728fdc2b0e1e1d8/datafusion/physical-plan/src/lib.rs#L155 which is always an array of None

Describe the solution you'd like

Implement required_input_ordering in UnionExec to have it ask its inputs to keep their sort order if the Union has output_ordering

Describe alternatives you've considered

No response

Additional context

No response

[NGA-TRAN]

I am working on this and this is the in-progress PR #9867

[rgehan]

Is there any plan on moving forward with the fix suggested in #9867?

We've got a use case where we need to merge multiple similarly dimensioned, and mostly sorted, Parquet files.
To achieve this, we leverage unions, add explicit sorts on the few files that are not sorted, then aggregate.
We want to sort near the data, so that subsequent aggregation can leverage SortPreservingMergeExec, minimizing memory usage.

While the initial plan only sorts the unsorted files, before the union happens, the optimized plan would cause the entire result set to be resorted, increasing memory usage/spilling.

I'm testing with target_partitions: 1, repartition_joins: false.

Before optimization

OutputRequirementExec: order_by=[], dist_by=Unspecified
  AggregateExec: mode=Final, gby=[col1@0 as col1, col2@1 as col2], aggr=[], ordering_mode=Sorted
    AggregateExec: mode=Partial, gby=[col1@0 as col1, col2@1 as col2], aggr=[], ordering_mode=Sorted
      UnionExec
        SortExec: expr=[col1@0 ASC, col2@1 ASC], preserve_partitioning=[false]
          DataSourceExec: file_groups={1 group: [[file1.parquet]]}, projection=[col1, col2], file_type=parquet
        DataSourceExec: file_groups={1 group: [[file2.parquet]]}, projection=[col1, col2], output_ordering=[col1@0 ASC, col2@1 ASC], file_type=parquet
        DataSourceExec: file_groups={1 group: [[file3.parquet]]}, projection=[col1, col2], output_ordering=[col1@0 ASC, col2@1 ASC], file_type=parquet

Actual plan, after optimization

• Sort below the union is removed
• Sort is lifted higher up
• We don't use SortPreservingMergeExec

OutputRequirementExec: order_by=[], dist_by=Unspecified
  AggregateExec: mode=Final, gby=[col1@0 as col1, col2@1 as col2], aggr=[], ordering_mode=Sorted
    SortExec: expr=[col1@0 ASC NULLS LAST, col2@1 ASC NULLS LAST], preserve_partitioning=[false]
      CoalescePartitionsExec
        AggregateExec: mode=Partial, gby=[col1@0 as col1, col2@1 as col2], aggr=[]
          UnionExec
            DataSourceExec: file_groups={1 group: [[file1.parquet]]}, projection=[col1, col2], file_type=parquet
            DataSourceExec: file_groups={1 group: [[file2.parquet]]}, projection=[col1, col2], output_ordering=[col1@0 ASC, col2@1 ASC], file_type=parquet
            DataSourceExec: file_groups={1 group: [[file3.parquet]]}, projection=[col1, col2], output_ordering=[col1@0 ASC, col2@1 ASC], file_type=parquet

Expected plan, after optimization

• Sort below the union is preserved
• No top-level sort is added
• We use SortPreservingMergeExec
• Aggregation is streamed / minimizes spills (i.e. ordering_mode=Sorted)

OutputRequirementExec: order_by=[], dist_by=Unspecified
  AggregateExec: mode=Final, gby=[col1@0 as col1, col2@1 as col2], aggr=[], ordering_mode=Sorted
    SortPreservingMergeExec: [col1@0 ASC, col2@1 ASC]
      AggregateExec: mode=Partial, gby=[col1@0 as col1, col2@1 as col2], aggr=[], ordering_mode=Sorted
        UnionExec
          SortExec: expr=[col1@0 ASC, col2@1 ASC], preserve_partitioning=[false]
            DataSourceExec: file_groups={1 group: [[file1.parquet]]}, projection=[col1, col2], output_ordering=[col1@0 ASC, col2@1 ASC], file_type=parquet
          DataSourceExec: file_groups={1 group: [[file2.parquet]]}, projection=[col1, col2], output_ordering=[col1@0 ASC, col2@1 ASC], file_type=parquet
          DataSourceExec: file_groups={1 group: [[file3.parquet]]}, projection=[col1, col2], output_ordering=[col1@0 ASC, col2@1 ASC], file_type=parquet

I'd be happy to help bring the PR to completion if needs be, but might need some guidance

[NGA-TRAN]

@rgehan: Your use case is both different and more common, and I agree it’s something we should aim to support. The best way to move forward is to share a reproducer for your Before Optimization plan—specifically for the many partitions case—and describe the improvement you’d like to see in the optimized plan.

Once you have a reproducer, feel free to tag me. I believe I can propose a change to support it and will help recruit someone to implement it.

[rgehan]

@NGA-TRAN Thank you for your answer, I opened a PR with a few tests showcasing the issues we're observing: #18352 (I left it as a draft)

Note that with your fix (#9867), all these tests pass

[NGA-TRAN]

Thanks @rgehan . I have left a few comments in your reproducer PR.
I guess after you address the comments, you may want to create new feature request ticket for the work you described above and link to your reproducer. Then you can try to fix it using #9867. You may need some more work to have all the tests pass