You are viewing a plain text version of this content. The canonical link for it is here.
Posted to github@datafusion.apache.org by "Dandandan (via GitHub)" <gi...@apache.org> on 2025/04/18 19:52:47 UTC
[PR] User `interleave` in hash repartitioning [datafusion]
Dandandan opened a new pull request, #15768:
URL: https://github.com/apache/datafusion/pull/15768
## Which issue does this PR close?
<!--
We generally require a GitHub issue to be filed for all bug fixes and enhancements and this helps us generate change logs for our releases. You can link an issue to this PR using the GitHub syntax. For example `Closes #123` indicates that this PR will close issue #123.
-->
Closes https://github.com/apache/datafusion/issues/7957
## Rationale for this change
<!--
Why are you proposing this change? If this is already explained clearly in the issue then this section is not needed.
Explaining clearly why changes are proposed helps reviewers understand your changes and offer better suggestions for fixes.
-->
## What changes are included in this PR?
<!--
There is no need to duplicate the description in the issue here but it is sometimes worth providing a summary of the individual changes in this PR.
-->
## Are these changes tested?
<!--
We typically require tests for all PRs in order to:
1. Prevent the code from being accidentally broken by subsequent changes
2. Serve as another way to document the expected behavior of the code
If tests are not included in your PR, please explain why (for example, are they covered by existing tests)?
-->
## Are there any user-facing changes?
<!--
If there are user-facing changes then we may require documentation to be updated before approving the PR.
-->
<!--
If there are any breaking changes to public APIs, please add the `api change` label.
-->
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Use `interleave` in hash repartitioning [datafusion]
Posted by "ctsk (via GitHub)" <gi...@apache.org>.
ctsk commented on code in PR #15768:
URL: https://github.com/apache/datafusion/pull/15768#discussion_r2051758338
##########
datafusion/physical-plan/src/repartition/mod.rs:
##########
@@ -298,25 +299,15 @@ impl BatchPartitioner {
.into_iter()
.enumerate()
.filter_map(|(partition, indices)| {
- let indices: PrimitiveArray<UInt32Type> = indices.into();
(!indices.is_empty()).then_some((partition, indices))
})
.map(move |(partition, indices)| {
// Tracking time required for repartitioned batches construction
let _timer = partitioner_timer.timer();
+ let b: Vec<&RecordBatch> = batches.iter().collect();
// Produce batches based on indices
- let columns = take_arrays(batch.columns(), &indices, None)?;
-
- let mut options = RecordBatchOptions::new();
- options = options.with_row_count(Some(indices.len()));
- let batch = RecordBatch::try_new_with_options(
- batch.schema(),
- columns,
- &options,
- )
- .unwrap();
-
+ let batch = interleave_record_batch(&b, &indices)?;
Review Comment:
I see I had misunderstood this PR. It makes a lot of sense to do this. As part of prototyping the integration of a take_in in API in datafusion, I made a similar change - move the buffering *before* sending to their destination thread. I don't remeber see as much speedup when I benchmarked that change independently - I guess using interleave instead of a take/concat (like I did back then) combo makes a significant difference. Awesome!
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Use `interleave` in hash repartitioning [datafusion]
Posted by "ctsk (via GitHub)" <gi...@apache.org>.
ctsk commented on code in PR #15768:
URL: https://github.com/apache/datafusion/pull/15768#discussion_r2051665380
##########
datafusion/physical-plan/src/repartition/mod.rs:
##########
@@ -298,25 +299,15 @@ impl BatchPartitioner {
.into_iter()
.enumerate()
.filter_map(|(partition, indices)| {
- let indices: PrimitiveArray<UInt32Type> = indices.into();
(!indices.is_empty()).then_some((partition, indices))
})
.map(move |(partition, indices)| {
// Tracking time required for repartitioned batches construction
let _timer = partitioner_timer.timer();
+ let b: Vec<&RecordBatch> = batches.iter().collect();
// Produce batches based on indices
- let columns = take_arrays(batch.columns(), &indices, None)?;
-
- let mut options = RecordBatchOptions::new();
- options = options.with_row_count(Some(indices.len()));
- let batch = RecordBatch::try_new_with_options(
- batch.schema(),
- columns,
- &options,
- )
- .unwrap();
-
+ let batch = interleave_record_batch(&b, &indices)?;
Review Comment:
Nice change, and a much clearer performance win than https://github.com/apache/datafusion/pull/15479. I expect (without testing) that these two PRs interact negatively with one another - Removing coalesce will mean that the data is "more scattered" in memory and probably make interleave work worse - as well as the computation of the left join keys.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] User `interleave` in hash repartitioning [datafusion]
Posted by "Dandandan (via GitHub)" <gi...@apache.org>.
Dandandan commented on code in PR #15768:
URL: https://github.com/apache/datafusion/pull/15768#discussion_r2051062202
##########
datafusion/physical-plan/src/repartition/mod.rs:
##########
@@ -298,25 +299,15 @@ impl BatchPartitioner {
.into_iter()
.enumerate()
.filter_map(|(partition, indices)| {
- let indices: PrimitiveArray<UInt32Type> = indices.into();
(!indices.is_empty()).then_some((partition, indices))
})
.map(move |(partition, indices)| {
// Tracking time required for repartitioned batches construction
let _timer = partitioner_timer.timer();
+ let b: Vec<&RecordBatch> = batches.iter().collect();
// Produce batches based on indices
- let columns = take_arrays(batch.columns(), &indices, None)?;
-
- let mut options = RecordBatchOptions::new();
- options = options.with_row_count(Some(indices.len()));
- let batch = RecordBatch::try_new_with_options(
- batch.schema(),
- columns,
- &options,
- )
- .unwrap();
-
+ let batch = interleave_record_batch(&b, &indices)?;
Review Comment:
Probably an api like https://github.com/apache/arrow-rs/pull/7325 would be faster (avoiding one level of "trivial" indexing).
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Use `interleave` to speed up hash repartitioning [datafusion]
Posted by "Dandandan (via GitHub)" <gi...@apache.org>.
Dandandan commented on PR #15768:
URL: https://github.com/apache/datafusion/pull/15768#issuecomment-2836318864
closing for now
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Use `interleave` in hash repartitioning [datafusion]
Posted by "comphead (via GitHub)" <gi...@apache.org>.
comphead commented on code in PR #15768:
URL: https://github.com/apache/datafusion/pull/15768#discussion_r2053072665
##########
datafusion/physical-plan/src/repartition/mod.rs:
##########
@@ -260,33 +260,43 @@ impl BatchPartitioner {
} => {
let idx = *next_idx;
*next_idx = (*next_idx + 1) % *num_partitions;
- Box::new(std::iter::once(Ok((idx, batch))))
+ assert_eq!(batches.len(), 1);
Review Comment:
should we have
```
mut batches: Vec<RecordBatch>
```
if ```assert_eq!(batches.len(), 1)```
? 🤔
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Use `interleave` to speed up hash repartitioning [datafusion]
Posted by "Dandandan (via GitHub)" <gi...@apache.org>.
Dandandan commented on PR #15768:
URL: https://github.com/apache/datafusion/pull/15768#issuecomment-2836308324
Hmmm... I think the speedup actually doesn't come from using `interleave` but from yielding less often (once per partition_num * partition_num batches, rather than once per partition_num). Changing this in the original also makes this faster (thus `interleave` is actually still slower than `take` + `concat` ATM.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Use `interleave` in hash repartitioning [datafusion]
Posted by "Dandandan (via GitHub)" <gi...@apache.org>.
Dandandan commented on code in PR #15768:
URL: https://github.com/apache/datafusion/pull/15768#discussion_r2051497967
##########
datafusion/physical-plan/src/repartition/mod.rs:
##########
@@ -233,11 +233,11 @@ impl BatchPartitioner {
///
/// The time spent repartitioning, not including time spent in `f` will be recorded
/// to the [`metrics::Time`] provided on construction
- pub fn partition<F>(&mut self, batch: RecordBatch, mut f: F) -> Result<()>
+ pub fn partition<F>(&mut self, batches: Vec<RecordBatch>, mut f: F) -> Result<()>
Review Comment:
Good one
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Use `interleave` in hash repartitioning [datafusion]
Posted by "Dandandan (via GitHub)" <gi...@apache.org>.
Dandandan commented on code in PR #15768:
URL: https://github.com/apache/datafusion/pull/15768#discussion_r2051763714
##########
datafusion/physical-plan/src/repartition/mod.rs:
##########
@@ -298,25 +299,15 @@ impl BatchPartitioner {
.into_iter()
.enumerate()
.filter_map(|(partition, indices)| {
- let indices: PrimitiveArray<UInt32Type> = indices.into();
(!indices.is_empty()).then_some((partition, indices))
})
.map(move |(partition, indices)| {
// Tracking time required for repartitioned batches construction
let _timer = partitioner_timer.timer();
+ let b: Vec<&RecordBatch> = batches.iter().collect();
// Produce batches based on indices
- let columns = take_arrays(batch.columns(), &indices, None)?;
-
- let mut options = RecordBatchOptions::new();
- options = options.with_row_count(Some(indices.len()));
- let batch = RecordBatch::try_new_with_options(
- batch.schema(),
- columns,
- &options,
- )
- .unwrap();
-
+ let batch = interleave_record_batch(&b, &indices)?;
Review Comment:
Yeah the speedup comes from avoiding copying the data a second time in `concat` / CoalesceBatches. So when using `take_in` we should be careful to use it once (for a single destination batch) to avoid doing the concat on the small batches
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Use `interleave` in hash repartitioning [datafusion]
Posted by "Dandandan (via GitHub)" <gi...@apache.org>.
Dandandan commented on PR #15768:
URL: https://github.com/apache/datafusion/pull/15768#issuecomment-2818679075
FYI @Dandandan would be nice to have some benchmarks to be run. This improves on some long standing issue in DataFusion.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Use `interleave` to speed up hash repartitioning [datafusion]
Posted by "comphead (via GitHub)" <gi...@apache.org>.
comphead commented on code in PR #15768:
URL: https://github.com/apache/datafusion/pull/15768#discussion_r2063914703
##########
datafusion/physical-plan/src/repartition/mod.rs:
##########
@@ -260,33 +260,43 @@ impl BatchPartitioner {
} => {
let idx = *next_idx;
*next_idx = (*next_idx + 1) % *num_partitions;
- Box::new(std::iter::once(Ok((idx, batch))))
+ assert_eq!(batches.len(), 1);
Review Comment:
Thanks @Dandandan what I mean do we really need a vector of batches as input if there is only 1 batch always ?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] User `interleave` in hash repartitioning [datafusion]
Posted by "Dandandan (via GitHub)" <gi...@apache.org>.
Dandandan commented on code in PR #15768:
URL: https://github.com/apache/datafusion/pull/15768#discussion_r2051129739
##########
datafusion/physical-plan/src/repartition/mod.rs:
##########
@@ -298,25 +299,15 @@ impl BatchPartitioner {
.into_iter()
.enumerate()
.filter_map(|(partition, indices)| {
- let indices: PrimitiveArray<UInt32Type> = indices.into();
(!indices.is_empty()).then_some((partition, indices))
})
.map(move |(partition, indices)| {
// Tracking time required for repartitioned batches construction
let _timer = partitioner_timer.timer();
+ let b: Vec<&RecordBatch> = batches.iter().collect();
// Produce batches based on indices
- let columns = take_arrays(batch.columns(), &indices, None)?;
-
- let mut options = RecordBatchOptions::new();
- options = options.with_row_count(Some(indices.len()));
- let batch = RecordBatch::try_new_with_options(
- batch.schema(),
- columns,
- &options,
- )
- .unwrap();
-
+ let batch = interleave_record_batch(&b, &indices)?;
Review Comment:
FYI @ctsk
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Use `interleave` in hash repartitioning [datafusion]
Posted by "ctsk (via GitHub)" <gi...@apache.org>.
ctsk commented on PR #15768:
URL: https://github.com/apache/datafusion/pull/15768#issuecomment-2817240670
Does the `batches_until_yield` threshold need to be adjusted?
https://github.com/apache/datafusion/blob/91870b626b184a901c713fada3a1adb0eda2de7c/datafusion/physical-plan/src/repartition/mod.rs#L921-L926
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Use `interleave` to speed up hash repartitioning [datafusion]
Posted by "Dandandan (via GitHub)" <gi...@apache.org>.
Dandandan closed pull request #15768: Use `interleave` to speed up hash repartitioning
URL: https://github.com/apache/datafusion/pull/15768
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Use `interleave` to speed up hash repartitioning [datafusion]
Posted by "Dandandan (via GitHub)" <gi...@apache.org>.
Dandandan closed pull request #15768: Use `interleave` to speed up hash repartitioning
URL: https://github.com/apache/datafusion/pull/15768
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] User `interleave` in hash repartitioning [datafusion]
Posted by "Dandandan (via GitHub)" <gi...@apache.org>.
Dandandan commented on code in PR #15768:
URL: https://github.com/apache/datafusion/pull/15768#discussion_r2051062202
##########
datafusion/physical-plan/src/repartition/mod.rs:
##########
@@ -298,25 +299,15 @@ impl BatchPartitioner {
.into_iter()
.enumerate()
.filter_map(|(partition, indices)| {
- let indices: PrimitiveArray<UInt32Type> = indices.into();
(!indices.is_empty()).then_some((partition, indices))
})
.map(move |(partition, indices)| {
// Tracking time required for repartitioned batches construction
let _timer = partitioner_timer.timer();
+ let b: Vec<&RecordBatch> = batches.iter().collect();
// Produce batches based on indices
- let columns = take_arrays(batch.columns(), &indices, None)?;
-
- let mut options = RecordBatchOptions::new();
- options = options.with_row_count(Some(indices.len()));
- let batch = RecordBatch::try_new_with_options(
- batch.schema(),
- columns,
- &options,
- )
- .unwrap();
-
+ let batch = interleave_record_batch(&b, &indices)?;
Review Comment:
Probably an api like https://github.com/apache/arrow-rs/pull/7325 would be even faster (avoiding one level of "trivial" indexing).
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Use `interleave` in hash repartitioning [datafusion]
Posted by "Dandandan (via GitHub)" <gi...@apache.org>.
Dandandan commented on code in PR #15768:
URL: https://github.com/apache/datafusion/pull/15768#discussion_r2051731884
##########
datafusion/physical-plan/src/repartition/mod.rs:
##########
@@ -298,25 +299,15 @@ impl BatchPartitioner {
.into_iter()
.enumerate()
.filter_map(|(partition, indices)| {
- let indices: PrimitiveArray<UInt32Type> = indices.into();
(!indices.is_empty()).then_some((partition, indices))
})
.map(move |(partition, indices)| {
// Tracking time required for repartitioned batches construction
let _timer = partitioner_timer.timer();
+ let b: Vec<&RecordBatch> = batches.iter().collect();
// Produce batches based on indices
- let columns = take_arrays(batch.columns(), &indices, None)?;
-
- let mut options = RecordBatchOptions::new();
- options = options.with_row_count(Some(indices.len()));
- let batch = RecordBatch::try_new_with_options(
- batch.schema(),
- columns,
- &options,
- )
- .unwrap();
-
+ let batch = interleave_record_batch(&b, &indices)?;
Review Comment:
I think removing coalesce after this change (for all hash repartitions) might be possible, as the output batch size will be roughly equal to input batch size (instead of roughly 1/partitions * batch_size). Unless hash values are somehow skewed (but this is currently also not good anyway).
A future api could use your `take_in` api maybe to only output rows once batch size has been reached.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Use `interleave` in hash repartitioning [datafusion]
Posted by "comphead (via GitHub)" <gi...@apache.org>.
comphead commented on code in PR #15768:
URL: https://github.com/apache/datafusion/pull/15768#discussion_r2053072665
##########
datafusion/physical-plan/src/repartition/mod.rs:
##########
@@ -260,33 +260,43 @@ impl BatchPartitioner {
} => {
let idx = *next_idx;
*next_idx = (*next_idx + 1) % *num_partitions;
- Box::new(std::iter::once(Ok((idx, batch))))
+ assert_eq!(batches.len(), 1);
Review Comment:
should we have
```
mut batches: Vec<RecordBatch>
```
if ```assert_eq!(batches.len(), 1)```
?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Use `interleave` in hash repartitioning [datafusion]
Posted by "xudong963 (via GitHub)" <gi...@apache.org>.
xudong963 commented on code in PR #15768:
URL: https://github.com/apache/datafusion/pull/15768#discussion_r2051494486
##########
datafusion/physical-plan/src/repartition/mod.rs:
##########
@@ -233,11 +233,11 @@ impl BatchPartitioner {
///
/// The time spent repartitioning, not including time spent in `f` will be recorded
/// to the [`metrics::Time`] provided on construction
- pub fn partition<F>(&mut self, batch: RecordBatch, mut f: F) -> Result<()>
+ pub fn partition<F>(&mut self, batches: Vec<RecordBatch>, mut f: F) -> Result<()>
Review Comment:
Since this is a *pub** API, it would be best if we could keep the old API, mark it as deprecated, and add a new API with the changed arguments
##########
datafusion/physical-plan/src/repartition/mod.rs:
##########
@@ -233,11 +233,11 @@ impl BatchPartitioner {
///
/// The time spent repartitioning, not including time spent in `f` will be recorded
/// to the [`metrics::Time`] provided on construction
- pub fn partition<F>(&mut self, batch: RecordBatch, mut f: F) -> Result<()>
+ pub fn partition<F>(&mut self, batches: Vec<RecordBatch>, mut f: F) -> Result<()>
Review Comment:
Since this is a **pub** API, it would be best if we could keep the old API, mark it as deprecated, and add a new API with the changed arguments
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Use `interleave` in hash repartitioning [datafusion]
Posted by "ctsk (via GitHub)" <gi...@apache.org>.
ctsk commented on code in PR #15768:
URL: https://github.com/apache/datafusion/pull/15768#discussion_r2051758338
##########
datafusion/physical-plan/src/repartition/mod.rs:
##########
@@ -298,25 +299,15 @@ impl BatchPartitioner {
.into_iter()
.enumerate()
.filter_map(|(partition, indices)| {
- let indices: PrimitiveArray<UInt32Type> = indices.into();
(!indices.is_empty()).then_some((partition, indices))
})
.map(move |(partition, indices)| {
// Tracking time required for repartitioned batches construction
let _timer = partitioner_timer.timer();
+ let b: Vec<&RecordBatch> = batches.iter().collect();
// Produce batches based on indices
- let columns = take_arrays(batch.columns(), &indices, None)?;
-
- let mut options = RecordBatchOptions::new();
- options = options.with_row_count(Some(indices.len()));
- let batch = RecordBatch::try_new_with_options(
- batch.schema(),
- columns,
- &options,
- )
- .unwrap();
-
+ let batch = interleave_record_batch(&b, &indices)?;
Review Comment:
I see I had misunderstood this PR. It makes a lot of sense to do this. As part of prototyping the integration of a take_in API in datafusion, I made a similar change - move the buffering *before* sending the small batches to their destination thread. I don't remember seeing as much speedup when I benchmarked that change independently - I guess using interleave instead of a take/concat combo (like I did back then) makes a significant difference. Awesome!
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Use `interleave` in hash repartitioning [datafusion]
Posted by "Copilot (via GitHub)" <gi...@apache.org>.
Copilot commented on code in PR #15768:
URL: https://github.com/apache/datafusion/pull/15768#discussion_r2051134690
##########
datafusion/physical-plan/src/repartition/mod.rs:
##########
@@ -260,33 +260,43 @@ impl BatchPartitioner {
} => {
let idx = *next_idx;
*next_idx = (*next_idx + 1) % *num_partitions;
- Box::new(std::iter::once(Ok((idx, batch))))
+ assert_eq!(batches.len(), 1);
+ Box::new(std::iter::once(Ok((idx, batches.swap_remove(0)))))
Review Comment:
In the RoundRobin branch, an assertion forces batches to contain exactly one RecordBatch. If multiple batches are ever passed, this assertion will panic. Consider either handling multiple batches or documenting the assumption to prevent future misuse.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Use `interleave` to speed up hash repartitioning [datafusion]
Posted by "Dandandan (via GitHub)" <gi...@apache.org>.
Dandandan commented on PR #15768:
URL: https://github.com/apache/datafusion/pull/15768#issuecomment-2835325709
@alamb could you run the benchmarks maybe?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Use `interleave` to speed up hash repartitioning [datafusion]
Posted by "Dandandan (via GitHub)" <gi...@apache.org>.
Dandandan commented on code in PR #15768:
URL: https://github.com/apache/datafusion/pull/15768#discussion_r2063781770
##########
datafusion/physical-plan/src/repartition/mod.rs:
##########
@@ -260,33 +260,43 @@ impl BatchPartitioner {
} => {
let idx = *next_idx;
*next_idx = (*next_idx + 1) % *num_partitions;
- Box::new(std::iter::once(Ok((idx, batch))))
+ assert_eq!(batches.len(), 1);
Review Comment:
For round robin, `batches.len` is always 1 (no buffering needed)/
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org
Re: [PR] Use `interleave` in hash repartitioning [datafusion]
Posted by "ctsk (via GitHub)" <gi...@apache.org>.
ctsk commented on code in PR #15768:
URL: https://github.com/apache/datafusion/pull/15768#discussion_r2051758338
##########
datafusion/physical-plan/src/repartition/mod.rs:
##########
@@ -298,25 +299,15 @@ impl BatchPartitioner {
.into_iter()
.enumerate()
.filter_map(|(partition, indices)| {
- let indices: PrimitiveArray<UInt32Type> = indices.into();
(!indices.is_empty()).then_some((partition, indices))
})
.map(move |(partition, indices)| {
// Tracking time required for repartitioned batches construction
let _timer = partitioner_timer.timer();
+ let b: Vec<&RecordBatch> = batches.iter().collect();
// Produce batches based on indices
- let columns = take_arrays(batch.columns(), &indices, None)?;
-
- let mut options = RecordBatchOptions::new();
- options = options.with_row_count(Some(indices.len()));
- let batch = RecordBatch::try_new_with_options(
- batch.schema(),
- columns,
- &options,
- )
- .unwrap();
-
+ let batch = interleave_record_batch(&b, &indices)?;
Review Comment:
I see I had misunderstood this PR. It makes a lot of sense to do this. As part of prototyping the integration of a take_in in API in datafusion, I made a similar change - move the buffering *before* sending the small batches to their destination thread. I don't remeber see as much speedup when I benchmarked that change independently - I guess using interleave instead of a take/concat (like I did back then) combo makes a significant difference. Awesome!
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: github-unsubscribe@datafusion.apache.org
For additional commands, e-mail: github-help@datafusion.apache.org