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.
(
&self,
target_partitions: usize,
output_ordering: LexOrdering,
)
| 475 | /// as having too few batches to fulfill the `target_partitions` or if unable |
| 476 | /// to preserve output ordering. |
| 477 | fn repartition_preserving_order( |
| 478 | &self, |
| 479 | target_partitions: usize, |
| 480 | output_ordering: LexOrdering, |
| 481 | ) -> Result<Option<Vec<Vec<RecordBatch>>>> { |
| 482 | if !self.eq_properties().ordering_satisfy(output_ordering)? { |
| 483 | Ok(None) |
| 484 | } else { |
| 485 | let total_num_batches = |
| 486 | self.partitions.iter().map(|b| b.len()).sum::<usize>(); |
| 487 | if total_num_batches < target_partitions { |
| 488 | // no way to create the desired repartitioning |
| 489 | return Ok(None); |
| 490 | } |
| 491 | |
| 492 | let cnt_to_repartition = target_partitions - self.partitions.len(); |
| 493 | |
| 494 | // Label the current partitions and their order. |
| 495 | // Such that when we later split up the partitions into smaller sizes, we are maintaining the order. |
| 496 | let to_repartition = self |
| 497 | .partitions |
| 498 | .iter() |
| 499 | .enumerate() |
| 500 | .map(|(idx, batches)| RePartition { |
| 501 | idx: idx + (cnt_to_repartition * idx), // make space in ordering for split partitions |
| 502 | row_count: batches.iter().map(|batch| batch.num_rows()).sum(), |
| 503 | batches: batches.clone(), |
| 504 | }) |
| 505 | .collect_vec(); |
| 506 | |
| 507 | // Put all of the partitions into a heap ordered by `RePartition::partial_cmp`, which sizes |
| 508 | // by count of rows. |
| 509 | let mut max_heap = BinaryHeap::with_capacity(target_partitions); |
| 510 | for rep in to_repartition { |
| 511 | max_heap.push(CompareByRowCount(rep)); |
| 512 | } |
| 513 | |
| 514 | // Split the largest partitions into smaller partitions. Maintaining the output |
| 515 | // order of the partitions & newly created partitions. |
| 516 | let mut cannot_split_further = Vec::with_capacity(target_partitions); |
| 517 | for _ in 0..cnt_to_repartition { |
| 518 | // triggers loop for the cnt_to_repartition. So if need another 4 partitions, it attempts to split 4 times. |
| 519 | loop { |
| 520 | // Take the largest item off the heap, and attempt to split. |
| 521 | let Some(to_split) = max_heap.pop() else { |
| 522 | // Nothing left to attempt repartition. Break inner loop. |
| 523 | break; |
| 524 | }; |
| 525 | |
| 526 | // Split the partition. The new partitions will be ordered with idx and idx+1. |
| 527 | let mut new_partitions = to_split.into_inner().split(); |
| 528 | if new_partitions.len() > 1 { |
| 529 | for new_partition in new_partitions { |
| 530 | max_heap.push(CompareByRowCount(new_partition)); |
| 531 | } |
| 532 | // Successful repartition. Break inner loop, and return to outer `cnt_to_repartition` loop. |
| 533 | break; |
| 534 | } else { |