diff --git a/iceberg/README.md b/iceberg/README.md index 6fe6e00ad..dbdda2198 100644 --- a/iceberg/README.md +++ b/iceberg/README.md @@ -78,9 +78,12 @@ Iceberg's `FileIO` handles storage in both cases. Eager planning waits for metad discovery and increases worker-plan payloads; lazy scans overlap discovery with reading and deliver descriptors through the coordinator's feed. -Both modes retain the existing snapshot statistics and task-size estimates. -Selected-file task sizing, greedy assignment, and selected-file column statistics -are separate changes. Task-local eager variants and decoded plans report unknown +Eager scans derive row counts and bytes from selected files and cap task estimates +by the number of selected files. Predicates or deletes make row counts inexact; +empty selections have exactly zero rows. Whole-snapshot column statistics are +skipped when they would describe a different row set. Lazy scans retain snapshot +statistics and task-size estimates. Greedy assignment and selected-file column +statistics are separate changes. Task-local eager variants and decoded plans report unknown statistics rather than applying whole-snapshot aggregates to a subset. ```bash diff --git a/iceberg/src/distributed_desired_task_count_handler.rs b/iceberg/src/distributed_desired_task_count_handler.rs index 27b10c27f..e4a15ebe4 100644 --- a/iceberg/src/distributed_desired_task_count_handler.rs +++ b/iceberg/src/distributed_desired_task_count_handler.rs @@ -8,7 +8,8 @@ use datafusion_distributed::{ use crate::{EagerIcebergDataSource, LazyIcebergDataSource}; -/// Estimates scan parallelism from the selected Iceberg snapshot's total file size. +/// Estimates scan parallelism from selected file bytes, bounded by available file work. +/// Lazy scans continue to use snapshot totals. pub fn iceberg_desired_task_count( ev: DesiredTaskCountEvent, ) -> Option> { @@ -29,7 +30,14 @@ pub fn iceberg_desired_task_count( config.file_scan_config_bytes_per_partition, ev.session_config.target_partitions(), ) - .map(DesiredTaskCountEventResponse::soft), + .map(|count| { + let count = node + .downcast_ref::() + .map_or(count, |source| { + count.min(source.planned_files().count() as f64) + }); + DesiredTaskCountEventResponse::soft(count) + }), ) } diff --git a/iceberg/src/eager/planned_files.rs b/iceberg/src/eager/planned_files.rs index 46a34c5e1..eb7f7c9cf 100644 --- a/iceberg/src/eager/planned_files.rs +++ b/iceberg/src/eager/planned_files.rs @@ -1,6 +1,8 @@ use std::mem::{size_of, size_of_val}; -use datafusion::common::Result; +use datafusion::arrow::datatypes::SchemaRef; +use datafusion::common::stats::Precision; +use datafusion::common::{Result, Statistics}; use datafusion::execution::TaskContext; use datafusion::execution::memory_pool::{MemoryConsumer, MemoryReservation}; use futures::TryStreamExt; @@ -21,6 +23,29 @@ impl PlannedFiles { &self.tasks } + /// Summarize selected file work; predicates and deletes can reduce output rows. + pub(crate) fn statistics(&self, schema: &SchemaRef, may_filter_rows: bool) -> Statistics { + let bytes = self.tasks.iter().fold(0usize, |bytes, task| { + bytes.saturating_add(task.length() as usize) + }); + let rows = self.tasks.iter().try_fold(0usize, |rows, task| { + rows.checked_add(usize::try_from(task.record_count()?).ok()?) + }); + let mut statistics = Statistics::new_unknown(schema); + statistics.num_rows = match rows { + Some(0) if self.tasks.is_empty() => Precision::Exact(0), + Some(rows) if !may_filter_rows => Precision::Exact(rows), + Some(rows) => Precision::Inexact(rows), + None => Precision::Absent, + }; + statistics.total_byte_size = if may_filter_rows { + Precision::Inexact(bytes) + } else { + Precision::Exact(bytes) + }; + statistics + } + fn new(context: &TaskContext) -> Self { Self { tasks: Vec::new(), diff --git a/iceberg/src/table_provider/scan.rs b/iceberg/src/table_provider/scan.rs index b9b0d35ed..612342ffa 100644 --- a/iceberg/src/table_provider/scan.rs +++ b/iceberg/src/table_provider/scan.rs @@ -1,4 +1,4 @@ -use crate::common::scan_config::{ScanConfig, stats_from_snapshot}; +use crate::common::scan_config::ScanConfig; use crate::common::{df_err, file_partition}; use crate::eager::planned_files::PlannedFiles; use crate::table_provider::column_statistics::compute_column_stats; @@ -60,31 +60,10 @@ pub(crate) async fn plan_scan( .iceberg_runtime .unwrap_or_else(iceberg::Runtime::current), }; - let column_stats = if IcebergConfig::from_task_context(&context).column_stats_enabled { - let schema = match &table_snapshot { - Some(snap) => snap.schema(table.metadata()).map_err(df_err)?, - // empty table - None => table.metadata().current_schema().clone(), - }; - let fields = schema.as_struct().fields().to_vec(); - let field_ids: Vec = match opts.projection { - Some(projection) => projection.iter().map(|&idx| fields[idx].id).collect(), - None => fields.iter().map(|f| f.id).collect(), - }; - Some(compute_column_stats(table.clone(), field_ids, table_snapshot.clone()).await?) - } else { - None - }; - if IcebergConfig::from_task_context(&context).eager_file_planning { + let planned = if IcebergConfig::from_task_context(&context).eager_file_planning { if partitioning.partition_count() == 0 { return config_err!("Iceberg scan requires at least one output partition"); } - // TODO: derive eager statistics from selected files in PR #787: - // https://github.com/datafusion-contrib/datafusion-distributed/pull/787 - let mut statistics = stats_from_snapshot(table_snapshot.clone(), &scan.schema)?; - if let Some(column_stats) = column_stats { - statistics.column_statistics = column_stats; - } let mut builder = table.scan(); if let Some(snapshot_id) = opts.snapshot_id { builder = builder.snapshot_id(snapshot_id); @@ -92,19 +71,49 @@ pub(crate) async fn plan_scan( if let Some(columns) = &projection { builder = builder.select(columns); } - if let Some(predicates) = predicates { - builder = builder.with_filter(predicates); + if let Some(predicates) = &predicates { + builder = builder.with_filter(predicates.clone()); } let table_scan = builder.build().map_err(df_err)?; // File discovery runs on Iceberg's I/O runtime and is canceled if planning is dropped. let (_cancel_on_drop, canceled) = oneshot::channel::<()>(); + let planning_context = Arc::clone(&context); let task = scan.iceberg_runtime.io().spawn(async move { tokio::select! { _ = canceled => exec_err!("Iceberg scan planning canceled"), - planned = PlannedFiles::collect(table_scan, &context) => planned, + planned = PlannedFiles::collect(table_scan, &planning_context) => planned, } }); - let planned = Arc::new(task.await.map_err(df_err)??); + Some(Arc::new(task.await.map_err(df_err)??)) + } else { + None + }; + let may_filter_rows = predicates.is_some() + || planned + .as_ref() + .is_some_and(|files| files.tasks().iter().any(|task| !task.deletes().is_empty())); + let column_stats = if IcebergConfig::from_task_context(&context).column_stats_enabled + && !(planned.is_some() && may_filter_rows) + { + let schema = match &table_snapshot { + Some(snap) => snap.schema(table.metadata()).map_err(df_err)?, + // empty table + None => table.metadata().current_schema().clone(), + }; + let fields = schema.as_struct().fields().to_vec(); + let field_ids: Vec = match opts.projection { + Some(projection) => projection.iter().map(|&idx| fields[idx].id).collect(), + None => fields.iter().map(|f| f.id).collect(), + }; + Some(compute_column_stats(table.clone(), field_ids, table_snapshot.clone()).await?) + } else { + None + }; + if let Some(planned) = planned { + let mut statistics = planned.statistics(&scan.schema, may_filter_rows); + if let Some(column_stats) = column_stats { + statistics.column_statistics = column_stats; + } let mut partitions = vec![Vec::new(); partitioning.partition_count()]; for index in 0..planned.tasks().len() { partitions[file_partition(index, 1, partitioning.partition_count())].push(index); diff --git a/iceberg/tests/desired_task_count.rs b/iceberg/tests/desired_task_count.rs index 222d20c8b..2a03be262 100644 --- a/iceberg/tests/desired_task_count.rs +++ b/iceberg/tests/desired_task_count.rs @@ -1,7 +1,10 @@ #[cfg(all(feature = "integration", test))] mod tests { use datafusion::common::Result; + use datafusion::common::tree_node::{TreeNode, TreeNodeRecursion}; + use datafusion::datasource::source::DataSourceExec; use datafusion_distributed::DistributedExt; + use datafusion_distributed_iceberg::EagerIcebergDataSource; use datafusion_distributed_iceberg::test_utils::{ IcebergTestHarness, empty_taxi_metadata_builder, taxi_metadata, }; @@ -57,6 +60,37 @@ mod tests { Ok(()) } + #[tokio::test] + async fn planned_scan_tasks_are_bounded_by_selected_file_work() -> Result<()> { + let harness = IcebergTestHarness::builder() + .configure_iceberg(|config| config.eager_file_planning = true) + .configure_session(|mut state| { + let session = state.config().get_or_insert_default(); + session.options_mut().execution.target_partitions = 2; + state.with_distributed_file_scan_config_bytes_per_partition(1) + })? + .build() + .await?; + assert_eq!( + harness.estimate_task_count(&harness.scan().await?)?, + Some(7) + ); + let plan = harness + .physical_plan("SELECT pickup_date FROM taxi WHERE pickup_date = DATE '2024-01-10'") + .await?; + let mut estimates = Vec::new(); + plan.apply(|node| { + if let Some(exec) = node.downcast_ref::() + && exec.data_source().is::() + { + estimates.push(harness.estimate_task_count(node)?); + } + Ok(TreeNodeRecursion::Continue) + })?; + assert_eq!(estimates, vec![Some(1)]); + Ok(()) + } + #[tokio::test] async fn estimates_current_snapshot() -> Result<()> { assert_task_count(metadata_with_file_size(Some("24000000")), None, Some(12)).await diff --git a/iceberg/tests/snapshots/assigned_file_routing__tests__pruned_files_choose_the_same_worker_across_union_branches_and_queries.snap b/iceberg/tests/snapshots/assigned_file_routing__tests__pruned_files_choose_the_same_worker_across_union_branches_and_queries.snap index 8599b0d89..496fdf61c 100644 --- a/iceberg/tests/snapshots/assigned_file_routing__tests__pruned_files_choose_the_same_worker_across_union_branches_and_queries.snap +++ b/iceberg/tests/snapshots/assigned_file_routing__tests__pruned_files_choose_the_same_worker_across_union_branches_and_queries.snap @@ -5,7 +5,7 @@ expression: output Date: 2024-01-08 Task 0 -> http://url-1/ s3://iceberg-test/warehouse/taxi/data/pickup_date=2024-01-08/data_0.parquet -Task 2 -> http://url-1/ +Task 1 -> http://url-1/ s3://iceberg-test/warehouse/taxi/data/pickup_date=2024-01-08/data_0.parquet +------+----------+ | kind | count(*) | @@ -16,7 +16,7 @@ Task 2 -> http://url-1/ Date: 2024-01-09 Task 0 -> http://url-2/ s3://iceberg-test/warehouse/taxi/data/pickup_date=2024-01-09/data_0.parquet -Task 2 -> http://url-2/ +Task 1 -> http://url-2/ s3://iceberg-test/warehouse/taxi/data/pickup_date=2024-01-09/data_0.parquet +------+----------+ | kind | count(*) | @@ -27,7 +27,7 @@ Task 2 -> http://url-2/ Date: 2024-01-08 Task 0 -> http://url-1/ s3://iceberg-test/warehouse/taxi/data/pickup_date=2024-01-08/data_0.parquet -Task 2 -> http://url-1/ +Task 1 -> http://url-1/ s3://iceberg-test/warehouse/taxi/data/pickup_date=2024-01-08/data_0.parquet +------+----------+ | kind | count(*) | diff --git a/iceberg/tests/statistics.rs b/iceberg/tests/statistics.rs index 5e8e779d6..7997139e5 100644 --- a/iceberg/tests/statistics.rs +++ b/iceberg/tests/statistics.rs @@ -11,7 +11,8 @@ mod tests { use datafusion::scalar::ScalarValue; use datafusion_distributed_iceberg::IcebergExt; use datafusion_distributed_iceberg::test_utils::{ - FIXTURE_URI, IcebergTestHarness, empty_taxi_metadata_builder, taxi_metadata, + FIXTURE_URI, IcebergTestHarness, IcebergTestHarnessBuilder, empty_taxi_metadata_builder, + taxi_metadata, }; use iceberg::io::{MemoryStorage, Storage}; use iceberg::spec::{ @@ -37,6 +38,53 @@ mod tests { Ok(()) } + #[tokio::test] + async fn planned_file_statistics_use_manifest_totals_when_snapshot_summary_is_missing() + -> Result<()> { + let harness = IcebergTestHarness::builder() + .with_table_metadata(metadata_without_summary_statistics()) + .configure_iceberg(|config| config.eager_file_planning = true) + .build() + .await?; + let stats = query_statistics(&harness, "SELECT * FROM taxi").await?; + assert_eq!(stats.num_rows, Precision::Exact(TAXI_ROWS)); + assert_eq!(stats.total_byte_size, Precision::Exact(TAXI_BYTES)); + + let filtered = query_statistics( + &harness, + "SELECT pickup_date FROM taxi WHERE pickup_date = DATE '2024-01-10'", + ) + .await?; + assert!(matches!(filtered.num_rows, Precision::Inexact(_))); + assert!(matches!(filtered.total_byte_size, Precision::Inexact(_))); + Ok(()) + } + + #[tokio::test] + async fn eager_filtered_scan_skips_unused_column_statistics() -> Result<(), Box> { + // The manifest's partition summary excludes this query. Poison its contents + // so a redundant whole-snapshot column-statistics read fails the test. + let harness = manifest_metrics_builder() + .await? + .with_file( + format!("{FIXTURE_URI}/metadata/column-metrics.avro"), + vec![], + ) + .configure_iceberg(|config| { + config.eager_file_planning = true; + config.column_stats_enabled = true; + }) + .build() + .await?; + let stats = query_statistics( + &harness, + "SELECT * FROM taxi WHERE pickup_date = DATE '1900-01-01'", + ) + .await?; + assert_eq!(stats.num_rows, Precision::Exact(0)); + Ok(()) + } + #[tokio::test] async fn full_scan_without_column_stats() -> Result<(), Box> { assert_manifest_statistics(false, false).await @@ -253,6 +301,14 @@ mod tests { async fn harness_with_manifest_metrics( enabled: bool, ) -> Result> { + Ok(manifest_metrics_builder() + .await? + .configure_session(|state| Ok(state.with_iceberg_column_stats_enabled(enabled)))? + .build() + .await?) + } + + async fn manifest_metrics_builder() -> Result> { let metadata = taxi_metadata(); let snapshot = metadata.current_snapshot().expect("taxi has a snapshot"); let storage = MemoryStorage::new(); @@ -309,9 +365,6 @@ mod tests { storage.read(snapshot.manifest_list()).await?.to_vec(), ) .with_table_option("iceberg.snapshot_id", snapshot.snapshot_id().to_string()) - .with_table_metadata(metadata) - .configure_session(|state| Ok(state.with_iceberg_column_stats_enabled(enabled)))? - .build() - .await?) + .with_table_metadata(metadata)) } }