[SPARK-46378][SQL][FOLLOWUP] Do not rely on TreeNodeTag in Project #44429

cloud-fan · 2023-12-20T15:55:38Z

What changes were proposed in this pull request?

This is a followup of #44310 . It turns out that TreeNodeTag in Project is way too fragile. Project is a very basic node and very easy to get removed/transformed during plan optimization.

This PR switches to a different approach: since we can't retain the information (input data order doesn't matter) from Aggregate, let's leverage this information immediately. We pull out the expensive part of EliminateSorts to a new rule, so that we can safely call EliminateSorts right before we turn Aggregate into Project.

Why are the changes needed?

to make the optimizer more robust.

Does this PR introduce any user-facing change?

no

How was this patch tested?

existing tests

Was this patch authored or co-authored using generative AI tooling?

no

cloud-fan · 2023-12-20T15:56:30Z

sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/Optimizer.scala

      } else {
        s.copy(order = newOrders)
      }
-    case Sort(orders, false, child) if SortOrder.orderingSatisfies(child.outputOrdering, orders) =>


This is the expensive part as it need to calculate the ordering of children.

cloud-fan · 2023-12-20T16:01:34Z

sql/core/src/test/scala/org/apache/spark/sql/execution/datasources/V1WriteCommandSuite.scala

-        }, plan)
+        val sort = plan.collectFirst { case s: SortExec => s }
+        if (enabled) {
+          // With planned write, optimizer is more efficient and can eliminate `SORT BY value, key`.


This is a good side effect of this change. Before this PR, there is a conflict in EliminateSorts: ideally we remove the bottom Sort and keep the top Sort, but if the child sort ordering satisfies the top Sort, we remove top Sort. This is inconsistent and also suboptimal as we sort by more keys.

Now we have fixed the conflict. We always remove bottom Sort first.

cloud-fan · 2023-12-20T16:02:21Z

cc @dongjoon-hyun @viirya @ulysses-you

cloud-fan · 2023-12-20T16:03:53Z

sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/Optimizer.scala

 *    RepartitionByExpression, RebalancePartitions (with deterministic expressions) operators only
 *    and the Join condition is deterministic
- * 5) if the Sort operator is within GroupBy separated by 0...n Project, Filter, Repartition or
+ * 4) if the Sort operator is within GroupBy separated by 0...n Project, Filter, Repartition or


This part is still in EliminateSorts, so EliminateSorts is good enough for LimitPushDown

dongjoon-hyun

+1, LGTM.

dongjoon-hyun

Please fix the remaining failures.

[info] - SPARK-41914: v1 write with AQE and in-partition sorted - non-string partition column *** FAILED *** (188 milliseconds)

ulysses-you · 2023-12-21T02:51:02Z

sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/RemoveRedundantSorts.scala

+  }
+
+  private def recursiveRemoveSort(plan: LogicalPlan, optimizeGlobalSort: Boolean): LogicalPlan = {
+    if (!plan.containsPattern(SORT)) {


shall we pull out this to apply method ?

we should put it here to skip some children of a plan node.

plan.containsPattern contains the bitset of children..

When we traverse down a tree, we still need to apply the skipping for each plan node that has more than one children.

oh I see, make sense. Here we traverse the tree manually

viirya · 2023-12-21T03:01:29Z

sql/catalyst/src/main/scala/org/apache/spark/sql/catalyst/optimizer/Optimizer.scala

-        // so here we only preserve its output partitioning using `RepartitionByExpression`.
-        // We should use `None` as the optNumPartitions so AQE can coalesce shuffle partitions.
-        // This behavior is same with original global sort.
-        RepartitionByExpression(sortOrder, recursiveRemoveSort(child, true), None)


Hmm, previously this rule looks into this global Sort's child to remove local and global Sort recursively without condition. But in the new RemoveRedundantSorts rule:

case s @ Sort(orders, true, child) => val newChild = recursiveRemoveSort(child, optimizeGlobalSort = false)

recursiveRemoveSort in RemoveRedundantSorts only removes local Sort if its child is already sorted. Do we miss this optimization?

EliminateSorts still does this job: https://github.com/apache/spark/pull/44429/files#diff-11264d807efa58054cca2d220aae8fba644ee0f0f2a4722c46d52828394846efR1577

Said there are Sorts like

- Sort (local) - Sort (global) - Sort (local)

We reach:

case s @ Sort(_, global, child) => s.copy(child = recursiveRemoveSort(child, global))

Previously we can get rid of the middle global Sort and the bottom local Sort by RepartitionByExpression(sortOrder, recursiveRemoveSort(child, true), None) and:

case Sort(_, global, child) if canRemoveGlobalSort || !global => recursiveRemoveSort(child, canRemoveGlobalSort)

How does EliminateSorts still do it?
The code you point is same (not changed in this PR):

case s @ Sort(_, global, child) => s.copy(child = recursiveRemoveSort(child, global))

But in recursiveRemoveSort, as canRemoveGlobalSort is false, we don't get rid of the middle global Sort now (it will be done in RemoveRedundantSorts now).

Then the bottom local Sort under the rewritten RepartitionByExpression won't be optimized as it requires its child is sorted.

Do I miss or misread something?

After running EliminateSorts, the bottom sort is removed, then we run RemoveRedundantSorts which will turn the middle sort to local sort.

These two rules are in the same batch

viirya · 2023-12-21T08:19:31Z

Test failure looks unrelated?

cloud-fan · 2023-12-21T08:22:21Z

Yea the pyspark failure is unrelated. Thanks for the review, merging to master!

EnricoMi · 2023-12-22T16:03:04Z

This once again breaks writing sorted partitioned files, last time broken with 3.0.0 and fixed in 3.3.2: #38358, #39431.

When the user calls

ds.repartition(partitionColumns: _*)
  .sortWithinPartitions((partitionColumns ++ sortColumns): _*)
  .write
  .partitionBy(partitionColumns: _*)
  .parquet(...)

then the in-partition sort by sortColumns is not redundant but desired.

Is that meant to be fixed in #44458?

cloud-fan · 2023-12-22T16:07:51Z

@EnricoMi yes, will be fixed soon.

### What changes were proposed in this pull request? In `V1Writes`, we try to avoid adding Sort if the output ordering always satisfies. However, the code is completely broken with two issues: - we put `SortOrder` as the child of another `SortOrder` and compare, which always returns false. - once we add a project to do `empty2null`, we change the query output attribute id and the sort order never matches. It's not a big issue as we still have QO rules to eliminate useless sorts, but #44429 exposes this problem because the way we optimize sort is a bit different. For `V1Writes`, we should always avoid adding sort even if the number of ordering key is less, to not change the user query. ### Why are the changes needed? fix code mistakes. ### Does this PR introduce _any_ user-facing change? no ### How was this patch tested? updated test ### Was this patch authored or co-authored using generative AI tooling? no Closes #44458 from cloud-fan/sort. Authored-by: Wenchen Fan <[email protected]> Signed-off-by: Wenchen Fan <[email protected]>

EnricoMi · 2023-12-22T16:13:22Z

Thanks!

### What changes were proposed in this pull request? In `V1Writes`, we try to avoid adding Sort if the output ordering always satisfies. However, the code is completely broken with two issues: - we put `SortOrder` as the child of another `SortOrder` and compare, which always returns false. - once we add a project to do `empty2null`, we change the query output attribute id and the sort order never matches. It's not a big issue as we still have QO rules to eliminate useless sorts, but apache#44429 exposes this problem because the way we optimize sort is a bit different. For `V1Writes`, we should always avoid adding sort even if the number of ordering key is less, to not change the user query. ### Why are the changes needed? fix code mistakes. ### Does this PR introduce _any_ user-facing change? no ### How was this patch tested? updated test ### Was this patch authored or co-authored using generative AI tooling? no Closes apache#44458 from cloud-fan/sort. Authored-by: Wenchen Fan <[email protected]> Signed-off-by: Wenchen Fan <[email protected]>

Backport #44458 to branch-3.5. Justification: it fixes a hidden bug (until exposed by #44429) that has existed since 3.4. ### What changes were proposed in this pull request? In `V1Writes`, we try to avoid adding Sort if the output ordering always satisfies. However, the code is completely broken with two issues: - we put `SortOrder` as the child of another `SortOrder` and compare, which always returns false. - once we add a project to do `empty2null`, we change the query output attribute id and the sort order never matches. It's not a big issue as we still have QO rules to eliminate useless sorts, but #44429 exposes this problem because the way we optimize sort is a bit different. For `V1Writes`, we should always avoid adding sort even if the number of ordering key is less, to not change the user query. ### Why are the changes needed? fix code mistakes. ### Does this PR introduce _any_ user-facing change? no ### How was this patch tested? updated test ### Was this patch authored or co-authored using generative AI tooling? no Closes #52692 from pan3793/SPARK-46485-3.5. Authored-by: Wenchen Fan <[email protected]> Signed-off-by: Peter Toth <[email protected]>

do not rely on TreeNodeTag in Project

e550314

github-actions bot added the SQL label Dec 20, 2023

cloud-fan commented Dec 20, 2023

View reviewed changes

dongjoon-hyun approved these changes Dec 20, 2023

View reviewed changes

dongjoon-hyun reviewed Dec 20, 2023

View reviewed changes

update test

75f24b8

ulysses-you reviewed Dec 21, 2023

View reviewed changes

viirya reviewed Dec 21, 2023

View reviewed changes

ulysses-you approved these changes Dec 21, 2023

View reviewed changes

viirya approved these changes Dec 21, 2023

View reviewed changes

cloud-fan closed this in 0e94f34 Dec 21, 2023

cloud-fan mentioned this pull request Dec 22, 2023

[SPARK-46485][SQL] V1Write should not add Sort when not needed #44458

Closed

maytasm mentioned this pull request Dec 22, 2023

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

Closed

pan3793 mentioned this pull request Oct 22, 2025

[SPARK-46485][SQL][3.5] V1Write should not add Sort when not needed #52692

Closed

[SPARK-46378][SQL][FOLLOWUP] Do not rely on TreeNodeTag in Project #44429

[SPARK-46378][SQL][FOLLOWUP] Do not rely on TreeNodeTag in Project #44429

Uh oh!

Conversation

cloud-fan commented Dec 20, 2023

What changes were proposed in this pull request?

Why are the changes needed?

Does this PR introduce any user-facing change?

How was this patch tested?

Was this patch authored or co-authored using generative AI tooling?

Uh oh!

Choose a reason for hiding this comment

Uh oh!

cloud-fan Dec 20, 2023 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

cloud-fan commented Dec 20, 2023

Uh oh!

cloud-fan Dec 20, 2023 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

dongjoon-hyun left a comment

Choose a reason for hiding this comment

Uh oh!

dongjoon-hyun left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

viirya Dec 21, 2023 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

cloud-fan Dec 21, 2023 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

viirya commented Dec 21, 2023

Uh oh!

cloud-fan commented Dec 21, 2023

Uh oh!

EnricoMi commented Dec 22, 2023 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

cloud-fan commented Dec 22, 2023

Uh oh!

EnricoMi commented Dec 22, 2023

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

5 participants

cloud-fan Dec 20, 2023 •

edited

Loading

cloud-fan Dec 20, 2023 •

edited

Loading

viirya Dec 21, 2023 •

edited

Loading

cloud-fan Dec 21, 2023 •

edited

Loading

EnricoMi commented Dec 22, 2023 •

edited

Loading