Unbounded SortExec (and Top-K) Implementation When Req's Are Satisfied #12174

berkaysynnada · 2024-08-26T13:24:55Z

Which issue does this PR close?

Closes #.

Rationale for this change

SortExec (with or w/out fetch) can work without an actual sort if the existing input order is required.

What changes are included in this PR?

When we create a new SortExec, its execution mode is set according to its input order.
Adding a fetch can update its execution mode.
During execution, we consult the information of matching the requirements and having a fetch. If we already satisfy the requirements, SortExec either short-circuits itself (if it does not have a fetch), or behaves as a limit operator (if it has a fetch).

Are these changes tested?

Yes. I cannot practice the newly added stream types via an .slt test, but a unit test is added.

Are there any user-facing changes?

ozankabak

LGTM thank you

alamb

I recommend we add a test for this change (so that we don't accidentally break it in some future refactor).

alamb · 2024-08-26T19:24:04Z

datafusion/physical-plan/src/sorts/sort.rs

+    pub fn with_fetch(&self, fetch: Option<usize>) -> Self {
+        let mut cache = self.cache.clone();
+        if fetch.is_some() {
+            // When a theoretically unnecessary sort becomes a top-K (which


I don't fully understand how a top-k sort would become bounded. I may misundersrtand what the ExecutionMode trait means, but it seems like TopK could not complete until its input completed, but if its input was unbounded the sort itself therefore would also be unbounded

In general you are right -- however, you are missing that we turn it into Bounded only when the sort requirement is already satisfied. This happens when a sort "becomes" unnecessary during one of the plan optimization steps (and it will eventually get removed).

I don't fully understand how a top-k sort would become bounded. I may misundersrtand what the ExecutionMode trait means, but it seems like TopK could not complete until its input completed, but if its input was unbounded the sort itself therefore would also be unbounded

I misassumed the implementation of top-k. Could you please take a look to the new idea? I will update the PR title and body

FWIW I think the comments help explain what is going on here well. Thank you

ozankabak

Thanks for improving the corner cases to match the theory. IMO this is ready to go -- @alamb a quick look would be appreciated

datafusion/physical-plan/src/sorts/sort.rs

alamb

Thanks @berkaysynnada and @ozankabak

I am sorry, but I am still struggling to understand this PR

It seems like this PR changes the execution plan for a Sort(limit = 5) into a Limit if the data is already sorted correctly (aka doesn't need an additional sort).

But in this case I would have expected the Sort to have been removed by one of the optimizer passes rather than the SortExec implementing a limit. 🤔

alamb · 2024-08-28T12:43:34Z

datafusion/physical-plan/src/sorts/sort.rs

-                execution_options.sort_in_place_threshold_bytes,
-                &self.metrics_set,
-                context.runtime_env(),
+        let sort_satisfied = self


This is the same calculation as self.execution_mode(), right? Maybe we could call self.execution_mode here instead to be more efficient and ensure the calculations remained in sync

It is similar but not exactly the same (i.e. execution mode is derived from sort_satisfied but AFAIK the reverse is not possible). I think @berkaysynnada tried this but it didn't work

I have tried that but didn't work. Knowing the sort is bounded or unbounded does not mean sort is satisfied.

alamb · 2024-08-28T12:47:12Z

datafusion/physical-plan/src/sorts/sort.rs

    }

    fn fetch(&self) -> Option<usize> {
        self.fetch
    }
 }

+struct TopKStream {


I think it would help to add documentation to this struct, specifically that explains how it is different than TopK

alamb · 2024-08-28T12:47:31Z

datafusion/physical-plan/src/sorts/sort.rs

+    fetch: usize,
+}
+
+impl Stream for TopKStream {


This looks very similar to LimitStream -- https://docs.rs/datafusion-physical-plan/41.0.0/src/datafusion_physical_plan/limit.rs.html#434

though limit stream has metrics and some other features

Yes, maybe we can reuse it. We will take a look

You are right -- sent a commit to reuse LimitStream

ozankabak · 2024-08-28T12:59:20Z

But in this case I would have expected the Sort to have been removed by one of the optimizer passes rather than the SortExec implementing a limit. 🤔

Right. The sort will actually be entirely removed in later passes. However, in the meantime (before its removal), any rule looking at the execution mode will still think this TopK is pipeline-breaking (while it is not), which results in them behaving incorrectly. We caught this behavior downstream in the context of a custom rule.

Basically, this PR does two things:

It makes the execution_mode API give a correct result when data is already sorted, so intermediate rules that look at execution modes do not falsely think that the pipeline breaks somewhere.
It makes TopK's execution logic behave like a limit if it somehow ends up in the plan without being removed. This is not possible in upstream DataFusion (sans bugs), but people using custom optimizers can end up in such a state. In such cases, it will also execute more efficiently.

datafusion/physical-plan/src/sorts/sort.rs

ozankabak · 2024-08-28T17:59:24Z

Thanks for taking another look -- incorporated your feedback to leverage LimitStream

ozankabak · 2024-08-29T05:14:21Z

I will go ahead and merge this soon since it is a small, localized change that improves corner cases without interfering with the main use case of SortExec.

If there are any lingering concerns, we will address with a quick follow-on PR

alamb · 2024-09-01T12:28:37Z

Thank you for responding to the feedback @berkaysynnada and @ozankabak -- sorry for my delay -- I have been out. This PR now looks good to me

Sort fetch updates execution mode

4e9424b

github-actions bot added the physical-expr Physical Expressions label Aug 26, 2024

ozankabak approved these changes Aug 26, 2024

View reviewed changes

alamb reviewed Aug 26, 2024

View reviewed changes

berkaysynnada added 4 commits August 27, 2024 15:02

Update sort.rs

61ca06c

Update sort.rs

5f80cdc

Update sort.rs

65679db

Update sort.rs

35bfb10

berkaysynnada changed the title ~~TopK Should Update ExecutionMode~~ Unbounded SortExec (and Top-K) Implementation When Req's Are Satisfied Aug 27, 2024

Update sort.rs

7715255

ozankabak approved these changes Aug 28, 2024

View reviewed changes

datafusion/physical-plan/src/sorts/sort.rs Outdated Show resolved Hide resolved

datafusion/physical-plan/src/sorts/sort.rs Outdated Show resolved Hide resolved

ozankabak and others added 2 commits August 28, 2024 11:27

Apply suggestions from code review

7d713ff

Update sort.rs

8f43cf6

alamb reviewed Aug 28, 2024

View reviewed changes

alamb mentioned this pull request Aug 28, 2024

Minor: Improve ExecutionMode documentation #12214

Merged

Merge branch 'main' into feature/sort-fetch-exec-mode

9d1cdea

ozankabak reviewed Aug 28, 2024

View reviewed changes

datafusion/physical-plan/src/sorts/sort.rs Outdated Show resolved Hide resolved

Update datafusion/physical-plan/src/sorts/sort.rs

23643cc

ozankabak reviewed Aug 28, 2024

View reviewed changes

datafusion/physical-plan/src/sorts/sort.rs Outdated Show resolved Hide resolved

ozankabak added 2 commits August 28, 2024 18:17

Update datafusion/physical-plan/src/sorts/sort.rs

2d8954e

Reuse LimitStream

d2208b6

ozankabak merged commit f5dcdf0 into apache:main Aug 29, 2024
24 checks passed

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Unbounded SortExec (and Top-K) Implementation When Req's Are Satisfied #12174

Unbounded SortExec (and Top-K) Implementation When Req's Are Satisfied #12174

berkaysynnada commented Aug 26, 2024 •

edited

Loading

ozankabak left a comment

alamb left a comment

alamb Aug 26, 2024

ozankabak Aug 27, 2024

berkaysynnada Aug 27, 2024 •

edited

Loading

alamb Sep 1, 2024

ozankabak left a comment

alamb left a comment

alamb Aug 28, 2024

ozankabak Aug 28, 2024

berkaysynnada Aug 29, 2024

alamb Aug 28, 2024

alamb Aug 28, 2024

ozankabak Aug 28, 2024

ozankabak Aug 28, 2024

ozankabak commented Aug 28, 2024 •

edited

Loading

ozankabak commented Aug 28, 2024

ozankabak commented Aug 29, 2024 •

edited

Loading

alamb commented Sep 1, 2024

Unbounded SortExec (and Top-K) Implementation When Req's Are Satisfied #12174

Unbounded SortExec (and Top-K) Implementation When Req's Are Satisfied #12174

Conversation

berkaysynnada commented Aug 26, 2024 • edited Loading

Which issue does this PR close?

Rationale for this change

What changes are included in this PR?

Are these changes tested?

Are there any user-facing changes?

ozankabak left a comment

Choose a reason for hiding this comment

alamb left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

berkaysynnada Aug 27, 2024 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

ozankabak left a comment

Choose a reason for hiding this comment

alamb left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

ozankabak commented Aug 28, 2024 • edited Loading

ozankabak commented Aug 28, 2024

ozankabak commented Aug 29, 2024 • edited Loading

alamb commented Sep 1, 2024

berkaysynnada commented Aug 26, 2024 •

edited

Loading

berkaysynnada Aug 27, 2024 •

edited

Loading

ozankabak commented Aug 28, 2024 •

edited

Loading

ozankabak commented Aug 29, 2024 •

edited

Loading