use crate::{DeltaFunnelError, error::DeltaScanFileTaskPartitionPlanningSnafu};
use super::file_task::DeltaScanFileTask;
#[allow(dead_code)]
pub(crate) struct DeltaScanFileTaskPartitionPlanRequest {
pub(crate) source_name: String,
pub(crate) table_uri: String,
pub(crate) snapshot_version: u64,
pub(crate) scan_metadata_exhausted: bool,
pub(crate) file_tasks: Vec<DeltaScanFileTask>,
pub(crate) options: DeltaScanFileTaskPartitionOptions,
}
#[allow(dead_code)]
pub(crate) struct DeltaScanFileTaskPartitionPlan {
pub(crate) source_name: String,
pub(crate) table_uri: String,
pub(crate) snapshot_version: u64,
pub(crate) partitions: Vec<DeltaScanFileTaskPartition>,
pub(crate) scan_metadata_exhausted: bool,
pub(crate) estimated_bytes: Option<u64>,
pub(crate) estimated_rows: Option<u64>,
}
#[allow(dead_code)]
pub(crate) struct DeltaScanFileTaskPartition {
pub(crate) file_tasks: Vec<DeltaScanFileTask>,
pub(crate) estimated_bytes: Option<u64>,
pub(crate) estimated_rows: Option<u64>,
}
#[allow(dead_code)]
pub(crate) struct DeltaScanFileTaskPartitionOptions {
pub(crate) target_partitions: usize,
}
impl DeltaScanFileTaskPartitionOptions {
#[allow(dead_code)]
pub(crate) fn validate_for_scan_context(
&self,
source_name: &str,
table_uri: &str,
snapshot_version: u64,
) -> Result<(), DeltaFunnelError> {
validate_partition_options(source_name, table_uri, snapshot_version, self)
}
}
impl DeltaScanFileTaskPartitionPlan {
#[allow(dead_code)]
pub(crate) fn try_new(
request: DeltaScanFileTaskPartitionPlanRequest,
) -> Result<Self, DeltaFunnelError> {
let DeltaScanFileTaskPartitionPlanRequest {
source_name,
table_uri,
snapshot_version,
scan_metadata_exhausted,
file_tasks,
options,
} = request;
options.validate_for_scan_context(&source_name, &table_uri, snapshot_version)?;
validate_file_task_context(&source_name, &table_uri, snapshot_version, &file_tasks)?;
let estimated_bytes = sum_task_estimate(
&source_name,
&table_uri,
snapshot_version,
"estimated bytes",
file_tasks.iter().map(|file_task| file_task.estimated_bytes),
)?;
let estimated_rows = sum_task_estimate(
&source_name,
&table_uri,
snapshot_version,
"estimated rows",
file_tasks.iter().map(|file_task| file_task.estimated_rows),
)?;
let partitions = if file_tasks.is_empty() {
Vec::new()
} else if matches!(estimated_bytes, Some(0) | None) {
group_by_file_count(
&source_name,
&table_uri,
snapshot_version,
file_tasks,
options.target_partitions,
)?
} else {
group_by_estimated_bytes(
&source_name,
&table_uri,
snapshot_version,
file_tasks,
options.target_partitions,
)?
};
Ok(Self {
source_name,
table_uri,
snapshot_version,
partitions,
scan_metadata_exhausted,
estimated_bytes,
estimated_rows,
})
}
}
fn validate_partition_options(
source_name: &str,
table_uri: &str,
snapshot_version: u64,
options: &DeltaScanFileTaskPartitionOptions,
) -> Result<(), DeltaFunnelError> {
if options.target_partitions == 0 {
return partition_planning_error(
source_name,
table_uri,
snapshot_version,
"target_partitions must be greater than zero",
);
}
Ok(())
}
fn validate_file_task_context(
source_name: &str,
table_uri: &str,
snapshot_version: u64,
file_tasks: &[DeltaScanFileTask],
) -> Result<(), DeltaFunnelError> {
for file_task in file_tasks {
if file_task.source_name != source_name {
return partition_planning_error(
source_name,
table_uri,
snapshot_version,
format!(
"file task `{}` belongs to source `{}`, not `{source_name}`",
file_task.path, file_task.source_name
),
);
}
if file_task.table_uri != table_uri {
return partition_planning_error(
source_name,
table_uri,
snapshot_version,
format!(
"file task `{}` belongs to table URI `{}`, not `{table_uri}`",
file_task.path, file_task.table_uri
),
);
}
if file_task.snapshot_version != snapshot_version {
return partition_planning_error(
source_name,
table_uri,
snapshot_version,
format!(
"file task `{}` belongs to snapshot version {}, not {snapshot_version}",
file_task.path, file_task.snapshot_version
),
);
}
}
Ok(())
}
fn group_by_estimated_bytes(
source_name: &str,
table_uri: &str,
snapshot_version: u64,
file_tasks: Vec<DeltaScanFileTask>,
target_partitions: usize,
) -> Result<Vec<DeltaScanFileTaskPartition>, DeltaFunnelError> {
let output_limit = target_partitions.min(file_tasks.len());
let Some(total_bytes) = sum_task_estimate(
source_name,
table_uri,
snapshot_version,
"estimated bytes",
file_tasks.iter().map(|file_task| file_task.estimated_bytes),
)?
else {
return partition_planning_error(
source_name,
table_uri,
snapshot_version,
"known-size grouping requires every file task to have estimated bytes",
);
};
let target_bytes = total_bytes.div_ceil(output_limit as u64);
let mut partitions = Vec::new();
let mut current_file_tasks = Vec::new();
let mut current_bytes = 0_u64;
for file_task in file_tasks {
let Some(file_bytes) = file_task.estimated_bytes else {
return partition_planning_error(
source_name,
table_uri,
snapshot_version,
"known-size grouping requires every file task to have estimated bytes",
);
};
let can_start_next_partition = !current_file_tasks.is_empty()
&& partitions.len() + 1 < output_limit
&& current_bytes.saturating_add(file_bytes) > target_bytes;
if can_start_next_partition {
partitions.push(build_partition(
source_name,
table_uri,
snapshot_version,
current_file_tasks,
)?);
current_file_tasks = Vec::new();
current_bytes = 0;
}
current_bytes = match current_bytes.checked_add(file_bytes) {
Some(bytes) => bytes,
None => {
return partition_planning_error(
source_name,
table_uri,
snapshot_version,
"partition estimated bytes overflowed u64",
);
}
};
current_file_tasks.push(file_task);
}
if !current_file_tasks.is_empty() {
partitions.push(build_partition(
source_name,
table_uri,
snapshot_version,
current_file_tasks,
)?);
}
Ok(partitions)
}
fn group_by_file_count(
source_name: &str,
table_uri: &str,
snapshot_version: u64,
file_tasks: Vec<DeltaScanFileTask>,
target_partitions: usize,
) -> Result<Vec<DeltaScanFileTaskPartition>, DeltaFunnelError> {
let output_limit = target_partitions.min(file_tasks.len());
let mut partitions = Vec::new();
let mut file_tasks = file_tasks.into_iter();
let mut remaining_files = file_tasks.len();
for partition_index in 0..output_limit {
let remaining_partitions = output_limit - partition_index;
let take_count = remaining_files.div_ceil(remaining_partitions);
let mut partition_file_tasks = Vec::with_capacity(take_count);
for _ in 0..take_count {
let Some(file_task) = file_tasks.next() else {
return partition_planning_error(
source_name,
table_uri,
snapshot_version,
"file-count grouping exhausted file tasks unexpectedly",
);
};
partition_file_tasks.push(file_task);
}
remaining_files -= take_count;
partitions.push(build_partition(
source_name,
table_uri,
snapshot_version,
partition_file_tasks,
)?);
}
Ok(partitions)
}
fn build_partition(
source_name: &str,
table_uri: &str,
snapshot_version: u64,
file_tasks: Vec<DeltaScanFileTask>,
) -> Result<DeltaScanFileTaskPartition, DeltaFunnelError> {
let estimated_bytes = sum_task_estimate(
source_name,
table_uri,
snapshot_version,
"estimated bytes",
file_tasks.iter().map(|file_task| file_task.estimated_bytes),
)?;
let estimated_rows = sum_task_estimate(
source_name,
table_uri,
snapshot_version,
"estimated rows",
file_tasks.iter().map(|file_task| file_task.estimated_rows),
)?;
Ok(DeltaScanFileTaskPartition {
file_tasks,
estimated_bytes,
estimated_rows,
})
}
fn sum_task_estimate(
source_name: &str,
table_uri: &str,
snapshot_version: u64,
estimate_name: &str,
estimates: impl IntoIterator<Item = Option<u64>>,
) -> Result<Option<u64>, DeltaFunnelError> {
let mut total = 0_u64;
for estimate in estimates {
let Some(estimate) = estimate else {
return Ok(None);
};
total = match total.checked_add(estimate) {
Some(total) => total,
None => {
return partition_planning_error(
source_name,
table_uri,
snapshot_version,
format!("{estimate_name} overflowed u64"),
);
}
};
}
Ok(Some(total))
}
fn partition_planning_error<T>(
source_name: &str,
table_uri: &str,
snapshot_version: u64,
reason: impl Into<String>,
) -> Result<T, DeltaFunnelError> {
DeltaScanFileTaskPartitionPlanningSnafu {
source_name: source_name.to_owned(),
table_uri: table_uri.to_owned(),
snapshot_version,
reason: reason.into(),
}
.fail()
}
#[cfg(test)]
mod tests {
use std::collections::HashMap;
use crate::table_formats::{
KernelPhysicalToLogicalTransform, KernelScanDeletionVectorMetadata, KernelScanFileMetadata,
KernelScanFileStats,
};
use super::*;
fn plan_request(
file_tasks: Vec<DeltaScanFileTask>,
target_partitions: usize,
) -> DeltaScanFileTaskPartitionPlanRequest {
DeltaScanFileTaskPartitionPlanRequest {
source_name: "orders".to_owned(),
table_uri: "file:///tmp/table".to_owned(),
snapshot_version: 42,
scan_metadata_exhausted: true,
file_tasks,
options: DeltaScanFileTaskPartitionOptions { target_partitions },
}
}
fn file_task(
path: &str,
estimated_bytes: Option<u64>,
estimated_rows: Option<u64>,
) -> Result<DeltaScanFileTask, DeltaFunnelError> {
file_task_with_metadata(
path,
estimated_bytes,
estimated_rows,
KernelScanDeletionVectorMetadata::NotPresent,
KernelPhysicalToLogicalTransform::NotRequired,
)
}
fn file_task_with_metadata(
path: &str,
estimated_bytes: Option<u64>,
estimated_rows: Option<u64>,
deletion_vector: KernelScanDeletionVectorMetadata,
physical_to_logical_transform: KernelPhysicalToLogicalTransform,
) -> Result<DeltaScanFileTask, DeltaFunnelError> {
let size = match i64::try_from(estimated_bytes.unwrap_or(10)) {
Ok(size) => size,
Err(_) => {
return partition_planning_error(
"orders",
"file:///tmp/table",
42,
"test file size does not fit i64",
);
}
};
let mut task = DeltaScanFileTask::from_kernel_metadata(
"orders",
"file:///tmp/table",
42,
KernelScanFileMetadata {
path: path.to_owned(),
size,
modification_time: 1587968586000,
stats: Some(KernelScanFileStats {
num_records: estimated_rows.unwrap_or(1),
}),
deletion_vector,
physical_to_logical_transform,
partition_values: HashMap::from([("region".to_owned(), "us-west".to_owned())]),
},
)?;
task.estimated_bytes = estimated_bytes;
task.estimated_rows = estimated_rows;
if estimated_rows.is_none() {
task.stats = None;
}
Ok(task)
}
fn plan_partition_paths(plan: &DeltaScanFileTaskPartitionPlan) -> Vec<Vec<&str>> {
plan.partitions
.iter()
.map(|partition| {
partition
.file_tasks
.iter()
.map(|file_task| file_task.path.as_str())
.collect()
})
.collect()
}
#[test]
fn partition_plan_rejects_zero_target_partitions() -> Result<(), Box<dyn std::error::Error>> {
let error = match DeltaScanFileTaskPartitionPlan::try_new(plan_request(Vec::new(), 0)) {
Ok(_) => return Err("expected target_partitions validation to fail".into()),
Err(error) => error,
};
assert!(error.to_string().contains("target_partitions"));
Ok(())
}
#[test]
fn partition_plan_rejects_mismatched_file_task_context()
-> Result<(), Box<dyn std::error::Error>> {
let mut mismatched_task = file_task("part-0.parquet", Some(10), Some(1))?;
mismatched_task.snapshot_version = 43;
let error =
match DeltaScanFileTaskPartitionPlan::try_new(plan_request(vec![mismatched_task], 1)) {
Ok(_) => return Err("expected file task context validation to fail".into()),
Err(error) => error,
};
assert!(error.to_string().contains("snapshot version"));
assert!(error.to_string().contains("part-0.parquet"));
Ok(())
}
#[test]
fn oversized_known_size_file_stays_whole() -> Result<(), Box<dyn std::error::Error>> {
let plan = DeltaScanFileTaskPartitionPlan::try_new(plan_request(
vec![
file_task("huge.parquet", Some(1_000), Some(100))?,
file_task("small-0.parquet", Some(10), Some(1))?,
file_task("small-1.parquet", Some(10), Some(1))?,
],
2,
))?;
assert_eq!(
plan_partition_paths(&plan),
vec![
vec!["huge.parquet"],
vec!["small-0.parquet", "small-1.parquet"]
]
);
assert_eq!(plan.partitions[0].estimated_bytes, Some(1_000));
assert_eq!(plan.partitions[1].estimated_bytes, Some(20));
Ok(())
}
#[test]
fn empty_file_task_list_returns_empty_partition_plan() -> Result<(), Box<dyn std::error::Error>>
{
let plan = DeltaScanFileTaskPartitionPlan::try_new(plan_request(Vec::new(), 4))?;
assert_eq!(plan.source_name, "orders");
assert_eq!(plan.table_uri, "file:///tmp/table");
assert_eq!(plan.snapshot_version, 42);
assert!(plan.scan_metadata_exhausted);
assert!(plan.partitions.is_empty());
assert_eq!(plan.estimated_bytes, Some(0));
assert_eq!(plan.estimated_rows, Some(0));
Ok(())
}
#[test]
fn known_size_files_group_by_estimated_bytes() -> Result<(), Box<dyn std::error::Error>> {
let plan = DeltaScanFileTaskPartitionPlan::try_new(plan_request(
vec![
file_task("large.parquet", Some(90), Some(9))?,
file_task("small-1.parquet", Some(10), Some(1))?,
file_task("small-2.parquet", Some(10), Some(1))?,
file_task("small-3.parquet", Some(10), Some(1))?,
],
2,
))?;
assert_eq!(
plan_partition_paths(&plan),
vec![
vec!["large.parquet"],
vec!["small-1.parquet", "small-2.parquet", "small-3.parquet"],
]
);
assert_eq!(
plan.partitions
.iter()
.map(|partition| partition.estimated_bytes)
.collect::<Vec<_>>(),
vec![Some(90), Some(30)]
);
assert_eq!(plan.estimated_bytes, Some(120));
assert_eq!(plan.estimated_rows, Some(12));
Ok(())
}
#[test]
fn known_size_grouping_can_emit_fewer_partitions_than_requested()
-> Result<(), Box<dyn std::error::Error>> {
let plan = DeltaScanFileTaskPartitionPlan::try_new(plan_request(
vec![
file_task("part-0.parquet", Some(10), Some(1))?,
file_task("part-1.parquet", Some(10), Some(1))?,
],
8,
))?;
assert_eq!(plan.partitions.len(), 2);
assert_eq!(
plan_partition_paths(&plan),
vec![vec!["part-0.parquet"], vec!["part-1.parquet"]]
);
Ok(())
}
#[test]
fn unknown_size_files_fallback_to_file_count_balancing()
-> Result<(), Box<dyn std::error::Error>> {
let plan = DeltaScanFileTaskPartitionPlan::try_new(plan_request(
vec![
file_task("part-0.parquet", None, Some(1))?,
file_task("part-1.parquet", Some(10), Some(1))?,
file_task("part-2.parquet", Some(10), Some(1))?,
file_task("part-3.parquet", Some(10), Some(1))?,
file_task("part-4.parquet", Some(10), Some(1))?,
],
2,
))?;
assert_eq!(
plan_partition_paths(&plan),
vec![
vec!["part-0.parquet", "part-1.parquet", "part-2.parquet"],
vec!["part-3.parquet", "part-4.parquet"],
]
);
assert_eq!(plan.estimated_bytes, None);
assert_eq!(plan.partitions[0].estimated_bytes, None);
assert_eq!(plan.partitions[1].estimated_bytes, Some(20));
Ok(())
}
#[test]
fn zero_byte_files_are_preserved() -> Result<(), Box<dyn std::error::Error>> {
let plan = DeltaScanFileTaskPartitionPlan::try_new(plan_request(
vec![
file_task("zero-0.parquet", Some(0), Some(0))?,
file_task("zero-1.parquet", Some(0), Some(0))?,
],
4,
))?;
assert_eq!(
plan_partition_paths(&plan),
vec![vec!["zero-0.parquet"], vec!["zero-1.parquet"]]
);
assert_eq!(
plan.partitions
.iter()
.map(|partition| partition.estimated_bytes)
.collect::<Vec<_>>(),
vec![Some(0), Some(0)]
);
assert_eq!(plan.estimated_bytes, Some(0));
assert_eq!(plan.estimated_rows, Some(0));
Ok(())
}
#[test]
fn each_input_file_task_appears_exactly_once() -> Result<(), Box<dyn std::error::Error>> {
let plan = DeltaScanFileTaskPartitionPlan::try_new(plan_request(
vec![
file_task("part-0.parquet", Some(10), Some(1))?,
file_task("part-1.parquet", Some(10), Some(1))?,
file_task("part-2.parquet", Some(10), Some(1))?,
file_task("part-3.parquet", Some(10), Some(1))?,
],
2,
))?;
let flattened_paths = plan
.partitions
.iter()
.flat_map(|partition| partition.file_tasks.iter())
.map(|file_task| file_task.path.as_str())
.collect::<Vec<_>>();
assert_eq!(
flattened_paths,
vec![
"part-0.parquet",
"part-1.parquet",
"part-2.parquet",
"part-3.parquet",
]
);
Ok(())
}
#[test]
fn grouped_plan_remains_delta_aware_file_task_plan() -> Result<(), Box<dyn std::error::Error>> {
let plan = DeltaScanFileTaskPartitionPlan::try_new(plan_request(
vec![file_task_with_metadata(
"part-with-delta-metadata.parquet",
Some(10),
Some(1),
KernelScanDeletionVectorMetadata::NotPresent,
KernelPhysicalToLogicalTransform::test_required_column_transform("physical_id"),
)?],
1,
))?;
let file_task = &plan.partitions[0].file_tasks[0];
let path_only = file_task.path.clone();
assert_eq!(path_only, "part-with-delta-metadata.parquet");
assert_eq!(file_task.source_name, "orders");
assert_eq!(file_task.table_uri, "file:///tmp/table");
assert_eq!(file_task.snapshot_version, 42);
assert_eq!(file_task.estimated_bytes, Some(10));
assert_eq!(file_task.estimated_rows, Some(1));
assert_eq!(
file_task.partition_values.get("region").map(String::as_str),
Some("us-west")
);
assert!(matches!(
file_task.deletion_vector,
KernelScanDeletionVectorMetadata::NotPresent
));
assert!(file_task.transform.is_required());
Ok(())
}
#[test]
fn unknown_rows_keep_row_estimates_unknown() -> Result<(), Box<dyn std::error::Error>> {
let plan = DeltaScanFileTaskPartitionPlan::try_new(plan_request(
vec![
file_task("part-0.parquet", Some(10), None)?,
file_task("part-1.parquet", Some(10), Some(1))?,
],
1,
))?;
assert_eq!(plan.estimated_rows, None);
assert_eq!(plan.partitions[0].estimated_rows, None);
Ok(())
}
}