You are viewing a plain text version of this content. The canonical link for it is here.
Posted to github@datafusion.apache.org by "wiedld (via GitHub)" <gi...@apache.org> on 2025/03/25 07:35:25 UTC
[PR] Enable repartitioning on MemTable. [datafusion]
wiedld opened a new pull request, #15409:
URL: https://github.com/apache/datafusion/pull/15409
**Draft PR.**
Still have 1 test to fix, and then I'll run benchmarks.
## Which issue does this PR close?
Closes #15088
## Rationale for this change
As noted in the issue, the MemTable->Sort does not have parallelized sorts. The reason is because the `MemorySourceConfig` data source does not implement [repartitioning](https://docs.rs/datafusion/latest/datafusion/datasource/source/trait.DataSource.html#method.repartitioned) (instead uses the default implementation which does not repartition).
## What changes are included in this PR?
* Reproducer for the [issue here](https://github.com/apache/datafusion/issues/15088#issue-2904662974).
* Add repartitioning to `MemorySourceConfig`, along with unit tests.
* Demonstrate the fix to the reproducer.
## Are these changes tested?
Yes.
## Are there any user-facing changes?
Only in possible EXPLAIN plans seen.
--
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] Enable repartitioning on MemTable. [datafusion]
Posted by "wiedld (via GitHub)" <gi...@apache.org>.
wiedld commented on code in PR #15409:
URL: https://github.com/apache/datafusion/pull/15409#discussion_r2053131713
##########
datafusion/datasource/src/memory.rs:
##########
@@ -440,6 +443,35 @@ impl DataSource for MemorySourceConfig {
}
}
+ fn repartitioned(
Review Comment:
I added a bunch of docs to clarify how the bits fits together, as well as what None means. https://github.com/apache/datafusion/pull/15409/commits/d8929dc9d179e89868964f5f33dd2e387e54c299
--
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] Enable repartitioning on MemTable. [datafusion]
Posted by "wiedld (via GitHub)" <gi...@apache.org>.
wiedld commented on PR #15409:
URL: https://github.com/apache/datafusion/pull/15409#issuecomment-2777243569
Thanks for the review.
Haven't had time to do the updates.
Converting to draft, and will mark ready again after updates.
--
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] Enable repartitioning on MemTable. [datafusion]
Posted by "wiedld (via GitHub)" <gi...@apache.org>.
wiedld commented on code in PR #15409:
URL: https://github.com/apache/datafusion/pull/15409#discussion_r2011766884
##########
datafusion/core/tests/physical_optimizer/enforce_distribution.rs:
##########
@@ -3461,3 +3467,102 @@ fn optimize_away_unnecessary_repartition2() -> Result<()> {
Ok(())
}
+
+#[tokio::test]
+async fn test_distribute_sort_parquet() -> Result<()> {
+ let test_config: TestConfig =
+ TestConfig::default().with_prefer_repartition_file_scans(1000);
+ assert!(
+ test_config.config.optimizer.repartition_file_scans,
+ "should enable scans to be repartitioned"
+ );
+
+ let schema = schema();
+ let sort_key = LexOrdering::new(vec![PhysicalSortExpr {
+ expr: col("c", &schema).unwrap(),
+ options: SortOptions::default(),
+ }]);
+ let physical_plan = sort_exec(sort_key, parquet_exec_with_stats(10000 * 8192), false);
+
+ // prior to optimization, this is the starting plan
+ let starting = &[
+ "SortExec: expr=[c@2 ASC], preserve_partitioning=[false]",
+ " DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], file_type=parquet",
+ ];
+ plans_matches_expected!(starting, physical_plan.clone());
+
+ // what the enforce distribution run does.
+ let expected = &[
+ "SortExec: expr=[c@2 ASC], preserve_partitioning=[false]",
+ " CoalescePartitionsExec",
+ " DataSourceExec: file_groups={10 groups: [[x:0..8192000], [x:8192000..16384000], [x:16384000..24576000], [x:24576000..32768000], [x:32768000..40960000], [x:40960000..49152000], [x:49152000..57344000], [x:57344000..65536000], [x:65536000..73728000], [x:73728000..81920000]]}, projection=[a, b, c, d, e], file_type=parquet",
+ ];
+ test_config.run(expected, physical_plan.clone(), &[Run::Distribution])?;
+
+ // what the sort parallelization (in enforce sorting), does after the enforce distribution changes
+ let expected = &[
+ "SortPreservingMergeExec: [c@2 ASC]",
+ " SortExec: expr=[c@2 ASC], preserve_partitioning=[true]",
+ " DataSourceExec: file_groups={10 groups: [[x:0..8192000], [x:8192000..16384000], [x:16384000..24576000], [x:24576000..32768000], [x:32768000..40960000], [x:40960000..49152000], [x:49152000..57344000], [x:57344000..65536000], [x:65536000..73728000], [x:73728000..81920000]]}, projection=[a, b, c, d, e], file_type=parquet",
+ ];
+ test_config.run(expected, physical_plan, &[Run::Distribution, Run::Sorting])?;
+ Ok(())
+}
+
+#[tokio::test]
+async fn test_distribute_sort_memtable() -> Result<()> {
Review Comment:
Reproducer from the issue.
This compares with the test above (`test_distribute_sort_parquet`), which demonstrates how it (already) works in oarquet.
--
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] Enable repartitioning on MemTable. [datafusion]
Posted by "2010YOUY01 (via GitHub)" <gi...@apache.org>.
2010YOUY01 commented on code in PR #15409:
URL: https://github.com/apache/datafusion/pull/15409#discussion_r2018202146
##########
datafusion/datasource/src/memory.rs:
##########
@@ -440,6 +443,35 @@ impl DataSource for MemorySourceConfig {
}
}
+ fn repartitioned(
Review Comment:
I don't quite understand how `repartitioned()` should be implemented with `Some(...)` `output_ordering` (I think this is not documented yet, I'll update the comments afterward)
Apparently, within each partition all entries should follow the specified ordering, but is there any cross-partition requirements? For example, should all entries in partition 2 required to be ordered after those in partition 1?
##########
datafusion/datasource/src/memory.rs:
##########
@@ -902,4 +1130,319 @@ mod tests {
Ok(())
}
+
+ fn batch(row_size: usize) -> RecordBatch {
+ let a: ArrayRef = Arc::new(Int32Array::from(vec![1; row_size]));
+ let b: ArrayRef = Arc::new(StringArray::from_iter(vec![Some("foo"); row_size]));
+ let c: ArrayRef = Arc::new(Int64Array::from_iter(vec![1; row_size]));
+ RecordBatch::try_from_iter(vec![("a", a), ("b", b), ("c", c)]).unwrap()
+ }
+
+ fn schema() -> SchemaRef {
+ batch(1).schema()
+ }
+
+ fn memorysrcconfig_no_partitions(
+ sort_information: Vec<LexOrdering>,
+ ) -> Result<MemorySourceConfig> {
+ let partitions = vec![];
+ MemorySourceConfig::try_new(&partitions, schema(), None)?
+ .try_with_sort_information(sort_information)
+ }
+
+ fn memorysrcconfig_1_partition_1_batch(
+ sort_information: Vec<LexOrdering>,
+ ) -> Result<MemorySourceConfig> {
+ let partitions = vec![vec![batch(100)]];
+ MemorySourceConfig::try_new(&partitions, schema(), None)?
+ .try_with_sort_information(sort_information)
+ }
+
+ fn memorysrcconfig_3_partitions_1_batch_each(
+ sort_information: Vec<LexOrdering>,
+ ) -> Result<MemorySourceConfig> {
+ let partitions = vec![vec![batch(100)], vec![batch(100)], vec![batch(100)]];
+ MemorySourceConfig::try_new(&partitions, schema(), None)?
+ .try_with_sort_information(sort_information)
+ }
+
+ fn memorysrcconfig_3_partitions_with_2_batches_each(
+ sort_information: Vec<LexOrdering>,
+ ) -> Result<MemorySourceConfig> {
+ let partitions = vec![
+ vec![batch(100), batch(100)],
+ vec![batch(100), batch(100)],
+ vec![batch(100), batch(100)],
+ ];
+ MemorySourceConfig::try_new(&partitions, schema(), None)?
+ .try_with_sort_information(sort_information)
+ }
+
+ fn memorysrcconfig_1_partition_with_different_sized_batches(
+ sort_information: Vec<LexOrdering>,
+ ) -> Result<MemorySourceConfig> {
+ let partitions = vec![vec![batch(100_000), batch(10_000), batch(100), batch(1)]];
+ MemorySourceConfig::try_new(&partitions, schema(), None)?
+ .try_with_sort_information(sort_information)
+ }
+
+ fn memorysrcconfig_2_partition_with_different_sized_batches(
+ sort_information: Vec<LexOrdering>,
+ ) -> Result<MemorySourceConfig> {
+ let partitions = vec![
+ vec![batch(100_000), batch(10_000), batch(1_000)],
Review Comment:
I recommend to include a test with size 1 and size 0 batches for extreme edge cases.
##########
datafusion/datasource/src/memory.rs:
##########
@@ -440,6 +443,35 @@ impl DataSource for MemorySourceConfig {
}
}
+ fn repartitioned(
+ &self,
+ target_partitions: usize,
+ _repartition_file_min_size: usize,
+ output_ordering: Option<LexOrdering>,
+ ) -> Result<Option<Arc<dyn DataSource>>> {
+ if self.partitions.is_empty() || self.partitions.len() >= target_partitions
+ // if already have more partitions than desired, do not merge
+ {
+ return Ok(None);
+ }
+
+ let maybe_repartitioned = if let Some(output_ordering) = output_ordering {
Review Comment:
It would be great to briefly explain under what circumstance this repartition will fail (with returning `None`)
##########
datafusion/core/tests/fuzz_cases/aggregate_fuzz.rs:
##########
@@ -520,7 +520,9 @@ async fn group_by_string_test(
let expected = compute_counts(&input, column_name);
let schema = input[0].schema();
- let session_config = SessionConfig::new().with_batch_size(50);
+ let session_config = SessionConfig::new()
+ .with_batch_size(50)
+ .with_repartition_file_scans(false);
Review Comment:
Probably a better idea is to use a separate config to control whether to repartition `MemTable`, instead of using on config option to control all of csv/parquet/memtable
##########
datafusion/datasource/src/memory.rs:
##########
@@ -902,4 +1130,319 @@ mod tests {
Ok(())
}
+
+ fn batch(row_size: usize) -> RecordBatch {
+ let a: ArrayRef = Arc::new(Int32Array::from(vec![1; row_size]));
+ let b: ArrayRef = Arc::new(StringArray::from_iter(vec![Some("foo"); row_size]));
+ let c: ArrayRef = Arc::new(Int64Array::from_iter(vec![1; row_size]));
+ RecordBatch::try_from_iter(vec![("a", a), ("b", b), ("c", c)]).unwrap()
+ }
+
+ fn schema() -> SchemaRef {
+ batch(1).schema()
+ }
+
+ fn memorysrcconfig_no_partitions(
+ sort_information: Vec<LexOrdering>,
+ ) -> Result<MemorySourceConfig> {
+ let partitions = vec![];
+ MemorySourceConfig::try_new(&partitions, schema(), None)?
+ .try_with_sort_information(sort_information)
+ }
+
+ fn memorysrcconfig_1_partition_1_batch(
+ sort_information: Vec<LexOrdering>,
+ ) -> Result<MemorySourceConfig> {
+ let partitions = vec![vec![batch(100)]];
+ MemorySourceConfig::try_new(&partitions, schema(), None)?
+ .try_with_sort_information(sort_information)
+ }
+
+ fn memorysrcconfig_3_partitions_1_batch_each(
+ sort_information: Vec<LexOrdering>,
+ ) -> Result<MemorySourceConfig> {
+ let partitions = vec![vec![batch(100)], vec![batch(100)], vec![batch(100)]];
+ MemorySourceConfig::try_new(&partitions, schema(), None)?
+ .try_with_sort_information(sort_information)
+ }
+
+ fn memorysrcconfig_3_partitions_with_2_batches_each(
+ sort_information: Vec<LexOrdering>,
+ ) -> Result<MemorySourceConfig> {
+ let partitions = vec![
+ vec![batch(100), batch(100)],
+ vec![batch(100), batch(100)],
+ vec![batch(100), batch(100)],
+ ];
+ MemorySourceConfig::try_new(&partitions, schema(), None)?
+ .try_with_sort_information(sort_information)
+ }
+
+ fn memorysrcconfig_1_partition_with_different_sized_batches(
+ sort_information: Vec<LexOrdering>,
+ ) -> Result<MemorySourceConfig> {
+ let partitions = vec![vec![batch(100_000), batch(10_000), batch(100), batch(1)]];
+ MemorySourceConfig::try_new(&partitions, schema(), None)?
+ .try_with_sort_information(sort_information)
+ }
+
+ fn memorysrcconfig_2_partition_with_different_sized_batches(
+ sort_information: Vec<LexOrdering>,
+ ) -> Result<MemorySourceConfig> {
+ let partitions = vec![
+ vec![batch(100_000), batch(10_000), batch(1_000)],
Review Comment:
I recommend to include a test with size 1 and size 0 batches for extreme edge cases.
--
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] Enable repartitioning on MemTable. [datafusion]
Posted by "alamb (via GitHub)" <gi...@apache.org>.
alamb commented on PR #15409:
URL: https://github.com/apache/datafusion/pull/15409#issuecomment-2868765074
Thank you @wiedld and @2010YOUY01
--
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] Enable repartitioning on MemTable. [datafusion]
Posted by "2010YOUY01 (via GitHub)" <gi...@apache.org>.
2010YOUY01 commented on code in PR #15409:
URL: https://github.com/apache/datafusion/pull/15409#discussion_r2018210999
##########
datafusion/datasource/src/memory.rs:
##########
@@ -902,4 +1130,319 @@ mod tests {
Ok(())
}
+
+ fn batch(row_size: usize) -> RecordBatch {
+ let a: ArrayRef = Arc::new(Int32Array::from(vec![1; row_size]));
+ let b: ArrayRef = Arc::new(StringArray::from_iter(vec![Some("foo"); row_size]));
+ let c: ArrayRef = Arc::new(Int64Array::from_iter(vec![1; row_size]));
+ RecordBatch::try_from_iter(vec![("a", a), ("b", b), ("c", c)]).unwrap()
+ }
+
+ fn schema() -> SchemaRef {
+ batch(1).schema()
+ }
+
+ fn memorysrcconfig_no_partitions(
+ sort_information: Vec<LexOrdering>,
+ ) -> Result<MemorySourceConfig> {
+ let partitions = vec![];
+ MemorySourceConfig::try_new(&partitions, schema(), None)?
+ .try_with_sort_information(sort_information)
+ }
+
+ fn memorysrcconfig_1_partition_1_batch(
+ sort_information: Vec<LexOrdering>,
+ ) -> Result<MemorySourceConfig> {
+ let partitions = vec![vec![batch(100)]];
+ MemorySourceConfig::try_new(&partitions, schema(), None)?
+ .try_with_sort_information(sort_information)
+ }
+
+ fn memorysrcconfig_3_partitions_1_batch_each(
+ sort_information: Vec<LexOrdering>,
+ ) -> Result<MemorySourceConfig> {
+ let partitions = vec![vec![batch(100)], vec![batch(100)], vec![batch(100)]];
+ MemorySourceConfig::try_new(&partitions, schema(), None)?
+ .try_with_sort_information(sort_information)
+ }
+
+ fn memorysrcconfig_3_partitions_with_2_batches_each(
+ sort_information: Vec<LexOrdering>,
+ ) -> Result<MemorySourceConfig> {
+ let partitions = vec![
+ vec![batch(100), batch(100)],
+ vec![batch(100), batch(100)],
+ vec![batch(100), batch(100)],
+ ];
+ MemorySourceConfig::try_new(&partitions, schema(), None)?
+ .try_with_sort_information(sort_information)
+ }
+
+ fn memorysrcconfig_1_partition_with_different_sized_batches(
+ sort_information: Vec<LexOrdering>,
+ ) -> Result<MemorySourceConfig> {
+ let partitions = vec![vec![batch(100_000), batch(10_000), batch(100), batch(1)]];
+ MemorySourceConfig::try_new(&partitions, schema(), None)?
+ .try_with_sort_information(sort_information)
+ }
+
+ fn memorysrcconfig_2_partition_with_different_sized_batches(
+ sort_information: Vec<LexOrdering>,
+ ) -> Result<MemorySourceConfig> {
+ let partitions = vec![
+ vec![batch(100_000), batch(10_000), batch(1_000)],
Review Comment:
I recommend to include a test with size 1 and size 0 batches for extreme edge cases.
--
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] Enable repartitioning on MemTable. [datafusion]
Posted by "wiedld (via GitHub)" <gi...@apache.org>.
wiedld commented on code in PR #15409:
URL: https://github.com/apache/datafusion/pull/15409#discussion_r2011720446
##########
datafusion/core/tests/fuzz_cases/aggregate_fuzz.rs:
##########
@@ -520,7 +520,9 @@ async fn group_by_string_test(
let expected = compute_counts(&input, column_name);
let schema = input[0].schema();
- let session_config = SessionConfig::new().with_batch_size(50);
+ let session_config = SessionConfig::new()
+ .with_batch_size(50)
+ .with_repartition_file_scans(false);
Review Comment:
On the other hand, the fact that I had to disable mem table repartitioning (via config) to get some of the test cases to pass, makes me would if the should-repartition-files vs should-repartition-memtable decision should be made separately. 🤔
--
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] Enable repartitioning on MemTable. [datafusion]
Posted by "2010YOUY01 (via GitHub)" <gi...@apache.org>.
2010YOUY01 commented on code in PR #15409:
URL: https://github.com/apache/datafusion/pull/15409#discussion_r2065594544
##########
datafusion/datasource/src/memory.rs:
##########
@@ -723,6 +761,222 @@ impl MemorySourceConfig {
pub fn original_schema(&self) -> SchemaRef {
Arc::clone(&self.schema)
}
+
+ /// Repartition while preserving order.
+ ///
+ /// Returns `Ok(None)` if cannot fulfill the requested repartitioning, such
+ /// as having too few batches to fulfill the `target_partitions` or if unable
+ /// to preserve output ordering.
+ fn repartition_preserving_order(
+ &self,
+ target_partitions: usize,
+ output_ordering: LexOrdering,
+ ) -> Result<Option<Vec<Vec<RecordBatch>>>> {
+ if !self.eq_properties().ordering_satisfy(&output_ordering) {
+ Ok(None)
+ } else if self.partitions.len() == 1 {
+ self.repartition_evenly_by_size(target_partitions)
+ } else {
Review Comment:
It would be great to include a simple example to explain this code block, like how is an input table be repartitioned, what is the key in the max heap, and how the state changes during each loop step.
--
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] Enable repartitioning on MemTable. [datafusion]
Posted by "wiedld (via GitHub)" <gi...@apache.org>.
wiedld commented on PR #15409:
URL: https://github.com/apache/datafusion/pull/15409#issuecomment-2858928895
@2010YOUY01 -- comments addressed, and a conditional removed due to this reason: https://github.com/apache/datafusion/pull/15409#discussion_r2076020887
--
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] Enable repartitioning on MemTable. [datafusion]
Posted by "wiedld (via GitHub)" <gi...@apache.org>.
wiedld commented on code in PR #15409:
URL: https://github.com/apache/datafusion/pull/15409#discussion_r2053131373
##########
datafusion/core/tests/fuzz_cases/aggregate_fuzz.rs:
##########
@@ -520,7 +520,9 @@ async fn group_by_string_test(
let expected = compute_counts(&input, column_name);
let schema = input[0].schema();
- let session_config = SessionConfig::new().with_batch_size(50);
+ let session_config = SessionConfig::new()
+ .with_batch_size(50)
+ .with_repartition_file_scans(false);
Review Comment:
When updating the docs (see this commit: https://github.com/apache/datafusion/pull/15409/commits/d8929dc9d179e89868964f5f33dd2e387e54c299), it felt like we don't want to split out the 2 configurations IMO. Since the repartitioning config variable is used to both (a) repartition at the datasrc (which could be a file or memtable), as well as (b) insert a repartition operator later in the plan.
Please take a look at the added and docs, the lmk what you think.
--
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] Enable repartitioning on MemTable. [datafusion]
Posted by "wiedld (via GitHub)" <gi...@apache.org>.
wiedld commented on PR #15409:
URL: https://github.com/apache/datafusion/pull/15409#issuecomment-2750626791
TODO: I'll run benchmarks on this tmrw.
--
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] Enable repartitioning on MemTable. [datafusion]
Posted by "wiedld (via GitHub)" <gi...@apache.org>.
wiedld commented on code in PR #15409:
URL: https://github.com/apache/datafusion/pull/15409#discussion_r2011495377
##########
datafusion/datasource/src/memory.rs:
##########
@@ -902,4 +1108,319 @@ mod tests {
Ok(())
}
+
+ fn batch(row_size: usize) -> RecordBatch {
+ let a: ArrayRef = Arc::new(Int32Array::from(vec![1; row_size]));
+ let b: ArrayRef = Arc::new(StringArray::from_iter(vec![Some("foo"); row_size]));
+ let c: ArrayRef = Arc::new(Int64Array::from_iter(vec![1; row_size]));
+ RecordBatch::try_from_iter(vec![("a", a), ("b", b), ("c", c)]).unwrap()
+ }
+
+ fn schema() -> SchemaRef {
+ batch(1).schema()
+ }
+
+ fn memorysrcconfig_no_partitions(
+ sort_information: Vec<LexOrdering>,
+ ) -> Result<MemorySourceConfig> {
+ let partitions = vec![];
+ MemorySourceConfig::try_new(&partitions, schema(), None)?
+ .try_with_sort_information(sort_information)
+ }
+
+ fn memorysrcconfig_1_partition_1_batch(
+ sort_information: Vec<LexOrdering>,
+ ) -> Result<MemorySourceConfig> {
+ let partitions = vec![vec![batch(100)]];
+ MemorySourceConfig::try_new(&partitions, schema(), None)?
+ .try_with_sort_information(sort_information)
+ }
+
+ fn memorysrcconfig_3_partitions_1_batch_each(
+ sort_information: Vec<LexOrdering>,
+ ) -> Result<MemorySourceConfig> {
+ let partitions = vec![vec![batch(100)], vec![batch(100)], vec![batch(100)]];
+ MemorySourceConfig::try_new(&partitions, schema(), None)?
+ .try_with_sort_information(sort_information)
+ }
+
+ fn memorysrcconfig_3_partitions_with_2_batches_each(
+ sort_information: Vec<LexOrdering>,
+ ) -> Result<MemorySourceConfig> {
+ let partitions = vec![
+ vec![batch(100), batch(100)],
+ vec![batch(100), batch(100)],
+ vec![batch(100), batch(100)],
+ ];
+ MemorySourceConfig::try_new(&partitions, schema(), None)?
+ .try_with_sort_information(sort_information)
+ }
+
+ fn memorysrcconfig_1_partition_with_different_sized_batches(
+ sort_information: Vec<LexOrdering>,
+ ) -> Result<MemorySourceConfig> {
+ let partitions = vec![vec![batch(100_000), batch(10_000), batch(100), batch(1)]];
+ MemorySourceConfig::try_new(&partitions, schema(), None)?
+ .try_with_sort_information(sort_information)
+ }
+
+ fn memorysrcconfig_2_partition_with_different_sized_batches(
+ sort_information: Vec<LexOrdering>,
+ ) -> Result<MemorySourceConfig> {
+ let partitions = vec![
+ vec![batch(100_000), batch(10_000), batch(1_000)],
+ vec![batch(2_000), batch(20)],
+ ];
+ MemorySourceConfig::try_new(&partitions, schema(), None)?
+ .try_with_sort_information(sort_information)
+ }
+
+ /// Assert that we get the expected count of partitions after repartitioning.
+ ///
+ /// If None, then we expected the [`DataSource::repartitioned`] to return None.
+ fn assert_partitioning(
+ partitioned_datasrc: Option<Arc<dyn DataSource>>,
+ partition_cnt: Option<usize>,
+ ) {
+ let should_exist = if let Some(partition_cnt) = partition_cnt {
+ format!(
+ "new datasource should exist and have {:?} partitions",
+ partition_cnt
+ )
+ } else {
+ "new datasource should not exist".into()
+ };
+
+ let actual = partitioned_datasrc
+ .map(|datasrc| datasrc.output_partitioning().partition_count());
+ assert_eq!(
+ actual,
+ partition_cnt,
+ "partitioned datasrc does not match expected, we expected {}, instead found {:?}",
+ should_exist,
+ actual
+ );
+ }
+
+ fn run_all_test_scenarios(
+ output_ordering: Option<LexOrdering>,
+ sort_information_on_config: Vec<LexOrdering>,
+ ) -> Result<()> {
+ let not_used = usize::MAX;
+
+ // src has no partitions
+ let mem_src_config =
+ memorysrcconfig_no_partitions(sort_information_on_config.clone())?;
+ let partitioned_datasrc =
+ mem_src_config.repartitioned(1, not_used, output_ordering.clone())?;
+ assert_partitioning(partitioned_datasrc, None);
+
+ // src has partitions == target partitions (=1)
+ let target_partitions = 1;
+ let mem_src_config =
+ memorysrcconfig_1_partition_1_batch(sort_information_on_config.clone())?;
+ let partitioned_datasrc = mem_src_config.repartitioned(
+ target_partitions,
+ not_used,
+ output_ordering.clone(),
+ )?;
+ assert_partitioning(partitioned_datasrc, None);
+
+ // src has partitions == target partitions (=3)
+ let target_partitions = 3;
+ let mem_src_config = memorysrcconfig_3_partitions_1_batch_each(
+ sort_information_on_config.clone(),
+ )?;
+ let partitioned_datasrc = mem_src_config.repartitioned(
+ target_partitions,
+ not_used,
+ output_ordering.clone(),
+ )?;
+ assert_partitioning(partitioned_datasrc, None);
+
+ // src has partitions > target partitions, but we don't merge them
+ let target_partitions = 2;
+ let mem_src_config = memorysrcconfig_3_partitions_1_batch_each(
+ sort_information_on_config.clone(),
+ )?;
+ let partitioned_datasrc = mem_src_config.repartitioned(
+ target_partitions,
+ not_used,
+ output_ordering.clone(),
+ )?;
+ assert_partitioning(partitioned_datasrc, None);
+
+ // src has partitions < target partitions, but not enough batches (per partition) to split into more partitions
+ let target_partitions = 4;
+ let mem_src_config = memorysrcconfig_3_partitions_1_batch_each(
+ sort_information_on_config.clone(),
+ )?;
+ let partitioned_datasrc = mem_src_config.repartitioned(
+ target_partitions,
+ not_used,
+ output_ordering.clone(),
+ )?;
+ assert_partitioning(partitioned_datasrc, None);
+
+ // src has partitions < target partitions, and can split to sufficient amount
+ // has 6 batches across 3 partitions. Will need to split 2 of it's partitions.
+ let target_partitions = 5;
+ let mem_src_config = memorysrcconfig_3_partitions_with_2_batches_each(
+ sort_information_on_config.clone(),
+ )?;
+ let partitioned_datasrc = mem_src_config.repartitioned(
+ target_partitions,
+ not_used,
+ output_ordering.clone(),
+ )?;
+ assert_partitioning(partitioned_datasrc, Some(5));
+
+ // src has partitions < target partitions, and can split to sufficient amount
+ // has 6 batches across 3 partitions. Will need to split all of it's partitions.
+ let target_partitions = 6;
+ let mem_src_config = memorysrcconfig_3_partitions_with_2_batches_each(
+ sort_information_on_config.clone(),
+ )?;
+ let partitioned_datasrc = mem_src_config.repartitioned(
+ target_partitions,
+ not_used,
+ output_ordering.clone(),
+ )?;
+ assert_partitioning(partitioned_datasrc, Some(6));
+
+ // src has partitions < target partitions, but not enough total batches to fulfill the split (desired target_partitions)
+ let target_partitions = 3 * 2 + 1;
+ let mem_src_config = memorysrcconfig_3_partitions_with_2_batches_each(
+ sort_information_on_config.clone(),
+ )?;
+ let partitioned_datasrc = mem_src_config.repartitioned(
+ target_partitions,
+ not_used,
+ output_ordering.clone(),
+ )?;
+ assert_partitioning(partitioned_datasrc, None);
+
+ // src has 1 partition with many batches of lopsided sizes
+ // make sure it handles the split properly
+ let target_partitions = 2;
+ let mem_src_config = memorysrcconfig_1_partition_with_different_sized_batches(
+ sort_information_on_config,
+ )?;
+ let partitioned_datasrc = mem_src_config.clone().repartitioned(
+ target_partitions,
+ not_used,
+ output_ordering,
+ )?;
+ assert_partitioning(partitioned_datasrc.clone(), Some(2));
+ // Starting = batch(100_000), batch(10_000), batch(100), batch(1).
+ // It should have split as p1=batch(100_000), p2=[batch(10_000), batch(100), batch(1)]
+ let repartitioned_raw_batches = mem_src_config
+ .repartition_evenly_by_size(target_partitions)?
+ .unwrap();
+ assert_eq!(repartitioned_raw_batches.len(), 2);
+ let [ref p1, ref p2] = repartitioned_raw_batches[..] else {
+ unreachable!()
+ };
+ // p1=batch(100_000)
+ assert_eq!(p1.len(), 1);
+ assert_eq!(p1[0].num_rows(), 100_000);
+ // p2=[batch(10_000), batch(100), batch(1)]
+ assert_eq!(p2.len(), 3);
+ assert_eq!(p2[0].num_rows(), 10_000);
+ assert_eq!(p2[1].num_rows(), 100);
+ assert_eq!(p2[2].num_rows(), 1);
+
+ Ok(())
+ }
+
+ #[test]
+ fn test_repartition_no_sort_information_no_output_ordering() -> Result<()> {
+ let no_sort = vec![];
+ let no_output_ordering = None;
+
+ // Test: common set of functionality
+ run_all_test_scenarios(no_output_ordering.clone(), no_sort.clone())?;
+
+ // Test: does not preserve separate partitions (with own internal ordering) on even split
Review Comment:
This defines what I believe the different behavior should be when not preserving sort ordering. I looked at how [FileScanConfig::repartition](https://github.com/apache/datafusion/blob/908ca0e684f3184075ef799128ee3f58f1e9db86/datafusion/datasource/src/file_scan_config.rs#L237) was implemented as a guide. The key branch point in it's implementation is here:
https://github.com/apache/datafusion/blob/908ca0e684f3184075ef799128ee3f58f1e9db86/datafusion/datasource/src/file_groups.rs#L194-L199
Please feel free to redirect. 🙏🏼
--
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] Enable repartitioning on MemTable. [datafusion]
Posted by "wiedld (via GitHub)" <gi...@apache.org>.
wiedld commented on code in PR #15409:
URL: https://github.com/apache/datafusion/pull/15409#discussion_r2076020887
##########
datafusion/datasource/src/memory.rs:
##########
@@ -723,6 +761,222 @@ impl MemorySourceConfig {
pub fn original_schema(&self) -> SchemaRef {
Arc::clone(&self.schema)
}
+
+ /// Repartition while preserving order.
+ ///
+ /// Returns `Ok(None)` if cannot fulfill the requested repartitioning, such
+ /// as having too few batches to fulfill the `target_partitions` or if unable
+ /// to preserve output ordering.
+ fn repartition_preserving_order(
+ &self,
+ target_partitions: usize,
+ output_ordering: LexOrdering,
+ ) -> Result<Option<Vec<Vec<RecordBatch>>>> {
+ if !self.eq_properties().ordering_satisfy(&output_ordering) {
+ Ok(None)
+ } else if self.partitions.len() == 1 {
+ self.repartition_evenly_by_size(target_partitions)
+ } else {
Review Comment:
Thank you for the code review.
After adding an additional test case targeting the `self.partitions.len() == 1`, and ended up removing this conditional due to this: https://github.com/apache/datafusion/pull/15409/commits/57147add2bb1d540334b08639504b80930c1c982
Additionally, I added code comments about how the max heap is being used: https://github.com/apache/datafusion/pull/15409/commits/12e6dfb1dec6038b47a63019ee6965ed626bc68e
I believe the `RePartition` struct already documents what prop is the max heap cmp, vs ordering.
Please let me know if you spot anything else. Thank you.
--
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] Enable repartitioning on MemTable. [datafusion]
Posted by "alamb (via GitHub)" <gi...@apache.org>.
alamb commented on PR #15409:
URL: https://github.com/apache/datafusion/pull/15409#issuecomment-2759400185
@2010YOUY01 and @alan910127 -- do you have time to review this PR as you filed / expressed interested on https://github.com/apache/datafusion/issues/15088#top
--
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] Enable repartitioning on MemTable. [datafusion]
Posted by "wiedld (via GitHub)" <gi...@apache.org>.
wiedld commented on code in PR #15409:
URL: https://github.com/apache/datafusion/pull/15409#discussion_r2011766884
##########
datafusion/core/tests/physical_optimizer/enforce_distribution.rs:
##########
@@ -3461,3 +3467,102 @@ fn optimize_away_unnecessary_repartition2() -> Result<()> {
Ok(())
}
+
+#[tokio::test]
+async fn test_distribute_sort_parquet() -> Result<()> {
+ let test_config: TestConfig =
+ TestConfig::default().with_prefer_repartition_file_scans(1000);
+ assert!(
+ test_config.config.optimizer.repartition_file_scans,
+ "should enable scans to be repartitioned"
+ );
+
+ let schema = schema();
+ let sort_key = LexOrdering::new(vec![PhysicalSortExpr {
+ expr: col("c", &schema).unwrap(),
+ options: SortOptions::default(),
+ }]);
+ let physical_plan = sort_exec(sort_key, parquet_exec_with_stats(10000 * 8192), false);
+
+ // prior to optimization, this is the starting plan
+ let starting = &[
+ "SortExec: expr=[c@2 ASC], preserve_partitioning=[false]",
+ " DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], file_type=parquet",
+ ];
+ plans_matches_expected!(starting, physical_plan.clone());
+
+ // what the enforce distribution run does.
+ let expected = &[
+ "SortExec: expr=[c@2 ASC], preserve_partitioning=[false]",
+ " CoalescePartitionsExec",
+ " DataSourceExec: file_groups={10 groups: [[x:0..8192000], [x:8192000..16384000], [x:16384000..24576000], [x:24576000..32768000], [x:32768000..40960000], [x:40960000..49152000], [x:49152000..57344000], [x:57344000..65536000], [x:65536000..73728000], [x:73728000..81920000]]}, projection=[a, b, c, d, e], file_type=parquet",
+ ];
+ test_config.run(expected, physical_plan.clone(), &[Run::Distribution])?;
+
+ // what the sort parallelization (in enforce sorting), does after the enforce distribution changes
+ let expected = &[
+ "SortPreservingMergeExec: [c@2 ASC]",
+ " SortExec: expr=[c@2 ASC], preserve_partitioning=[true]",
+ " DataSourceExec: file_groups={10 groups: [[x:0..8192000], [x:8192000..16384000], [x:16384000..24576000], [x:24576000..32768000], [x:32768000..40960000], [x:40960000..49152000], [x:49152000..57344000], [x:57344000..65536000], [x:65536000..73728000], [x:73728000..81920000]]}, projection=[a, b, c, d, e], file_type=parquet",
+ ];
+ test_config.run(expected, physical_plan, &[Run::Distribution, Run::Sorting])?;
+ Ok(())
+}
+
+#[tokio::test]
+async fn test_distribute_sort_memtable() -> Result<()> {
Review Comment:
Reproducer from the issue.
This compares with the test above (`test_distribute_sort_parquet`), which demonstrates how it (already) works in parquet.
--
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] Enable repartitioning on MemTable. [datafusion]
Posted by "wiedld (via GitHub)" <gi...@apache.org>.
wiedld commented on PR #15409:
URL: https://github.com/apache/datafusion/pull/15409#issuecomment-2750536557
TODO: I'll run benchmarks on this tmrw.
--
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] Enable repartitioning on MemTable. [datafusion]
Posted by "wiedld (via GitHub)" <gi...@apache.org>.
wiedld commented on code in PR #15409:
URL: https://github.com/apache/datafusion/pull/15409#discussion_r2011599235
##########
datafusion/core/tests/memory_limit/mod.rs:
##########
@@ -455,7 +455,9 @@ async fn test_stringview_external_sort() {
.with_memory_pool(Arc::new(FairSpillPool::new(60 * 1024 * 1024)));
let runtime = builder.build_arc().unwrap();
- let config = SessionConfig::new().with_sort_spill_reservation_bytes(40 * 1024 * 1024);
+ let config = SessionConfig::new()
+ .with_sort_spill_reservation_bytes(40 * 1024 * 1024)
+ .with_repartition_file_scans(false);
Review Comment:
Update needed for test to pass.
--
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] Enable repartitioning on MemTable. [datafusion]
Posted by "wiedld (via GitHub)" <gi...@apache.org>.
wiedld commented on code in PR #15409:
URL: https://github.com/apache/datafusion/pull/15409#discussion_r2011688314
##########
datafusion/core/tests/fuzz_cases/aggregate_fuzz.rs:
##########
@@ -520,7 +520,9 @@ async fn group_by_string_test(
let expected = compute_counts(&input, column_name);
let schema = input[0].schema();
- let session_config = SessionConfig::new().with_batch_size(50);
+ let session_config = SessionConfig::new()
+ .with_batch_size(50)
+ .with_repartition_file_scans(false);
Review Comment:
Also note the existing configuration `.repartition_file_scans` that is [used here](https://github.com/apache/datafusion/blob/908ca0e684f3184075ef799128ee3f58f1e9db86/datafusion/physical-optimizer/src/enforce_distribution.rs#L1171) to determine repartitioning.
It's naming doesn't quite match a mem table (because not a file)? Should I be updating the naming (and the config docs as well) in a follow up PR?
--
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] Enable repartitioning on MemTable. [datafusion]
Posted by "wiedld (via GitHub)" <gi...@apache.org>.
wiedld commented on PR #15409:
URL: https://github.com/apache/datafusion/pull/15409#issuecomment-2820043803
Took me awhile to circle back to this PR. I believe it's now ready for re-review @2010YOUY01
--
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] Enable repartitioning on MemTable. [datafusion]
Posted by "wiedld (via GitHub)" <gi...@apache.org>.
wiedld commented on code in PR #15409:
URL: https://github.com/apache/datafusion/pull/15409#discussion_r2011691796
##########
datafusion/datasource/src/memory.rs:
##########
@@ -718,6 +750,181 @@ impl MemorySourceConfig {
pub fn original_schema(&self) -> SchemaRef {
Arc::clone(&self.schema)
}
+
+ /// Repartition while preserving order.
+ ///
+ /// Returns None if cannot fulfill requested repartitioning.
+ fn repartition_preserving_order(
+ &self,
+ target_partitions: usize,
+ output_ordering: LexOrdering,
+ ) -> Result<Option<Vec<Vec<RecordBatch>>>> {
+ if !self.eq_properties().ordering_satisfy(&output_ordering) {
+ Ok(None)
+ } else if self.partitions.len() == 1 {
+ self.repartition_evenly_by_size(target_partitions)
+ } else {
+ let total_num_batches =
+ self.partitions.iter().map(|b| b.len()).sum::<usize>();
+ if total_num_batches < target_partitions {
+ // no way to create the desired repartitioning
+ return Ok(None);
+ }
+
+ let cnt_to_repartition = target_partitions - self.partitions.len();
+
+ let to_repartition = self
+ .partitions
+ .iter()
+ .enumerate()
+ .map(|(idx, batches)| RePartition {
+ idx: idx + (cnt_to_repartition * idx), // make space in ordering for split partitions
+ row_count: batches.iter().map(|batch| batch.num_rows()).sum(),
+ batches: batches.clone(),
+ })
+ .collect_vec();
+
+ // split the largest partitions
+ let mut max_heap = BinaryHeap::with_capacity(target_partitions);
+ for rep in to_repartition {
+ max_heap.push(rep);
+ }
+ for _ in 0..cnt_to_repartition {
+ let Some(to_split) = max_heap.pop() else {
+ unreachable!()
+ };
+ for new_partition in to_split.split() {
+ max_heap.push(new_partition);
+ }
+ }
+ let mut partitions = max_heap.drain().collect_vec();
+ partitions.sort_by_key(|p| p.idx);
+ let partitions = partitions.into_iter().map(|rep| rep.batches).collect_vec();
+
+ Ok(Some(partitions))
+ }
+ }
+
+ /// Repartition into evenly sized chunks (as much as possible without batch splitting),
+ /// disregarding any ordering.
+ ///
+ /// Returns None if cannot fulfill requested repartitioning.
+ fn repartition_evenly_by_size(
+ &self,
+ target_partitions: usize,
+ ) -> Result<Option<Vec<Vec<RecordBatch>>>> {
+ // determine if we have enough total batches to fulfill request
+ let flatten_batches = self.partitions.clone().into_iter().flatten().collect_vec();
+ if flatten_batches.len() < target_partitions {
+ return Ok(None);
+ }
+
+ // repartition evenly
+ let total_num_rows = flatten_batches.iter().map(|b| b.num_rows()).sum::<usize>();
+ let target_partition_size = total_num_rows.div_ceil(target_partitions);
+ let mut partitions =
+ vec![Vec::with_capacity(flatten_batches.len()); target_partitions];
+ let mut curr_row_count = 0;
+ let mut next_idx = 0;
+ for batch in flatten_batches {
+ let row_cnt = batch.num_rows();
+
+ // handle very lopsided batch sizing
+ if partitions[next_idx].is_empty() {
+ partitions[next_idx].push(batch);
+ } else {
+ // have at least 1 batch per partition
+ let idx =
+ std::cmp::min(next_idx + 1, curr_row_count / target_partition_size);
+ if let Some(partition) = partitions.get_mut(idx) {
+ partition.push(batch);
+ } else {
+ partitions[target_partitions - 1].push(batch);
+ }
+ next_idx = idx;
+ }
Review Comment:
Note to self; this feels ugly and I should fix it 😅 .
--
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] Enable repartitioning on MemTable. [datafusion]
Posted by "wiedld (via GitHub)" <gi...@apache.org>.
wiedld commented on code in PR #15409:
URL: https://github.com/apache/datafusion/pull/15409#discussion_r2011688314
##########
datafusion/core/tests/fuzz_cases/aggregate_fuzz.rs:
##########
@@ -520,7 +520,9 @@ async fn group_by_string_test(
let expected = compute_counts(&input, column_name);
let schema = input[0].schema();
- let session_config = SessionConfig::new().with_batch_size(50);
+ let session_config = SessionConfig::new()
+ .with_batch_size(50)
+ .with_repartition_file_scans(false);
Review Comment:
Also note the existing configuration `repartition_file_scans` that is [used here](https://github.com/apache/datafusion/blob/908ca0e684f3184075ef799128ee3f58f1e9db86/datafusion/physical-optimizer/src/enforce_distribution.rs#L1171) to determine repartitioning.
It's naming doesn't quite match a mem table (because not a file)? Should I be updating the naming (and the config docs as well) in a follow up PR?
--
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] Enable repartitioning on MemTable. [datafusion]
Posted by "wiedld (via GitHub)" <gi...@apache.org>.
wiedld commented on code in PR #15409:
URL: https://github.com/apache/datafusion/pull/15409#discussion_r2076020887
##########
datafusion/datasource/src/memory.rs:
##########
@@ -723,6 +761,222 @@ impl MemorySourceConfig {
pub fn original_schema(&self) -> SchemaRef {
Arc::clone(&self.schema)
}
+
+ /// Repartition while preserving order.
+ ///
+ /// Returns `Ok(None)` if cannot fulfill the requested repartitioning, such
+ /// as having too few batches to fulfill the `target_partitions` or if unable
+ /// to preserve output ordering.
+ fn repartition_preserving_order(
+ &self,
+ target_partitions: usize,
+ output_ordering: LexOrdering,
+ ) -> Result<Option<Vec<Vec<RecordBatch>>>> {
+ if !self.eq_properties().ordering_satisfy(&output_ordering) {
+ Ok(None)
+ } else if self.partitions.len() == 1 {
+ self.repartition_evenly_by_size(target_partitions)
+ } else {
Review Comment:
Thank you for the code review.
After adding an additional test case targeting the `self.partitions.len() == 1`, and ended up removing this conditional due to this: https://github.com/apache/datafusion/pull/15409/commits/e901482c84caca72c2105979b826221463ea6e0c
Additionally, I added code comments about how the max heap is being used: https://github.com/apache/datafusion/pull/15409/commits/3d5d50fcf5a9b065084c0da453cf967472928033
I believe the `RePartition` struct already documents what prop is the max heap cmp, vs ordering.
Please let me know if you spot anything else. Thank you.
--
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] Enable repartitioning on MemTable. [datafusion]
Posted by "wiedld (via GitHub)" <gi...@apache.org>.
wiedld commented on code in PR #15409:
URL: https://github.com/apache/datafusion/pull/15409#discussion_r2076020887
##########
datafusion/datasource/src/memory.rs:
##########
@@ -723,6 +761,222 @@ impl MemorySourceConfig {
pub fn original_schema(&self) -> SchemaRef {
Arc::clone(&self.schema)
}
+
+ /// Repartition while preserving order.
+ ///
+ /// Returns `Ok(None)` if cannot fulfill the requested repartitioning, such
+ /// as having too few batches to fulfill the `target_partitions` or if unable
+ /// to preserve output ordering.
+ fn repartition_preserving_order(
+ &self,
+ target_partitions: usize,
+ output_ordering: LexOrdering,
+ ) -> Result<Option<Vec<Vec<RecordBatch>>>> {
+ if !self.eq_properties().ordering_satisfy(&output_ordering) {
+ Ok(None)
+ } else if self.partitions.len() == 1 {
+ self.repartition_evenly_by_size(target_partitions)
+ } else {
Review Comment:
Thank you for the code review.
After adding an additional test case targeting the `self.partitions.len() == 1`, and ended up removing this conditional due to this: https://github.com/apache/datafusion/pull/15409/commits/57147add2bb1d540334b08639504b80930c1c982
Additionally, I added code comments about how the max heap is being used: https://github.com/apache/datafusion/pull/15409/commits/12e6dfb1dec6038b47a63019ee6965ed626bc68e
I believe the `RePartition` struct already documents what prop is the max heap cmp, vs ordering.
Please let me know if anything else is still not clear.
--
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] Enable repartitioning on MemTable. [datafusion]
Posted by "wiedld (via GitHub)" <gi...@apache.org>.
wiedld commented on code in PR #15409:
URL: https://github.com/apache/datafusion/pull/15409#discussion_r2011495377
##########
datafusion/datasource/src/memory.rs:
##########
@@ -902,4 +1108,319 @@ mod tests {
Ok(())
}
+
+ fn batch(row_size: usize) -> RecordBatch {
+ let a: ArrayRef = Arc::new(Int32Array::from(vec![1; row_size]));
+ let b: ArrayRef = Arc::new(StringArray::from_iter(vec![Some("foo"); row_size]));
+ let c: ArrayRef = Arc::new(Int64Array::from_iter(vec![1; row_size]));
+ RecordBatch::try_from_iter(vec![("a", a), ("b", b), ("c", c)]).unwrap()
+ }
+
+ fn schema() -> SchemaRef {
+ batch(1).schema()
+ }
+
+ fn memorysrcconfig_no_partitions(
+ sort_information: Vec<LexOrdering>,
+ ) -> Result<MemorySourceConfig> {
+ let partitions = vec![];
+ MemorySourceConfig::try_new(&partitions, schema(), None)?
+ .try_with_sort_information(sort_information)
+ }
+
+ fn memorysrcconfig_1_partition_1_batch(
+ sort_information: Vec<LexOrdering>,
+ ) -> Result<MemorySourceConfig> {
+ let partitions = vec![vec![batch(100)]];
+ MemorySourceConfig::try_new(&partitions, schema(), None)?
+ .try_with_sort_information(sort_information)
+ }
+
+ fn memorysrcconfig_3_partitions_1_batch_each(
+ sort_information: Vec<LexOrdering>,
+ ) -> Result<MemorySourceConfig> {
+ let partitions = vec![vec![batch(100)], vec![batch(100)], vec![batch(100)]];
+ MemorySourceConfig::try_new(&partitions, schema(), None)?
+ .try_with_sort_information(sort_information)
+ }
+
+ fn memorysrcconfig_3_partitions_with_2_batches_each(
+ sort_information: Vec<LexOrdering>,
+ ) -> Result<MemorySourceConfig> {
+ let partitions = vec![
+ vec![batch(100), batch(100)],
+ vec![batch(100), batch(100)],
+ vec![batch(100), batch(100)],
+ ];
+ MemorySourceConfig::try_new(&partitions, schema(), None)?
+ .try_with_sort_information(sort_information)
+ }
+
+ fn memorysrcconfig_1_partition_with_different_sized_batches(
+ sort_information: Vec<LexOrdering>,
+ ) -> Result<MemorySourceConfig> {
+ let partitions = vec![vec![batch(100_000), batch(10_000), batch(100), batch(1)]];
+ MemorySourceConfig::try_new(&partitions, schema(), None)?
+ .try_with_sort_information(sort_information)
+ }
+
+ fn memorysrcconfig_2_partition_with_different_sized_batches(
+ sort_information: Vec<LexOrdering>,
+ ) -> Result<MemorySourceConfig> {
+ let partitions = vec![
+ vec![batch(100_000), batch(10_000), batch(1_000)],
+ vec![batch(2_000), batch(20)],
+ ];
+ MemorySourceConfig::try_new(&partitions, schema(), None)?
+ .try_with_sort_information(sort_information)
+ }
+
+ /// Assert that we get the expected count of partitions after repartitioning.
+ ///
+ /// If None, then we expected the [`DataSource::repartitioned`] to return None.
+ fn assert_partitioning(
+ partitioned_datasrc: Option<Arc<dyn DataSource>>,
+ partition_cnt: Option<usize>,
+ ) {
+ let should_exist = if let Some(partition_cnt) = partition_cnt {
+ format!(
+ "new datasource should exist and have {:?} partitions",
+ partition_cnt
+ )
+ } else {
+ "new datasource should not exist".into()
+ };
+
+ let actual = partitioned_datasrc
+ .map(|datasrc| datasrc.output_partitioning().partition_count());
+ assert_eq!(
+ actual,
+ partition_cnt,
+ "partitioned datasrc does not match expected, we expected {}, instead found {:?}",
+ should_exist,
+ actual
+ );
+ }
+
+ fn run_all_test_scenarios(
+ output_ordering: Option<LexOrdering>,
+ sort_information_on_config: Vec<LexOrdering>,
+ ) -> Result<()> {
+ let not_used = usize::MAX;
+
+ // src has no partitions
+ let mem_src_config =
+ memorysrcconfig_no_partitions(sort_information_on_config.clone())?;
+ let partitioned_datasrc =
+ mem_src_config.repartitioned(1, not_used, output_ordering.clone())?;
+ assert_partitioning(partitioned_datasrc, None);
+
+ // src has partitions == target partitions (=1)
+ let target_partitions = 1;
+ let mem_src_config =
+ memorysrcconfig_1_partition_1_batch(sort_information_on_config.clone())?;
+ let partitioned_datasrc = mem_src_config.repartitioned(
+ target_partitions,
+ not_used,
+ output_ordering.clone(),
+ )?;
+ assert_partitioning(partitioned_datasrc, None);
+
+ // src has partitions == target partitions (=3)
+ let target_partitions = 3;
+ let mem_src_config = memorysrcconfig_3_partitions_1_batch_each(
+ sort_information_on_config.clone(),
+ )?;
+ let partitioned_datasrc = mem_src_config.repartitioned(
+ target_partitions,
+ not_used,
+ output_ordering.clone(),
+ )?;
+ assert_partitioning(partitioned_datasrc, None);
+
+ // src has partitions > target partitions, but we don't merge them
+ let target_partitions = 2;
+ let mem_src_config = memorysrcconfig_3_partitions_1_batch_each(
+ sort_information_on_config.clone(),
+ )?;
+ let partitioned_datasrc = mem_src_config.repartitioned(
+ target_partitions,
+ not_used,
+ output_ordering.clone(),
+ )?;
+ assert_partitioning(partitioned_datasrc, None);
+
+ // src has partitions < target partitions, but not enough batches (per partition) to split into more partitions
+ let target_partitions = 4;
+ let mem_src_config = memorysrcconfig_3_partitions_1_batch_each(
+ sort_information_on_config.clone(),
+ )?;
+ let partitioned_datasrc = mem_src_config.repartitioned(
+ target_partitions,
+ not_used,
+ output_ordering.clone(),
+ )?;
+ assert_partitioning(partitioned_datasrc, None);
+
+ // src has partitions < target partitions, and can split to sufficient amount
+ // has 6 batches across 3 partitions. Will need to split 2 of it's partitions.
+ let target_partitions = 5;
+ let mem_src_config = memorysrcconfig_3_partitions_with_2_batches_each(
+ sort_information_on_config.clone(),
+ )?;
+ let partitioned_datasrc = mem_src_config.repartitioned(
+ target_partitions,
+ not_used,
+ output_ordering.clone(),
+ )?;
+ assert_partitioning(partitioned_datasrc, Some(5));
+
+ // src has partitions < target partitions, and can split to sufficient amount
+ // has 6 batches across 3 partitions. Will need to split all of it's partitions.
+ let target_partitions = 6;
+ let mem_src_config = memorysrcconfig_3_partitions_with_2_batches_each(
+ sort_information_on_config.clone(),
+ )?;
+ let partitioned_datasrc = mem_src_config.repartitioned(
+ target_partitions,
+ not_used,
+ output_ordering.clone(),
+ )?;
+ assert_partitioning(partitioned_datasrc, Some(6));
+
+ // src has partitions < target partitions, but not enough total batches to fulfill the split (desired target_partitions)
+ let target_partitions = 3 * 2 + 1;
+ let mem_src_config = memorysrcconfig_3_partitions_with_2_batches_each(
+ sort_information_on_config.clone(),
+ )?;
+ let partitioned_datasrc = mem_src_config.repartitioned(
+ target_partitions,
+ not_used,
+ output_ordering.clone(),
+ )?;
+ assert_partitioning(partitioned_datasrc, None);
+
+ // src has 1 partition with many batches of lopsided sizes
+ // make sure it handles the split properly
+ let target_partitions = 2;
+ let mem_src_config = memorysrcconfig_1_partition_with_different_sized_batches(
+ sort_information_on_config,
+ )?;
+ let partitioned_datasrc = mem_src_config.clone().repartitioned(
+ target_partitions,
+ not_used,
+ output_ordering,
+ )?;
+ assert_partitioning(partitioned_datasrc.clone(), Some(2));
+ // Starting = batch(100_000), batch(10_000), batch(100), batch(1).
+ // It should have split as p1=batch(100_000), p2=[batch(10_000), batch(100), batch(1)]
+ let repartitioned_raw_batches = mem_src_config
+ .repartition_evenly_by_size(target_partitions)?
+ .unwrap();
+ assert_eq!(repartitioned_raw_batches.len(), 2);
+ let [ref p1, ref p2] = repartitioned_raw_batches[..] else {
+ unreachable!()
+ };
+ // p1=batch(100_000)
+ assert_eq!(p1.len(), 1);
+ assert_eq!(p1[0].num_rows(), 100_000);
+ // p2=[batch(10_000), batch(100), batch(1)]
+ assert_eq!(p2.len(), 3);
+ assert_eq!(p2[0].num_rows(), 10_000);
+ assert_eq!(p2[1].num_rows(), 100);
+ assert_eq!(p2[2].num_rows(), 1);
+
+ Ok(())
+ }
+
+ #[test]
+ fn test_repartition_no_sort_information_no_output_ordering() -> Result<()> {
+ let no_sort = vec![];
+ let no_output_ordering = None;
+
+ // Test: common set of functionality
+ run_all_test_scenarios(no_output_ordering.clone(), no_sort.clone())?;
+
+ // Test: does not preserve separate partitions (with own internal ordering) on even split
Review Comment:
This defines what I believe the different behavior should be when not preserving sort ordering. I looked at how [FileScanConfig::repartition](https://github.com/apache/datafusion/blob/908ca0e684f3184075ef799128ee3f58f1e9db86/datafusion/datasource/src/file_scan_config.rs#L237) was implemented as a guide. The key branch point in it's implementation is here:
https://github.com/apache/datafusion/blob/908ca0e684f3184075ef799128ee3f58f1e9db86/datafusion/datasource/src/file_groups.rs#L194-L199
--
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] Enable repartitioning on MemTable. [datafusion]
Posted by "wiedld (via GitHub)" <gi...@apache.org>.
wiedld commented on code in PR #15409:
URL: https://github.com/apache/datafusion/pull/15409#discussion_r2053132567
##########
datafusion/datasource/src/memory.rs:
##########
@@ -902,4 +1130,319 @@ mod tests {
Ok(())
}
+
+ fn batch(row_size: usize) -> RecordBatch {
+ let a: ArrayRef = Arc::new(Int32Array::from(vec![1; row_size]));
+ let b: ArrayRef = Arc::new(StringArray::from_iter(vec![Some("foo"); row_size]));
+ let c: ArrayRef = Arc::new(Int64Array::from_iter(vec![1; row_size]));
+ RecordBatch::try_from_iter(vec![("a", a), ("b", b), ("c", c)]).unwrap()
+ }
+
+ fn schema() -> SchemaRef {
+ batch(1).schema()
+ }
+
+ fn memorysrcconfig_no_partitions(
+ sort_information: Vec<LexOrdering>,
+ ) -> Result<MemorySourceConfig> {
+ let partitions = vec![];
+ MemorySourceConfig::try_new(&partitions, schema(), None)?
+ .try_with_sort_information(sort_information)
+ }
+
+ fn memorysrcconfig_1_partition_1_batch(
+ sort_information: Vec<LexOrdering>,
+ ) -> Result<MemorySourceConfig> {
+ let partitions = vec![vec![batch(100)]];
+ MemorySourceConfig::try_new(&partitions, schema(), None)?
+ .try_with_sort_information(sort_information)
+ }
+
+ fn memorysrcconfig_3_partitions_1_batch_each(
+ sort_information: Vec<LexOrdering>,
+ ) -> Result<MemorySourceConfig> {
+ let partitions = vec![vec![batch(100)], vec![batch(100)], vec![batch(100)]];
+ MemorySourceConfig::try_new(&partitions, schema(), None)?
+ .try_with_sort_information(sort_information)
+ }
+
+ fn memorysrcconfig_3_partitions_with_2_batches_each(
+ sort_information: Vec<LexOrdering>,
+ ) -> Result<MemorySourceConfig> {
+ let partitions = vec![
+ vec![batch(100), batch(100)],
+ vec![batch(100), batch(100)],
+ vec![batch(100), batch(100)],
+ ];
+ MemorySourceConfig::try_new(&partitions, schema(), None)?
+ .try_with_sort_information(sort_information)
+ }
+
+ fn memorysrcconfig_1_partition_with_different_sized_batches(
+ sort_information: Vec<LexOrdering>,
+ ) -> Result<MemorySourceConfig> {
+ let partitions = vec![vec![batch(100_000), batch(10_000), batch(100), batch(1)]];
+ MemorySourceConfig::try_new(&partitions, schema(), None)?
+ .try_with_sort_information(sort_information)
+ }
+
+ fn memorysrcconfig_2_partition_with_different_sized_batches(
+ sort_information: Vec<LexOrdering>,
+ ) -> Result<MemorySourceConfig> {
+ let partitions = vec![
+ vec![batch(100_000), batch(10_000), batch(1_000)],
Review Comment:
Excellent idea, and it caught a bug. I've updated the algo & added the test: https://github.com/apache/datafusion/pull/15409/commits/c661d97221f12b53520690a7d707439fc3100205
thank you.
--
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] Enable repartitioning on MemTable. [datafusion]
Posted by "wiedld (via GitHub)" <gi...@apache.org>.
wiedld commented on code in PR #15409:
URL: https://github.com/apache/datafusion/pull/15409#discussion_r2053132567
##########
datafusion/datasource/src/memory.rs:
##########
@@ -902,4 +1130,319 @@ mod tests {
Ok(())
}
+
+ fn batch(row_size: usize) -> RecordBatch {
+ let a: ArrayRef = Arc::new(Int32Array::from(vec![1; row_size]));
+ let b: ArrayRef = Arc::new(StringArray::from_iter(vec![Some("foo"); row_size]));
+ let c: ArrayRef = Arc::new(Int64Array::from_iter(vec![1; row_size]));
+ RecordBatch::try_from_iter(vec![("a", a), ("b", b), ("c", c)]).unwrap()
+ }
+
+ fn schema() -> SchemaRef {
+ batch(1).schema()
+ }
+
+ fn memorysrcconfig_no_partitions(
+ sort_information: Vec<LexOrdering>,
+ ) -> Result<MemorySourceConfig> {
+ let partitions = vec![];
+ MemorySourceConfig::try_new(&partitions, schema(), None)?
+ .try_with_sort_information(sort_information)
+ }
+
+ fn memorysrcconfig_1_partition_1_batch(
+ sort_information: Vec<LexOrdering>,
+ ) -> Result<MemorySourceConfig> {
+ let partitions = vec![vec![batch(100)]];
+ MemorySourceConfig::try_new(&partitions, schema(), None)?
+ .try_with_sort_information(sort_information)
+ }
+
+ fn memorysrcconfig_3_partitions_1_batch_each(
+ sort_information: Vec<LexOrdering>,
+ ) -> Result<MemorySourceConfig> {
+ let partitions = vec![vec![batch(100)], vec![batch(100)], vec![batch(100)]];
+ MemorySourceConfig::try_new(&partitions, schema(), None)?
+ .try_with_sort_information(sort_information)
+ }
+
+ fn memorysrcconfig_3_partitions_with_2_batches_each(
+ sort_information: Vec<LexOrdering>,
+ ) -> Result<MemorySourceConfig> {
+ let partitions = vec![
+ vec![batch(100), batch(100)],
+ vec![batch(100), batch(100)],
+ vec![batch(100), batch(100)],
+ ];
+ MemorySourceConfig::try_new(&partitions, schema(), None)?
+ .try_with_sort_information(sort_information)
+ }
+
+ fn memorysrcconfig_1_partition_with_different_sized_batches(
+ sort_information: Vec<LexOrdering>,
+ ) -> Result<MemorySourceConfig> {
+ let partitions = vec![vec![batch(100_000), batch(10_000), batch(100), batch(1)]];
+ MemorySourceConfig::try_new(&partitions, schema(), None)?
+ .try_with_sort_information(sort_information)
+ }
+
+ fn memorysrcconfig_2_partition_with_different_sized_batches(
+ sort_information: Vec<LexOrdering>,
+ ) -> Result<MemorySourceConfig> {
+ let partitions = vec![
+ vec![batch(100_000), batch(10_000), batch(1_000)],
Review Comment:
Excellent idea, and it caught a bug. I've updated the algo & added the test: https://github.com/apache/datafusion/pull/15409/commits/c661d97221f12b53520690a7d707439fc3100205. Thank you.
--
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] Enable repartitioning on MemTable. [datafusion]
Posted by "wiedld (via GitHub)" <gi...@apache.org>.
wiedld commented on code in PR #15409:
URL: https://github.com/apache/datafusion/pull/15409#discussion_r2053131373
##########
datafusion/core/tests/fuzz_cases/aggregate_fuzz.rs:
##########
@@ -520,7 +520,9 @@ async fn group_by_string_test(
let expected = compute_counts(&input, column_name);
let schema = input[0].schema();
- let session_config = SessionConfig::new().with_batch_size(50);
+ let session_config = SessionConfig::new()
+ .with_batch_size(50)
+ .with_repartition_file_scans(false);
Review Comment:
When updating the docs (see this commit: https://github.com/apache/datafusion/pull/15409/commits/d8929dc9d179e89868964f5f33dd2e387e54c299), it felt like we don't want to split out the 2 configurations IMO. Since the repartitioning config variable is used to both (a) repartition at the datasrc (which could be a file or memtable), as well as (b) insert a repartition operator later in the plan.
Please take a look at the added docs (in the commit linked above), the lmk what you think.
--
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] Enable repartitioning on MemTable. [datafusion]
Posted by "2010YOUY01 (via GitHub)" <gi...@apache.org>.
2010YOUY01 commented on code in PR #15409:
URL: https://github.com/apache/datafusion/pull/15409#discussion_r2065566166
##########
datafusion/core/tests/physical_optimizer/enforce_distribution.rs:
##########
@@ -3471,3 +3477,102 @@ fn optimize_away_unnecessary_repartition2() -> Result<()> {
Ok(())
}
+
+#[tokio::test]
+async fn test_distribute_sort_parquet() -> Result<()> {
+ let test_config: TestConfig =
+ TestConfig::default().with_prefer_repartition_file_scans(1000);
+ assert!(
+ test_config.config.optimizer.repartition_datasource_scans,
+ "should enable scans to be repartitioned"
+ );
+
+ let schema = schema();
+ let sort_key = LexOrdering::new(vec![PhysicalSortExpr {
+ expr: col("c", &schema).unwrap(),
+ options: SortOptions::default(),
+ }]);
+ let physical_plan = sort_exec(sort_key, parquet_exec_with_stats(10000 * 8192), false);
+
+ // prior to optimization, this is the starting plan
+ let starting = &[
+ "SortExec: expr=[c@2 ASC], preserve_partitioning=[false]",
+ " DataSourceExec: file_groups={1 group: [[x]]}, projection=[a, b, c, d, e], file_type=parquet",
+ ];
+ plans_matches_expected!(starting, physical_plan.clone());
+
+ // what the enforce distribution run does.
+ let expected = &[
+ "SortExec: expr=[c@2 ASC], preserve_partitioning=[false]",
+ " CoalescePartitionsExec",
+ " DataSourceExec: file_groups={10 groups: [[x:0..8192000], [x:8192000..16384000], [x:16384000..24576000], [x:24576000..32768000], [x:32768000..40960000], [x:40960000..49152000], [x:49152000..57344000], [x:57344000..65536000], [x:65536000..73728000], [x:73728000..81920000]]}, projection=[a, b, c, d, e], file_type=parquet",
+ ];
+ test_config.run(expected, physical_plan.clone(), &[Run::Distribution])?;
+
+ // what the sort parallelization (in enforce sorting), does after the enforce distribution changes
+ let expected = &[
+ "SortPreservingMergeExec: [c@2 ASC]",
+ " SortExec: expr=[c@2 ASC], preserve_partitioning=[true]",
+ " DataSourceExec: file_groups={10 groups: [[x:0..8192000], [x:8192000..16384000], [x:16384000..24576000], [x:24576000..32768000], [x:32768000..40960000], [x:40960000..49152000], [x:49152000..57344000], [x:57344000..65536000], [x:65536000..73728000], [x:73728000..81920000]]}, projection=[a, b, c, d, e], file_type=parquet",
+ ];
+ test_config.run(expected, physical_plan, &[Run::Distribution, Run::Sorting])?;
+ Ok(())
+}
+
+#[tokio::test]
Review Comment:
```suggestion
/// Ensures that `DataSourceExec` has been repartitioned into `target_partitions` file groups
#[tokio::test]
```
##########
datafusion/common/src/config.rs:
##########
@@ -632,14 +632,21 @@ config_namespace! {
/// long runner execution, all types of joins may encounter out-of-memory errors.
pub allow_symmetric_joins_without_pruning: bool, default = true
- /// When set to `true`, file groups will be repartitioned to achieve maximum parallelism.
- /// Currently Parquet and CSV formats are supported.
+ /// When set to `true`, datasource partitions will be repartitioned to achieve maximum parallelism.
+ /// This applies to both in-memory partitions and FileSource's file groups (1 group is 1 partition).
///
- /// If set to `true`, all files will be repartitioned evenly (i.e., a single large file
+ /// For FileSources, only Parquet and CSV formats are currently supported.
+ ///
+ /// If set to `true` for a FileSource, all files will be repartitioned evenly (i.e., a single large file
/// might be partitioned into smaller chunks) for parallel scanning.
- /// If set to `false`, different files will be read in parallel, but repartitioning won't
+ /// If set to `false` for a FileSource, different files will be read in parallel, but repartitioning won't
/// happen within a single file.
- pub repartition_file_scans: bool, default = true
+ ///
+ /// If set to `true` for an in-memory source, all memtable's partitions will have their batches
+ /// repartitioned evenly to the desired number of `target_partitions`. Repartitioning can change
+ /// the total number of partitions and batches per partition, but does not slice the initial
+ /// record tables provided to the MemTable on creation.
+ pub repartition_datasource_scans: bool, default = true
Review Comment:
This name is better, however, I recommend using the old name to avoid changing the config API.
##########
datafusion/datasource/src/memory.rs:
##########
@@ -723,6 +761,222 @@ impl MemorySourceConfig {
pub fn original_schema(&self) -> SchemaRef {
Arc::clone(&self.schema)
}
+
+ /// Repartition while preserving order.
+ ///
+ /// Returns `Ok(None)` if cannot fulfill the requested repartitioning, such
+ /// as having too few batches to fulfill the `target_partitions` or if unable
+ /// to preserve output ordering.
+ fn repartition_preserving_order(
+ &self,
+ target_partitions: usize,
+ output_ordering: LexOrdering,
+ ) -> Result<Option<Vec<Vec<RecordBatch>>>> {
+ if !self.eq_properties().ordering_satisfy(&output_ordering) {
+ Ok(None)
+ } else if self.partitions.len() == 1 {
+ self.repartition_evenly_by_size(target_partitions)
+ } else {
Review Comment:
It would be great to include a simple example to explain this code block, like how is an input table be repartitioned, what is the key in the max heap, and how the state changes during each loop step.
##########
datafusion/core/tests/physical_optimizer/enforce_distribution.rs:
##########
@@ -3471,3 +3477,102 @@ fn optimize_away_unnecessary_repartition2() -> Result<()> {
Ok(())
}
+
+#[tokio::test]
Review Comment:
```suggestion
/// Ensures that `DataSourceExec` has been repartitioned into `target_partitions` file groups
#[tokio::test]
```
##########
datafusion/datasource/src/memory.rs:
##########
@@ -723,6 +761,222 @@ impl MemorySourceConfig {
pub fn original_schema(&self) -> SchemaRef {
Arc::clone(&self.schema)
}
+
+ /// Repartition while preserving order.
+ ///
+ /// Returns `Ok(None)` if cannot fulfill the requested repartitioning, such
+ /// as having too few batches to fulfill the `target_partitions` or if unable
+ /// to preserve output ordering.
+ fn repartition_preserving_order(
+ &self,
+ target_partitions: usize,
+ output_ordering: LexOrdering,
+ ) -> Result<Option<Vec<Vec<RecordBatch>>>> {
+ if !self.eq_properties().ordering_satisfy(&output_ordering) {
+ Ok(None)
+ } else if self.partitions.len() == 1 {
+ self.repartition_evenly_by_size(target_partitions)
+ } else {
+ let total_num_batches =
+ self.partitions.iter().map(|b| b.len()).sum::<usize>();
+ if total_num_batches < target_partitions {
+ // no way to create the desired repartitioning
+ return Ok(None);
+ }
+
+ let cnt_to_repartition = target_partitions - self.partitions.len();
+
+ // Label the current partitions and their order.
+ // Such that when we later split up the partitions into smaller sizes, we are maintaining the order.
+ let to_repartition = self
+ .partitions
+ .iter()
+ .enumerate()
+ .map(|(idx, batches)| RePartition {
+ idx: idx + (cnt_to_repartition * idx), // make space in ordering for split partitions
+ row_count: batches.iter().map(|batch| batch.num_rows()).sum(),
+ batches: batches.clone(),
+ })
+ .collect_vec();
+
+ // Put all of the partitions into a heap ordered by `RePartition::partial_cmp`, which sizes
+ // by count of rows.
+ let mut max_heap = BinaryHeap::with_capacity(target_partitions);
+ for rep in to_repartition {
+ max_heap.push(rep);
+ }
+
+ // Split the largest partitions into smaller partitions. Maintaining the output
+ // order of the partitions & newly created partitions.
+ let mut cannot_split_further = Vec::with_capacity(target_partitions);
+ for _ in 0..cnt_to_repartition {
+ loop {
+ let Some(to_split) = max_heap.pop() else {
+ break;
+ };
+
+ let mut new_partitions = to_split.split();
+ if new_partitions.len() > 1 {
+ for new_partition in new_partitions {
+ max_heap.push(new_partition);
+ }
+ break;
+ } else {
+ cannot_split_further.push(new_partitions.remove(0));
+ }
+ }
+ }
+ let mut partitions = max_heap.drain().collect_vec();
+ partitions.extend(cannot_split_further);
+
+ // Finally, sort all partitions by the output ordering.
+ partitions.sort_by_key(|p| p.idx);
Review Comment:
I think currently datafusion only uses partial ordering (the arg `output_ordering` of this function means the memory table before partitioning has such order within each partition, not necessarily mean it maintains order across partitions), see:
https://github.com/apache/datafusion/blob/5e1214c55e37d198d732667b770943cfba4fe5c3/datafusion/physical-plan/src/execution_plan.rs#L533
If it is the case, this final sorting step is not necessary. However, I think it can be kept, since it's efficient.
We might need additional API to do such repartitioning that further maintains global order, and this implementation can be reused.
##########
datafusion/datasource/src/memory.rs:
##########
@@ -723,6 +761,222 @@ impl MemorySourceConfig {
pub fn original_schema(&self) -> SchemaRef {
Arc::clone(&self.schema)
}
+
+ /// Repartition while preserving order.
+ ///
+ /// Returns `Ok(None)` if cannot fulfill the requested repartitioning, such
+ /// as having too few batches to fulfill the `target_partitions` or if unable
+ /// to preserve output ordering.
+ fn repartition_preserving_order(
+ &self,
+ target_partitions: usize,
+ output_ordering: LexOrdering,
+ ) -> Result<Option<Vec<Vec<RecordBatch>>>> {
+ if !self.eq_properties().ordering_satisfy(&output_ordering) {
+ Ok(None)
+ } else if self.partitions.len() == 1 {
+ self.repartition_evenly_by_size(target_partitions)
+ } else {
Review Comment:
It would be great to include a simple example to explain this code block, like how is an input table be repartitioned, what is the key in the max heap, and how the state changes during each loop step.
--
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] Enable repartitioning on MemTable. [datafusion]
Posted by "alamb (via GitHub)" <gi...@apache.org>.
alamb merged PR #15409:
URL: https://github.com/apache/datafusion/pull/15409
--
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] Enable repartitioning on MemTable. [datafusion]
Posted by "wiedld (via GitHub)" <gi...@apache.org>.
wiedld commented on code in PR #15409:
URL: https://github.com/apache/datafusion/pull/15409#discussion_r2011496396
##########
datafusion/datasource/src/memory.rs:
##########
@@ -902,4 +1108,319 @@ mod tests {
Ok(())
}
+
+ fn batch(row_size: usize) -> RecordBatch {
+ let a: ArrayRef = Arc::new(Int32Array::from(vec![1; row_size]));
+ let b: ArrayRef = Arc::new(StringArray::from_iter(vec![Some("foo"); row_size]));
+ let c: ArrayRef = Arc::new(Int64Array::from_iter(vec![1; row_size]));
+ RecordBatch::try_from_iter(vec![("a", a), ("b", b), ("c", c)]).unwrap()
+ }
+
+ fn schema() -> SchemaRef {
+ batch(1).schema()
+ }
+
+ fn memorysrcconfig_no_partitions(
+ sort_information: Vec<LexOrdering>,
+ ) -> Result<MemorySourceConfig> {
+ let partitions = vec![];
+ MemorySourceConfig::try_new(&partitions, schema(), None)?
+ .try_with_sort_information(sort_information)
+ }
+
+ fn memorysrcconfig_1_partition_1_batch(
+ sort_information: Vec<LexOrdering>,
+ ) -> Result<MemorySourceConfig> {
+ let partitions = vec![vec![batch(100)]];
+ MemorySourceConfig::try_new(&partitions, schema(), None)?
+ .try_with_sort_information(sort_information)
+ }
+
+ fn memorysrcconfig_3_partitions_1_batch_each(
+ sort_information: Vec<LexOrdering>,
+ ) -> Result<MemorySourceConfig> {
+ let partitions = vec![vec![batch(100)], vec![batch(100)], vec![batch(100)]];
+ MemorySourceConfig::try_new(&partitions, schema(), None)?
+ .try_with_sort_information(sort_information)
+ }
+
+ fn memorysrcconfig_3_partitions_with_2_batches_each(
+ sort_information: Vec<LexOrdering>,
+ ) -> Result<MemorySourceConfig> {
+ let partitions = vec![
+ vec![batch(100), batch(100)],
+ vec![batch(100), batch(100)],
+ vec![batch(100), batch(100)],
+ ];
+ MemorySourceConfig::try_new(&partitions, schema(), None)?
+ .try_with_sort_information(sort_information)
+ }
+
+ fn memorysrcconfig_1_partition_with_different_sized_batches(
+ sort_information: Vec<LexOrdering>,
+ ) -> Result<MemorySourceConfig> {
+ let partitions = vec![vec![batch(100_000), batch(10_000), batch(100), batch(1)]];
+ MemorySourceConfig::try_new(&partitions, schema(), None)?
+ .try_with_sort_information(sort_information)
+ }
+
+ fn memorysrcconfig_2_partition_with_different_sized_batches(
+ sort_information: Vec<LexOrdering>,
+ ) -> Result<MemorySourceConfig> {
+ let partitions = vec![
+ vec![batch(100_000), batch(10_000), batch(1_000)],
+ vec![batch(2_000), batch(20)],
+ ];
+ MemorySourceConfig::try_new(&partitions, schema(), None)?
+ .try_with_sort_information(sort_information)
+ }
+
+ /// Assert that we get the expected count of partitions after repartitioning.
+ ///
+ /// If None, then we expected the [`DataSource::repartitioned`] to return None.
+ fn assert_partitioning(
+ partitioned_datasrc: Option<Arc<dyn DataSource>>,
+ partition_cnt: Option<usize>,
+ ) {
+ let should_exist = if let Some(partition_cnt) = partition_cnt {
+ format!(
+ "new datasource should exist and have {:?} partitions",
+ partition_cnt
+ )
+ } else {
+ "new datasource should not exist".into()
+ };
+
+ let actual = partitioned_datasrc
+ .map(|datasrc| datasrc.output_partitioning().partition_count());
+ assert_eq!(
+ actual,
+ partition_cnt,
+ "partitioned datasrc does not match expected, we expected {}, instead found {:?}",
+ should_exist,
+ actual
+ );
+ }
+
+ fn run_all_test_scenarios(
+ output_ordering: Option<LexOrdering>,
+ sort_information_on_config: Vec<LexOrdering>,
+ ) -> Result<()> {
+ let not_used = usize::MAX;
+
+ // src has no partitions
+ let mem_src_config =
+ memorysrcconfig_no_partitions(sort_information_on_config.clone())?;
+ let partitioned_datasrc =
+ mem_src_config.repartitioned(1, not_used, output_ordering.clone())?;
+ assert_partitioning(partitioned_datasrc, None);
+
+ // src has partitions == target partitions (=1)
+ let target_partitions = 1;
+ let mem_src_config =
+ memorysrcconfig_1_partition_1_batch(sort_information_on_config.clone())?;
+ let partitioned_datasrc = mem_src_config.repartitioned(
+ target_partitions,
+ not_used,
+ output_ordering.clone(),
+ )?;
+ assert_partitioning(partitioned_datasrc, None);
+
+ // src has partitions == target partitions (=3)
+ let target_partitions = 3;
+ let mem_src_config = memorysrcconfig_3_partitions_1_batch_each(
+ sort_information_on_config.clone(),
+ )?;
+ let partitioned_datasrc = mem_src_config.repartitioned(
+ target_partitions,
+ not_used,
+ output_ordering.clone(),
+ )?;
+ assert_partitioning(partitioned_datasrc, None);
+
+ // src has partitions > target partitions, but we don't merge them
+ let target_partitions = 2;
+ let mem_src_config = memorysrcconfig_3_partitions_1_batch_each(
+ sort_information_on_config.clone(),
+ )?;
+ let partitioned_datasrc = mem_src_config.repartitioned(
+ target_partitions,
+ not_used,
+ output_ordering.clone(),
+ )?;
+ assert_partitioning(partitioned_datasrc, None);
+
+ // src has partitions < target partitions, but not enough batches (per partition) to split into more partitions
+ let target_partitions = 4;
+ let mem_src_config = memorysrcconfig_3_partitions_1_batch_each(
+ sort_information_on_config.clone(),
+ )?;
+ let partitioned_datasrc = mem_src_config.repartitioned(
+ target_partitions,
+ not_used,
+ output_ordering.clone(),
+ )?;
+ assert_partitioning(partitioned_datasrc, None);
+
+ // src has partitions < target partitions, and can split to sufficient amount
+ // has 6 batches across 3 partitions. Will need to split 2 of it's partitions.
+ let target_partitions = 5;
+ let mem_src_config = memorysrcconfig_3_partitions_with_2_batches_each(
+ sort_information_on_config.clone(),
+ )?;
+ let partitioned_datasrc = mem_src_config.repartitioned(
+ target_partitions,
+ not_used,
+ output_ordering.clone(),
+ )?;
+ assert_partitioning(partitioned_datasrc, Some(5));
+
+ // src has partitions < target partitions, and can split to sufficient amount
+ // has 6 batches across 3 partitions. Will need to split all of it's partitions.
+ let target_partitions = 6;
+ let mem_src_config = memorysrcconfig_3_partitions_with_2_batches_each(
+ sort_information_on_config.clone(),
+ )?;
+ let partitioned_datasrc = mem_src_config.repartitioned(
+ target_partitions,
+ not_used,
+ output_ordering.clone(),
+ )?;
+ assert_partitioning(partitioned_datasrc, Some(6));
+
+ // src has partitions < target partitions, but not enough total batches to fulfill the split (desired target_partitions)
+ let target_partitions = 3 * 2 + 1;
+ let mem_src_config = memorysrcconfig_3_partitions_with_2_batches_each(
+ sort_information_on_config.clone(),
+ )?;
+ let partitioned_datasrc = mem_src_config.repartitioned(
+ target_partitions,
+ not_used,
+ output_ordering.clone(),
+ )?;
+ assert_partitioning(partitioned_datasrc, None);
+
+ // src has 1 partition with many batches of lopsided sizes
+ // make sure it handles the split properly
+ let target_partitions = 2;
+ let mem_src_config = memorysrcconfig_1_partition_with_different_sized_batches(
+ sort_information_on_config,
+ )?;
+ let partitioned_datasrc = mem_src_config.clone().repartitioned(
+ target_partitions,
+ not_used,
+ output_ordering,
+ )?;
+ assert_partitioning(partitioned_datasrc.clone(), Some(2));
+ // Starting = batch(100_000), batch(10_000), batch(100), batch(1).
+ // It should have split as p1=batch(100_000), p2=[batch(10_000), batch(100), batch(1)]
+ let repartitioned_raw_batches = mem_src_config
+ .repartition_evenly_by_size(target_partitions)?
+ .unwrap();
+ assert_eq!(repartitioned_raw_batches.len(), 2);
+ let [ref p1, ref p2] = repartitioned_raw_batches[..] else {
+ unreachable!()
+ };
+ // p1=batch(100_000)
+ assert_eq!(p1.len(), 1);
+ assert_eq!(p1[0].num_rows(), 100_000);
+ // p2=[batch(10_000), batch(100), batch(1)]
+ assert_eq!(p2.len(), 3);
+ assert_eq!(p2[0].num_rows(), 10_000);
+ assert_eq!(p2[1].num_rows(), 100);
+ assert_eq!(p2[2].num_rows(), 1);
+
+ Ok(())
+ }
+
+ #[test]
+ fn test_repartition_no_sort_information_no_output_ordering() -> Result<()> {
+ let no_sort = vec![];
+ let no_output_ordering = None;
+
+ // Test: common set of functionality
+ run_all_test_scenarios(no_output_ordering.clone(), no_sort.clone())?;
+
+ // Test: does not preserve separate partitions (with own internal ordering) on even split
+ let target_partitions = 3;
+ let mem_src_config =
+ memorysrcconfig_2_partition_with_different_sized_batches(no_sort)?;
+ let partitioned_datasrc = mem_src_config.clone().repartitioned(
+ target_partitions,
+ usize::MAX,
+ no_output_ordering,
+ )?;
+ assert_partitioning(partitioned_datasrc.clone(), Some(3));
+ // Starting = batch(100_000), batch(10_000), batch(1_000), batch(2_000), batch(20)
+ // It should have split as p1=batch(100_000), p2=batch(10_000), p3=rest(mixed across original partitions)
+ let repartitioned_raw_batches = mem_src_config
+ .repartition_evenly_by_size(target_partitions)?
+ .unwrap();
+ assert_eq!(repartitioned_raw_batches.len(), 3);
+ let [ref p1, ref p2, ref p3] = repartitioned_raw_batches[..] else {
+ unreachable!()
+ };
+ // p1=batch(100_000)
+ assert_eq!(p1.len(), 1);
+ assert_eq!(p1[0].num_rows(), 100_000);
+ // p2=batch(10_000)
+ assert_eq!(p2.len(), 1);
+ assert_eq!(p2[0].num_rows(), 10_000);
+ // p3= batch(1_000), batch(2_000), batch(20)
+ assert_eq!(p3.len(), 3);
+ assert_eq!(p3[0].num_rows(), 1_000);
+ assert_eq!(p3[1].num_rows(), 2_000);
+ assert_eq!(p3[2].num_rows(), 20);
+
+ Ok(())
+ }
+
+ #[test]
+ fn test_repartition_with_sort_information() -> Result<()> {
+ let schema = schema();
+ let sort_key = LexOrdering::new(vec![PhysicalSortExpr {
+ expr: col("c", &schema).unwrap(),
+ options: SortOptions::default(),
+ }]);
+ let has_sort = vec![sort_key.clone()];
+ let output_ordering = Some(sort_key);
+
+ // Test: common set of functionality
+ run_all_test_scenarios(output_ordering.clone(), has_sort.clone())?;
+
+ // Test: DOES preserve separate partitions (with own internal ordering)
Review Comment:
Here is the preservation of sort ordering. Again, using the [FileSourceConfig implementation](https://github.com/apache/datafusion/pull/15409/files#r2011495377) as a guide.
--
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] Enable repartitioning on MemTable. [datafusion]
Posted by "wiedld (via GitHub)" <gi...@apache.org>.
wiedld commented on code in PR #15409:
URL: https://github.com/apache/datafusion/pull/15409#discussion_r2011720446
##########
datafusion/core/tests/fuzz_cases/aggregate_fuzz.rs:
##########
@@ -520,7 +520,9 @@ async fn group_by_string_test(
let expected = compute_counts(&input, column_name);
let schema = input[0].schema();
- let session_config = SessionConfig::new().with_batch_size(50);
+ let session_config = SessionConfig::new()
+ .with_batch_size(50)
+ .with_repartition_file_scans(false);
Review Comment:
On the other hand, the fact that I had to disable mem table repartitioning (via config) to get some of the test cases to pass, makes me would if the should-repartition-files vs should-repartition-memtable decision should be made separately in the enforce distribution optimizer. 🤔
--
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] Enable repartitioning on MemTable. [datafusion]
Posted by "wiedld (via GitHub)" <gi...@apache.org>.
wiedld commented on code in PR #15409:
URL: https://github.com/apache/datafusion/pull/15409#discussion_r2011678484
##########
datafusion/core/tests/fuzz_cases/aggregate_fuzz.rs:
##########
@@ -520,7 +520,9 @@ async fn group_by_string_test(
let expected = compute_counts(&input, column_name);
let schema = input[0].schema();
- let session_config = SessionConfig::new().with_batch_size(50);
+ let session_config = SessionConfig::new()
+ .with_batch_size(50)
+ .with_repartition_file_scans(false);
Review Comment:
Update needed for test to pass.
--
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] Enable repartitioning on MemTable. [datafusion]
Posted by "wiedld (via GitHub)" <gi...@apache.org>.
wiedld commented on code in PR #15409:
URL: https://github.com/apache/datafusion/pull/15409#discussion_r2011691796
##########
datafusion/datasource/src/memory.rs:
##########
@@ -718,6 +750,181 @@ impl MemorySourceConfig {
pub fn original_schema(&self) -> SchemaRef {
Arc::clone(&self.schema)
}
+
+ /// Repartition while preserving order.
+ ///
+ /// Returns None if cannot fulfill requested repartitioning.
+ fn repartition_preserving_order(
+ &self,
+ target_partitions: usize,
+ output_ordering: LexOrdering,
+ ) -> Result<Option<Vec<Vec<RecordBatch>>>> {
+ if !self.eq_properties().ordering_satisfy(&output_ordering) {
+ Ok(None)
+ } else if self.partitions.len() == 1 {
+ self.repartition_evenly_by_size(target_partitions)
+ } else {
+ let total_num_batches =
+ self.partitions.iter().map(|b| b.len()).sum::<usize>();
+ if total_num_batches < target_partitions {
+ // no way to create the desired repartitioning
+ return Ok(None);
+ }
+
+ let cnt_to_repartition = target_partitions - self.partitions.len();
+
+ let to_repartition = self
+ .partitions
+ .iter()
+ .enumerate()
+ .map(|(idx, batches)| RePartition {
+ idx: idx + (cnt_to_repartition * idx), // make space in ordering for split partitions
+ row_count: batches.iter().map(|batch| batch.num_rows()).sum(),
+ batches: batches.clone(),
+ })
+ .collect_vec();
+
+ // split the largest partitions
+ let mut max_heap = BinaryHeap::with_capacity(target_partitions);
+ for rep in to_repartition {
+ max_heap.push(rep);
+ }
+ for _ in 0..cnt_to_repartition {
+ let Some(to_split) = max_heap.pop() else {
+ unreachable!()
+ };
+ for new_partition in to_split.split() {
+ max_heap.push(new_partition);
+ }
+ }
+ let mut partitions = max_heap.drain().collect_vec();
+ partitions.sort_by_key(|p| p.idx);
+ let partitions = partitions.into_iter().map(|rep| rep.batches).collect_vec();
+
+ Ok(Some(partitions))
+ }
+ }
+
+ /// Repartition into evenly sized chunks (as much as possible without batch splitting),
+ /// disregarding any ordering.
+ ///
+ /// Returns None if cannot fulfill requested repartitioning.
+ fn repartition_evenly_by_size(
+ &self,
+ target_partitions: usize,
+ ) -> Result<Option<Vec<Vec<RecordBatch>>>> {
+ // determine if we have enough total batches to fulfill request
+ let flatten_batches = self.partitions.clone().into_iter().flatten().collect_vec();
+ if flatten_batches.len() < target_partitions {
+ return Ok(None);
+ }
+
+ // repartition evenly
+ let total_num_rows = flatten_batches.iter().map(|b| b.num_rows()).sum::<usize>();
+ let target_partition_size = total_num_rows.div_ceil(target_partitions);
+ let mut partitions =
+ vec![Vec::with_capacity(flatten_batches.len()); target_partitions];
+ let mut curr_row_count = 0;
+ let mut next_idx = 0;
+ for batch in flatten_batches {
+ let row_cnt = batch.num_rows();
+
+ // handle very lopsided batch sizing
+ if partitions[next_idx].is_empty() {
+ partitions[next_idx].push(batch);
+ } else {
+ // have at least 1 batch per partition
+ let idx =
+ std::cmp::min(next_idx + 1, curr_row_count / target_partition_size);
+ if let Some(partition) = partitions.get_mut(idx) {
+ partition.push(batch);
+ } else {
+ partitions[target_partitions - 1].push(batch);
+ }
+ next_idx = idx;
+ }
Review Comment:
Note to self; this feels ugly and I should fix it 😅 .
--
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