[SPARK-39911][SQL] Optimize global Sort to RepartitionByExpression

[ulysses-you]

What changes were proposed in this pull request?

Optimize Global sort to RepartitionByExpression, for example:

Sort local             Sort local 
  Sort global    =>      RepartitionByExpression

Why are the changes needed?

If a global sort below a local sort, the only meaningful thing is it's distribution. So this pr optimizes that global sort to RepartitionByExpression to save a local sort.

Does this PR introduce any user-facing change?

no, only improve performance

How was this patch tested?

add test

[cloud-fan]

cc @sigmod

[sigmod]

Thanks @ulysses-you !

[cloud-fan]

thanks, merging to master!

[cloud-fan]

@ulysses-you can you open a backport PR for 3.3? I think this is a necessary followup of #37250 to avoid perf regression.

[maytasm]

@ulysses-you @cloud-fan
Not sure if I am missing something but can this causes performance degradation if my sort order is on a key with few/single values?
For example, if I have 500 shuffle partitions...
Without this change:

Sort local    
  Sort global   

both of the above stages would run with 500 tasks
With this change: and say the RepartitionByExpression is on the column date and there is only a single value for this in my dataset

Sort local 
   RepartitionByExpression

RepartitionByExpression will run with 500 tasks and create a single partition
then the Sort would run with a single task (as there is only a single partition)

[ulysses-you]

hi @maytasm , Sort global is semantics equal to Sort local + RepartitionByExpression. For your case, if there is only single value on the sort column, then Sort global would also introduce a range partitioning shuffle and make a single partition task at the end. Can you show a real case if you find any behavior change ?

[maytasm]

@ulysses-you
Sorry for the late reply. That is true unless the following condition is True for the Sort global:

   case s @ Sort(orders, _, child) if orders.isEmpty || orders.exists(_.child.foldable) =>
      val newOrders = orders.filterNot(_.child.foldable)
      if (newOrders.isEmpty) {
        applyLocally.lift(child).getOrElse(child)
      }

In the case that the above condition is True, then this Sort global will be removed. However, with this change (in this PR), the Sort global will not be removed as it would have been turn into a RepartitionByExpression

Consider the following case:

+- Sort [date#0 ASC NULLS FIRST], false
   +- Project [...]
      +- Sort [20231221 ASC NULLS FIRST], true
         +- Aggregate [...]
               +- RelationV2[...] some_table

without the change from this PR, the Sort [20231221 ASC NULLS FIRST], true would be remove. We would end up with

+- Sort [date#0 ASC NULLS FIRST], false
   +- Project [...]
         +- Aggregate [...]
               +- RelationV2[...] some_table

This ran fast as we have high parallelism in all stages (many partitions -> many tasks running)
However, with the change in this PR:

+- Sort [date#0 ASC NULLS FIRST], false
   +- Project [...]
      +- Exchange rangepartition (20231221 ASC NULL FIRST)
         +- Aggregate [...]
               +- RelationV2[...] some_table

The problem here is that Exchange rangepartition reduces the partition to 1. Then the stages after that like the Sort on date#0 ran with a single task.

[cloud-fan]

without the change from this PR, the Sort [20231221 ASC NULLS FIRST], true would be remove.

Are you sure? It looks wrong to do so

[maytasm]

@cloud-fan
Can you expand on why you think the Sort [20231221 ASC NULLS FIRST], true looks wrong to remove? It is a sort that has no reference and should be fine to remove.
Sort [20231221 ASC NULLS FIRST], true would look something like:

Since Literal is foldable, newOrders would be empty. As newOrders is empty, the sort is then removed.
This is explained in #11840

I believe this PR prevents the optimization added in #11840 to work as intended since it convert the no-op sort into RepartitionByExpression, which can no longer be removed

[cloud-fan]

Oh now I get it. This is sort by a constant. Can you try the latest master branch? I think https://github.com/apache/spark/pull/44429/files#diff-11264d807efa58054cca2d220aae8fba644ee0f0f2a4722c46d52828394846efR214 has solved it.

[maytasm]

Ah! I haven't tried running yet but took a look at PR #44429 and I think it does solve this issue. So basically the logic for converting global sort to Exchange rangepartition is moved into another rule, RemoveRedundantSorts, and EliminateSorts is run before RemoveRedundantSorts. Meaning the no-op sort like Sort [20231221 ASC NULLS FIRST], true would be removed by EliminateSorts already by the time RemoveRedundantSorts is run.

[maytasm]

@cloud-fan Confirmed that PR #44429 fixed the issue I was having with the no-op sort!