feat: physical execution for range partitioning

[saadtajwar]

Which issue does this PR close?

• Closes Support physical execution of Range repartitioning #23137

Rationale for this change

Range repartitioning was already planned and serialized into physical plans, but RepartitionExec could not execute it. This PR completes the core execution path so rows in an input batch are routed to the correct output partition based on range split points and the ordering defined on the partitioning scheme.

What changes are included in this PR?

This PR adds a Range variant to BatchPartitioner that evaluates the ordering expressions on each input batch, compares each row's key against split points using compare_rows (respecting ASC/DESC and null ordering), and assigns row indices to output partitions via binary search. The partitioned row indices are then materialized into sub-batches using the same partition_grouped_take path as hash repartitioning. pull_from_input is wired to construct a range partitioner for Partitioning::Range, replacing the previous not_impl_err! at execution time.

Optimizer-related paths remain intentionally unimplemented and are tracked in #23230: projection pushdown through RepartitionExec (try_swapping_with_projection), sort pushdown (try_pushdown_sort), and changing partition counts via repartitioned().

Are these changes tested?

Yes!

Are there any user-facing changes?

No public API changes

[saadtajwar]

cc @gene-bordegaray ! 🎉

[gene-bordegaray]

Thank you for the work @saadtajwar , I think this will be very useful in upcoming efforts 😄

Before really diving into this we shoudl step back and plan how repartitioning will work from a high level first before diving into the nitty gritty. Per descussions here #23236 it seems that we will be working toward deprecating HashPartitioned and move to KeyPartitioned distribution variant.

So essentially we are going to have operators that require a KeyPartitioned distirbution with two options to achieve this. Repartition via Hash or repartition via Range. It is unclear to me exactly the best way to make this decision and if / how we can recognize to use one or the other. Should this be something that users specify as a config? Is there some way to dtect this? Should we only repartition to range if it is to a superset of the current range partitioning (example: data partitioned on day -> repartition to hour)?

These are some things I would like to discuss with other before we decide to implement anything regarding repartitioning (as of now we just preserve it from a DataSourceExec)

cc: @alamb @gabotechs @stuhood @2010YOUY01

[saadtajwar]

Hey @gene-bordegaray - that makes sense, thanks! I just posted some thoughts in #23236 just to help us keep the discussion centralized in one spot - looking forward to working on this all together!

[asolimando]

So essentially we are going to have operators that require a KeyPartitioned distirbution with two options to achieve this. Repartition via Hash or repartition via Range. [...] Should we only repartition to range if it is to a superset of the current range partitioning (example: data partitioned on day -> repartition to hour)?

Filters involving ranges (at least BETWEEN, <=, <, >, >=, involving literals at first, but possibly we can do something smart for columns and more complex expressions too) could benefit from range partitioning too, as it would allow partition pruning of entire partitions without evaluation.

[stuhood]

It is unclear to me exactly the best way to make this decision and if / how we can recognize to use one or the other. Should this be something that users specify as a config? Is there some way to dtect this? Should we only repartition to range if it is to a superset of the current range partitioning (example: data partitioned on day -> repartition to hour)?

From briefly looking around, I only see a few cases where a logical optimizer might want to request Range rather than Hash (things like non-equi joins, re-organizing data for output that preserves partitioning, global window functions.)

Filters involving ranges (at least BETWEEN, <=, <, >, >=, involving literals at first, but possibly we can do something smart for columns and more complex expressions too) could benefit from range partitioning too, as it would allow partition pruning of entire partitions without evaluation.

They benefit from existing Range partitioning, but it probably wouldn't make sense to repartition data using Range for that purpose: Hash will get you better balance more cheaply (rather than via binary search), and then each partition can directly evaluate the filter.

Before really diving into this we shoudl step back and plan how repartitioning will work from a high level first before diving into the nitty gritty.

But the choice to introduce Range partitioning would be a logical decision, right? So, while I agree that changing logical optimizers to request Range would take a lot of thought and design, implementing the physical side (this PR) doesn't seem to be blocked on that? Or are you concerned that the API might still shift, or that it won't have enough test-coverage?

[saadtajwar]

Agree with @stuhood on the above, especially on the below - while I'm still trying to understand the Distribution options and when range partitioning would be chosen as the distribution scheme, as long as the current Range partitioning API wouldn't change (using ordering, split points, etc), especially the physical executors, it makes sense to me that the physical execution here should remain the same

But the choice to introduce Range partitioning would be a logical decision, right? So, while I agree that changing logical optimizers to request Range would take a lot of thought and design, implementing the physical side (this PR) doesn't seem to be blocked on that? Or are you concerned that the API might still shift, or that it won't have enough test-coverage?

[gene-bordegaray]

But the choice to introduce Range partitioning would be a logical decision, right? So, while I agree that changing logical optimizers to request Range would take a lot of thought and design, implementing the physical side (this PR) doesn't seem to be blocked on that? Or are you concerned that the API might still shift, or that it won't have enough test-coverage?

@stuhood I am most concerned with implementing physical layer behavior before having a real use for it that we can represent. What would the use case of being able to repartiution on range right now be? Do you have a use case where you would like to phsyically insert a repartition on range? maybe this is a good place to start the conversation on where and how this should be decided 🤔

[alamb]

@stuhood I am most concerned with implementing physical layer behavior before having a real use for it that we can represent. What would the use case of being able to repartiution on range right now be? Do you have a use case where you would like to phsyically insert a repartition on range? maybe this is a good place to start the conversation on where and how this should be decided 🤔

The main usecase we have at the moment for range partitioning is when the input source data is already range partitioned and the point of the work in this epic is for DataFusion to know about that (pre-existing) partitioning and take advantage of it

I think you guys are talking about having hte optimizer decide to repartition data into ranges (e.g. when it wants to add more parallelism to the plan). That would probably need to be a cost based decision based on statistics (like value distributions) that we don't yet have in DataFusion (and maybe never will have).

[alamb]

TLDR is I agree with @stuhood

They benefit from existing Range partitioning, but it probably wouldn't make sense to repartition data using Range for that purpose: Hash will get you better balance more cheaply (rather than via binary search), and then each partition can directly evaluate the filter.

👍

[stuhood]

But the choice to introduce Range partitioning would be a logical decision, right? So, while I agree that changing logical optimizers to request Range would take a lot of thought and design, implementing the physical side (this PR) doesn't seem to be blocked on that? Or are you concerned that the API might still shift, or that it won't have enough test-coverage?

@stuhood I am most concerned with implementing physical layer behavior before having a real use for it that we can represent. What would the use case of being able to repartiution on range right now be? Do you have a use case where you would like to phsyically insert a repartition on range? maybe this is a good place to start the conversation on where and how this should be decided 🤔

Understood.

Yea, I don't feel strongly about it either way... but I don't really love the idea of leaving in todo!()s that would panic if actually used. Even if the DF repository itself does not contain optimizer rules which introduce Range, I could imagine consuming repos attempting to introduce it, only to have it fail? But perhaps the DF project would prefer those folks to then come and talk about their use cases, to see whether they could be upstreamed.

[alamb]

Yea, I don't feel strongly about it either way... but I don't really love the idea of leaving in todo!()s that would panic if actually used. Even if the DF repository itself does not contain optimizer rules which introduce Range, I could imagine consuming repos attempting to introduce it, only to have it fail? But perhaps the DF project would prefer those folks to then come and talk about their use cases, to see whether they could be upstreamed.

I agree panics are not great -- returing an NotYetImplemned error would be better.

[gene-bordegaray]

I agree panics are not good and to keep the existing not_impl_err()s but for this specific PR are we agreeing that it is best to not introduce this until DF can explicitly make a decision like with CBO as mentioned or a user has a need for this so we can better understand the use case before implementing?

[alamb]

I agree panics are not good and to keep the existing not_impl_err()s but for this specific PR are we agreeing that it is best to not introduce this until DF can explicitly make a decision like with CBO as mentioned or a user has a need for this so we can better understand the use case before implementing?

One question I would have is:

1. What type of query would produce a range repartitioning? Can we write such a think in SQL today?

I can envision it being used when one side of a join is unpartitioned and one side is range partitioned -- the optimizer would have to put a range repartition to match up the keys

If that is the case, adding range repartitioning in preparation seems ok to me

[gene-bordegaray]

I can envision it being used when one side of a join is unpartitioned and one side is range partitioned -- the optimizer would have to put a range repartition to match up the keys

If that is the case, adding range repartitioning in preparation seems ok to me

@alamb yes it seems I was getting to far ahead of myself 😅 . This would be a great first use case and a good place to start. I immediately jumped to some fantasy of choosing a partitioning based on the stats or some complex planning

[saadtajwar]

I can envision it being used when one side of a join is unpartitioned and one side is range partitioned -- the optimizer would have to put a range repartition to match up the keys
If that is the case, adding range repartitioning in preparation seems ok to me

@alamb yes it seems I was getting to far ahead of myself 😅 . This would be a great first use case and a good place to start. I immediately jumped to some fantasy of choosing a partitioning based on the stats or some complex planning

If we're aligned on moving forward with this change set, then please let me know if you have any thoughts from reviewing! 😁

[gene-bordegaray]

thank you @saadtajwar!

overall looks solid, love seeing a binary search in the wild 💯

I had some ideas for optimizations and saw you made a follow up issue for partitioning, could we do similar for repartitioning (if you could link the Range partitioning epic: #22395 that would be awesome), I think some of those can be addressed there if they are more nuanced and marked clearly.

I would also appreciate some more test cses covered and small nits 👍

[saadtajwar]

thank you @saadtajwar!

overall looks solid, love seeing a binary search in the wild 💯

I had some ideas for optimizations and saw you made a follow up issue for that (if you could link the Range partitioning epic: #22395 that would be awesome), I think some of those can be addressed there if they are more nuanced and marked clearly.

I would also appreciate some more test cses covered and small nits 👍

Awesome, thanks @gene-bordegaray ! Will look over these and hoping to make changes start of next week!

[saadtajwar]

@gene-bordegaray thanks again for the thorough review here! Have pushed changes to address your comments, please let me know your thoughts!

The only open question I had was regarding the follow-ups for optimizations - just to clarify, would you like another GH issue opened to track changes like the one described here, beyond the existing #23230 issue that we've already opened? Or can #23230 encapsulate all of the follow-up work needed here?

[gene-bordegaray]

looking much better 💯 few more suggestions

[saadtajwar]

looking much better 💯 few more suggestions

Thank you! Hopefully will be able to address these later tonight!

[saadtajwar]

@gene-bordegaray - just pushed a couple commits to address the new comments! 🙏

[gene-bordegaray]

this looks good from my end

@stuhood I know you were involve in this as well 🙇

[gene-bordegaray]

wowzer thats pretty good 👍

[gene-bordegaray]

also a nit, I dont know what the preference for DF is but we could squash commit history a bit

[stuhood]

Thanks to @alamb for reminding us of the elephant-in-the-room use case for Range repartitioning.

[gabotechs]

also a nit, I dont know what the preference for DF is but we could squash commit history a bit

Should not be an issue, the merge queue will squash it automatically

[gabotechs]

Thanks @saadtajwar for the PR, and @gene-bordegaray and @stuhood for the reviews!

Looks good to me, as soon as CI is green, I think this is ready to go.

[gene-bordegaray]

Had one last omment to be addressed on pushdowns 👍

https://github.com/apache/datafusion/pull/23231/changes#r3543211552

[saadtajwar]

@gene-bordegaray pushed a commit to address your comment! Thank you all folks @gene-bordegaray & @gabotechs & @stuhood for taking the time to review! The cargo fmt CI test should pass now too 🙏

Quick bump on the above question, this probably just got lost in the shuffle of all of the comments & changes haha - @gene-bordegaray lmk if you have thoughts on this below!

The only open question I had was regarding the follow-ups for optimizations - just to clarify, would you like another GH issue opened to track changes like the one described #23231 (comment), beyond the existing #23230 issue that we've already opened? Or can #23230 encapsulate all of the follow-up work needed here?

[gene-bordegaray]

@gene-bordegaray pushed a commit to address your comment! Thank you all folks @gene-bordegaray & @gabotechs & @stuhood for taking the time to review! The cargo fmt CI test should pass now too 🙏

Quick bump on the above question, this probably just got lost in the shuffle of all of the comments & changes haha - @gene-bordegaray lmk if you have thoughts on this below!

The only open question I had was regarding the follow-ups for optimizations - just to clarify, would you like another GH issue opened to track changes like the one described #23231 (comment), beyond the existing #23230 issue that we've already opened? Or can #23230 encapsulate all of the follow-up work needed here?

@saadtajwar Lets hold off on doing optimizations here before e profile, then can make issues after 👍

[saadtajwar]

@gene-bordegaray pushed a commit to address your comment! Thank you all folks @gene-bordegaray & @gabotechs & @stuhood for taking the time to review! The cargo fmt CI test should pass now too 🙏
Quick bump on the above question, this probably just got lost in the shuffle of all of the comments & changes haha - @gene-bordegaray lmk if you have thoughts on this below!

The only open question I had was regarding the follow-ups for optimizations - just to clarify, would you like another GH issue opened to track changes like the one described #23231 (comment), beyond the existing #23230 issue that we've already opened? Or can #23230 encapsulate all of the follow-up work needed here?

@saadtajwar Lets hold off on doing optimizations here before e profile, then can make issues after 👍

Sounds good! Should I keep #23230 open and have a follow-up PR for those parts we had labeled as TODOs? Or hold off for now?

[alamb]

Is this PR ready to merge?

[gene-bordegaray]

Is this PR ready to merge?

@alamb yes