// SPDX-License-Identifier: Apache-2.0

// SPDX-FileCopyrightText: See the NOTICE file distributed with this work for additional information regarding copyright ownership

// SPDX-FileContributor: Licensed to the Apache Software Foundation (ASF) under one or more contributor license agreements

//! Logic for managing groups of [`PartitionedFile`]s in DataFusion

use crate::{FileRange, PartitionedFile};

use arrow::compute::SortOptions;

use datafusion_common::utils::compare_rows;

use datafusion_common::{ScalarValue, SplitPoint, Statistics};

use itertools::Itertools;

use std::cmp::{Ordering, min};

use std::collections::{BinaryHeap, HashMap};

use std::iter::repeat_with;

use std::mem;

use std::ops::{Deref, DerefMut, Index, IndexMut};

use std::sync::Arc;

/// Repartition input files into `target_partitions` partitions, if total file size exceed

/// `repartition_file_min_size`

///

/// This partitions evenly by file byte range, and does not have any knowledge

/// of how data is laid out in specific files. The specific `FileOpener` are

/// responsible for the actual partitioning on specific data source type. (e.g.

/// the `CsvOpener` will read lines overlap with byte range as well as

/// handle boundaries to ensure all lines will be read exactly once)

///

/// # Example

///

/// For example, if there are two files `A` and `B` that we wish to read with 4

/// partitions (with 4 threads) they will be divided as follows:

///

/// ```text

/// ┌ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ┐

/// ┌─────────────────┐

/// │ │ │ │

/// │ File A │

/// │ │ Range: 0-2MB │ │

/// │ │

/// │ └─────────────────┘ │

/// ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─

/// ┌─────────────────┐ ┌ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ┐

/// │ │ ┌─────────────────┐

/// │ │ │ │ │ │

/// │ │ │ File A │

/// │ │ │ │ Range 2-4MB │ │

/// │ │ │ │

/// │ │ │ └─────────────────┘ │

/// │ File A (7MB) │ ────────▶ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─

/// │ │ ┌ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ┐

/// │ │ ┌─────────────────┐

/// │ │ │ │ │ │

/// │ │ │ File A │

/// │ │ │ │ Range: 4-6MB │ │

/// │ │ │ │

/// │ │ │ └─────────────────┘ │

/// └─────────────────┘ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─

/// ┌─────────────────┐ ┌ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ┐

/// │ File B (1MB) │ ┌─────────────────┐

/// │ │ │ │ File A │ │

/// └─────────────────┘ │ Range: 6-7MB │

/// │ └─────────────────┘ │

/// ┌─────────────────┐

/// │ │ File B (1MB) │ │

/// │ │

/// │ └─────────────────┘ │

/// ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─

///

/// If target_partitions = 4,

/// divides into 4 groups

/// ```

///

/// # Maintaining Order

///

/// Within each group files are read sequentially. Thus, if the overall order of

/// tuples must be preserved, multiple files can not be mixed in the same group.

///

/// In this case, the code will split the largest files evenly into any

/// available empty groups, but the overall distribution may not be as even

/// as if the order did not need to be preserved.

///

/// ```text

/// ┌ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ┐

/// ┌─────────────────┐

/// │ │ │ │

/// │ File A │

/// │ │ Range: 0-2MB │ │

/// │ │

/// ┌─────────────────┐ │ └─────────────────┘ │

/// │ │ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─

/// │ │ ┌ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ┐

/// │ │ ┌─────────────────┐

/// │ │ │ │ │ │

/// │ │ │ File A │

/// │ │ │ │ Range 2-4MB │ │

/// │ File A (6MB) │ ────────▶ │ │

/// │ (ordered) │ │ └─────────────────┘ │

/// │ │ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─

/// │ │ ┌ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ┐

/// │ │ ┌─────────────────┐

/// │ │ │ │ │ │

/// │ │ │ File A │

/// │ │ │ │ Range: 4-6MB │ │

/// └─────────────────┘ │ │

/// ┌─────────────────┐ │ └─────────────────┘ │

/// │ File B (1MB) │ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─

/// │ (ordered) │ ┌ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ┐

/// └─────────────────┘ ┌─────────────────┐

/// │ │ File B (1MB) │ │

/// │ │

/// │ └─────────────────┘ │

/// ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─

///

/// If target_partitions = 4,

/// divides into 4 groups

/// ```

#[derive(Debug, Clone, Copy)]

pub struct FileGroupPartitioner {

/// how many partitions should be created

target_partitions: usize,

/// the minimum size for a file to be repartitioned.

repartition_file_min_size: usize,

/// if the order when reading the files must be preserved

preserve_order_within_groups: bool,

}

impl Default for FileGroupPartitioner {

fn default() -> Self {

Self::new()

}

}

impl FileGroupPartitioner {

/// Creates a new [`FileGroupPartitioner`] with default values:

/// 1. `target_partitions = 1`

/// 2. `repartition_file_min_size = 10MB`

/// 3. `preserve_order_within_groups = false`

pub fn new() -> Self {

Self {

target_partitions: 1,

repartition_file_min_size: 10 * 1024 * 1024,

preserve_order_within_groups: false,

}

}

/// Set the target partitions

pub fn with_target_partitions(mut self, target_partitions: usize) -> Self {

self.target_partitions = target_partitions;

self

}

/// Set the minimum size at which to repartition a file

pub fn with_repartition_file_min_size(

mut self,

repartition_file_min_size: usize,

) -> Self {

self.repartition_file_min_size = repartition_file_min_size;

self

}

/// Set whether the order of tuples within a file must be preserved

pub fn with_preserve_order_within_groups(

mut self,

preserve_order_within_groups: bool,

) -> Self {

self.preserve_order_within_groups = preserve_order_within_groups;

self

}

/// Repartition input files according to the settings on this [`FileGroupPartitioner`].

///

/// If no repartitioning is needed or possible, return `None`.

pub fn repartition_file_groups(

&self,

file_groups: &[FileGroup],

) -> Option<Vec<FileGroup>> {

if file_groups.is_empty() {

return None;

}

// special case when order must be preserved

if self.preserve_order_within_groups {

self.repartition_preserving_order(file_groups)

} else {

self.repartition_evenly_by_size(file_groups)

}

}

/// Evenly repartition files across partitions by size, ignoring any

/// existing grouping / ordering

fn repartition_evenly_by_size(

&self,

file_groups: &[FileGroup],

) -> Option<Vec<FileGroup>> {

let target_partitions = self.target_partitions;

let repartition_file_min_size = self.repartition_file_min_size;

let flattened_files = file_groups.iter().flat_map(FileGroup::iter).collect_vec();

let total_size = flattened_files

.iter()

.map(|f| f.effective_size())

.sum::<u64>();

if total_size < (repartition_file_min_size as u64) || total_size == 0 {

return None;

}

let target_partition_size = total_size.div_ceil(target_partitions as u64);

let current_partition_index: usize = 0;

let current_partition_size: u64 = 0;

// Partition byte range evenly for all `PartitionedFile`s

let repartitioned_files = flattened_files

.into_iter()

.scan(

(current_partition_index, current_partition_size),

|(current_partition_index, current_partition_size), source_file| {

let mut produced_files = vec![];

let (mut range_start, file_end) = source_file.range();

while range_start < file_end {

let range_end = min(

range_start

+ (target_partition_size - *current_partition_size),

file_end,

);

let mut produced_file = source_file.clone();

produced_file.range = Some(FileRange {

start: range_start as i64,

end: range_end as i64,

});

produced_files.push((*current_partition_index, produced_file));

if *current_partition_size + (range_end - range_start)

>= target_partition_size

{

*current_partition_index += 1;

*current_partition_size = 0;

} else {

*current_partition_size += range_end - range_start;

}

range_start = range_end;

}

Some(produced_files)

},

)

.flatten()

.chunk_by(|(partition_idx, _)| *partition_idx)

.into_iter()

.map(|(_, group)| FileGroup::new(group.map(|(_, vals)| vals).collect_vec()))

.collect_vec();

Some(repartitioned_files)

}

/// Redistribute file groups across size preserving order

fn repartition_preserving_order(

&self,

file_groups: &[FileGroup],

) -> Option<Vec<FileGroup>> {

// Can't repartition and preserve order if there are more groups

// than partitions

if file_groups.len() >= self.target_partitions {

return None;

}

let num_new_groups = self.target_partitions - file_groups.len();

// If there is only a single file

if file_groups.len() == 1 && file_groups[0].len() == 1 {

return self.repartition_evenly_by_size(file_groups);

}

// Find which files could be split (single file groups)

let mut heap: BinaryHeap<_> = file_groups

.iter()

.enumerate()

.filter_map(|(group_index, group)| {

// ignore groups that do not have exactly 1 file

if group.len() == 1 {

Some(ToRepartition {

source_index: group_index,

file_size: group[0].effective_size(),

new_groups: vec![group_index],

})

} else {

None

}

})

.map(CompareByRangeSize)

.collect();

// No files can be redistributed

if heap.is_empty() {

return None;

}

// Add new empty groups to which we will redistribute ranges of existing files

// Add new empty groups to which we will redistribute ranges of existing files

let mut file_groups: Vec<_> = file_groups

.iter()

.cloned()

.chain(repeat_with(|| FileGroup::new(Vec::new())).take(num_new_groups))

.collect();

// Divide up empty groups

for (group_index, group) in file_groups.iter().enumerate() {

if !group.is_empty() {

continue;

}

// Pick the file that has the largest ranges to read so far

let mut largest_group = heap.pop().unwrap();

largest_group.new_groups.push(group_index);

heap.push(largest_group);

}

// Distribute files to their newly assigned groups

while let Some(to_repartition) = heap.pop() {

let range_size = to_repartition.range_size();

let ToRepartition {

source_index,

file_size: _,

new_groups,

} = to_repartition.into_inner();

assert_eq!(file_groups[source_index].len(), 1);

let original_file = file_groups[source_index].pop().unwrap();

let last_group = new_groups.len() - 1;

let (mut range_start, file_end) = original_file.range();

let mut range_end = range_start + range_size;

for (i, group_index) in new_groups.into_iter().enumerate() {

let target_group = &mut file_groups[group_index];

assert!(target_group.is_empty());

// adjust last range to include the entire file

if i == last_group {

range_end = file_end;

}

target_group.push(

original_file

.clone()

.with_range(range_start as i64, range_end as i64),

);

range_start = range_end;

range_end += range_size;

}

}

Some(file_groups)

}

}

/// Represents a group of partitioned files that'll be processed by a single thread.

/// Maintains optional statistics across all files in the group.

///

/// # Statistics

///

/// The group-level [`FileGroup::file_statistics`] field contains merged statistics from all files

/// in the group for the **full table schema** (file columns + partition columns).

///

/// Partition column statistics are derived from the individual file partition values:

/// - `min` = minimum partition value across all files in the group, or absent if any

/// file in the group has a null partition value (see #26748)

/// - `max` = maximum partition value across all files in the group, or absent if any

/// file in the group has a null partition value

/// - `null_count` = the number of rows in files with a null partition value

///

/// This allows query optimizers to prune entire file groups based on partition bounds.

#[derive(Debug, Clone)]

pub struct FileGroup {

/// The files in this group

files: Vec<PartitionedFile>,

/// Optional statistics for the data across all files in the group.

///

/// These statistics cover the full table schema: file columns plus partition columns.

/// Partition column statistics are merged from individual [`PartitionedFile::statistics`],

/// which compute exact values from [`PartitionedFile::partition_values`].

statistics: Option<Arc<Statistics>>,

}

impl FileGroup {

/// Creates a new FileGroup from a vector of PartitionedFile objects

pub fn new(files: Vec<PartitionedFile>) -> Self {

Self {

files,

statistics: None,

}

}

/// Returns the number of files in this group

pub fn len(&self) -> usize {

self.files.len()

}

/// Set the statistics for this group

pub fn with_statistics(mut self, statistics: Arc<Statistics>) -> Self {

self.statistics = Some(statistics);

self

}

/// Returns a slice of the files in this group

pub fn files(&self) -> &[PartitionedFile] {

&self.files

}

pub fn iter(&self) -> impl Iterator<Item = &PartitionedFile> {

self.files.iter()

}

pub fn into_inner(self) -> Vec<PartitionedFile> {

self.files

}

pub fn is_empty(&self) -> bool {

self.files.is_empty()

}

/// Removes the last element from the files vector and returns it, or None if empty

pub fn pop(&mut self) -> Option<PartitionedFile> {

self.files.pop()

}

/// Adds a file to the group

pub fn push(&mut self, partitioned_file: PartitionedFile) {

self.files.push(partitioned_file);

}

/// Get the specific file statistics for the given index

/// If the index is None, return the `FileGroup` statistics

pub fn file_statistics(&self, index: Option<usize>) -> Option<&Statistics> {

if let Some(index) = index {

self.files.get(index).and_then(|f| f.statistics.as_deref())

} else {

self.statistics.as_deref()

}

}

/// Get the mutable reference to the statistics for this group

pub fn statistics_mut(&mut self) -> Option<&mut Statistics> {

self.statistics.as_mut().map(Arc::make_mut)

}

/// Partition the list of files into `n` groups

pub fn split_files(mut self, n: usize) -> Vec<FileGroup> {

if self.is_empty() {

return vec![];

}

// ObjectStore::list does not guarantee any consistent order and for some

// implementations such as LocalFileSystem, it may be inconsistent. Thus

// Sort files by path to ensure consistent plans when run more than once.

self.files.sort_by(|a, b| a.path().cmp(b.path()));

// effectively this is div with rounding up instead of truncating

let chunk_size = self.len().div_ceil(n);

let mut chunks = Vec::with_capacity(n);

let mut current_chunk = Vec::with_capacity(chunk_size);

for file in self.files.drain(..) {

current_chunk.push(file);

if current_chunk.len() == chunk_size {

let full_chunk = FileGroup::new(mem::replace(

&mut current_chunk,

Vec::with_capacity(chunk_size),

));

chunks.push(full_chunk);

}

}

if !current_chunk.is_empty() {

chunks.push(FileGroup::new(current_chunk))

}

chunks

}

/// Groups files by their partition values, ensuring all files with same

/// partition values are in the same group.

///

/// Note: May return fewer groups than `max_target_partitions` when the

/// number of unique partition values is less than the target.

///

/// Wrapper around [`Self::group_by_partition_values_with_split_points`] that discards

/// the split points.

pub fn group_by_partition_values(

self,

max_target_partitions: usize,

) -> Vec<FileGroup> {

self.group_by_partition_values_with_split_points(max_target_partitions)

.0

}

/// Groups files by partition value and returns the [`SplitPoint`]s between groups.

///

/// Values are sorted with [`SortOptions::default()`] and cut into at most

/// `max_target_partitions` contiguous chunks so routing a file's `partition_values`

/// through the split points yields the group it was placed in.

#[allow(clippy::allow_attributes, clippy::mutable_key_type)]

pub fn group_by_partition_values_with_split_points(

self,

max_target_partitions: usize,

) -> (Vec<FileGroup>, Vec<SplitPoint>) {

if self.is_empty() || max_target_partitions == 0 {

return (vec![], vec![]);

}

let mut partition_groups: HashMap<Vec<ScalarValue>, Vec<PartitionedFile>> =

HashMap::new();

for file in self.files {

partition_groups

.entry(file.partition_values.clone())

.or_default()

.push(file);

}

let num_unique_partitions = partition_groups.len();

// Sort for deterministic bucket assignment across query executions. Must match

// the ordering the caller declares, otherwise the split points would not

// describe the groups.

let mut sorted_partitions: Vec<_> = partition_groups.into_iter().collect();

let sort_options =

vec![

SortOptions::default();

sorted_partitions.first().map(|(k, _)| k.len()).unwrap_or(0)

];

sorted_partitions.sort_by(|a, b| {

compare_rows(&a.0, &b.0, &sort_options).unwrap_or(Ordering::Equal)

});

let bucket_count = min(num_unique_partitions, max_target_partitions);

let mut groups = Vec::with_capacity(bucket_count);

let mut split_points = Vec::with_capacity(bucket_count.saturating_sub(1));

// Every chunk is non-empty because `bucket_count <= num_unique_partitions`, so

// its first tuple is a well defined lower bound.

let mut iter = sorted_partitions.into_iter();

for bucket in 0..bucket_count {

let start = bucket * num_unique_partitions / bucket_count;

let end = (bucket + 1) * num_unique_partitions / bucket_count;

let mut files = Vec::new();

for offset in start..end {

let (values, group_files) = iter

.next()

.expect("contiguous chunking never runs past the sorted partitions");

if offset == start && bucket > 0 {

split_points.push(SplitPoint::new(values));

}

files.extend(group_files);

}

groups.push(FileGroup::new(files));

}

debug_assert_eq!(split_points.len() + 1, groups.len());

(groups, split_points)

}

}

impl Index<usize> for FileGroup {

type Output = PartitionedFile;

fn index(&self, index: usize) -> &Self::Output {

&self.files[index]

}

}

impl IndexMut<usize> for FileGroup {

fn index_mut(&mut self, index: usize) -> &mut Self::Output {

&mut self.files[index]

}

}

impl FromIterator<PartitionedFile> for FileGroup {

fn from_iter<I: IntoIterator<Item = PartitionedFile>>(iter: I) -> Self {

let files = iter.into_iter().collect();

FileGroup::new(files)

}

}

impl From<Vec<PartitionedFile>> for FileGroup {

fn from(files: Vec<PartitionedFile>) -> Self {

FileGroup::new(files)

}

}

impl Default for FileGroup {

fn default() -> Self {

Self::new(Vec::new())

}

}

/// Tracks how a individual file will be repartitioned

#[derive(Debug, Clone)]

struct ToRepartition {

/// the index from which the original file will be taken

source_index: usize,

/// the size of the original file

file_size: u64,

/// indexes of which group(s) will this be distributed to (including `source_index`)

new_groups: Vec<usize>,

}

impl ToRepartition {

/// How big will each file range be when this file is read in its new groups?

fn range_size(&self) -> u64 {

self.file_size / (self.new_groups.len() as u64)

}

}

struct CompareByRangeSize(ToRepartition);

impl CompareByRangeSize {

fn into_inner(self) -> ToRepartition {

self.0

}

}

impl Ord for CompareByRangeSize {

fn cmp(&self, other: &Self) -> Ordering {

self.0.range_size().cmp(&other.0.range_size())

}

}

impl PartialOrd for CompareByRangeSize {

fn partial_cmp(&self, other: &Self) -> Option<Ordering> {

Some(self.cmp(other))

}

}

impl PartialEq for CompareByRangeSize {

fn eq(&self, other: &Self) -> bool {

// PartialEq must be consistent with PartialOrd

self.cmp(other) == Ordering::Equal

}

}

impl Eq for CompareByRangeSize {}

impl Deref for CompareByRangeSize {

type Target = ToRepartition;

fn deref(&self) -> &Self::Target {

&self.0

}

}

impl DerefMut for CompareByRangeSize {

fn deref_mut(&mut self) -> &mut Self::Target {

&mut self.0

}

}

#[cfg(test)]

mod test {

use super::*;

use datafusion_physical_expr::RangePartitioning;

use datafusion_physical_expr::expressions::Column;

use datafusion_physical_expr_common::sort_expr::{LexOrdering, PhysicalSortExpr};

/// Empty file won't get partitioned

#[test]

fn repartition_empty_file_only() {

let partitioned_file_empty = pfile("empty", 0);

let file_group = vec![FileGroup::new(vec![partitioned_file_empty])];

let partitioned_files = FileGroupPartitioner::new()

.with_target_partitions(4)

.with_repartition_file_min_size(0)

.repartition_file_groups(&file_group);

assert_partitioned_files(None, partitioned_files);

}

/// Repartition when there is a empty file in file groups

#[test]

fn repartition_empty_files() {

let pfile_a = pfile("a", 10);

let pfile_b = pfile("b", 10);

let pfile_empty = pfile("empty", 0);

let empty_first = vec![

FileGroup::new(vec![pfile_empty.clone()]),

FileGroup::new(vec![pfile_a.clone()]),

FileGroup::new(vec![pfile_b.clone()]),

];

let empty_middle = vec![

FileGroup::new(vec![pfile_a.clone()]),

FileGroup::new(vec![pfile_empty.clone()]),

FileGroup::new(vec![pfile_b.clone()]),

];

let empty_last = vec![

FileGroup::new(vec![pfile_a]),

FileGroup::new(vec![pfile_b]),

FileGroup::new(vec![pfile_empty]),

];

// Repartition file groups into x partitions

let expected_2 = vec![

FileGroup::new(vec![pfile("a", 10).with_range(0, 10)]),

FileGroup::new(vec![pfile("b", 10).with_range(0, 10)]),

];

let expected_3 = vec![

FileGroup::new(vec![pfile("a", 10).with_range(0, 7)]),

FileGroup::new(vec![

pfile("a", 10).with_range(7, 10),

pfile("b", 10).with_range(0, 4),

]),

FileGroup::new(vec![pfile("b", 10).with_range(4, 10)]),

];

let file_groups_tests = [empty_first, empty_middle, empty_last];

for fg in file_groups_tests {

let all_expected = [(2, expected_2.clone()), (3, expected_3.clone())];

for (n_partition, expected) in all_expected {

let actual = FileGroupPartitioner::new()

.with_target_partitions(n_partition)

.with_repartition_file_min_size(10)

.repartition_file_groups(&fg);

assert_partitioned_files(Some(expected), actual);

}

}

}

#[test]

fn repartition_single_file() {

// Single file, single partition into multiple partitions

let single_partition = vec![FileGroup::new(vec![pfile("a", 123)])];

let actual = FileGroupPartitioner::new()

.with_target_partitions(4)

.with_repartition_file_min_size(10)

.repartition_file_groups(&single_partition);

let expected = Some(vec![

FileGroup::new(vec![pfile("a", 123).with_range(0, 31)]),

FileGroup::new(vec![pfile("a", 123).with_range(31, 62)]),

FileGroup::new(vec![pfile("a", 123).with_range(62, 93)]),

FileGroup::new(vec![pfile("a", 123).with_range(93, 123)]),

]);

assert_partitioned_files(expected, actual);

}

#[test]

fn repartition_single_file_with_range() {

// Single file, single partition into multiple partitions

let single_partition =

vec![FileGroup::new(vec![pfile("a", 123).with_range(0, 123)])];

let actual = FileGroupPartitioner::new()

.with_target_partitions(4)

.with_repartition_file_min_size(10)

.repartition_file_groups(&single_partition);

let expected = Some(vec![

FileGroup::new(vec![pfile("a", 123).with_range(0, 31)]),

FileGroup::new(vec![pfile("a", 123).with_range(31, 62)]),

FileGroup::new(vec![pfile("a", 123).with_range(62, 93)]),

FileGroup::new(vec![pfile("a", 123).with_range(93, 123)]),

]);

assert_partitioned_files(expected, actual);

}

#[test]

fn repartition_single_file_with_incomplete_range() {

// Single file, single partition into multiple partitions

let single_partition =

vec![FileGroup::new(vec![pfile("a", 123).with_range(10, 100)])];

let actual = FileGroupPartitioner::new()

.with_target_partitions(4)

.with_repartition_file_min_size(10)

.repartition_file_groups(&single_partition);

let expected = Some(vec![

FileGroup::new(vec![pfile("a", 123).with_range(10, 33)]),

FileGroup::new(vec![pfile("a", 123).with_range(33, 56)]),

FileGroup::new(vec![pfile("a", 123).with_range(56, 79)]),

FileGroup::new(vec![pfile("a", 123).with_range(79, 100)]),

]);

assert_partitioned_files(expected, actual);

}

#[test]

fn repartition_single_file_duplicated_with_range() {

// Single file, two partitions into multiple partitions

let single_partition = vec![FileGroup::new(vec![

pfile("a", 100).with_range(0, 50),

pfile("a", 100).with_range(50, 100),

])];

let actual = FileGroupPartitioner::new()

.with_target_partitions(4)

.with_repartition_file_min_size(10)

.repartition_file_groups(&single_partition);

let expected = Some(vec![

FileGroup::new(vec![pfile("a", 100).with_range(0, 25)]),

FileGroup::new(vec![pfile("a", 100).with_range(25, 50)]),

FileGroup::new(vec![pfile("a", 100).with_range(50, 75)]),

FileGroup::new(vec![pfile("a", 100).with_range(75, 100)]),

]);

assert_partitioned_files(expected, actual);

}

#[test]

fn repartition_too_much_partitions() {

// Single file, single partition into 96 partitions

let partitioned_file = pfile("a", 8);

let single_partition = vec![FileGroup::new(vec![partitioned_file])];

let actual = FileGroupPartitioner::new()

.with_target_partitions(96)

.with_repartition_file_min_size(5)

.repartition_file_groups(&single_partition);

let expected = Some(vec![

FileGroup::new(vec![pfile("a", 8).with_range(0, 1)]),

FileGroup::new(vec![pfile("a", 8).with_range(1, 2)]),

FileGroup::new(vec![pfile("a", 8).with_range(2, 3)]),

FileGroup::new(vec![pfile("a", 8).with_range(3, 4)]),

FileGroup::new(vec![pfile("a", 8).with_range(4, 5)]),

FileGroup::new(vec![pfile("a", 8).with_range(5, 6)]),

FileGroup::new(vec![pfile("a", 8).with_range(6, 7)]),

FileGroup::new(vec![pfile("a", 8).with_range(7, 8)]),

]);

assert_partitioned_files(expected, actual);

}

#[test]

fn repartition_multiple_partitions() {

// Multiple files in single partition after redistribution

let source_partitions = vec![

FileGroup::new(vec![pfile("a", 40)]),

FileGroup::new(vec![pfile("b", 60)]),

];

let actual = FileGroupPartitioner::new()

.with_target_partitions(3)

.with_repartition_file_min_size(10)

.repartition_file_groups(&source_partitions);

let expected = Some(vec![

FileGroup::new(vec![pfile("a", 40).with_range(0, 34)]),

FileGroup::new(vec![

pfile("a", 40).with_range(34, 40),

pfile("b", 60).with_range(0, 28),

]),

FileGroup::new(vec![pfile("b", 60).with_range(28, 60)]),

]);

assert_partitioned_files(expected, actual);

}

#[test]

fn repartition_same_num_partitions() {

// "Rebalance" files across partitions

let source_partitions = vec![

FileGroup::new(vec![pfile("a", 40)]),

FileGroup::new(vec![pfile("b", 60)]),

];

let actual = FileGroupPartitioner::new()

.with_target_partitions(2)

.with_repartition_file_min_size(10)

.repartition_file_groups(&source_partitions);

let expected = Some(vec![

FileGroup::new(vec![

pfile("a", 40).with_range(0, 40),

pfile("b", 60).with_range(0, 10),

]),

FileGroup::new(vec![pfile("b", 60).with_range(10, 60)]),

]);

assert_partitioned_files(expected, actual);

}

#[test]

fn repartition_no_action_min_size() {

// No action due to target_partition_size

let single_partition = vec![FileGroup::new(vec![pfile("a", 123)])];

let actual = FileGroupPartitioner::new()

.with_target_partitions(65)

.with_repartition_file_min_size(500)

.repartition_file_groups(&single_partition);

assert_partitioned_files(None, actual)

}

#[test]

fn repartition_no_action_zero_files() {

// No action due to no files

let empty_partition = vec![];

let partitioner = FileGroupPartitioner::new()

.with_target_partitions(65)

.with_repartition_file_min_size(500);

assert_partitioned_files(None, repartition_test(partitioner, empty_partition))

}

#[test]

fn repartition_ordered_no_action_too_few_partitions() {

// No action as there are no new groups to redistribute to

let input_partitions = vec![

FileGroup::new(vec![pfile("a", 100)]),

FileGroup::new(vec![pfile("b", 200)]),

];

let actual = FileGroupPartitioner::new()

.with_preserve_order_within_groups(true)

.with_target_partitions(2)

.with_repartition_file_min_size(10)

.repartition_file_groups(&input_partitions);

assert_partitioned_files(None, actual)

}

#[test]

fn repartition_ordered_no_action_file_too_small() {

// No action as there are no new groups to redistribute to

let single_partition = vec![FileGroup::new(vec![pfile("a", 100)])];

let actual = FileGroupPartitioner::new()

.with_preserve_order_within_groups(true)

.with_target_partitions(2)

// file is too small to repartition

.with_repartition_file_min_size(1000)

.repartition_file_groups(&single_partition);

assert_partitioned_files(None, actual)

}

#[test]

fn repartition_ordered_one_large_file() {

// "Rebalance" the single large file across partitions

let source_partitions = vec![FileGroup::new(vec![pfile("a", 100)])];

let actual = FileGroupPartitioner::new()

.with_preserve_order_within_groups(true)

.with_target_partitions(3)

.with_repartition_file_min_size(10)

.repartition_file_groups(&source_partitions);

let expected = Some(vec![

FileGroup::new(vec![pfile("a", 100).with_range(0, 34)]),

FileGroup::new(vec![pfile("a", 100).with_range(34, 68)]),

FileGroup::new(vec![pfile("a", 100).with_range(68, 100)]),

]);

assert_partitioned_files(expected, actual);

}

#[test]

fn repartition_ordered_one_large_file_with_range() {

// "Rebalance" the single large file across partitions

let source_partitions =

vec![FileGroup::new(vec![pfile("a", 100).with_range(0, 100)])];

let actual = FileGroupPartitioner::new()

.with_preserve_order_within_groups(true)

.with_target_partitions(3)

.with_repartition_file_min_size(10)

.repartition_file_groups(&source_partitions);

let expected = Some(vec![

FileGroup::new(vec![pfile("a", 100).with_range(0, 34)]),

FileGroup::new(vec![pfile("a", 100).with_range(34, 68)]),

FileGroup::new(vec![pfile("a", 100).with_range(68, 100)]),

]);

assert_partitioned_files(expected, actual);

}

#[test]

fn repartition_ordered_one_large_one_small_file() {

// "Rebalance" the single large file across empty partitions, but can't split

// small file

let source_partitions = vec![

FileGroup::new(vec![pfile("a", 100)]),

FileGroup::new(vec![pfile("b", 30)]),

];

let actual = FileGroupPartitioner::new()

.with_preserve_order_within_groups(true)

.with_target_partitions(4)

.with_repartition_file_min_size(10)

.repartition_file_groups(&source_partitions);

let expected = Some(vec![

// scan first third of "a"

FileGroup::new(vec![pfile("a", 100).with_range(0, 33)]),

// only b in this group (can't do this)

FileGroup::new(vec![pfile("b", 30).with_range(0, 30)]),

// second third of "a"

FileGroup::new(vec![pfile("a", 100).with_range(33, 66)]),

// final third of "a"

FileGroup::new(vec![pfile("a", 100).with_range(66, 100)]),

]);

assert_partitioned_files(expected, actual);

}

#[test]

fn repartition_ordered_one_large_one_small_file_with_full_range() {

// "Rebalance" the single large file across empty partitions, but can't split

// small file

let source_partitions = vec![

FileGroup::new(vec![pfile("a", 100).with_range(0, 100)]),

FileGroup::new(vec![pfile("b", 30)]),