From 016d097717e513f60bf79aa1a2e079c8ec11b0f2 Mon Sep 17 00:00:00 2001 From: "alex.bianchi" Date: Sun, 11 Oct 2026 23:38:11 -0400 Subject: [PATCH 1/7] feat: plan Iceberg file tasks before execution with public APIs --- crates/datafusion/README.md | 45 ++++ crates/datafusion/src/config.rs | 34 +++ crates/datafusion/src/lib.rs | 3 + crates/datafusion/src/read/file_planning.rs | 97 +++++++ crates/datafusion/src/read/mod.rs | 2 + crates/datafusion/src/read/scan.rs | 155 +++++++++-- .../table_provider/catalog_table_provider.rs | 24 +- .../table_provider/static_table_provider.rs | 33 ++- crates/datafusion/src/test_utils/harness.rs | 37 ++- .../datafusion/tests/eager_file_planning.rs | 247 ++++++++++++++++++ 10 files changed, 635 insertions(+), 42 deletions(-) create mode 100644 crates/datafusion/src/config.rs create mode 100644 crates/datafusion/src/read/file_planning.rs create mode 100644 crates/datafusion/tests/eager_file_planning.rs diff --git a/crates/datafusion/README.md b/crates/datafusion/README.md index 134a8ef..66f4f00 100644 --- a/crates/datafusion/README.md +++ b/crates/datafusion/README.md @@ -20,3 +20,48 @@ # Apache Iceberg DataFusion Integration This crate contains the integration of Apache DataFusion and Apache Iceberg. + +## Optional file planning before execution + +By default, the Iceberg reader discovers files when execution starts. To select +`FileScanTask`s during `TableProvider::scan()` instead, register the planning +options on the DataFusion session before planning a query: + +```rust +use datafusion::prelude::SessionConfig; +use iceberg_datafusion::IcebergPlanningConfig; + +let config = SessionConfig::new() + .with_option_extension(IcebergPlanningConfig::default()) + .set_bool("iceberg_local.plan_files", true); +``` + +The option defaults to `false`. When enabled, each physical scan retains its +selected tasks within the query's memory pool, failing rather than truncating +if more than `iceberg_local.planning_max_files` (default: 100,000) are selected. +Tasks are assigned round-robin to at most the session's `target_partitions`; +an empty scan exposes one empty partition. This does **not** read Parquet data +during planning. The physical scan continues to report unknown statistics. +`iceberg_local` is distinct from DataFusion Distributed's `iceberg` options so +both can be registered in the same session. + +The local reader is constructed with the existing public +`Table::reader_builder()` API, which uses the table's `FileIO` and runtime. +This matches the current integration's reader defaults because it does not +set reader options on `TableScan`: the snapshot, projection, and predicate are +carried by the tasks. For multiple output partitions, each partition reads one +file at a time instead of using the lazy scan's per-reader file concurrency. +**Reader settings are not transferred separately from the scan.** +If scan-specific batch size, data-file concurrency, row-group, row-selection, +or bloom-filter settings are introduced later, the lazy and planned readers +could diverge. A follow-up should expose portable scan-level reader options +and apply them to a locally constructed builder in iceberg-rust, with +`TableScan::to_arrow()` using the same mapping. Distributed consumers should +construct readers using worker-local IO and runtimes rather than transport a +builder or reader. + +Planning the whole file list earlier adds metadata I/O and retained-task memory +even when a query stops early (for example, `LIMIT`), and equivalent provider +`scan()` calls currently repeat planning. No performance improvement is claimed +by this option alone. Delete-bearing eager scans and cross-partition delete-cache +reuse also need dedicated validation before enabling it by default. diff --git a/crates/datafusion/src/config.rs b/crates/datafusion/src/config.rs new file mode 100644 index 0000000..4d71716 --- /dev/null +++ b/crates/datafusion/src/config.rs @@ -0,0 +1,34 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +use datafusion::common::extensions_options; +use datafusion::config::ConfigExtension; + +extensions_options! { + /// Opt-in file planning during physical planning for local Iceberg scans. + pub struct IcebergPlanningConfig { + /// Discover and retain file tasks in TableProvider::scan instead of at execution. + pub plan_files: bool, default = false + /// Fail rather than retain more than this many tasks in a single scan. + pub planning_max_files: usize, default = 100_000 + } +} + +impl ConfigExtension for IcebergPlanningConfig { + // Distinct from DFD's `iceberg` options so both integrations can share a session. + const PREFIX: &'static str = "iceberg_local"; +} diff --git a/crates/datafusion/src/lib.rs b/crates/datafusion/src/lib.rs index 41cbc8b..80b6371 100644 --- a/crates/datafusion/src/lib.rs +++ b/crates/datafusion/src/lib.rs @@ -18,6 +18,9 @@ mod catalog; pub use catalog::*; +mod config; +pub use config::IcebergPlanningConfig; + mod error; pub use error::*; diff --git a/crates/datafusion/src/read/file_planning.rs b/crates/datafusion/src/read/file_planning.rs new file mode 100644 index 0000000..4f6709c --- /dev/null +++ b/crates/datafusion/src/read/file_planning.rs @@ -0,0 +1,97 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Query-local file work selected before execution, independently of task assignment. + +use std::mem::{size_of, size_of_val}; +use std::sync::Arc; + +use datafusion::common::{Result, exec_err}; +use datafusion::execution::TaskContext; +use datafusion::execution::memory_pool::{MemoryConsumer, MemoryReservation}; +use futures::TryStreamExt; +use iceberg::scan::{FileScanTask, TableScan}; + +use crate::to_datafusion_error; + +/// File work selected by one Iceberg scan before assignment to output partitions. +/// +/// The tasks retain their bound predicates, projection, and delete references. Keeping +/// this value alive also keeps the query memory reservation for the retained tasks. +/// This is not a reader and contains no coordinator-local reader state. +#[derive(Debug)] +pub struct PlannedFileScan { + tasks: Vec>, + _reservation: MemoryReservation, +} + +impl PlannedFileScan { + /// Tasks in Iceberg's planning order. Each must be assigned exactly once by a consumer. + pub fn tasks(&self) -> &[Arc] { + &self.tasks + } + + pub(crate) async fn collect( + scan: TableScan, + context: &TaskContext, + max_files: usize, + ) -> Result { + let reservation = + MemoryConsumer::new("Iceberg file planning").register(context.memory_pool()); + let mut planned = Self { + tasks: Vec::new(), + _reservation: reservation, + }; + let mut stream = scan.plan_files().await.map_err(to_datafusion_error)?; + while let Some(task) = stream.try_next().await.map_err(to_datafusion_error)? { + if planned.tasks.len() >= max_files { + return exec_err!( + "Iceberg file planning exceeded planning_max_files={max_files}" + ); + } + // Allow for spare Vec capacity and allocator overhead; schemas and specs + // are shared Arcs, but paths and delete references are owned by each task. + planned + ._reservation + .try_grow(retained_size(&task).saturating_mul(2))?; + planned.tasks.push(Arc::new(task)); + } + Ok(planned) + } +} + +fn retained_size(task: &FileScanTask) -> usize { + let mut bytes = size_of::() + + size_of::>() + + task.data_file_path().len() + + size_of_val(task.project_field_ids()) + + size_of_val(task.deletes()) + + task.key_metadata().map_or(0, <[u8]>::len); + for delete in task.deletes() { + bytes = bytes + .saturating_add(delete.file_path.len()) + .saturating_add(delete.referenced_data_file.as_ref().map_or(0, String::len)) + .saturating_add( + delete + .equality_ids + .as_ref() + .map_or(0, |ids| size_of_val(ids.as_slice())), + ) + .saturating_add(delete.key_metadata.as_ref().map_or(0, |key| key.len())); + } + bytes +} diff --git a/crates/datafusion/src/read/mod.rs b/crates/datafusion/src/read/mod.rs index 4c0157d..ceb2817 100644 --- a/crates/datafusion/src/read/mod.rs +++ b/crates/datafusion/src/read/mod.rs @@ -15,8 +15,10 @@ // specific language governing permissions and limitations // under the License. +mod file_planning; pub(crate) mod predicate; pub(crate) mod scan; +pub use file_planning::PlannedFileScan; pub use predicate::convert_filters_to_predicate; pub use scan::{IcebergDataSource, IcebergDataSourceBuilder}; diff --git a/crates/datafusion/src/read/scan.rs b/crates/datafusion/src/read/scan.rs index afcf3ee..b5048c0 100644 --- a/crates/datafusion/src/read/scan.rs +++ b/crates/datafusion/src/read/scan.rs @@ -20,8 +20,8 @@ use std::sync::Arc; use datafusion::arrow::array::RecordBatch; use datafusion::arrow::datatypes::SchemaRef as ArrowSchemaRef; -use datafusion::common::Statistics; use datafusion::common::tree_node::TreeNodeRecursion; +use datafusion::common::{Statistics, internal_err}; use datafusion::datasource::source::DataSource; use datafusion::error::Result; use datafusion::execution::{SendableRecordBatchStream, TaskContext}; @@ -30,13 +30,27 @@ use datafusion::physical_expr::{EquivalenceProperties, PhysicalExpr}; use datafusion::physical_plan::stream::RecordBatchStreamAdapter; use datafusion::physical_plan::{DisplayFormatType, Partitioning}; use datafusion::prelude::Expr; -use futures::{Stream, TryStreamExt}; +use futures::{Stream, StreamExt, TryStreamExt}; +use iceberg::arrow::ArrowReader; use iceberg::expr::Predicate; +use iceberg::scan::TableScan; use iceberg::table::Table; +use super::file_planning::PlannedFileScan; use super::predicate::convert_filters_to_predicate; use crate::to_datafusion_error; +// A reader is created once per eager scan. Its clones share the delete loader's +// cache while each read() call still owns its own scan metrics. +#[derive(Clone)] +struct ScanReader(ArrowReader); + +impl std::fmt::Debug for ScanReader { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.write_str("ScanReader") + } +} + /// Manages the scanning process of an Iceberg [`Table`], encapsulating the /// necessary details and computed properties required for execution planning. #[derive(Debug, Clone)] @@ -53,6 +67,9 @@ pub struct IcebergDataSource { predicate: Option, /// Optional limit on the number of rows to return fetch: Option, + planned_files: Option>, + output_partitions: usize, + reader: Option>, } impl IcebergDataSource { @@ -70,6 +87,45 @@ impl IcebergDataSource { } } + /// Plan file tasks during physical planning; read them only at execution. + pub(crate) async fn with_planned_files( + mut self, + context: Arc, + max_files: usize, + target_partitions: usize, + ) -> Result { + let scan = build_table_scan( + &self.table, + self.snapshot_id, + Some(&self.projection), + self.predicate.as_ref(), + )?; + let planned = PlannedFileScan::collect(scan, &context, max_files).await?; + let partitions = planned.tasks().len().max(1).min(target_partitions.max(1)); + + // Table::reader_builder retains the table's FileIO and custom runtime. + // This scan currently changes only snapshot, projection, and predicate; + // those values are already carried by FileScanTask. A future scan-specific + // reader setting must be applied here and in TableScan::to_arrow(). + let builder = self.table.reader_builder(); + let builder = if partitions > 1 { + builder.with_data_file_concurrency_limit(1) + } else { + builder + }; + self.reader = Some(Arc::new(ScanReader(builder.build()))); + self.output_partitions = partitions; + self.planned_files = Some(Arc::new(planned)); + Ok(self) + } + + /// Planned, ungrouped tasks, if file discovery was enabled for this scan. + /// A distributed consumer may assign these tasks to its own workers instead + /// of using the local output partitions of this source. + pub fn planned_files(&self) -> Option<&Arc> { + self.planned_files.as_ref() + } + pub fn table(&self) -> &Table { &self.table } @@ -100,7 +156,8 @@ impl IcebergDataSource { /// A table provider passes the table's schema and selects columns with /// [`Self::with_projection`]. To rebuild a source, such as in another process, /// pass its [`IcebergDataSource::schema`], which is already projected, and no -/// projection. +/// projection. This rebuilds a lazy single-partition source; planned file tasks +/// must be transferred separately when reconstructing a distributed scan. #[derive(Debug, Clone)] pub struct IcebergDataSourceBuilder { table: Table, @@ -167,6 +224,9 @@ impl IcebergDataSourceBuilder { projection, predicate: self.predicate, fetch: self.fetch, + planned_files: None, + output_partitions: 1, + reader: None, }) } } @@ -181,16 +241,45 @@ impl DataSource for IcebergDataSource { fn open( &self, - _partition: usize, + partition: usize, _context: Arc, ) -> Result { - let fut = get_batch_stream( - self.table.clone(), - self.snapshot_id, - Some(self.projection.clone()), - self.predicate.clone(), - ); - let stream = futures::stream::once(fut).try_flatten(); + let stream: Pin> + Send>> = + if let Some(planned) = &self.planned_files { + if partition >= self.output_partitions { + return internal_err!( + "Iceberg scan partition {partition} does not exist" + ); + } + let Some(reader) = &self.reader else { + return internal_err!("Planned Iceberg scan has no reader"); + }; + let planned = Arc::clone(planned); + let partitions = self.output_partitions; + let tasks = futures::stream::iter( + (partition..planned.tasks().len()).step_by(partitions), + ) + .map(move |index| { + Ok(Arc::unwrap_or_clone(Arc::clone(&planned.tasks()[index]))) + }); + Box::pin( + reader + .0 + .clone() + .read(Box::pin(tasks)) + .map_err(to_datafusion_error)? + .stream() + .map_err(to_datafusion_error), + ) + } else { + let fut = get_batch_stream( + self.table.clone(), + self.snapshot_id, + Some(self.projection.clone()), + self.predicate.clone(), + ); + Box::pin(futures::stream::once(fut).try_flatten()) + }; // Apply the fetch if specified. let limited_stream: Pin> + Send>> = @@ -231,11 +320,14 @@ impl DataSource for IcebergDataSource { if let Some(fetch) = self.fetch { write!(f, ", fetch={fetch}")?; } + if let Some(planned) = &self.planned_files { + write!(f, ", planned_files={}", planned.tasks().len())?; + } Ok(()) } fn output_partitioning(&self) -> Partitioning { - Partitioning::UnknownPartitioning(1) + Partitioning::UnknownPartitioning(self.output_partitions) } fn eq_properties(&self) -> EquivalenceProperties { @@ -275,19 +367,12 @@ async fn get_batch_stream( column_names: Option>, predicate: Option, ) -> Result> + Send>>> { - let scan_builder = match snapshot_id { - Some(snapshot_id) => table.scan().snapshot_id(snapshot_id), - None => table.scan(), - }; - - let mut scan_builder = match column_names { - Some(column_names) => scan_builder.select(column_names), - None => scan_builder.select_all(), - }; - if let Some(predicate) = predicate { - scan_builder = scan_builder.with_filter(predicate); - } - let table_scan = scan_builder.build().map_err(to_datafusion_error)?; + let table_scan = build_table_scan( + &table, + snapshot_id, + column_names.as_deref(), + predicate.as_ref(), + )?; let stream = table_scan .to_arrow() @@ -296,3 +381,23 @@ async fn get_batch_stream( .map_err(to_datafusion_error); Ok(Box::pin(stream)) } + +fn build_table_scan( + table: &Table, + snapshot_id: Option, + column_names: Option<&[String]>, + predicate: Option<&Predicate>, +) -> Result { + let builder = match snapshot_id { + Some(snapshot_id) => table.scan().snapshot_id(snapshot_id), + None => table.scan(), + }; + let mut builder = match column_names { + Some(columns) => builder.select(columns), + None => builder.select_all(), + }; + if let Some(predicate) = predicate { + builder = builder.with_filter(predicate.clone()); + } + builder.build().map_err(to_datafusion_error) +} diff --git a/crates/datafusion/src/table_provider/catalog_table_provider.rs b/crates/datafusion/src/table_provider/catalog_table_provider.rs index b804f56..a6a0619 100644 --- a/crates/datafusion/src/table_provider/catalog_table_provider.rs +++ b/crates/datafusion/src/table_provider/catalog_table_provider.rs @@ -34,6 +34,7 @@ use iceberg::inspect::MetadataTableType; use iceberg::spec::TableProperties; use iceberg::{Catalog, NamespaceIdent, TableIdent}; +use crate::IcebergPlanningConfig; use crate::error::to_datafusion_error; use crate::metadata::table::IcebergMetadataTableProvider; use crate::read::IcebergDataSource; @@ -126,24 +127,39 @@ impl TableProvider for IcebergCatalogTableProvider { async fn scan( &self, - _state: &dyn Session, + state: &dyn Session, projection: Option<&Vec>, filters: &[Expr], limit: Option, ) -> Result> { - // Load fresh table metadata from catalog + // Load fresh table metadata from catalog. let table = self .catalog .load_table(&self.table_ident) .await .map_err(to_datafusion_error)?; - // Create scan with fresh metadata (always use current snapshot) - let source = IcebergDataSource::builder(table, self.schema.clone()) + // Create scan with fresh metadata (always use current snapshot). + let mut source = IcebergDataSource::builder(table, self.schema.clone()) .with_projection(projection) .with_filters(filters) .with_fetch(limit) .build()?; + if let Some(config) = state + .config() + .options() + .extensions + .get::() + .filter(|config| config.plan_files) + { + source = source + .with_planned_files( + state.task_ctx(), + config.planning_max_files, + state.config().target_partitions(), + ) + .await?; + } Ok(DataSourceExec::from_data_source(source)) } diff --git a/crates/datafusion/src/table_provider/static_table_provider.rs b/crates/datafusion/src/table_provider/static_table_provider.rs index 5ebabc8..bda652f 100644 --- a/crates/datafusion/src/table_provider/static_table_provider.rs +++ b/crates/datafusion/src/table_provider/static_table_provider.rs @@ -30,6 +30,7 @@ use datafusion::physical_plan::ExecutionPlan; use iceberg::arrow::schema_to_arrow_schema; use iceberg::table::Table; +use crate::IcebergPlanningConfig; use crate::error::to_datafusion_error; use crate::read::IcebergDataSource; @@ -120,18 +121,34 @@ impl TableProvider for IcebergStaticTableProvider { async fn scan( &self, - _state: &dyn Session, + state: &dyn Session, projection: Option<&Vec>, filters: &[Expr], limit: Option, ) -> Result> { - // Use cached table (no refresh) - let source = IcebergDataSource::builder(self.table.clone(), self.schema.clone()) - .with_snapshot_id(self.snapshot_id) - .with_projection(projection) - .with_filters(filters) - .with_fetch(limit) - .build()?; + // Use cached table (no refresh). + let mut source = + IcebergDataSource::builder(self.table.clone(), self.schema.clone()) + .with_snapshot_id(self.snapshot_id) + .with_projection(projection) + .with_filters(filters) + .with_fetch(limit) + .build()?; + if let Some(config) = state + .config() + .options() + .extensions + .get::() + .filter(|config| config.plan_files) + { + source = source + .with_planned_files( + state.task_ctx(), + config.planning_max_files, + state.config().target_partitions(), + ) + .await?; + } Ok(DataSourceExec::from_data_source(source)) } diff --git a/crates/datafusion/src/test_utils/harness.rs b/crates/datafusion/src/test_utils/harness.rs index d106467..13bea7f 100644 --- a/crates/datafusion/src/test_utils/harness.rs +++ b/crates/datafusion/src/test_utils/harness.rs @@ -25,6 +25,7 @@ use datafusion::arrow::util::pretty::pretty_format_batches; use datafusion::dataframe::DataFrame; use datafusion::error::Result; use datafusion::execution::SessionStateBuilder; +use datafusion::execution::runtime_env::RuntimeEnv; use datafusion::physical_plan::displayable; use datafusion::prelude::{SessionConfig, SessionContext}; use futures::StreamExt; @@ -53,12 +54,34 @@ pub struct IcebergTestHarness { impl IcebergTestHarness { pub async fn new() -> Result { - let state = SessionStateBuilder::new() + Self::with_config(SessionConfig::new().with_target_partitions(16)).await + } + + pub async fn with_config(config: SessionConfig) -> Result { + Self::build(config, None).await + } + + pub async fn with_config_and_runtime( + config: SessionConfig, + runtime_env: Arc, + ) -> Result { + Self::build(config, Some(runtime_env)).await + } + + async fn build( + config: SessionConfig, + runtime_env: Option>, + ) -> Result { + let builder = SessionStateBuilder::new() .with_default_features() - // Plan snapshots must not vary with the runner's CPU count. - .with_config(SessionConfig::new().with_target_partitions(16)) - // We register a factory in addition to the catalog below - // so that tests can be written against the factory's behaviours + .with_config(config); + let builder = match runtime_env { + Some(runtime_env) => builder.with_runtime_env(runtime_env), + None => builder, + }; + // Register a factory in addition to the catalog below so tests can + // exercise the factory's behavior. + let state = builder .with_table_factory( "ICEBERG".to_string(), Arc::new(IcebergTableProviderFactory::new_with_storage_factory( @@ -99,6 +122,10 @@ impl IcebergTestHarness { Ok(Self { ctx }) } + pub fn context(&self) -> &SessionContext { + &self.ctx + } + pub async fn query(&self, sql: &str) -> Result<(String, String)> { let dataframe: DataFrame = self.ctx.sql(sql).await?; let plan = dataframe.create_physical_plan().await?; diff --git a/crates/datafusion/tests/eager_file_planning.rs b/crates/datafusion/tests/eager_file_planning.rs new file mode 100644 index 0000000..05e2fce --- /dev/null +++ b/crates/datafusion/tests/eager_file_planning.rs @@ -0,0 +1,247 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +use std::sync::Arc; + +use datafusion::common::{Result, exec_datafusion_err, stats::Precision}; +use datafusion::datasource::TableProvider; +use datafusion::datasource::source::{DataSource, DataSourceExec}; +use datafusion::execution::memory_pool::GreedyMemoryPool; +use datafusion::execution::runtime_env::RuntimeEnvBuilder; +use datafusion::physical_plan::ExecutionPlanProperties; +use datafusion::physical_plan::common::collect; +use datafusion::prelude::{SessionConfig, col, lit}; +use datafusion::scalar::ScalarValue; +use iceberg::Runtime; +use iceberg::table::Table; +use iceberg_datafusion::read::IcebergDataSource; +use iceberg_datafusion::test_utils::IcebergTestHarness; +use iceberg_datafusion::{ + IcebergCatalogTableProvider, IcebergPlanningConfig, IcebergStaticTableProvider, +}; + +fn eager_config(max_files: usize) -> SessionConfig { + SessionConfig::new() + .with_target_partitions(4) + .with_option_extension(IcebergPlanningConfig::default()) + .set_bool("iceberg_local.plan_files", true) + .set_usize("iceberg_local.planning_max_files", max_files) +} + +fn iceberg_source( + plan: &dyn datafusion::physical_plan::ExecutionPlan, +) -> &IcebergDataSource { + plan.downcast_ref::() + .expect("the Iceberg provider creates a DataSourceExec") + .data_source() + .downcast_ref::() + .expect("the plan contains the Iceberg data source") +} + +#[tokio::test] +async fn plans_file_tasks_before_execution_without_changing_statistics() -> Result<()> { + let eager = IcebergTestHarness::with_config(eager_config(100)).await?; + let lazy = IcebergTestHarness::new().await?; + let eager_provider = eager.context().table_provider("taxi").await?; + let lazy_provider = lazy.context().table_provider("taxi").await?; + let eager_plan = eager_provider + .scan(&eager.context().state(), None, &[], None) + .await?; + let lazy_plan = lazy_provider + .scan(&lazy.context().state(), None, &[], None) + .await?; + let eager_source = iceberg_source(eager_plan.as_ref()); + let lazy_source = iceberg_source(lazy_plan.as_ref()); + + assert!(lazy_source.planned_files().is_none()); + assert_eq!(lazy_plan.output_partitioning().partition_count(), 1); + assert_eq!(eager_source.planned_files().unwrap().tasks().len(), 7); + assert_eq!(eager_plan.output_partitioning().partition_count(), 4); + assert_eq!( + eager_source.partition_statistics(None)?.num_rows, + Precision::Absent + ); + assert_eq!( + eager_source.partition_statistics(None)?.total_byte_size, + lazy_source.partition_statistics(None)?.total_byte_size + ); + + let empty = eager_provider + .scan( + &eager.context().state(), + None, + &[col("pickup_date").eq(lit(ScalarValue::Date32(Some(0))))], + None, + ) + .await?; + assert!( + iceberg_source(empty.as_ref()) + .planned_files() + .unwrap() + .tasks() + .is_empty() + ); + assert_eq!(empty.output_partitioning().partition_count(), 1); + + Ok(()) +} + +#[tokio::test] +async fn eager_partition_count_is_bounded_by_selected_files() -> Result<()> { + let eager = + IcebergTestHarness::with_config(eager_config(100).with_target_partitions(16)) + .await?; + let provider = eager.context().table_provider("taxi").await?; + let plan = provider + .scan(&eager.context().state(), None, &[], None) + .await?; + assert_eq!(plan.output_partitioning().partition_count(), 7); + Ok(()) +} + +#[tokio::test] +async fn filtered_and_empty_eager_scans_match_lazy_results() -> Result<()> { + let eager = IcebergTestHarness::with_config(eager_config(100)).await?; + let lazy = IcebergTestHarness::new().await?; + for (query, planned_files) in [ + ( + "SELECT pickup_date, COUNT(*) AS trips FROM taxi \ + WHERE pickup_date = DATE '2024-01-10' \ + GROUP BY pickup_date ORDER BY pickup_date", + 1, + ), + ( + "SELECT COUNT(*) AS trips FROM taxi \ + WHERE pickup_date = DATE '2025-01-01'", + 0, + ), + ] { + let (plan, eager_batches) = eager.query(query).await?; + let (_, lazy_batches) = lazy.query(query).await?; + assert!( + plan.contains(&format!("planned_files={planned_files}")), + "eager scan not visible in plan: {plan}" + ); + assert_eq!(eager_batches, lazy_batches); + } + Ok(()) +} + +#[tokio::test] +async fn eager_multi_partition_limit_is_global() -> Result<()> { + let eager = IcebergTestHarness::with_config(eager_config(100)).await?; + let (plan, batches) = eager.query("SELECT pickup_date FROM taxi LIMIT 3").await?; + assert!(plan.contains("planned_files=7"), "unexpected plan: {plan}"); + assert_eq!(batches.matches("2024-01-").count(), 3, "{batches}"); + Ok(()) +} + +#[tokio::test] +async fn eager_ordered_limit_matches_lazy_results() -> Result<()> { + let eager = IcebergTestHarness::with_config(eager_config(100)).await?; + let lazy = IcebergTestHarness::new().await?; + let sql = "SELECT vendor_id, pickup_at FROM taxi \ + ORDER BY pickup_at, vendor_id, pickup_location_id LIMIT 3"; + let (plan, eager_batches) = eager.query(sql).await?; + let (_, lazy_batches) = lazy.query(sql).await?; + assert!(plan.contains("planned_files=7"), "unexpected plan: {plan}"); + assert_eq!(eager_batches, lazy_batches); + Ok(()) +} + +#[tokio::test] +async fn eager_scan_uses_the_tables_custom_runtime() -> Result<()> { + let harness = IcebergTestHarness::with_config(eager_config(100)).await?; + let provider = harness.context().table_provider("taxi").await?; + let catalog_provider = provider + .downcast_ref::() + .expect("the fixture creates a catalog Iceberg provider"); + let table = catalog_provider + .catalog() + .load_table(catalog_provider.table_ident()) + .await + .map_err(iceberg_datafusion::to_datafusion_error)?; + let io_runtime = tokio::runtime::Builder::new_multi_thread() + .enable_all() + .build()?; + let table = Table::builder() + .file_io(table.file_io().clone()) + .metadata(table.metadata_ref()) + .identifier(table.identifier().clone()) + .runtime(Runtime::new(&io_runtime)) + .build() + .map_err(iceberg_datafusion::to_datafusion_error)?; + let provider = IcebergStaticTableProvider::try_new_from_table(table).await?; + let state = harness.context().state(); + let plan = provider.scan(&state, None, &[], None).await?; + assert_eq!( + iceberg_source(plan.as_ref()) + .planned_files() + .unwrap() + .tasks() + .len(), + 7 + ); + let mut rows = 0; + for partition in 0..plan.output_partitioning().partition_count() { + let batches = collect(plan.execute(partition, state.task_ctx())?).await?; + rows += batches.iter().map(|batch| batch.num_rows()).sum::(); + } + assert_eq!(rows, 175_000); + tokio::task::spawn_blocking(move || drop(io_runtime)) + .await + .map_err(|error| { + exec_datafusion_err!("failed to shut down test runtime: {error}") + })?; + Ok(()) +} + +#[tokio::test] +async fn eager_tasks_respect_the_query_memory_pool() -> Result<()> { + let runtime = Arc::new( + RuntimeEnvBuilder::new() + .with_memory_pool(Arc::new(GreedyMemoryPool::new(1))) + .build()?, + ); + let eager = + IcebergTestHarness::with_config_and_runtime(eager_config(100), runtime).await?; + let provider = eager.context().table_provider("taxi").await?; + let error = provider + .scan(&eager.context().state(), None, &[], None) + .await + .expect_err("file metadata cannot fit in one byte"); + assert!( + error.to_string().contains("Iceberg file planning"), + "{error}" + ); + Ok(()) +} + +#[tokio::test] +async fn planning_file_limit_rejects_instead_of_truncating() -> Result<()> { + let eager = IcebergTestHarness::with_config(eager_config(2)).await?; + let provider = eager.context().table_provider("taxi").await?; + let error = provider + .scan(&eager.context().state(), None, &[], None) + .await + .expect_err("seven files exceed the two-file planning limit"); + assert!( + error.to_string().contains("planning_max_files=2"), + "{error}" + ); + Ok(()) +} From 1f3982e9cb181a1577f25a349282110d531278ac Mon Sep 17 00:00:00 2001 From: "alex.bianchi" Date: Wed, 21 Oct 2026 20:20:31 -0400 Subject: [PATCH 2/7] refactor: align eager and lazy Iceberg scans --- Cargo.lock | 1 + crates/datafusion/Cargo.toml | 1 + crates/datafusion/README.md | 25 +- crates/datafusion/eager-planning.md | 158 +++++++++ crates/datafusion/src/config.rs | 4 +- crates/datafusion/src/read/file_planning.rs | 36 +- crates/datafusion/src/read/scan.rs | 172 +++++----- .../table_provider/catalog_table_provider.rs | 19 +- .../table_provider/static_table_provider.rs | 30 +- crates/datafusion/src/test_utils/harness.rs | 28 +- .../datafusion/tests/eager_file_planning.rs | 320 +++++++++++++++--- .../tests/integration_datafusion_test.rs | 2 +- 12 files changed, 577 insertions(+), 219 deletions(-) create mode 100644 crates/datafusion/eager-planning.md diff --git a/Cargo.lock b/Cargo.lock index 5182cdf..4c19a68 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -3247,6 +3247,7 @@ dependencies = [ "insta", "parquet", "serde", + "serde_json", "tempfile", "tokio", "typetag", diff --git a/crates/datafusion/Cargo.toml b/crates/datafusion/Cargo.toml index cd6edd0..4320866 100644 --- a/crates/datafusion/Cargo.toml +++ b/crates/datafusion/Cargo.toml @@ -44,6 +44,7 @@ typetag = "0.2" [dev-dependencies] insta = "1" parquet = "59.2" +serde_json = "1" tempfile = "3.18" [lints] diff --git a/crates/datafusion/README.md b/crates/datafusion/README.md index 66f4f00..c278fc9 100644 --- a/crates/datafusion/README.md +++ b/crates/datafusion/README.md @@ -33,12 +33,13 @@ use iceberg_datafusion::IcebergPlanningConfig; let config = SessionConfig::new() .with_option_extension(IcebergPlanningConfig::default()) - .set_bool("iceberg_local.plan_files", true); + .set_bool("iceberg_local.eager_file_planning", true); ``` The option defaults to `false`. When enabled, each physical scan retains its -selected tasks within the query's memory pool, failing rather than truncating -if more than `iceberg_local.planning_max_files` (default: 100,000) are selected. +selected tasks with an estimated reservation in the query's memory pool. There +is no file-count cap: exceeding the memory pool returns an error, never a +partial file list. Tasks are assigned round-robin to at most the session's `target_partitions`; an empty scan exposes one empty partition. This does **not** read Parquet data during planning. The physical scan continues to report unknown statistics. @@ -60,8 +61,22 @@ and apply them to a locally constructed builder in iceberg-rust, with construct readers using worker-local IO and runtimes rather than transport a builder or reader. +Both modes feed tasks to the same reader implementation. One reader is built +per physical scan, and its clones share the delete loader across output +partitions. Dropping eager planning drops collection and its memory reservation. + +`TableScan::plan_files()` performs its initial manifest-list read on the calling +runtime before spawning manifest work on the table's runtime. The planning +caller must therefore support metadata I/O, even with a separately configured +table runtime. The table's runtime accessor is private; making `plan_files()` +offload that initial read is an upstream follow-up. + Planning the whole file list earlier adds metadata I/O and retained-task memory even when a query stops early (for example, `LIMIT`), and equivalent provider `scan()` calls currently repeat planning. No performance improvement is claimed -by this option alone. Delete-bearing eager scans and cross-partition delete-cache -reuse also need dedicated validation before enabling it by default. +by this option alone. A position-delete fixture verifies eager/lazy row counts; +equality deletes and cross-partition delete-cache reuse still need dedicated +validation before enabling it by default. + +See [the planning API discussion](eager-planning.md) for the relationship to +distributed planning and the proposed upstream metadata and reader-settings APIs. diff --git a/crates/datafusion/eager-planning.md b/crates/datafusion/eager-planning.md new file mode 100644 index 0000000..b43bcce --- /dev/null +++ b/crates/datafusion/eager-planning.md @@ -0,0 +1,158 @@ + + +# Eager planning across Iceberg, DataFusion, and distributed execution + +This document describes the boundary proposed by +[datafusion-iceberg#67](https://github.com/apache/datafusion-iceberg/pull/67). +API names below are illustrative, not APIs provided by this change. + +## One selection pipeline, two consumption times + +```text +Table metadata → snapshot → manifest list → selected manifests + → live entries → partition/metrics pruning → file tasks + delete references + +Lazy: execution starts → stream selected tasks → reader +Eager: planning starts → collect selected tasks → assign partitions + → execution starts → reader +``` + +The catalog does not return the complete file list. Reading the manifest list +and manifests is metadata I/O, even when Parquet reads are deferred. Eager +planning pays that cost before execution and retains tasks, allowing the +optimizer to see the selected work. Lazy planning can overlap discovery and +execution and avoid retaining the full task list. Both use Iceberg's pruning +and delete matching, and both retain residual DataFusion filters. + +#67 keeps both modes, defaults to lazy, and uses the same local reader flow. +It exposes borrowed access to retained tasks, accounts for their estimated +memory, and assigns them round-robin to local partitions. It does not add +column statistics, distributed routing, or a file-count limit. Statistics +remain unknown in both modes. + +## What the distributed stack adds + +| Change | Responsibility | +| --- | --- | +| [DFD#773](https://github.com/datafusion-contrib/datafusion-distributed/pull/773) | Select and retain work during planning while preserving lazy execution. | +| [DFD#786](https://github.com/datafusion-contrib/datafusion-distributed/pull/786) | Distribute feed work across tasks before filling their local partitions. | +| [DFD#787](https://github.com/datafusion-contrib/datafusion-distributed/pull/787) | Use selected-file counts and whole-file bytes to size distributed work. | +| [DFD#788](https://github.com/datafusion-contrib/datafusion-distributed/pull/788) | Balance eager assignments by file size. | +| [DFD#789](https://github.com/datafusion-contrib/datafusion-distributed/pull/789) | Reconstruct statistics for selected files. | +| [DFD#790](https://github.com/datafusion-contrib/datafusion-distributed/pull/790) | Extend placement with affinity. | + +Selection belongs in Iceberg; interpretation as DataFusion statistics belongs +in the integration; worker counts and placement belong in DFD. #67 carries +over the shared flow and ownership lessons without imposing distributed +scheduling on local readers. + +## Upstream need 1: preserve access to selected-file metadata + +At Iceberg revision `5b82ab99`, `ManifestEntryContext` holds the parsed manifest +entry, including `DataFile` column sizes, value/null/NaN counts, and lower/upper +bounds. `process_data_manifest_entry` applies pruning before +[`into_file_scan_task`](https://github.com/apache/iceberg-rust/blob/5b82ab99aeaf50bf4f5b3b776ac9cd6380a5f3e3/crates/iceberg/src/scan/context.rs#L140) constructs the public task. That conversion omits the +column metrics. `plan_files()` exposes only the resulting tasks. + +The metadata may remain in the table's internal object cache, but consumers +cannot obtain it from the selected tasks. DFD's current workaround traverses +manifests again and matches selected paths; its direct manifest reader also +bypasses that cache. Exposing the cache alone would not eliminate the repeated +traversal and matching. + +An additive API could return a stream of planning results: + +```rust,ignore +// Illustrative API, not a definition in this crate. +impl PlannedFile { + pub fn scan_task(&self) -> &FileScanTask; + pub fn data_file(&self) -> &DataFile; + pub fn into_scan_task(self) -> FileScanTask; +} +// TableScan::plan_files_with_metadata() -> stream of Result +``` + +Internally, a result can retain the existing `Arc` beside its +task. The generic [`plan_data_files`](https://github.com/apache/iceberg-rust/blob/5b82ab99aeaf50bf4f5b3b776ac9cd6380a5f3e3/crates/iceberg/src/scan/mod.rs#L412) result builder already supports alternative +outputs; copy-on-write planning already preserves a `DataFile` alongside a +task. Reuse that pipeline rather than expose planning contexts or ask callers +to reproduce pruning. Keep `plan_files()` and task serialization unchanged. +The internal result builder would have this shape: + +```rust,ignore +self.plan_data_files(|ctx| async move { + let entry = Arc::clone(&ctx.manifest_entry); + let task = ctx.into_file_scan_task().await?; + Ok(PlannedFile { entry, task }) +}).await +``` + +Document whether a future split produces multiple tasks for one file: consumers +must aggregate whole-file metrics once per file, not once per task. + +A consumer can aggregate metrics as results arrive, retain only tasks and +aggregate totals, then release each metadata reference. This avoids the second +pass without forcing every scan to retain all metrics. Missing metrics must +remain unknown; file bounds/counts are not exact post-filter or post-delete row +statistics. Snapshot NDV cannot be summed across selected files and remains a +separate, inexact fallback. Iceberg should expose metadata; DataFusion should +own precision and aggregation semantics. + +## Upstream need 2: portable reader settings + +The [#2671 discussion](https://github.com/apache/iceberg-rust/pull/2671#issuecomment-6116801448) +identifies a separate risk: building a reader outside `TableScan::to_arrow()` +can lose scan-level batch size, file concurrency, row-group, row-selection, +and bloom-filter settings. + +Expose a lightweight `ScanReaderOptions` value and an Iceberg-owned method to +apply it to a local `ArrowReaderBuilder`. Have `to_arrow()` use that same +mapping. Distributed consumers transport settings and tasks, then supply +worker-local `FileIO` and runtimes. They should not serialize a builder or +retain an entire `TableScan` merely to reconstruct a reader. + +#67 currently uses default reader settings in both modes, with a deliberate +one-file-at-a-time setting per eager output partition. A single built reader +is cloned across those partitions so its delete loader is shared. That local +cache lifetime is distinct from portable options. + +## Upstream runtime boundary + +`plan_files()` awaits the initial manifest-list read on its caller before +spawning the remaining work onto the table's runtime. `Table::runtime()` is +private. Prefer making `plan_files()` execute that initial read on the table's +I/O runtime too, with cancellation on drop, rather than requiring each consumer +to duplicate runtime dispatch. Until then, #67 requires an I/O-capable planning +caller. Dropping direct collection releases its retained-task reservation; +this does not promise immediate cancellation of all Iceberg background work. + +## Delivery order and validation + +#67 can land using existing public Iceberg APIs. The metadata-result API then +lets selected-file statistics replace DFD's second-pass workaround. Portable +reader options prevent future configuration divergence and support remote +reader construction; they need not block the default-settings foundation. + +Integration coverage should execute the same physical plan it inspects, +compare eager/lazy filtered results, verify empty scans and multi-partition +LIMIT, cover delete files, and check runtime and memory-pool behavior. The +current position-delete fixture checks correctness; equality deletes and +cross-partition delete-cache reuse need dedicated coverage before changing +the default. This proposal makes no measured performance claim. diff --git a/crates/datafusion/src/config.rs b/crates/datafusion/src/config.rs index 4d71716..dbc3011 100644 --- a/crates/datafusion/src/config.rs +++ b/crates/datafusion/src/config.rs @@ -22,9 +22,7 @@ extensions_options! { /// Opt-in file planning during physical planning for local Iceberg scans. pub struct IcebergPlanningConfig { /// Discover and retain file tasks in TableProvider::scan instead of at execution. - pub plan_files: bool, default = false - /// Fail rather than retain more than this many tasks in a single scan. - pub planning_max_files: usize, default = 100_000 + pub eager_file_planning: bool, default = false } } diff --git a/crates/datafusion/src/read/file_planning.rs b/crates/datafusion/src/read/file_planning.rs index 4f6709c..c1ccb89 100644 --- a/crates/datafusion/src/read/file_planning.rs +++ b/crates/datafusion/src/read/file_planning.rs @@ -18,9 +18,8 @@ //! Query-local file work selected before execution, independently of task assignment. use std::mem::{size_of, size_of_val}; -use std::sync::Arc; -use datafusion::common::{Result, exec_err}; +use datafusion::common::Result; use datafusion::execution::TaskContext; use datafusion::execution::memory_pool::{MemoryConsumer, MemoryReservation}; use futures::TryStreamExt; @@ -35,21 +34,17 @@ use crate::to_datafusion_error; /// This is not a reader and contains no coordinator-local reader state. #[derive(Debug)] pub struct PlannedFileScan { - tasks: Vec>, + tasks: Vec, _reservation: MemoryReservation, } impl PlannedFileScan { - /// Tasks in Iceberg's planning order. Each must be assigned exactly once by a consumer. - pub fn tasks(&self) -> &[Arc] { + /// Tasks in Iceberg's planning order. Assign each once per execution. + pub fn tasks(&self) -> &[FileScanTask] { &self.tasks } - pub(crate) async fn collect( - scan: TableScan, - context: &TaskContext, - max_files: usize, - ) -> Result { + pub(crate) async fn collect(scan: TableScan, context: &TaskContext) -> Result { let reservation = MemoryConsumer::new("Iceberg file planning").register(context.memory_pool()); let mut planned = Self { @@ -58,17 +53,12 @@ impl PlannedFileScan { }; let mut stream = scan.plan_files().await.map_err(to_datafusion_error)?; while let Some(task) = stream.try_next().await.map_err(to_datafusion_error)? { - if planned.tasks.len() >= max_files { - return exec_err!( - "Iceberg file planning exceeded planning_max_files={max_files}" - ); - } // Allow for spare Vec capacity and allocator overhead; schemas and specs // are shared Arcs, but paths and delete references are owned by each task. planned ._reservation .try_grow(retained_size(&task).saturating_mul(2))?; - planned.tasks.push(Arc::new(task)); + planned.tasks.push(task); } Ok(planned) } @@ -76,22 +66,16 @@ impl PlannedFileScan { fn retained_size(task: &FileScanTask) -> usize { let mut bytes = size_of::() - + size_of::>() + task.data_file_path().len() + size_of_val(task.project_field_ids()) + size_of_val(task.deletes()) + task.key_metadata().map_or(0, <[u8]>::len); for delete in task.deletes() { bytes = bytes - .saturating_add(delete.file_path.len()) - .saturating_add(delete.referenced_data_file.as_ref().map_or(0, String::len)) - .saturating_add( - delete - .equality_ids - .as_ref() - .map_or(0, |ids| size_of_val(ids.as_slice())), - ) - .saturating_add(delete.key_metadata.as_ref().map_or(0, |key| key.len())); + .saturating_add(delete.file_path().len()) + .saturating_add(delete.referenced_data_file().map_or(0, str::len)) + .saturating_add(delete.equality_ids().map_or(0, size_of_val)) + .saturating_add(delete.key_metadata().map_or(0, <[u8]>::len)); } bytes } diff --git a/crates/datafusion/src/read/scan.rs b/crates/datafusion/src/read/scan.rs index b5048c0..1612442 100644 --- a/crates/datafusion/src/read/scan.rs +++ b/crates/datafusion/src/read/scan.rs @@ -33,18 +33,24 @@ use datafusion::prelude::Expr; use futures::{Stream, StreamExt, TryStreamExt}; use iceberg::arrow::ArrowReader; use iceberg::expr::Predicate; -use iceberg::scan::TableScan; +use iceberg::scan::{FileScanTaskStream, TableScan}; use iceberg::table::Table; use super::file_planning::PlannedFileScan; use super::predicate::convert_filters_to_predicate; -use crate::to_datafusion_error; +use crate::{IcebergPlanningConfig, to_datafusion_error}; -// A reader is created once per eager scan. Its clones share the delete loader's +// A reader is created once per physical scan. Its clones share the delete loader's // cache while each read() call still owns its own scan metrics. #[derive(Clone)] struct ScanReader(ArrowReader); +#[derive(Debug, Clone)] +struct PlannedScan { + files: Arc, + partitions: usize, +} + impl std::fmt::Debug for ScanReader { fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { f.write_str("ScanReader") @@ -67,9 +73,8 @@ pub struct IcebergDataSource { predicate: Option, /// Optional limit on the number of rows to return fetch: Option, - planned_files: Option>, - output_partitions: usize, - reader: Option>, + planned: Option, + reader: ScanReader, } impl IcebergDataSource { @@ -91,39 +96,47 @@ impl IcebergDataSource { pub(crate) async fn with_planned_files( mut self, context: Arc, - max_files: usize, - target_partitions: usize, ) -> Result { + if !context + .session_config() + .options() + .extensions + .get::() + .is_some_and(|config| config.eager_file_planning) + { + return Ok(self); + } + let target_partitions = context.session_config().target_partitions(); let scan = build_table_scan( &self.table, self.snapshot_id, Some(&self.projection), self.predicate.as_ref(), )?; - let planned = PlannedFileScan::collect(scan, &context, max_files).await?; + // Dropping physical planning drops collection and its memory reservation. + let planned = PlannedFileScan::collect(scan, &context).await?; let partitions = planned.tasks().len().max(1).min(target_partitions.max(1)); - // Table::reader_builder retains the table's FileIO and custom runtime. - // This scan currently changes only snapshot, projection, and predicate; - // those values are already carried by FileScanTask. A future scan-specific - // reader setting must be applied here and in TableScan::to_arrow(). - let builder = self.table.reader_builder(); - let builder = if partitions > 1 { - builder.with_data_file_concurrency_limit(1) - } else { - builder - }; - self.reader = Some(Arc::new(ScanReader(builder.build()))); - self.output_partitions = partitions; - self.planned_files = Some(Arc::new(planned)); + if partitions > 1 { + self.reader = ScanReader( + self.table + .reader_builder() + .with_data_file_concurrency_limit(1) + .build(), + ); + } + self.planned = Some(PlannedScan { + files: Arc::new(planned), + partitions, + }); Ok(self) } /// Planned, ungrouped tasks, if file discovery was enabled for this scan. /// A distributed consumer may assign these tasks to its own workers instead /// of using the local output partitions of this source. - pub fn planned_files(&self) -> Option<&Arc> { - self.planned_files.as_ref() + pub fn planned_files(&self) -> Option<&PlannedFileScan> { + self.planned.as_ref().map(|planned| planned.files.as_ref()) } pub fn table(&self) -> &Table { @@ -217,6 +230,7 @@ impl IcebergDataSourceBuilder { .iter() .map(|field| field.name().clone()) .collect(); + let reader = ScanReader(self.table.reader_builder().build()); Ok(IcebergDataSource { table: self.table, snapshot_id: self.snapshot_id, @@ -224,9 +238,8 @@ impl IcebergDataSourceBuilder { projection, predicate: self.predicate, fetch: self.fetch, - planned_files: None, - output_partitions: 1, - reader: None, + planned: None, + reader, }) } } @@ -244,42 +257,47 @@ impl DataSource for IcebergDataSource { partition: usize, _context: Arc, ) -> Result { - let stream: Pin> + Send>> = - if let Some(planned) = &self.planned_files { - if partition >= self.output_partitions { - return internal_err!( - "Iceberg scan partition {partition} does not exist" - ); - } - let Some(reader) = &self.reader else { - return internal_err!("Planned Iceberg scan has no reader"); - }; - let planned = Arc::clone(planned); - let partitions = self.output_partitions; - let tasks = futures::stream::iter( - (partition..planned.tasks().len()).step_by(partitions), - ) - .map(move |index| { - Ok(Arc::unwrap_or_clone(Arc::clone(&planned.tasks()[index]))) - }); + let partitions = self + .planned + .as_ref() + .map_or(1, |planned| planned.partitions); + if partition >= partitions { + return internal_err!("Iceberg scan partition {partition} does not exist"); + } + let tasks: FileScanTaskStream = match &self.planned { + Some(planned) => { + // The reader owns each task; cloning keeps the physical plan reusable. + let files = Arc::clone(&planned.files); Box::pin( - reader - .0 - .clone() - .read(Box::pin(tasks)) - .map_err(to_datafusion_error)? - .stream() - .map_err(to_datafusion_error), + futures::stream::iter( + (partition..files.tasks().len()).step_by(partitions), + ) + .map(move |index| Ok(files.tasks()[index].clone())), ) - } else { - let fut = get_batch_stream( - self.table.clone(), + } + None => { + let scan = build_table_scan( + &self.table, self.snapshot_id, - Some(self.projection.clone()), - self.predicate.clone(), - ); - Box::pin(futures::stream::once(fut).try_flatten()) - }; + Some(&self.projection), + self.predicate.as_ref(), + )?; + Box::pin( + futures::stream::once(async move { scan.plan_files().await }) + .try_flatten(), + ) + } + }; + // Both modes consume a file-task stream through the same reader implementation. + // Reader clones share their delete loader; no per-partition cache is rebuilt. + let stream = self + .reader + .0 + .clone() + .read(tasks) + .map_err(to_datafusion_error)? + .stream() + .map_err(to_datafusion_error); // Apply the fetch if specified. let limited_stream: Pin> + Send>> = @@ -320,14 +338,18 @@ impl DataSource for IcebergDataSource { if let Some(fetch) = self.fetch { write!(f, ", fetch={fetch}")?; } - if let Some(planned) = &self.planned_files { + if let Some(planned) = self.planned_files() { write!(f, ", planned_files={}", planned.tasks().len())?; } Ok(()) } fn output_partitioning(&self) -> Partitioning { - Partitioning::UnknownPartitioning(self.output_partitions) + Partitioning::UnknownPartitioning( + self.planned + .as_ref() + .map_or(1, |planned| planned.partitions), + ) } fn eq_properties(&self) -> EquivalenceProperties { @@ -356,32 +378,6 @@ impl DataSource for IcebergDataSource { } } -/// Asynchronously retrieves a stream of [`RecordBatch`] instances -/// from a given table. -/// -/// This function initializes a [`TableScan`], builds it, -/// and then converts it into a stream of Arrow [`RecordBatch`]es. -async fn get_batch_stream( - table: Table, - snapshot_id: Option, - column_names: Option>, - predicate: Option, -) -> Result> + Send>>> { - let table_scan = build_table_scan( - &table, - snapshot_id, - column_names.as_deref(), - predicate.as_ref(), - )?; - - let stream = table_scan - .to_arrow() - .await - .map_err(to_datafusion_error)? - .map_err(to_datafusion_error); - Ok(Box::pin(stream)) -} - fn build_table_scan( table: &Table, snapshot_id: Option, diff --git a/crates/datafusion/src/table_provider/catalog_table_provider.rs b/crates/datafusion/src/table_provider/catalog_table_provider.rs index a6a0619..a34fe20 100644 --- a/crates/datafusion/src/table_provider/catalog_table_provider.rs +++ b/crates/datafusion/src/table_provider/catalog_table_provider.rs @@ -34,7 +34,6 @@ use iceberg::inspect::MetadataTableType; use iceberg::spec::TableProperties; use iceberg::{Catalog, NamespaceIdent, TableIdent}; -use crate::IcebergPlanningConfig; use crate::error::to_datafusion_error; use crate::metadata::table::IcebergMetadataTableProvider; use crate::read::IcebergDataSource; @@ -140,26 +139,12 @@ impl TableProvider for IcebergCatalogTableProvider { .map_err(to_datafusion_error)?; // Create scan with fresh metadata (always use current snapshot). - let mut source = IcebergDataSource::builder(table, self.schema.clone()) + let source = IcebergDataSource::builder(table, self.schema.clone()) .with_projection(projection) .with_filters(filters) .with_fetch(limit) .build()?; - if let Some(config) = state - .config() - .options() - .extensions - .get::() - .filter(|config| config.plan_files) - { - source = source - .with_planned_files( - state.task_ctx(), - config.planning_max_files, - state.config().target_partitions(), - ) - .await?; - } + let source = source.with_planned_files(state.task_ctx()).await?; Ok(DataSourceExec::from_data_source(source)) } diff --git a/crates/datafusion/src/table_provider/static_table_provider.rs b/crates/datafusion/src/table_provider/static_table_provider.rs index bda652f..09491fb 100644 --- a/crates/datafusion/src/table_provider/static_table_provider.rs +++ b/crates/datafusion/src/table_provider/static_table_provider.rs @@ -30,7 +30,6 @@ use datafusion::physical_plan::ExecutionPlan; use iceberg::arrow::schema_to_arrow_schema; use iceberg::table::Table; -use crate::IcebergPlanningConfig; use crate::error::to_datafusion_error; use crate::read::IcebergDataSource; @@ -127,28 +126,13 @@ impl TableProvider for IcebergStaticTableProvider { limit: Option, ) -> Result> { // Use cached table (no refresh). - let mut source = - IcebergDataSource::builder(self.table.clone(), self.schema.clone()) - .with_snapshot_id(self.snapshot_id) - .with_projection(projection) - .with_filters(filters) - .with_fetch(limit) - .build()?; - if let Some(config) = state - .config() - .options() - .extensions - .get::() - .filter(|config| config.plan_files) - { - source = source - .with_planned_files( - state.task_ctx(), - config.planning_max_files, - state.config().target_partitions(), - ) - .await?; - } + let source = IcebergDataSource::builder(self.table.clone(), self.schema.clone()) + .with_snapshot_id(self.snapshot_id) + .with_projection(projection) + .with_filters(filters) + .with_fetch(limit) + .build()?; + let source = source.with_planned_files(state.task_ctx()).await?; Ok(DataSourceExec::from_data_source(source)) } diff --git a/crates/datafusion/src/test_utils/harness.rs b/crates/datafusion/src/test_utils/harness.rs index 13bea7f..17dc446 100644 --- a/crates/datafusion/src/test_utils/harness.rs +++ b/crates/datafusion/src/test_utils/harness.rs @@ -26,7 +26,7 @@ use datafusion::dataframe::DataFrame; use datafusion::error::Result; use datafusion::execution::SessionStateBuilder; use datafusion::execution::runtime_env::RuntimeEnv; -use datafusion::physical_plan::displayable; +use datafusion::physical_plan::{collect, displayable}; use datafusion::prelude::{SessionConfig, SessionContext}; use futures::StreamExt; use futures::stream::BoxStream; @@ -58,20 +58,30 @@ impl IcebergTestHarness { } pub async fn with_config(config: SessionConfig) -> Result { - Self::build(config, None).await + Self::build(config, None, None).await + } + + /// Load the fixture from a separate directory so a test can write metadata safely. + pub async fn with_fixture_root(config: SessionConfig, root: PathBuf) -> Result { + Self::build(config, None, Some(root)).await } pub async fn with_config_and_runtime( config: SessionConfig, runtime_env: Arc, ) -> Result { - Self::build(config, Some(runtime_env)).await + Self::build(config, Some(runtime_env), None).await } async fn build( config: SessionConfig, runtime_env: Option>, + root: Option, ) -> Result { + let storage_factory: Arc = Arc::new(match root { + Some(root) => FixtureStorageFactory { root }, + None => FixtureStorageFactory::default(), + }); let builder = SessionStateBuilder::new() .with_default_features() .with_config(config); @@ -85,13 +95,13 @@ impl IcebergTestHarness { .with_table_factory( "ICEBERG".to_string(), Arc::new(IcebergTableProviderFactory::new_with_storage_factory( - Arc::new(FixtureStorageFactory::default()), + storage_factory.clone(), )), ) .build(); let catalog = MemoryCatalogBuilder::default() - .with_storage_factory(Arc::new(FixtureStorageFactory::default())) + .with_storage_factory(storage_factory) .load( CATALOG_NAME, HashMap::from([( @@ -129,12 +139,10 @@ impl IcebergTestHarness { pub async fn query(&self, sql: &str) -> Result<(String, String)> { let dataframe: DataFrame = self.ctx.sql(sql).await?; let plan = dataframe.create_physical_plan().await?; - let batches = dataframe.collect().await?; + let plan_display = displayable(plan.as_ref()).indent(true).to_string(); + let batches = collect(plan, Arc::new(dataframe.task_ctx())).await?; - Ok(( - displayable(plan.as_ref()).indent(true).to_string(), - pretty_format_batches(&batches)?.to_string(), - )) + Ok((plan_display, pretty_format_batches(&batches)?.to_string())) } } diff --git a/crates/datafusion/tests/eager_file_planning.rs b/crates/datafusion/tests/eager_file_planning.rs index 05e2fce..ec16269 100644 --- a/crates/datafusion/tests/eager_file_planning.rs +++ b/crates/datafusion/tests/eager_file_planning.rs @@ -15,9 +15,17 @@ // specific language governing permissions and limitations // under the License. +use std::collections::HashMap; +use std::error::Error; +use std::fs; +use std::io; +use std::path::Path; +use std::result::Result as StdResult; use std::sync::Arc; -use datafusion::common::{Result, exec_datafusion_err, stats::Precision}; +use datafusion::arrow::array::{ArrayRef, Int64Array, RecordBatch, StringArray}; +use datafusion::arrow::datatypes::{DataType, Field, Schema}; +use datafusion::common::{Result, stats::Precision}; use datafusion::datasource::TableProvider; use datafusion::datasource::source::{DataSource, DataSourceExec}; use datafusion::execution::memory_pool::GreedyMemoryPool; @@ -27,19 +35,29 @@ use datafusion::physical_plan::common::collect; use datafusion::prelude::{SessionConfig, col, lit}; use datafusion::scalar::ScalarValue; use iceberg::Runtime; +use iceberg::metadata_columns::{ + RESERVED_COL_NAME_DELETE_FILE_PATH, RESERVED_COL_NAME_DELETE_FILE_POS, + RESERVED_FIELD_ID_DELETE_FILE_PATH, RESERVED_FIELD_ID_DELETE_FILE_POS, +}; +use iceberg::spec::{ + DataContentType, DataFileBuilder, DataFileFormat, ManifestListWriter, + ManifestWriterBuilder, TableMetadata, +}; use iceberg::table::Table; use iceberg_datafusion::read::IcebergDataSource; use iceberg_datafusion::test_utils::IcebergTestHarness; use iceberg_datafusion::{ IcebergCatalogTableProvider, IcebergPlanningConfig, IcebergStaticTableProvider, }; +use parquet::arrow::{ArrowWriter, PARQUET_FIELD_ID_META_KEY}; +use serde_json::{Value, json}; +use tempfile::TempDir; -fn eager_config(max_files: usize) -> SessionConfig { +fn eager_config() -> SessionConfig { SessionConfig::new() .with_target_partitions(4) .with_option_extension(IcebergPlanningConfig::default()) - .set_bool("iceberg_local.plan_files", true) - .set_usize("iceberg_local.planning_max_files", max_files) + .set_bool("iceberg_local.eager_file_planning", true) } fn iceberg_source( @@ -52,9 +70,23 @@ fn iceberg_source( .expect("the plan contains the Iceberg data source") } +fn copy_fixture(from: &Path, to: &Path) -> io::Result<()> { + fs::create_dir_all(to)?; + for entry in fs::read_dir(from)? { + let entry = entry?; + let destination = to.join(entry.file_name()); + if entry.file_type()?.is_dir() { + copy_fixture(&entry.path(), &destination)?; + } else { + fs::copy(entry.path(), destination)?; + } + } + Ok(()) +} + #[tokio::test] async fn plans_file_tasks_before_execution_without_changing_statistics() -> Result<()> { - let eager = IcebergTestHarness::with_config(eager_config(100)).await?; + let eager = IcebergTestHarness::with_config(eager_config()).await?; let lazy = IcebergTestHarness::new().await?; let eager_provider = eager.context().table_provider("taxi").await?; let lazy_provider = lazy.context().table_provider("taxi").await?; @@ -71,6 +103,8 @@ async fn plans_file_tasks_before_execution_without_changing_statistics() -> Resu assert_eq!(lazy_plan.output_partitioning().partition_count(), 1); assert_eq!(eager_source.planned_files().unwrap().tasks().len(), 7); assert_eq!(eager_plan.output_partitioning().partition_count(), 4); + assert!(eager_source.open(4, eager.context().task_ctx()).is_err()); + assert!(lazy_source.open(1, lazy.context().task_ctx()).is_err()); assert_eq!( eager_source.partition_statistics(None)?.num_rows, Precision::Absent @@ -103,7 +137,7 @@ async fn plans_file_tasks_before_execution_without_changing_statistics() -> Resu #[tokio::test] async fn eager_partition_count_is_bounded_by_selected_files() -> Result<()> { let eager = - IcebergTestHarness::with_config(eager_config(100).with_target_partitions(16)) + IcebergTestHarness::with_config(eager_config().with_target_partitions(16)) .await?; let provider = eager.context().table_provider("taxi").await?; let plan = provider @@ -115,7 +149,7 @@ async fn eager_partition_count_is_bounded_by_selected_files() -> Result<()> { #[tokio::test] async fn filtered_and_empty_eager_scans_match_lazy_results() -> Result<()> { - let eager = IcebergTestHarness::with_config(eager_config(100)).await?; + let eager = IcebergTestHarness::with_config(eager_config()).await?; let lazy = IcebergTestHarness::new().await?; for (query, planned_files) in [ ( @@ -124,6 +158,11 @@ async fn filtered_and_empty_eager_scans_match_lazy_results() -> Result<()> { GROUP BY pickup_date ORDER BY pickup_date", 1, ), + ( + "SELECT COUNT(*) AS trips FROM taxi \ + WHERE pickup_date = DATE '2024-01-08' AND trip_distance > 50", + 1, + ), ( "SELECT COUNT(*) AS trips FROM taxi \ WHERE pickup_date = DATE '2025-01-01'", @@ -143,7 +182,7 @@ async fn filtered_and_empty_eager_scans_match_lazy_results() -> Result<()> { #[tokio::test] async fn eager_multi_partition_limit_is_global() -> Result<()> { - let eager = IcebergTestHarness::with_config(eager_config(100)).await?; + let eager = IcebergTestHarness::with_config(eager_config()).await?; let (plan, batches) = eager.query("SELECT pickup_date FROM taxi LIMIT 3").await?; assert!(plan.contains("planned_files=7"), "unexpected plan: {plan}"); assert_eq!(batches.matches("2024-01-").count(), 3, "{batches}"); @@ -152,7 +191,7 @@ async fn eager_multi_partition_limit_is_global() -> Result<()> { #[tokio::test] async fn eager_ordered_limit_matches_lazy_results() -> Result<()> { - let eager = IcebergTestHarness::with_config(eager_config(100)).await?; + let eager = IcebergTestHarness::with_config(eager_config()).await?; let lazy = IcebergTestHarness::new().await?; let sql = "SELECT vendor_id, pickup_at FROM taxi \ ORDER BY pickup_at, vendor_id, pickup_location_id LIMIT 3"; @@ -164,52 +203,227 @@ async fn eager_ordered_limit_matches_lazy_results() -> Result<()> { } #[tokio::test] -async fn eager_scan_uses_the_tables_custom_runtime() -> Result<()> { - let harness = IcebergTestHarness::with_config(eager_config(100)).await?; - let provider = harness.context().table_provider("taxi").await?; +async fn eager_and_lazy_scans_apply_position_deletes() -> StdResult<(), Box> { + // Isolate the checked-in fixture: the test writes a new delete manifest and + // snapshot, but must not modify the files used by other tests. + let temporary = TempDir::new()?; + let root = temporary.path().join("iceberg"); + copy_fixture( + &Path::new(env!("CARGO_MANIFEST_DIR")).join("testdata/iceberg"), + &root, + )?; + let eager = IcebergTestHarness::with_fixture_root(eager_config(), root).await?; + let provider = eager.context().table_provider("taxi").await?; let catalog_provider = provider .downcast_ref::() .expect("the fixture creates a catalog Iceberg provider"); let table = catalog_provider .catalog() .load_table(catalog_provider.table_ident()) - .await - .map_err(iceberg_datafusion::to_datafusion_error)?; - let io_runtime = tokio::runtime::Builder::new_multi_thread() - .enable_all() - .build()?; + .await?; + let snapshot = table + .metadata() + .current_snapshot() + .expect("fixture has a snapshot"); + let manifest_list = table.manifest_list_reader(snapshot).load().await?; + let data_manifest = table + .manifest_reader() + .read(&manifest_list.entries()[0]) + .await?; + let data_file = &data_manifest.entries()[0].data_file; + + let delete_path = "s3://iceberg-test/warehouse/taxi/data/pos-delete.parquet"; + let schema = Arc::new(Schema::new(vec![ + Field::new(RESERVED_COL_NAME_DELETE_FILE_PATH, DataType::Utf8, false) + .with_metadata(HashMap::from([( + PARQUET_FIELD_ID_META_KEY.to_string(), + RESERVED_FIELD_ID_DELETE_FILE_PATH.to_string(), + )])), + Field::new(RESERVED_COL_NAME_DELETE_FILE_POS, DataType::Int64, false) + .with_metadata(HashMap::from([( + PARQUET_FIELD_ID_META_KEY.to_string(), + RESERVED_FIELD_ID_DELETE_FILE_POS.to_string(), + )])), + ])); + let batch = RecordBatch::try_new( + schema.clone(), + vec![ + Arc::new(StringArray::from(vec![data_file.file_path()])) as ArrayRef, + Arc::new(Int64Array::from(vec![0])) as ArrayRef, + ], + )?; + let delete_file_on_disk = temporary + .path() + .join("iceberg/taxi/data/pos-delete.parquet"); + let mut writer = + ArrowWriter::try_new(fs::File::create(&delete_file_on_disk)?, schema, None)?; + writer.write(&batch)?; + writer.close()?; + let delete_size = fs::metadata(&delete_file_on_disk)?.len(); + + let next_snapshot_id = snapshot.snapshot_id() + 1; + let mut delete_manifest = ManifestWriterBuilder::new( + table + .file_io() + .new_output("s3://iceberg-test/warehouse/taxi/metadata/pos-delete.avro")?, + Some(next_snapshot_id), + snapshot.schema(table.metadata())?, + table.metadata().default_partition_spec().as_ref().clone(), + ) + .build_v2_deletes(); + delete_manifest.add_file( + DataFileBuilder::default() + .partition_spec_id(table.metadata().default_partition_spec_id()) + .content(DataContentType::PositionDeletes) + .file_path(delete_path.to_string()) + .file_format(DataFileFormat::Parquet) + .file_size_in_bytes(delete_size) + .record_count(1) + .partition(data_file.partition().clone()) + .build()?, + -1, + )?; + let delete_manifest = delete_manifest.write_manifest_file().await?; + let manifest_list_path = + "s3://iceberg-test/warehouse/taxi/metadata/snap-with-delete.avro"; + let mut writer = ManifestListWriter::v2( + table + .file_io() + .new_output(manifest_list_path)? + .writer() + .await?, + next_snapshot_id, + Some(snapshot.snapshot_id()), + snapshot.sequence_number() + 1, + ); + writer.add_manifests( + manifest_list + .entries() + .iter() + .cloned() + .chain([delete_manifest]), + )?; + writer.close().await?; + + let mut metadata: Value = serde_json::from_slice(&fs::read( + temporary + .path() + .join("iceberg/taxi/metadata/v1.metadata.json"), + )?)?; + let mut new_snapshot = metadata["snapshots"][0].clone(); + new_snapshot["snapshot-id"] = json!(next_snapshot_id); + new_snapshot["parent-snapshot-id"] = json!(snapshot.snapshot_id()); + new_snapshot["sequence-number"] = json!(snapshot.sequence_number() + 1); + new_snapshot["manifest-list"] = json!(manifest_list_path); + new_snapshot["timestamp-ms"] = json!(snapshot.timestamp_ms() + 1); + new_snapshot["summary"]["total-delete-files"] = json!("1"); + new_snapshot["summary"]["total-position-deletes"] = json!("1"); + metadata["current-snapshot-id"] = json!(next_snapshot_id); + metadata["last-sequence-number"] = json!(snapshot.sequence_number() + 1); + metadata["last-updated-ms"] = json!(snapshot.timestamp_ms() + 1); + metadata["refs"]["main"]["snapshot-id"] = json!(next_snapshot_id); + metadata["snapshot-log"] + .as_array_mut() + .expect("fixture has a snapshot log") + .push(json!({ + "snapshot-id": next_snapshot_id, + "timestamp-ms": snapshot.timestamp_ms() + 1, + })); + metadata["snapshots"] + .as_array_mut() + .expect("fixture has snapshots") + .push(new_snapshot); let table = Table::builder() .file_io(table.file_io().clone()) - .metadata(table.metadata_ref()) .identifier(table.identifier().clone()) - .runtime(Runtime::new(&io_runtime)) - .build() - .map_err(iceberg_datafusion::to_datafusion_error)?; + .metadata(serde_json::from_value::(metadata)?) + .runtime(Runtime::try_current()?) + .build()?; let provider = IcebergStaticTableProvider::try_new_from_table(table).await?; - let state = harness.context().state(); - let plan = provider.scan(&state, None, &[], None).await?; + let lazy_state = IcebergTestHarness::with_config(SessionConfig::new()).await?; + + let state = eager.context().state(); + let eager_plan = provider.scan(&state, None, &[], None).await?; assert_eq!( - iceberg_source(plan.as_ref()) + iceberg_source(eager_plan.as_ref()) .planned_files() - .unwrap() + .expect("eager plan has files") + .tasks() + .iter() + .filter(|task| !task.deletes().is_empty()) + .count(), + 1 + ); + assert_eq!( + iceberg_source(eager_plan.as_ref()) + .planned_files() + .expect("eager plan has files") .tasks() .len(), 7 ); - let mut rows = 0; - for partition in 0..plan.output_partitioning().partition_count() { - let batches = collect(plan.execute(partition, state.task_ctx())?).await?; - rows += batches.iter().map(|batch| batch.num_rows()).sum::(); + let lazy_plan = provider + .scan(&lazy_state.context().state(), None, &[], None) + .await?; + let mut eager_rows = 0; + for partition in 0..eager_plan.output_partitioning().partition_count() { + for batch in collect(eager_plan.execute(partition, state.task_ctx())?).await? { + eager_rows += batch.num_rows(); + } } - assert_eq!(rows, 175_000); - tokio::task::spawn_blocking(move || drop(io_runtime)) - .await - .map_err(|error| { - exec_datafusion_err!("failed to shut down test runtime: {error}") - })?; + let mut lazy_rows = 0; + for batch in collect(lazy_plan.execute(0, lazy_state.context().task_ctx())?).await? { + lazy_rows += batch.num_rows(); + } + assert_eq!(lazy_rows, 174_999); + assert_eq!(eager_rows, lazy_rows); Ok(()) } +#[test] +fn eager_scan_uses_the_tables_custom_runtime() -> Result<()> { + let io_runtime = tokio::runtime::Builder::new_multi_thread() + .enable_all() + .build()?; + let (harness, provider) = io_runtime.block_on(async { + let harness = IcebergTestHarness::with_config(eager_config()).await?; + let provider = harness.context().table_provider("taxi").await?; + let catalog_provider = provider + .downcast_ref::() + .expect("the fixture creates a catalog Iceberg provider"); + let table = catalog_provider + .catalog() + .load_table(catalog_provider.table_ident()) + .await + .map_err(iceberg_datafusion::to_datafusion_error)?; + let provider = IcebergStaticTableProvider::try_new_from_table(table).await?; + Ok::<_, datafusion::error::DataFusionError>((harness, provider)) + })?; + // The table retains its own runtime while planning is called from another runtime. + let planning_runtime = tokio::runtime::Builder::new_current_thread() + .enable_all() + .build()?; + planning_runtime.block_on(async { + let state = harness.context().state(); + let plan = provider.scan(&state, None, &[], None).await?; + assert_eq!( + iceberg_source(plan.as_ref()) + .planned_files() + .unwrap() + .tasks() + .len(), + 7 + ); + let mut rows = 0; + for partition in 0..plan.output_partitioning().partition_count() { + let batches = collect(plan.execute(partition, state.task_ctx())?).await?; + rows += batches.iter().map(|batch| batch.num_rows()).sum::(); + } + assert_eq!(rows, 175_000); + Ok(()) + }) +} + #[tokio::test] async fn eager_tasks_respect_the_query_memory_pool() -> Result<()> { let runtime = Arc::new( @@ -218,7 +432,8 @@ async fn eager_tasks_respect_the_query_memory_pool() -> Result<()> { .build()?, ); let eager = - IcebergTestHarness::with_config_and_runtime(eager_config(100), runtime).await?; + IcebergTestHarness::with_config_and_runtime(eager_config(), runtime.clone()) + .await?; let provider = eager.context().table_provider("taxi").await?; let error = provider .scan(&eager.context().state(), None, &[], None) @@ -228,20 +443,33 @@ async fn eager_tasks_respect_the_query_memory_pool() -> Result<()> { error.to_string().contains("Iceberg file planning"), "{error}" ); + assert_eq!(runtime.memory_pool.reserved(), 0); Ok(()) } #[tokio::test] -async fn planning_file_limit_rejects_instead_of_truncating() -> Result<()> { - let eager = IcebergTestHarness::with_config(eager_config(2)).await?; - let provider = eager.context().table_provider("taxi").await?; - let error = provider - .scan(&eager.context().state(), None, &[], None) - .await - .expect_err("seven files exceed the two-file planning limit"); - assert!( - error.to_string().contains("planning_max_files=2"), - "{error}" +async fn retained_tasks_follow_plan_and_stream_lifetimes() -> Result<()> { + let runtime = Arc::new( + RuntimeEnvBuilder::new() + .with_memory_pool(Arc::new(GreedyMemoryPool::new(1024 * 1024))) + .build()?, ); + let harness = + IcebergTestHarness::with_config_and_runtime(eager_config(), runtime.clone()) + .await?; + let provider = harness.context().table_provider("taxi").await?; + let plan = provider + .scan(&harness.context().state(), None, &[], None) + .await?; + let reserved = runtime.memory_pool.reserved(); + assert!(reserved > 0); + let source = iceberg_source(plan.as_ref()).clone(); + drop(plan); + assert_eq!(runtime.memory_pool.reserved(), reserved); + let stream = source.open(0, harness.context().task_ctx())?; + drop(source); + assert_eq!(runtime.memory_pool.reserved(), reserved); + drop(stream); + assert_eq!(runtime.memory_pool.reserved(), 0); Ok(()) } diff --git a/crates/datafusion/tests/integration_datafusion_test.rs b/crates/datafusion/tests/integration_datafusion_test.rs index 42d350a..fd1ab67 100644 --- a/crates/datafusion/tests/integration_datafusion_test.rs +++ b/crates/datafusion/tests/integration_datafusion_test.rs @@ -163,7 +163,7 @@ async fn test_provider_plan_stream_schema() -> Result<(), Box> { let task_ctx = Arc::new(df.task_ctx()); let plan = df.create_physical_plan().await.unwrap(); - let stream = plan.execute(1, task_ctx).unwrap(); + let stream = plan.execute(0, task_ctx).unwrap(); // Ensure both the plan and the stream conform to the same schema assert_eq!(plan.schema(), stream.schema()); From 0a88cee3c660823f226585b18a7f0c76ad74d4e4 Mon Sep 17 00:00:00 2001 From: "alex.bianchi" Date: Wed, 21 Oct 2026 23:15:22 -0400 Subject: [PATCH 3/7] docs: streamline eager planning documentation --- crates/datafusion/README.md | 65 ++++-------- crates/datafusion/eager-planning.md | 158 ---------------------------- 2 files changed, 21 insertions(+), 202 deletions(-) delete mode 100644 crates/datafusion/eager-planning.md diff --git a/crates/datafusion/README.md b/crates/datafusion/README.md index c278fc9..fc5508e 100644 --- a/crates/datafusion/README.md +++ b/crates/datafusion/README.md @@ -36,47 +36,24 @@ let config = SessionConfig::new() .set_bool("iceberg_local.eager_file_planning", true); ``` -The option defaults to `false`. When enabled, each physical scan retains its -selected tasks with an estimated reservation in the query's memory pool. There -is no file-count cap: exceeding the memory pool returns an error, never a -partial file list. -Tasks are assigned round-robin to at most the session's `target_partitions`; -an empty scan exposes one empty partition. This does **not** read Parquet data -during planning. The physical scan continues to report unknown statistics. -`iceberg_local` is distinct from DataFusion Distributed's `iceberg` options so -both can be registered in the same session. - -The local reader is constructed with the existing public -`Table::reader_builder()` API, which uses the table's `FileIO` and runtime. -This matches the current integration's reader defaults because it does not -set reader options on `TableScan`: the snapshot, projection, and predicate are -carried by the tasks. For multiple output partitions, each partition reads one -file at a time instead of using the lazy scan's per-reader file concurrency. -**Reader settings are not transferred separately from the scan.** -If scan-specific batch size, data-file concurrency, row-group, row-selection, -or bloom-filter settings are introduced later, the lazy and planned readers -could diverge. A follow-up should expose portable scan-level reader options -and apply them to a locally constructed builder in iceberg-rust, with -`TableScan::to_arrow()` using the same mapping. Distributed consumers should -construct readers using worker-local IO and runtimes rather than transport a -builder or reader. - -Both modes feed tasks to the same reader implementation. One reader is built -per physical scan, and its clones share the delete loader across output -partitions. Dropping eager planning drops collection and its memory reservation. - -`TableScan::plan_files()` performs its initial manifest-list read on the calling -runtime before spawning manifest work on the table's runtime. The planning -caller must therefore support metadata I/O, even with a separately configured -table runtime. The table's runtime accessor is private; making `plan_files()` -offload that initial read is an upstream follow-up. - -Planning the whole file list earlier adds metadata I/O and retained-task memory -even when a query stops early (for example, `LIMIT`), and equivalent provider -`scan()` calls currently repeat planning. No performance improvement is claimed -by this option alone. A position-delete fixture verifies eager/lazy row counts; -equality deletes and cross-partition delete-cache reuse still need dedicated -validation before enabling it by default. - -See [the planning API discussion](eager-planning.md) for the relationship to -distributed planning and the proposed upstream metadata and reader-settings APIs. +### Behavior + +- Lazy planning is the default. Eager planning selects files before execution; + it reads metadata, not Parquet data. +- Selected tasks are assigned round-robin to at most `target_partitions`. + An empty scan has one empty partition. Each eager partition reads one file + at a time when there are multiple output partitions. +- Retained tasks have an estimated reservation in the query memory pool. + Exceeding the pool fails planning; there is no file-count cap or truncation. +- Both modes use the same reader implementation. Reader clones share the + delete loader across partitions. Scan statistics remain unknown. + +### Trade-offs + +Eager planning completes file discovery and retains tasks even for queries +that stop early. Repeated provider scans repeat planning. The planning caller +must support metadata I/O: Iceberg's initial manifest-list read still runs on +that runtime, even when the table has a separate I/O runtime. + +The `iceberg_local` prefix allows these options to coexist with DataFusion +Distributed's `iceberg` options. diff --git a/crates/datafusion/eager-planning.md b/crates/datafusion/eager-planning.md deleted file mode 100644 index b43bcce..0000000 --- a/crates/datafusion/eager-planning.md +++ /dev/null @@ -1,158 +0,0 @@ - - -# Eager planning across Iceberg, DataFusion, and distributed execution - -This document describes the boundary proposed by -[datafusion-iceberg#67](https://github.com/apache/datafusion-iceberg/pull/67). -API names below are illustrative, not APIs provided by this change. - -## One selection pipeline, two consumption times - -```text -Table metadata → snapshot → manifest list → selected manifests - → live entries → partition/metrics pruning → file tasks + delete references - -Lazy: execution starts → stream selected tasks → reader -Eager: planning starts → collect selected tasks → assign partitions - → execution starts → reader -``` - -The catalog does not return the complete file list. Reading the manifest list -and manifests is metadata I/O, even when Parquet reads are deferred. Eager -planning pays that cost before execution and retains tasks, allowing the -optimizer to see the selected work. Lazy planning can overlap discovery and -execution and avoid retaining the full task list. Both use Iceberg's pruning -and delete matching, and both retain residual DataFusion filters. - -#67 keeps both modes, defaults to lazy, and uses the same local reader flow. -It exposes borrowed access to retained tasks, accounts for their estimated -memory, and assigns them round-robin to local partitions. It does not add -column statistics, distributed routing, or a file-count limit. Statistics -remain unknown in both modes. - -## What the distributed stack adds - -| Change | Responsibility | -| --- | --- | -| [DFD#773](https://github.com/datafusion-contrib/datafusion-distributed/pull/773) | Select and retain work during planning while preserving lazy execution. | -| [DFD#786](https://github.com/datafusion-contrib/datafusion-distributed/pull/786) | Distribute feed work across tasks before filling their local partitions. | -| [DFD#787](https://github.com/datafusion-contrib/datafusion-distributed/pull/787) | Use selected-file counts and whole-file bytes to size distributed work. | -| [DFD#788](https://github.com/datafusion-contrib/datafusion-distributed/pull/788) | Balance eager assignments by file size. | -| [DFD#789](https://github.com/datafusion-contrib/datafusion-distributed/pull/789) | Reconstruct statistics for selected files. | -| [DFD#790](https://github.com/datafusion-contrib/datafusion-distributed/pull/790) | Extend placement with affinity. | - -Selection belongs in Iceberg; interpretation as DataFusion statistics belongs -in the integration; worker counts and placement belong in DFD. #67 carries -over the shared flow and ownership lessons without imposing distributed -scheduling on local readers. - -## Upstream need 1: preserve access to selected-file metadata - -At Iceberg revision `5b82ab99`, `ManifestEntryContext` holds the parsed manifest -entry, including `DataFile` column sizes, value/null/NaN counts, and lower/upper -bounds. `process_data_manifest_entry` applies pruning before -[`into_file_scan_task`](https://github.com/apache/iceberg-rust/blob/5b82ab99aeaf50bf4f5b3b776ac9cd6380a5f3e3/crates/iceberg/src/scan/context.rs#L140) constructs the public task. That conversion omits the -column metrics. `plan_files()` exposes only the resulting tasks. - -The metadata may remain in the table's internal object cache, but consumers -cannot obtain it from the selected tasks. DFD's current workaround traverses -manifests again and matches selected paths; its direct manifest reader also -bypasses that cache. Exposing the cache alone would not eliminate the repeated -traversal and matching. - -An additive API could return a stream of planning results: - -```rust,ignore -// Illustrative API, not a definition in this crate. -impl PlannedFile { - pub fn scan_task(&self) -> &FileScanTask; - pub fn data_file(&self) -> &DataFile; - pub fn into_scan_task(self) -> FileScanTask; -} -// TableScan::plan_files_with_metadata() -> stream of Result -``` - -Internally, a result can retain the existing `Arc` beside its -task. The generic [`plan_data_files`](https://github.com/apache/iceberg-rust/blob/5b82ab99aeaf50bf4f5b3b776ac9cd6380a5f3e3/crates/iceberg/src/scan/mod.rs#L412) result builder already supports alternative -outputs; copy-on-write planning already preserves a `DataFile` alongside a -task. Reuse that pipeline rather than expose planning contexts or ask callers -to reproduce pruning. Keep `plan_files()` and task serialization unchanged. -The internal result builder would have this shape: - -```rust,ignore -self.plan_data_files(|ctx| async move { - let entry = Arc::clone(&ctx.manifest_entry); - let task = ctx.into_file_scan_task().await?; - Ok(PlannedFile { entry, task }) -}).await -``` - -Document whether a future split produces multiple tasks for one file: consumers -must aggregate whole-file metrics once per file, not once per task. - -A consumer can aggregate metrics as results arrive, retain only tasks and -aggregate totals, then release each metadata reference. This avoids the second -pass without forcing every scan to retain all metrics. Missing metrics must -remain unknown; file bounds/counts are not exact post-filter or post-delete row -statistics. Snapshot NDV cannot be summed across selected files and remains a -separate, inexact fallback. Iceberg should expose metadata; DataFusion should -own precision and aggregation semantics. - -## Upstream need 2: portable reader settings - -The [#2671 discussion](https://github.com/apache/iceberg-rust/pull/2671#issuecomment-6116801448) -identifies a separate risk: building a reader outside `TableScan::to_arrow()` -can lose scan-level batch size, file concurrency, row-group, row-selection, -and bloom-filter settings. - -Expose a lightweight `ScanReaderOptions` value and an Iceberg-owned method to -apply it to a local `ArrowReaderBuilder`. Have `to_arrow()` use that same -mapping. Distributed consumers transport settings and tasks, then supply -worker-local `FileIO` and runtimes. They should not serialize a builder or -retain an entire `TableScan` merely to reconstruct a reader. - -#67 currently uses default reader settings in both modes, with a deliberate -one-file-at-a-time setting per eager output partition. A single built reader -is cloned across those partitions so its delete loader is shared. That local -cache lifetime is distinct from portable options. - -## Upstream runtime boundary - -`plan_files()` awaits the initial manifest-list read on its caller before -spawning the remaining work onto the table's runtime. `Table::runtime()` is -private. Prefer making `plan_files()` execute that initial read on the table's -I/O runtime too, with cancellation on drop, rather than requiring each consumer -to duplicate runtime dispatch. Until then, #67 requires an I/O-capable planning -caller. Dropping direct collection releases its retained-task reservation; -this does not promise immediate cancellation of all Iceberg background work. - -## Delivery order and validation - -#67 can land using existing public Iceberg APIs. The metadata-result API then -lets selected-file statistics replace DFD's second-pass workaround. Portable -reader options prevent future configuration divergence and support remote -reader construction; they need not block the default-settings foundation. - -Integration coverage should execute the same physical plan it inspects, -compare eager/lazy filtered results, verify empty scans and multi-partition -LIMIT, cover delete files, and check runtime and memory-pool behavior. The -current position-delete fixture checks correctness; equality deletes and -cross-partition delete-cache reuse need dedicated coverage before changing -the default. This proposal makes no measured performance claim. From 15cc28ec1d882e89ce8064b0468c6a6299d94766 Mon Sep 17 00:00:00 2001 From: "alex.bianchi" Date: Wed, 21 Oct 2026 23:30:53 -0400 Subject: [PATCH 4/7] refactor: plan Iceberg files through the source builder --- crates/datafusion/src/read/mod.rs | 4 +- crates/datafusion/src/read/scan.rs | 150 ++++++++++++------ .../table_provider/catalog_table_provider.rs | 3 +- .../table_provider/static_table_provider.rs | 3 +- .../datafusion/tests/eager_file_planning.rs | 34 ++++ 5 files changed, 144 insertions(+), 50 deletions(-) diff --git a/crates/datafusion/src/read/mod.rs b/crates/datafusion/src/read/mod.rs index ceb2817..42626c3 100644 --- a/crates/datafusion/src/read/mod.rs +++ b/crates/datafusion/src/read/mod.rs @@ -21,4 +21,6 @@ pub(crate) mod scan; pub use file_planning::PlannedFileScan; pub use predicate::convert_filters_to_predicate; -pub use scan::{IcebergDataSource, IcebergDataSourceBuilder}; +pub use scan::{ + IcebergDataSource, IcebergDataSourceBuilder, PlannedIcebergDataSourceBuilder, +}; diff --git a/crates/datafusion/src/read/scan.rs b/crates/datafusion/src/read/scan.rs index 1612442..a6123e9 100644 --- a/crates/datafusion/src/read/scan.rs +++ b/crates/datafusion/src/read/scan.rs @@ -92,46 +92,6 @@ impl IcebergDataSource { } } - /// Plan file tasks during physical planning; read them only at execution. - pub(crate) async fn with_planned_files( - mut self, - context: Arc, - ) -> Result { - if !context - .session_config() - .options() - .extensions - .get::() - .is_some_and(|config| config.eager_file_planning) - { - return Ok(self); - } - let target_partitions = context.session_config().target_partitions(); - let scan = build_table_scan( - &self.table, - self.snapshot_id, - Some(&self.projection), - self.predicate.as_ref(), - )?; - // Dropping physical planning drops collection and its memory reservation. - let planned = PlannedFileScan::collect(scan, &context).await?; - let partitions = planned.tasks().len().max(1).min(target_partitions.max(1)); - - if partitions > 1 { - self.reader = ScanReader( - self.table - .reader_builder() - .with_data_file_concurrency_limit(1) - .build(), - ); - } - self.planned = Some(PlannedScan { - files: Arc::new(planned), - partitions, - }); - Ok(self) - } - /// Planned, ungrouped tasks, if file discovery was enabled for this scan. /// A distributed consumer may assign these tasks to its own workers instead /// of using the local output partitions of this source. @@ -167,7 +127,10 @@ impl IcebergDataSource { /// Builds an [`IcebergDataSource`]. Created by [`IcebergDataSource::builder`]. /// /// A table provider passes the table's schema and selects columns with -/// [`Self::with_projection`]. To rebuild a source, such as in another process, +/// [`Self::with_projection`]. Finish with [`Self::with_planned_files`] to honor +/// the session's eager-planning option before building. +/// +/// To rebuild a source, such as in another process, /// pass its [`IcebergDataSource::schema`], which is already projected, and no /// projection. This rebuilds a lazy single-partition source; planned file tasks /// must be transferred separately when reconstructing a distributed scan. @@ -215,22 +178,94 @@ impl IcebergDataSourceBuilder { self } + /// Finish configuring the scan and optionally select its files before execution. + /// + /// Honors [`IcebergPlanningConfig::eager_file_planning`]. The returned builder + /// only exposes `build`, so filters, projection, and snapshot cannot change + /// after files have been selected. + /// + /// # Errors + /// + /// Returns an error if the projection is invalid, file planning fails, or + /// the retained tasks exceed the query memory pool. + pub async fn with_planned_files( + self, + context: Arc, + ) -> Result { + if !context + .session_config() + .options() + .extensions + .get::() + .is_some_and(|config| config.eager_file_planning) + { + return Ok(PlannedIcebergDataSourceBuilder { + builder: self, + planned: None, + }); + } + let target_partitions = context.session_config().target_partitions(); + let schema = self.projected_schema()?; + let projection: Vec = schema + .fields() + .iter() + .map(|field| field.name().clone()) + .collect(); + let scan = build_table_scan( + &self.table, + self.snapshot_id, + Some(&projection), + self.predicate.as_ref(), + )?; + // Dropping physical planning drops collection and its memory reservation. + let planned = PlannedFileScan::collect(scan, &context).await?; + let partitions = planned.tasks().len().max(1).min(target_partitions.max(1)); + + Ok(PlannedIcebergDataSourceBuilder { + builder: self, + planned: Some(PlannedScan { + files: Arc::new(planned), + partitions, + }), + }) + } + /// Builds the source, which reads the columns of its projected schema by name. /// /// # Errors /// /// Returns an error if the projection is out of bounds for the schema. pub fn build(self) -> Result { - let schema = match &self.projection { - None => self.schema, - Some(projection) => Arc::new(self.schema.project(projection)?), - }; + self.build_with_planned_files(None) + } + + fn projected_schema(&self) -> Result { + match &self.projection { + None => Ok(self.schema.clone()), + Some(projection) => Ok(Arc::new(self.schema.project(projection)?)), + } + } + + fn build_with_planned_files( + self, + planned: Option, + ) -> Result { + let schema = self.projected_schema()?; let projection = schema .fields() .iter() .map(|field| field.name().clone()) .collect(); - let reader = ScanReader(self.table.reader_builder().build()); + let reader = self.table.reader_builder(); + let reader = if planned + .as_ref() + .is_some_and(|planned| planned.partitions > 1) + { + reader.with_data_file_concurrency_limit(1) + } else { + reader + }; + let reader = ScanReader(reader.build()); Ok(IcebergDataSource { table: self.table, snapshot_id: self.snapshot_id, @@ -238,12 +273,33 @@ impl IcebergDataSourceBuilder { projection, predicate: self.predicate, fetch: self.fetch, - planned: None, + planned, reader, }) } } +/// A configured scan ready to build after optional eager file planning. +/// +/// Created by [`IcebergDataSourceBuilder::with_planned_files`]. Configuration +/// is fixed so the retained tasks stay consistent with the resulting source. +#[derive(Debug, Clone)] +pub struct PlannedIcebergDataSourceBuilder { + builder: IcebergDataSourceBuilder, + planned: Option, +} + +impl PlannedIcebergDataSourceBuilder { + /// Construct the source and its shared reader from the configured scan. + /// + /// # Errors + /// + /// Returns an error if the projection is out of bounds for the schema. + pub fn build(self) -> Result { + self.builder.build_with_planned_files(self.planned) + } +} + impl DataSource for IcebergDataSource { fn apply_expressions( &self, diff --git a/crates/datafusion/src/table_provider/catalog_table_provider.rs b/crates/datafusion/src/table_provider/catalog_table_provider.rs index a34fe20..00d6145 100644 --- a/crates/datafusion/src/table_provider/catalog_table_provider.rs +++ b/crates/datafusion/src/table_provider/catalog_table_provider.rs @@ -143,8 +143,9 @@ impl TableProvider for IcebergCatalogTableProvider { .with_projection(projection) .with_filters(filters) .with_fetch(limit) + .with_planned_files(state.task_ctx()) + .await? .build()?; - let source = source.with_planned_files(state.task_ctx()).await?; Ok(DataSourceExec::from_data_source(source)) } diff --git a/crates/datafusion/src/table_provider/static_table_provider.rs b/crates/datafusion/src/table_provider/static_table_provider.rs index 09491fb..ef1857c 100644 --- a/crates/datafusion/src/table_provider/static_table_provider.rs +++ b/crates/datafusion/src/table_provider/static_table_provider.rs @@ -131,8 +131,9 @@ impl TableProvider for IcebergStaticTableProvider { .with_projection(projection) .with_filters(filters) .with_fetch(limit) + .with_planned_files(state.task_ctx()) + .await? .build()?; - let source = source.with_planned_files(state.task_ctx()).await?; Ok(DataSourceExec::from_data_source(source)) } diff --git a/crates/datafusion/tests/eager_file_planning.rs b/crates/datafusion/tests/eager_file_planning.rs index ec16269..2e18776 100644 --- a/crates/datafusion/tests/eager_file_planning.rs +++ b/crates/datafusion/tests/eager_file_planning.rs @@ -473,3 +473,37 @@ async fn retained_tasks_follow_plan_and_stream_lifetimes() -> Result<()> { assert_eq!(runtime.memory_pool.reserved(), 0); Ok(()) } + +#[tokio::test] +async fn public_builder_plans_projected_filtered_scans() -> Result<()> { + let harness = IcebergTestHarness::with_config(eager_config()).await?; + let provider = harness.context().table_provider("taxi").await?; + let provider = provider + .downcast_ref::() + .unwrap(); + let table = provider + .catalog() + .load_table(provider.table_ident()) + .await + .map_err(iceberg_datafusion::to_datafusion_error)?; + let schema = provider.schema(); + let projection = vec![schema.index_of("pickup_date")?]; + let pickup_date = ScalarValue::Date32(Some(19732)); // 2024-01-10 + let source = IcebergDataSource::builder(table, schema) + .with_projection(Some(&projection)) + .with_filters(&[col("pickup_date").eq(lit(pickup_date))]) + .with_planned_files(harness.context().task_ctx()) + .await? + .build()?; + assert_eq!(source.projection(), &["pickup_date"]); + assert_eq!(source.planned_files().unwrap().tasks().len(), 1); + assert_eq!(source.output_partitioning().partition_count(), 1); + let batches = collect(source.open(0, harness.context().task_ctx())?).await?; + assert!(!batches.is_empty()); + assert!( + batches + .iter() + .all(|batch| batch.schema() == source.schema()) + ); + Ok(()) +} From 58b8e63e2e94d07156b32b57a8413be4576f75d0 Mon Sep 17 00:00:00 2001 From: "alex.bianchi" Date: Wed, 21 Oct 2026 23:47:00 -0400 Subject: [PATCH 5/7] test: simplify eager planning fixtures with Iceberg builders --- Cargo.lock | 1 - crates/datafusion/Cargo.toml | 1 - crates/datafusion/tests/common/mod.rs | 18 ++ .../tests/common/position_deletes.rs | 197 ++++++++++++++++++ .../datafusion/tests/eager_file_planning.rs | 191 ++--------------- 5 files changed, 227 insertions(+), 181 deletions(-) create mode 100644 crates/datafusion/tests/common/mod.rs create mode 100644 crates/datafusion/tests/common/position_deletes.rs diff --git a/Cargo.lock b/Cargo.lock index 4c19a68..5182cdf 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -3247,7 +3247,6 @@ dependencies = [ "insta", "parquet", "serde", - "serde_json", "tempfile", "tokio", "typetag", diff --git a/crates/datafusion/Cargo.toml b/crates/datafusion/Cargo.toml index 4320866..cd6edd0 100644 --- a/crates/datafusion/Cargo.toml +++ b/crates/datafusion/Cargo.toml @@ -44,7 +44,6 @@ typetag = "0.2" [dev-dependencies] insta = "1" parquet = "59.2" -serde_json = "1" tempfile = "3.18" [lints] diff --git a/crates/datafusion/tests/common/mod.rs b/crates/datafusion/tests/common/mod.rs new file mode 100644 index 0000000..f34fe8d --- /dev/null +++ b/crates/datafusion/tests/common/mod.rs @@ -0,0 +1,18 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +pub mod position_deletes; diff --git a/crates/datafusion/tests/common/position_deletes.rs b/crates/datafusion/tests/common/position_deletes.rs new file mode 100644 index 0000000..ca3ff51 --- /dev/null +++ b/crates/datafusion/tests/common/position_deletes.rs @@ -0,0 +1,197 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +use std::collections::HashMap; +use std::error::Error; +use std::fs; +use std::io; +use std::path::Path; +use std::sync::Arc; + +use datafusion::arrow::array::{ArrayRef, Int64Array, RecordBatch, StringArray}; +use datafusion::arrow::datatypes::{DataType, Field, Schema}; +use datafusion::prelude::SessionConfig; +use iceberg::Runtime; +use iceberg::metadata_columns::{ + RESERVED_COL_NAME_DELETE_FILE_PATH, RESERVED_COL_NAME_DELETE_FILE_POS, + RESERVED_FIELD_ID_DELETE_FILE_PATH, RESERVED_FIELD_ID_DELETE_FILE_POS, +}; +use iceberg::spec::{ + DataContentType, DataFileBuilder, DataFileFormat, ManifestListWriter, + ManifestWriterBuilder, Operation, Snapshot, Summary, +}; +use iceberg::table::Table; +use iceberg_datafusion::IcebergCatalogTableProvider; +use iceberg_datafusion::test_utils::IcebergTestHarness; +use parquet::arrow::{ArrowWriter, PARQUET_FIELD_ID_META_KEY}; +use tempfile::TempDir; + +/// A taxi table with its first data file's first row deleted. +/// Owns the copied files for the lifetime of scans against this table. +pub struct PositionDeleteFixture { + pub table: Table, + _directory: TempDir, +} + +pub async fn taxi_with_position_delete() -> Result> +{ + // Isolate the checked-in fixture: the test writes a new delete manifest and + // snapshot, but must not modify the files used by other tests. + let temporary = TempDir::new()?; + let root = temporary.path().join("iceberg"); + copy_fixture( + &Path::new(env!("CARGO_MANIFEST_DIR")).join("testdata/iceberg"), + &root, + )?; + let harness = + IcebergTestHarness::with_fixture_root(SessionConfig::new(), root).await?; + let provider = harness.context().table_provider("taxi").await?; + let catalog_provider = provider + .downcast_ref::() + .expect("the fixture creates a catalog Iceberg provider"); + let table = catalog_provider + .catalog() + .load_table(catalog_provider.table_ident()) + .await?; + let snapshot = table + .metadata() + .current_snapshot() + .expect("fixture has a snapshot"); + let manifest_list = table.manifest_list_reader(snapshot).load().await?; + let data_manifest = table + .manifest_reader() + .read(&manifest_list.entries()[0]) + .await?; + let data_file = &data_manifest.entries()[0].data_file; + + let delete_path = "s3://iceberg-test/warehouse/taxi/data/pos-delete.parquet"; + let schema = Arc::new(Schema::new(vec![ + Field::new(RESERVED_COL_NAME_DELETE_FILE_PATH, DataType::Utf8, false) + .with_metadata(HashMap::from([( + PARQUET_FIELD_ID_META_KEY.to_string(), + RESERVED_FIELD_ID_DELETE_FILE_PATH.to_string(), + )])), + Field::new(RESERVED_COL_NAME_DELETE_FILE_POS, DataType::Int64, false) + .with_metadata(HashMap::from([( + PARQUET_FIELD_ID_META_KEY.to_string(), + RESERVED_FIELD_ID_DELETE_FILE_POS.to_string(), + )])), + ])); + let batch = RecordBatch::try_new( + schema.clone(), + vec![ + Arc::new(StringArray::from(vec![data_file.file_path()])) as ArrayRef, + Arc::new(Int64Array::from(vec![0])) as ArrayRef, + ], + )?; + let delete_file_on_disk = temporary + .path() + .join("iceberg/taxi/data/pos-delete.parquet"); + let mut writer = + ArrowWriter::try_new(fs::File::create(&delete_file_on_disk)?, schema, None)?; + writer.write(&batch)?; + writer.close()?; + let delete_size = fs::metadata(&delete_file_on_disk)?.len(); + + let next_snapshot_id = snapshot.snapshot_id() + 1; + let mut delete_manifest = ManifestWriterBuilder::new( + table + .file_io() + .new_output("s3://iceberg-test/warehouse/taxi/metadata/pos-delete.avro")?, + Some(next_snapshot_id), + snapshot.schema(table.metadata())?, + table.metadata().default_partition_spec().as_ref().clone(), + ) + .build_v2_deletes(); + delete_manifest.add_file( + DataFileBuilder::default() + .partition_spec_id(table.metadata().default_partition_spec_id()) + .content(DataContentType::PositionDeletes) + .file_path(delete_path.to_string()) + .file_format(DataFileFormat::Parquet) + .file_size_in_bytes(delete_size) + .record_count(1) + .partition(data_file.partition().clone()) + .build()?, + -1, + )?; + let delete_manifest = delete_manifest.write_manifest_file().await?; + let manifest_list_path = + "s3://iceberg-test/warehouse/taxi/metadata/snap-with-delete.avro"; + let mut writer = ManifestListWriter::v2( + table + .file_io() + .new_output(manifest_list_path)? + .writer() + .await?, + next_snapshot_id, + Some(snapshot.snapshot_id()), + snapshot.sequence_number() + 1, + ); + writer.add_manifests( + manifest_list + .entries() + .iter() + .cloned() + .chain([delete_manifest]), + )?; + writer.close().await?; + + let new_snapshot = Snapshot::builder() + .with_snapshot_id(next_snapshot_id) + .with_parent_snapshot_id(Some(snapshot.snapshot_id())) + .with_sequence_number(snapshot.sequence_number() + 1) + .with_timestamp_ms(snapshot.timestamp_ms() + 1) + .with_manifest_list(manifest_list_path) + .with_summary(Summary { + operation: Operation::Delete, + additional_properties: HashMap::new(), + }) + .with_schema_id(snapshot.schema(table.metadata())?.schema_id()) + .build(); + let metadata = table + .metadata() + .clone() + .into_builder(None) + .set_branch_snapshot(new_snapshot, "main")? + .build()? + .metadata; + let table = Table::builder() + .file_io(table.file_io().clone()) + .identifier(table.identifier().clone()) + .metadata(metadata) + .runtime(Runtime::try_current()?) + .build()?; + Ok(PositionDeleteFixture { + table, + _directory: temporary, + }) +} + +fn copy_fixture(from: &Path, to: &Path) -> io::Result<()> { + fs::create_dir_all(to)?; + for entry in fs::read_dir(from)? { + let entry = entry?; + let destination = to.join(entry.file_name()); + if entry.file_type()?.is_dir() { + copy_fixture(&entry.path(), &destination)?; + } else { + fs::copy(entry.path(), destination)?; + } + } + Ok(()) +} diff --git a/crates/datafusion/tests/eager_file_planning.rs b/crates/datafusion/tests/eager_file_planning.rs index 2e18776..f356356 100644 --- a/crates/datafusion/tests/eager_file_planning.rs +++ b/crates/datafusion/tests/eager_file_planning.rs @@ -15,43 +15,28 @@ // specific language governing permissions and limitations // under the License. -use std::collections::HashMap; use std::error::Error; -use std::fs; -use std::io; -use std::path::Path; use std::result::Result as StdResult; use std::sync::Arc; -use datafusion::arrow::array::{ArrayRef, Int64Array, RecordBatch, StringArray}; -use datafusion::arrow::datatypes::{DataType, Field, Schema}; use datafusion::common::{Result, stats::Precision}; use datafusion::datasource::TableProvider; use datafusion::datasource::source::{DataSource, DataSourceExec}; use datafusion::execution::memory_pool::GreedyMemoryPool; use datafusion::execution::runtime_env::RuntimeEnvBuilder; -use datafusion::physical_plan::ExecutionPlanProperties; use datafusion::physical_plan::common::collect; +use datafusion::physical_plan::{ExecutionPlanProperties, collect as collect_plan}; use datafusion::prelude::{SessionConfig, col, lit}; use datafusion::scalar::ScalarValue; -use iceberg::Runtime; -use iceberg::metadata_columns::{ - RESERVED_COL_NAME_DELETE_FILE_PATH, RESERVED_COL_NAME_DELETE_FILE_POS, - RESERVED_FIELD_ID_DELETE_FILE_PATH, RESERVED_FIELD_ID_DELETE_FILE_POS, -}; -use iceberg::spec::{ - DataContentType, DataFileBuilder, DataFileFormat, ManifestListWriter, - ManifestWriterBuilder, TableMetadata, -}; -use iceberg::table::Table; use iceberg_datafusion::read::IcebergDataSource; use iceberg_datafusion::test_utils::IcebergTestHarness; use iceberg_datafusion::{ IcebergCatalogTableProvider, IcebergPlanningConfig, IcebergStaticTableProvider, }; -use parquet::arrow::{ArrowWriter, PARQUET_FIELD_ID_META_KEY}; -use serde_json::{Value, json}; -use tempfile::TempDir; + +mod common; + +use common::position_deletes::taxi_with_position_delete; fn eager_config() -> SessionConfig { SessionConfig::new() @@ -70,20 +55,6 @@ fn iceberg_source( .expect("the plan contains the Iceberg data source") } -fn copy_fixture(from: &Path, to: &Path) -> io::Result<()> { - fs::create_dir_all(to)?; - for entry in fs::read_dir(from)? { - let entry = entry?; - let destination = to.join(entry.file_name()); - if entry.file_type()?.is_dir() { - copy_fixture(&entry.path(), &destination)?; - } else { - fs::copy(entry.path(), destination)?; - } - } - Ok(()) -} - #[tokio::test] async fn plans_file_tasks_before_execution_without_changing_statistics() -> Result<()> { let eager = IcebergTestHarness::with_config(eager_config()).await?; @@ -204,141 +175,9 @@ async fn eager_ordered_limit_matches_lazy_results() -> Result<()> { #[tokio::test] async fn eager_and_lazy_scans_apply_position_deletes() -> StdResult<(), Box> { - // Isolate the checked-in fixture: the test writes a new delete manifest and - // snapshot, but must not modify the files used by other tests. - let temporary = TempDir::new()?; - let root = temporary.path().join("iceberg"); - copy_fixture( - &Path::new(env!("CARGO_MANIFEST_DIR")).join("testdata/iceberg"), - &root, - )?; - let eager = IcebergTestHarness::with_fixture_root(eager_config(), root).await?; - let provider = eager.context().table_provider("taxi").await?; - let catalog_provider = provider - .downcast_ref::() - .expect("the fixture creates a catalog Iceberg provider"); - let table = catalog_provider - .catalog() - .load_table(catalog_provider.table_ident()) - .await?; - let snapshot = table - .metadata() - .current_snapshot() - .expect("fixture has a snapshot"); - let manifest_list = table.manifest_list_reader(snapshot).load().await?; - let data_manifest = table - .manifest_reader() - .read(&manifest_list.entries()[0]) - .await?; - let data_file = &data_manifest.entries()[0].data_file; - - let delete_path = "s3://iceberg-test/warehouse/taxi/data/pos-delete.parquet"; - let schema = Arc::new(Schema::new(vec![ - Field::new(RESERVED_COL_NAME_DELETE_FILE_PATH, DataType::Utf8, false) - .with_metadata(HashMap::from([( - PARQUET_FIELD_ID_META_KEY.to_string(), - RESERVED_FIELD_ID_DELETE_FILE_PATH.to_string(), - )])), - Field::new(RESERVED_COL_NAME_DELETE_FILE_POS, DataType::Int64, false) - .with_metadata(HashMap::from([( - PARQUET_FIELD_ID_META_KEY.to_string(), - RESERVED_FIELD_ID_DELETE_FILE_POS.to_string(), - )])), - ])); - let batch = RecordBatch::try_new( - schema.clone(), - vec![ - Arc::new(StringArray::from(vec![data_file.file_path()])) as ArrayRef, - Arc::new(Int64Array::from(vec![0])) as ArrayRef, - ], - )?; - let delete_file_on_disk = temporary - .path() - .join("iceberg/taxi/data/pos-delete.parquet"); - let mut writer = - ArrowWriter::try_new(fs::File::create(&delete_file_on_disk)?, schema, None)?; - writer.write(&batch)?; - writer.close()?; - let delete_size = fs::metadata(&delete_file_on_disk)?.len(); - - let next_snapshot_id = snapshot.snapshot_id() + 1; - let mut delete_manifest = ManifestWriterBuilder::new( - table - .file_io() - .new_output("s3://iceberg-test/warehouse/taxi/metadata/pos-delete.avro")?, - Some(next_snapshot_id), - snapshot.schema(table.metadata())?, - table.metadata().default_partition_spec().as_ref().clone(), - ) - .build_v2_deletes(); - delete_manifest.add_file( - DataFileBuilder::default() - .partition_spec_id(table.metadata().default_partition_spec_id()) - .content(DataContentType::PositionDeletes) - .file_path(delete_path.to_string()) - .file_format(DataFileFormat::Parquet) - .file_size_in_bytes(delete_size) - .record_count(1) - .partition(data_file.partition().clone()) - .build()?, - -1, - )?; - let delete_manifest = delete_manifest.write_manifest_file().await?; - let manifest_list_path = - "s3://iceberg-test/warehouse/taxi/metadata/snap-with-delete.avro"; - let mut writer = ManifestListWriter::v2( - table - .file_io() - .new_output(manifest_list_path)? - .writer() - .await?, - next_snapshot_id, - Some(snapshot.snapshot_id()), - snapshot.sequence_number() + 1, - ); - writer.add_manifests( - manifest_list - .entries() - .iter() - .cloned() - .chain([delete_manifest]), - )?; - writer.close().await?; - - let mut metadata: Value = serde_json::from_slice(&fs::read( - temporary - .path() - .join("iceberg/taxi/metadata/v1.metadata.json"), - )?)?; - let mut new_snapshot = metadata["snapshots"][0].clone(); - new_snapshot["snapshot-id"] = json!(next_snapshot_id); - new_snapshot["parent-snapshot-id"] = json!(snapshot.snapshot_id()); - new_snapshot["sequence-number"] = json!(snapshot.sequence_number() + 1); - new_snapshot["manifest-list"] = json!(manifest_list_path); - new_snapshot["timestamp-ms"] = json!(snapshot.timestamp_ms() + 1); - new_snapshot["summary"]["total-delete-files"] = json!("1"); - new_snapshot["summary"]["total-position-deletes"] = json!("1"); - metadata["current-snapshot-id"] = json!(next_snapshot_id); - metadata["last-sequence-number"] = json!(snapshot.sequence_number() + 1); - metadata["last-updated-ms"] = json!(snapshot.timestamp_ms() + 1); - metadata["refs"]["main"]["snapshot-id"] = json!(next_snapshot_id); - metadata["snapshot-log"] - .as_array_mut() - .expect("fixture has a snapshot log") - .push(json!({ - "snapshot-id": next_snapshot_id, - "timestamp-ms": snapshot.timestamp_ms() + 1, - })); - metadata["snapshots"] - .as_array_mut() - .expect("fixture has snapshots") - .push(new_snapshot); - let table = Table::builder() - .file_io(table.file_io().clone()) - .identifier(table.identifier().clone()) - .metadata(serde_json::from_value::(metadata)?) - .runtime(Runtime::try_current()?) - .build()?; + let fixture = taxi_with_position_delete().await?; + let eager = IcebergTestHarness::with_config(eager_config()).await?; + let table = fixture.table.clone(); let provider = IcebergStaticTableProvider::try_new_from_table(table).await?; let lazy_state = IcebergTestHarness::with_config(SessionConfig::new()).await?; @@ -365,16 +204,10 @@ async fn eager_and_lazy_scans_apply_position_deletes() -> StdResult<(), Box Date: Thu, 22 Oct 2026 00:03:20 -0400 Subject: [PATCH 6/7] refactor: expose assigned eager tasks before stream adaptation --- crates/datafusion/src/read/scan.rs | 60 ++++++++++++++----- .../datafusion/tests/eager_file_planning.rs | 20 +++++++ 2 files changed, 65 insertions(+), 15 deletions(-) diff --git a/crates/datafusion/src/read/scan.rs b/crates/datafusion/src/read/scan.rs index a6123e9..988d40c 100644 --- a/crates/datafusion/src/read/scan.rs +++ b/crates/datafusion/src/read/scan.rs @@ -30,10 +30,10 @@ use datafusion::physical_expr::{EquivalenceProperties, PhysicalExpr}; use datafusion::physical_plan::stream::RecordBatchStreamAdapter; use datafusion::physical_plan::{DisplayFormatType, Partitioning}; use datafusion::prelude::Expr; -use futures::{Stream, StreamExt, TryStreamExt}; +use futures::{Stream, TryStreamExt}; use iceberg::arrow::ArrowReader; use iceberg::expr::Predicate; -use iceberg::scan::{FileScanTaskStream, TableScan}; +use iceberg::scan::{FileScanTask, FileScanTaskStream, TableScan}; use iceberg::table::Table; use super::file_planning::PlannedFileScan; @@ -51,6 +51,12 @@ struct PlannedScan { partitions: usize, } +impl PlannedScan { + fn task_indices(&self, partition: usize) -> impl Iterator + use<> { + (partition..self.files.tasks().len()).step_by(self.partitions) + } +} + impl std::fmt::Debug for ScanReader { fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { f.write_str("ScanReader") @@ -99,6 +105,38 @@ impl IcebergDataSource { self.planned.as_ref().map(|planned| planned.files.as_ref()) } + /// Retained tasks assigned to a local output partition, in planning order. + /// Returns `None` for lazy scans, whose tasks are discovered during execution. + /// + /// This performs no I/O and does not clone tasks. Readers that outlive this + /// source must retain a source clone to keep the planning memory reservation. + /// + /// # Errors + /// + /// Returns an error if `partition` is outside this source's output partitions. + pub fn planned_partition_tasks( + &self, + partition: usize, + ) -> Result>> { + self.validate_partition(partition)?; + Ok(self.planned.as_ref().map(|planned| { + planned + .task_indices(partition) + .map(|index| &planned.files.tasks()[index]) + })) + } + + fn validate_partition(&self, partition: usize) -> Result<()> { + let partitions = self + .planned + .as_ref() + .map_or(1, |planned| planned.partitions); + if partition >= partitions { + return internal_err!("Iceberg scan partition {partition} does not exist"); + } + Ok(()) + } + pub fn table(&self) -> &Table { &self.table } @@ -313,23 +351,15 @@ impl DataSource for IcebergDataSource { partition: usize, _context: Arc, ) -> Result { - let partitions = self - .planned - .as_ref() - .map_or(1, |planned| planned.partitions); - if partition >= partitions { - return internal_err!("Iceberg scan partition {partition} does not exist"); - } + self.validate_partition(partition)?; let tasks: FileScanTaskStream = match &self.planned { Some(planned) => { // The reader owns each task; cloning keeps the physical plan reusable. let files = Arc::clone(&planned.files); - Box::pin( - futures::stream::iter( - (partition..files.tasks().len()).step_by(partitions), - ) - .map(move |index| Ok(files.tasks()[index].clone())), - ) + let tasks = planned + .task_indices(partition) + .map(move |index| Ok(files.tasks()[index].clone())); + Box::pin(futures::stream::iter(tasks)) } None => { let scan = build_table_scan( diff --git a/crates/datafusion/tests/eager_file_planning.rs b/crates/datafusion/tests/eager_file_planning.rs index f356356..d6feb28 100644 --- a/crates/datafusion/tests/eager_file_planning.rs +++ b/crates/datafusion/tests/eager_file_planning.rs @@ -71,6 +71,26 @@ async fn plans_file_tasks_before_execution_without_changing_statistics() -> Resu let lazy_source = iceberg_source(lazy_plan.as_ref()); assert!(lazy_source.planned_files().is_none()); + assert!(lazy_source.planned_partition_tasks(0)?.is_none()); + assert!(lazy_source.planned_partition_tasks(1).is_err()); + assert!(eager_source.planned_partition_tasks(4).is_err()); + let tasks = eager_source.planned_files().unwrap().tasks(); + for (partition, indices) in [vec![0, 4], vec![1, 5], vec![2, 6], vec![3]] + .into_iter() + .enumerate() + { + let assigned: Vec<_> = eager_source + .planned_partition_tasks(partition)? + .unwrap() + .map(std::ptr::from_ref) + .collect(); + let expected: Vec<_> = indices + .into_iter() + .map(|index| std::ptr::from_ref(&tasks[index])) + .collect(); + assert_eq!(assigned, expected); + } + assert_eq!(lazy_plan.output_partitioning().partition_count(), 1); assert_eq!(eager_source.planned_files().unwrap().tasks().len(), 7); assert_eq!(eager_plan.output_partitioning().partition_count(), 4); From 6a510dfde43092c288a42edde65b212f570876a7 Mon Sep 17 00:00:00 2001 From: "alex.bianchi" Date: Thu, 22 Oct 2026 10:02:17 -0400 Subject: [PATCH 7/7] refactor: simplify eager planning API and address review feedback --- crates/datafusion/README.md | 6 +- crates/datafusion/src/config.rs | 3 +- crates/datafusion/src/read/file_planning.rs | 4 +- crates/datafusion/src/read/mod.rs | 5 +- crates/datafusion/src/read/scan.rs | 161 +++++------- .../table_provider/catalog_table_provider.rs | 19 +- .../table_provider/static_table_provider.rs | 13 +- crates/datafusion/src/test_utils/harness.rs | 16 +- .../tests/common/position_deletes.rs | 243 +++++++++--------- .../datafusion/tests/eager_file_planning.rs | 97 +++---- .../tests/integration_datafusion_test.rs | 15 +- ...eager_multi_partition_limit_is_global.snap | 15 ++ ...er_ordered_limit_matches_lazy_results.snap | 16 ++ .../eager_file_planning__empty_scan.snap | 17 ++ ...ager_file_planning__partition_pruning.snap | 18 ++ .../eager_file_planning__residual_filter.snap | 17 ++ 16 files changed, 364 insertions(+), 301 deletions(-) create mode 100644 crates/datafusion/tests/snapshots/eager_file_planning__eager_multi_partition_limit_is_global.snap create mode 100644 crates/datafusion/tests/snapshots/eager_file_planning__eager_ordered_limit_matches_lazy_results.snap create mode 100644 crates/datafusion/tests/snapshots/eager_file_planning__empty_scan.snap create mode 100644 crates/datafusion/tests/snapshots/eager_file_planning__partition_pruning.snap create mode 100644 crates/datafusion/tests/snapshots/eager_file_planning__residual_filter.snap diff --git a/crates/datafusion/README.md b/crates/datafusion/README.md index fc5508e..4952ac3 100644 --- a/crates/datafusion/README.md +++ b/crates/datafusion/README.md @@ -33,7 +33,7 @@ use iceberg_datafusion::IcebergPlanningConfig; let config = SessionConfig::new() .with_option_extension(IcebergPlanningConfig::default()) - .set_bool("iceberg_local.eager_file_planning", true); + .set_bool("iceberg.eager_file_planning", true); ``` ### Behavior @@ -55,5 +55,5 @@ that stop early. Repeated provider scans repeat planning. The planning caller must support metadata I/O: Iceberg's initial manifest-list read still runs on that runtime, even when the table has a separate I/O runtime. -The `iceberg_local` prefix allows these options to coexist with DataFusion -Distributed's `iceberg` options. +Register only one configuration extension under the `iceberg` prefix per session. +DataFusion Distributed currently uses that same prefix for its own options. diff --git a/crates/datafusion/src/config.rs b/crates/datafusion/src/config.rs index dbc3011..361dfc4 100644 --- a/crates/datafusion/src/config.rs +++ b/crates/datafusion/src/config.rs @@ -27,6 +27,5 @@ extensions_options! { } impl ConfigExtension for IcebergPlanningConfig { - // Distinct from DFD's `iceberg` options so both integrations can share a session. - const PREFIX: &'static str = "iceberg_local"; + const PREFIX: &'static str = "iceberg"; } diff --git a/crates/datafusion/src/read/file_planning.rs b/crates/datafusion/src/read/file_planning.rs index c1ccb89..0c295f2 100644 --- a/crates/datafusion/src/read/file_planning.rs +++ b/crates/datafusion/src/read/file_planning.rs @@ -33,14 +33,14 @@ use crate::to_datafusion_error; /// this value alive also keeps the query memory reservation for the retained tasks. /// This is not a reader and contains no coordinator-local reader state. #[derive(Debug)] -pub struct PlannedFileScan { +pub(crate) struct PlannedFileScan { tasks: Vec, _reservation: MemoryReservation, } impl PlannedFileScan { /// Tasks in Iceberg's planning order. Assign each once per execution. - pub fn tasks(&self) -> &[FileScanTask] { + pub(crate) fn tasks(&self) -> &[FileScanTask] { &self.tasks } diff --git a/crates/datafusion/src/read/mod.rs b/crates/datafusion/src/read/mod.rs index 42626c3..095bcaa 100644 --- a/crates/datafusion/src/read/mod.rs +++ b/crates/datafusion/src/read/mod.rs @@ -19,8 +19,5 @@ mod file_planning; pub(crate) mod predicate; pub(crate) mod scan; -pub use file_planning::PlannedFileScan; pub use predicate::convert_filters_to_predicate; -pub use scan::{ - IcebergDataSource, IcebergDataSourceBuilder, PlannedIcebergDataSourceBuilder, -}; +pub use scan::{IcebergDataSource, IcebergDataSourceBuilder}; diff --git a/crates/datafusion/src/read/scan.rs b/crates/datafusion/src/read/scan.rs index 988d40c..15d04cb 100644 --- a/crates/datafusion/src/read/scan.rs +++ b/crates/datafusion/src/read/scan.rs @@ -38,12 +38,7 @@ use iceberg::table::Table; use super::file_planning::PlannedFileScan; use super::predicate::convert_filters_to_predicate; -use crate::{IcebergPlanningConfig, to_datafusion_error}; - -// A reader is created once per physical scan. Its clones share the delete loader's -// cache while each read() call still owns its own scan metrics. -#[derive(Clone)] -struct ScanReader(ArrowReader); +use crate::to_datafusion_error; #[derive(Debug, Clone)] struct PlannedScan { @@ -57,15 +52,9 @@ impl PlannedScan { } } -impl std::fmt::Debug for ScanReader { - fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - f.write_str("ScanReader") - } -} - /// Manages the scanning process of an Iceberg [`Table`], encapsulating the /// necessary details and computed properties required for execution planning. -#[derive(Debug, Clone)] +#[derive(Clone)] pub struct IcebergDataSource { /// A table in the catalog. table: Table, @@ -80,7 +69,21 @@ pub struct IcebergDataSource { /// Optional limit on the number of rows to return fetch: Option, planned: Option, - reader: ScanReader, + reader: ArrowReader, +} + +impl std::fmt::Debug for IcebergDataSource { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("IcebergDataSource") + .field("table", &self.table) + .field("snapshot_id", &self.snapshot_id) + .field("schema", &self.schema) + .field("projection", &self.projection) + .field("predicate", &self.predicate) + .field("fetch", &self.fetch) + .field("planned", &self.planned) + .finish_non_exhaustive() + } } impl IcebergDataSource { @@ -95,14 +98,15 @@ impl IcebergDataSource { snapshot_id: None, predicate: None, fetch: None, + eager_planning: false, } } /// Planned, ungrouped tasks, if file discovery was enabled for this scan. /// A distributed consumer may assign these tasks to its own workers instead /// of using the local output partitions of this source. - pub fn planned_files(&self) -> Option<&PlannedFileScan> { - self.planned.as_ref().map(|planned| planned.files.as_ref()) + pub fn planned_files(&self) -> Option<&[FileScanTask]> { + self.planned.as_ref().map(|planned| planned.files.tasks()) } /// Retained tasks assigned to a local output partition, in planning order. @@ -165,8 +169,8 @@ impl IcebergDataSource { /// Builds an [`IcebergDataSource`]. Created by [`IcebergDataSource::builder`]. /// /// A table provider passes the table's schema and selects columns with -/// [`Self::with_projection`]. Finish with [`Self::with_planned_files`] to honor -/// the session's eager-planning option before building. +/// [`Self::with_projection`]. Enable [`Self::with_eager_planning`] to select files +/// during [`Self::build`] instead of execution. /// /// To rebuild a source, such as in another process, /// pass its [`IcebergDataSource::schema`], which is already projected, and no @@ -180,6 +184,7 @@ pub struct IcebergDataSourceBuilder { snapshot_id: Option, predicate: Option, fetch: Option, + eager_planning: bool, } impl IcebergDataSourceBuilder { @@ -216,65 +221,49 @@ impl IcebergDataSourceBuilder { self } - /// Finish configuring the scan and optionally select its files before execution. - /// - /// Honors [`IcebergPlanningConfig::eager_file_planning`]. The returned builder - /// only exposes `build`, so filters, projection, and snapshot cannot change - /// after files have been selected. - /// - /// # Errors - /// - /// Returns an error if the projection is invalid, file planning fails, or - /// the retained tasks exceed the query memory pool. - pub async fn with_planned_files( - self, - context: Arc, - ) -> Result { - if !context - .session_config() - .options() - .extensions - .get::() - .is_some_and(|config| config.eager_file_planning) - { - return Ok(PlannedIcebergDataSourceBuilder { - builder: self, - planned: None, - }); - } - let target_partitions = context.session_config().target_partitions(); - let schema = self.projected_schema()?; - let projection: Vec = schema - .fields() - .iter() - .map(|field| field.name().clone()) - .collect(); - let scan = build_table_scan( - &self.table, - self.snapshot_id, - Some(&projection), - self.predicate.as_ref(), - )?; - // Dropping physical planning drops collection and its memory reservation. - let planned = PlannedFileScan::collect(scan, &context).await?; - let partitions = planned.tasks().len().max(1).min(target_partitions.max(1)); - - Ok(PlannedIcebergDataSourceBuilder { - builder: self, - planned: Some(PlannedScan { - files: Arc::new(planned), - partitions, - }), - }) + /// Select files during [`Self::build`] when enabled. Defaults to `false`. + pub fn with_eager_planning(mut self, enabled: bool) -> Self { + self.eager_planning = enabled; + self } - /// Builds the source, which reads the columns of its projected schema by name. + /// Builds the source and its shared reader. When eager planning is enabled, + /// discovers files using `context` for partition count and memory accounting. + /// Lazy builds perform no file-discovery I/O. /// /// # Errors /// - /// Returns an error if the projection is out of bounds for the schema. - pub fn build(self) -> Result { - self.build_with_planned_files(None) + /// Returns an error if the projection is invalid, file planning fails, or + /// retained tasks exceed the query memory pool. + pub async fn build(self, context: Arc) -> Result { + let planned = if self.eager_planning { + let schema = self.projected_schema()?; + let projection: Vec = schema + .fields() + .iter() + .map(|field| field.name().clone()) + .collect(); + let scan = build_table_scan( + &self.table, + self.snapshot_id, + Some(&projection), + self.predicate.as_ref(), + )?; + // Dropping physical planning drops collection and its reservation. + let files = PlannedFileScan::collect(scan, &context).await?; + let partitions = files + .tasks() + .len() + .max(1) + .min(context.session_config().target_partitions().max(1)); + Some(PlannedScan { + files: Arc::new(files), + partitions, + }) + } else { + None + }; + self.build_with_planned_files(planned) } fn projected_schema(&self) -> Result { @@ -303,7 +292,8 @@ impl IcebergDataSourceBuilder { } else { reader }; - let reader = ScanReader(reader.build()); + // Reader clones share the delete loader across output partitions. + let reader = reader.build(); Ok(IcebergDataSource { table: self.table, snapshot_id: self.snapshot_id, @@ -317,27 +307,6 @@ impl IcebergDataSourceBuilder { } } -/// A configured scan ready to build after optional eager file planning. -/// -/// Created by [`IcebergDataSourceBuilder::with_planned_files`]. Configuration -/// is fixed so the retained tasks stay consistent with the resulting source. -#[derive(Debug, Clone)] -pub struct PlannedIcebergDataSourceBuilder { - builder: IcebergDataSourceBuilder, - planned: Option, -} - -impl PlannedIcebergDataSourceBuilder { - /// Construct the source and its shared reader from the configured scan. - /// - /// # Errors - /// - /// Returns an error if the projection is out of bounds for the schema. - pub fn build(self) -> Result { - self.builder.build_with_planned_files(self.planned) - } -} - impl DataSource for IcebergDataSource { fn apply_expressions( &self, @@ -378,7 +347,6 @@ impl DataSource for IcebergDataSource { // Reader clones share their delete loader; no per-partition cache is rebuilt. let stream = self .reader - .0 .clone() .read(tasks) .map_err(to_datafusion_error)? @@ -425,12 +393,13 @@ impl DataSource for IcebergDataSource { write!(f, ", fetch={fetch}")?; } if let Some(planned) = self.planned_files() { - write!(f, ", planned_files={}", planned.tasks().len())?; + write!(f, ", planned_files={}", planned.len())?; } Ok(()) } fn output_partitioning(&self) -> Partitioning { + // Partitioned lazy discovery: https://github.com/apache/datafusion-iceberg/issues/84 Partitioning::UnknownPartitioning( self.planned .as_ref() diff --git a/crates/datafusion/src/table_provider/catalog_table_provider.rs b/crates/datafusion/src/table_provider/catalog_table_provider.rs index 00d6145..01aa191 100644 --- a/crates/datafusion/src/table_provider/catalog_table_provider.rs +++ b/crates/datafusion/src/table_provider/catalog_table_provider.rs @@ -34,6 +34,7 @@ use iceberg::inspect::MetadataTableType; use iceberg::spec::TableProperties; use iceberg::{Catalog, NamespaceIdent, TableIdent}; +use crate::IcebergPlanningConfig; use crate::error::to_datafusion_error; use crate::metadata::table::IcebergMetadataTableProvider; use crate::read::IcebergDataSource; @@ -143,9 +144,15 @@ impl TableProvider for IcebergCatalogTableProvider { .with_projection(projection) .with_filters(filters) .with_fetch(limit) - .with_planned_files(state.task_ctx()) - .await? - .build()?; + .with_eager_planning( + state + .config_options() + .extensions + .get::() + .is_some_and(|config| config.eager_file_planning), + ) + .build(state.task_ctx()) + .await?; Ok(DataSourceExec::from_data_source(source)) } @@ -459,7 +466,8 @@ mod tests { // Columns come back in the projection's order, not the table's. let source = IcebergDataSource::builder(table.clone(), schema.clone()) .with_projection(Some(&vec![1, 0])) - .build() + .build(SessionContext::new().task_ctx()) + .await .unwrap(); insta::assert_debug_snapshot!(source.projection(), @r#" [ @@ -470,7 +478,8 @@ mod tests { let err = IcebergDataSource::builder(table, schema) .with_projection(Some(&vec![0, out_of_range])) - .build() + .build(SessionContext::new().task_ctx()) + .await .unwrap_err(); assert_eq!( err.to_string(), diff --git a/crates/datafusion/src/table_provider/static_table_provider.rs b/crates/datafusion/src/table_provider/static_table_provider.rs index ef1857c..6336216 100644 --- a/crates/datafusion/src/table_provider/static_table_provider.rs +++ b/crates/datafusion/src/table_provider/static_table_provider.rs @@ -30,6 +30,7 @@ use datafusion::physical_plan::ExecutionPlan; use iceberg::arrow::schema_to_arrow_schema; use iceberg::table::Table; +use crate::IcebergPlanningConfig; use crate::error::to_datafusion_error; use crate::read::IcebergDataSource; @@ -131,9 +132,15 @@ impl TableProvider for IcebergStaticTableProvider { .with_projection(projection) .with_filters(filters) .with_fetch(limit) - .with_planned_files(state.task_ctx()) - .await? - .build()?; + .with_eager_planning( + state + .config_options() + .extensions + .get::() + .is_some_and(|config| config.eager_file_planning), + ) + .build(state.task_ctx()) + .await?; Ok(DataSourceExec::from_data_source(source)) } diff --git a/crates/datafusion/src/test_utils/harness.rs b/crates/datafusion/src/test_utils/harness.rs index 17dc446..bb8b9e5 100644 --- a/crates/datafusion/src/test_utils/harness.rs +++ b/crates/datafusion/src/test_utils/harness.rs @@ -58,30 +58,22 @@ impl IcebergTestHarness { } pub async fn with_config(config: SessionConfig) -> Result { - Self::build(config, None, None).await - } - - /// Load the fixture from a separate directory so a test can write metadata safely. - pub async fn with_fixture_root(config: SessionConfig, root: PathBuf) -> Result { - Self::build(config, None, Some(root)).await + Self::build(config, None).await } pub async fn with_config_and_runtime( config: SessionConfig, runtime_env: Arc, ) -> Result { - Self::build(config, Some(runtime_env), None).await + Self::build(config, Some(runtime_env)).await } async fn build( config: SessionConfig, runtime_env: Option>, - root: Option, ) -> Result { - let storage_factory: Arc = Arc::new(match root { - Some(root) => FixtureStorageFactory { root }, - None => FixtureStorageFactory::default(), - }); + let storage_factory: Arc = + Arc::new(FixtureStorageFactory::default()); let builder = SessionStateBuilder::new() .with_default_features() .with_config(config); diff --git a/crates/datafusion/tests/common/position_deletes.rs b/crates/datafusion/tests/common/position_deletes.rs index ca3ff51..b01ff09 100644 --- a/crates/datafusion/tests/common/position_deletes.rs +++ b/crates/datafusion/tests/common/position_deletes.rs @@ -17,181 +17,176 @@ use std::collections::HashMap; use std::error::Error; -use std::fs; -use std::io; -use std::path::Path; +use std::fs::File; use std::sync::Arc; -use datafusion::arrow::array::{ArrayRef, Int64Array, RecordBatch, StringArray}; -use datafusion::arrow::datatypes::{DataType, Field, Schema}; -use datafusion::prelude::SessionConfig; -use iceberg::Runtime; +use datafusion::arrow::array::{Int64Array, RecordBatch, StringArray}; +use datafusion::arrow::datatypes::{DataType, Field, Schema as ArrowSchema}; +use iceberg::io::FileIO; use iceberg::metadata_columns::{ RESERVED_COL_NAME_DELETE_FILE_PATH, RESERVED_COL_NAME_DELETE_FILE_POS, RESERVED_FIELD_ID_DELETE_FILE_PATH, RESERVED_FIELD_ID_DELETE_FILE_POS, }; use iceberg::spec::{ - DataContentType, DataFileBuilder, DataFileFormat, ManifestListWriter, - ManifestWriterBuilder, Operation, Snapshot, Summary, + DataContentType, DataFileBuilder, DataFileFormat, FormatVersion, ManifestListWriter, + ManifestWriterBuilder, NestedField, Operation, PrimitiveType, Schema, Snapshot, + SortOrder, Struct, Summary, TableMetadataBuilder, Type, UnboundPartitionSpec, }; use iceberg::table::Table; -use iceberg_datafusion::IcebergCatalogTableProvider; -use iceberg_datafusion::test_utils::IcebergTestHarness; +use iceberg::{Runtime, TableIdent}; use parquet::arrow::{ArrowWriter, PARQUET_FIELD_ID_META_KEY}; use tempfile::TempDir; -/// A taxi table with its first data file's first row deleted. -/// Owns the copied files for the lifetime of scans against this table. +/// Two data files containing IDs [0, 1] and [2, 3], with ID 0 position-deleted. +/// Keeps the generated files alive while the table is scanned. pub struct PositionDeleteFixture { pub table: Table, _directory: TempDir, } -pub async fn taxi_with_position_delete() -> Result> +pub async fn table_with_position_delete() -> Result> { - // Isolate the checked-in fixture: the test writes a new delete manifest and - // snapshot, but must not modify the files used by other tests. - let temporary = TempDir::new()?; - let root = temporary.path().join("iceberg"); - copy_fixture( - &Path::new(env!("CARGO_MANIFEST_DIR")).join("testdata/iceberg"), - &root, - )?; - let harness = - IcebergTestHarness::with_fixture_root(SessionConfig::new(), root).await?; - let provider = harness.context().table_provider("taxi").await?; - let catalog_provider = provider - .downcast_ref::() - .expect("the fixture creates a catalog Iceberg provider"); - let table = catalog_provider - .catalog() - .load_table(catalog_provider.table_ident()) - .await?; - let snapshot = table - .metadata() - .current_snapshot() - .expect("fixture has a snapshot"); - let manifest_list = table.manifest_list_reader(snapshot).load().await?; - let data_manifest = table - .manifest_reader() - .read(&manifest_list.entries()[0]) - .await?; - let data_file = &data_manifest.entries()[0].data_file; + let directory = TempDir::new()?; + let root = directory.path().to_str().unwrap(); + let io = FileIO::new_with_fs(); + let schema = Schema::builder() + .with_fields([NestedField::required( + 1, + "id", + Type::Primitive(PrimitiveType::Long), + ) + .into()]) + .build()?; + let metadata = TableMetadataBuilder::new( + schema, + UnboundPartitionSpec::default(), + SortOrder::unsorted_order(), + root.to_string(), + FormatVersion::V2, + HashMap::new(), + )? + .build()? + .metadata; + let schema = metadata.current_schema().clone(); + let spec = metadata.default_partition_spec().as_ref().clone(); + let mut data_manifest = ManifestWriterBuilder::new( + io.new_output(format!("{root}/data.avro"))?, + Some(2), + schema.clone(), + spec.clone(), + ) + .build_v2_data(); + let data_schema = Arc::new(ArrowSchema::new(vec![field("id", DataType::Int64, 1)])); + for file in 0..2 { + let path = format!("{root}/data-{file}.parquet"); + let batch = RecordBatch::try_new( + data_schema.clone(), + vec![Arc::new(Int64Array::from(vec![file * 2, file * 2 + 1]))], + )?; + let size = write_parquet(&path, batch)?; + data_manifest.add_file( + DataFileBuilder::default() + .partition_spec_id(0) + .content(DataContentType::Data) + .file_path(path) + .file_format(DataFileFormat::Parquet) + .file_size_in_bytes(size) + .record_count(2) + .partition(Struct::empty()) + .build()?, + 1, + )?; + } + let data_manifest = data_manifest.write_manifest_file().await?; - let delete_path = "s3://iceberg-test/warehouse/taxi/data/pos-delete.parquet"; - let schema = Arc::new(Schema::new(vec![ - Field::new(RESERVED_COL_NAME_DELETE_FILE_PATH, DataType::Utf8, false) - .with_metadata(HashMap::from([( - PARQUET_FIELD_ID_META_KEY.to_string(), - RESERVED_FIELD_ID_DELETE_FILE_PATH.to_string(), - )])), - Field::new(RESERVED_COL_NAME_DELETE_FILE_POS, DataType::Int64, false) - .with_metadata(HashMap::from([( - PARQUET_FIELD_ID_META_KEY.to_string(), - RESERVED_FIELD_ID_DELETE_FILE_POS.to_string(), - )])), + let delete_path = format!("{root}/delete.parquet"); + let delete_schema = Arc::new(ArrowSchema::new(vec![ + field( + RESERVED_COL_NAME_DELETE_FILE_PATH, + DataType::Utf8, + RESERVED_FIELD_ID_DELETE_FILE_PATH, + ), + field( + RESERVED_COL_NAME_DELETE_FILE_POS, + DataType::Int64, + RESERVED_FIELD_ID_DELETE_FILE_POS, + ), ])); let batch = RecordBatch::try_new( - schema.clone(), + delete_schema, vec![ - Arc::new(StringArray::from(vec![data_file.file_path()])) as ArrayRef, - Arc::new(Int64Array::from(vec![0])) as ArrayRef, + Arc::new(StringArray::from(vec![format!("{root}/data-0.parquet")])), + Arc::new(Int64Array::from(vec![0])), ], )?; - let delete_file_on_disk = temporary - .path() - .join("iceberg/taxi/data/pos-delete.parquet"); - let mut writer = - ArrowWriter::try_new(fs::File::create(&delete_file_on_disk)?, schema, None)?; - writer.write(&batch)?; - writer.close()?; - let delete_size = fs::metadata(&delete_file_on_disk)?.len(); - - let next_snapshot_id = snapshot.snapshot_id() + 1; + let size = write_parquet(&delete_path, batch)?; let mut delete_manifest = ManifestWriterBuilder::new( - table - .file_io() - .new_output("s3://iceberg-test/warehouse/taxi/metadata/pos-delete.avro")?, - Some(next_snapshot_id), - snapshot.schema(table.metadata())?, - table.metadata().default_partition_spec().as_ref().clone(), + io.new_output(format!("{root}/deletes.avro"))?, + Some(2), + schema.clone(), + spec, ) .build_v2_deletes(); delete_manifest.add_file( DataFileBuilder::default() - .partition_spec_id(table.metadata().default_partition_spec_id()) + .partition_spec_id(0) .content(DataContentType::PositionDeletes) - .file_path(delete_path.to_string()) + .file_path(delete_path) .file_format(DataFileFormat::Parquet) - .file_size_in_bytes(delete_size) + .file_size_in_bytes(size) .record_count(1) - .partition(data_file.partition().clone()) + .partition(Struct::empty()) .build()?, - -1, + 2, )?; let delete_manifest = delete_manifest.write_manifest_file().await?; - let manifest_list_path = - "s3://iceberg-test/warehouse/taxi/metadata/snap-with-delete.avro"; - let mut writer = ManifestListWriter::v2( - table - .file_io() - .new_output(manifest_list_path)? - .writer() - .await?, - next_snapshot_id, - Some(snapshot.snapshot_id()), - snapshot.sequence_number() + 1, + let manifest_list_path = format!("{root}/snapshot.avro"); + let mut list = ManifestListWriter::v2( + io.new_output(&manifest_list_path)?.writer().await?, + 2, + None, + 2, ); - writer.add_manifests( - manifest_list - .entries() - .iter() - .cloned() - .chain([delete_manifest]), - )?; - writer.close().await?; - - let new_snapshot = Snapshot::builder() - .with_snapshot_id(next_snapshot_id) - .with_parent_snapshot_id(Some(snapshot.snapshot_id())) - .with_sequence_number(snapshot.sequence_number() + 1) - .with_timestamp_ms(snapshot.timestamp_ms() + 1) + list.add_manifests([data_manifest, delete_manifest].into_iter())?; + list.close().await?; + let snapshot = Snapshot::builder() + .with_snapshot_id(2) + .with_sequence_number(2) + .with_timestamp_ms(metadata.last_updated_ms() + 1) .with_manifest_list(manifest_list_path) + .with_schema_id(schema.schema_id()) .with_summary(Summary { - operation: Operation::Delete, + operation: Operation::Overwrite, additional_properties: HashMap::new(), }) - .with_schema_id(snapshot.schema(table.metadata())?.schema_id()) .build(); - let metadata = table - .metadata() - .clone() + let metadata = metadata .into_builder(None) - .set_branch_snapshot(new_snapshot, "main")? + .set_branch_snapshot(snapshot, "main")? .build()? .metadata; let table = Table::builder() - .file_io(table.file_io().clone()) - .identifier(table.identifier().clone()) + .file_io(io) + .identifier(TableIdent::from_strs(["test", "deletes"])?) .metadata(metadata) .runtime(Runtime::try_current()?) .build()?; Ok(PositionDeleteFixture { table, - _directory: temporary, + _directory: directory, }) } -fn copy_fixture(from: &Path, to: &Path) -> io::Result<()> { - fs::create_dir_all(to)?; - for entry in fs::read_dir(from)? { - let entry = entry?; - let destination = to.join(entry.file_name()); - if entry.file_type()?.is_dir() { - copy_fixture(&entry.path(), &destination)?; - } else { - fs::copy(entry.path(), destination)?; - } - } - Ok(()) +fn field(name: &str, data_type: DataType, id: i32) -> Field { + Field::new(name, data_type, false).with_metadata(HashMap::from([( + PARQUET_FIELD_ID_META_KEY.to_string(), + id.to_string(), + )])) +} + +fn write_parquet(path: &str, batch: RecordBatch) -> Result> { + let mut writer = ArrowWriter::try_new(File::create(path)?, batch.schema(), None)?; + writer.write(&batch)?; + writer.close()?; + Ok(std::fs::metadata(path)?.len()) } diff --git a/crates/datafusion/tests/eager_file_planning.rs b/crates/datafusion/tests/eager_file_planning.rs index d6feb28..f751b82 100644 --- a/crates/datafusion/tests/eager_file_planning.rs +++ b/crates/datafusion/tests/eager_file_planning.rs @@ -19,6 +19,7 @@ use std::error::Error; use std::result::Result as StdResult; use std::sync::Arc; +use datafusion::arrow::array::Int64Array; use datafusion::common::{Result, stats::Precision}; use datafusion::datasource::TableProvider; use datafusion::datasource::source::{DataSource, DataSourceExec}; @@ -26,7 +27,7 @@ use datafusion::execution::memory_pool::GreedyMemoryPool; use datafusion::execution::runtime_env::RuntimeEnvBuilder; use datafusion::physical_plan::common::collect; use datafusion::physical_plan::{ExecutionPlanProperties, collect as collect_plan}; -use datafusion::prelude::{SessionConfig, col, lit}; +use datafusion::prelude::{SessionConfig, SessionContext, col, lit}; use datafusion::scalar::ScalarValue; use iceberg_datafusion::read::IcebergDataSource; use iceberg_datafusion::test_utils::IcebergTestHarness; @@ -36,13 +37,13 @@ use iceberg_datafusion::{ mod common; -use common::position_deletes::taxi_with_position_delete; +use common::position_deletes::table_with_position_delete; fn eager_config() -> SessionConfig { SessionConfig::new() .with_target_partitions(4) .with_option_extension(IcebergPlanningConfig::default()) - .set_bool("iceberg_local.eager_file_planning", true) + .set_bool("iceberg.eager_file_planning", true) } fn iceberg_source( @@ -74,7 +75,7 @@ async fn plans_file_tasks_before_execution_without_changing_statistics() -> Resu assert!(lazy_source.planned_partition_tasks(0)?.is_none()); assert!(lazy_source.planned_partition_tasks(1).is_err()); assert!(eager_source.planned_partition_tasks(4).is_err()); - let tasks = eager_source.planned_files().unwrap().tasks(); + let tasks = eager_source.planned_files().unwrap(); for (partition, indices) in [vec![0, 4], vec![1, 5], vec![2, 6], vec![3]] .into_iter() .enumerate() @@ -92,7 +93,7 @@ async fn plans_file_tasks_before_execution_without_changing_statistics() -> Resu } assert_eq!(lazy_plan.output_partitioning().partition_count(), 1); - assert_eq!(eager_source.planned_files().unwrap().tasks().len(), 7); + assert_eq!(eager_source.planned_files().unwrap().len(), 7); assert_eq!(eager_plan.output_partitioning().partition_count(), 4); assert!(eager_source.open(4, eager.context().task_ctx()).is_err()); assert!(lazy_source.open(1, lazy.context().task_ctx()).is_err()); @@ -117,7 +118,6 @@ async fn plans_file_tasks_before_execution_without_changing_statistics() -> Resu iceberg_source(empty.as_ref()) .planned_files() .unwrap() - .tasks() .is_empty() ); assert_eq!(empty.output_partitioning().partition_count(), 1); @@ -142,30 +142,27 @@ async fn eager_partition_count_is_bounded_by_selected_files() -> Result<()> { async fn filtered_and_empty_eager_scans_match_lazy_results() -> Result<()> { let eager = IcebergTestHarness::with_config(eager_config()).await?; let lazy = IcebergTestHarness::new().await?; - for (query, planned_files) in [ + for (name, query) in [ ( + "partition_pruning", "SELECT pickup_date, COUNT(*) AS trips FROM taxi \ WHERE pickup_date = DATE '2024-01-10' \ GROUP BY pickup_date ORDER BY pickup_date", - 1, ), ( + "residual_filter", "SELECT COUNT(*) AS trips FROM taxi \ WHERE pickup_date = DATE '2024-01-08' AND trip_distance > 50", - 1, ), ( + "empty_scan", "SELECT COUNT(*) AS trips FROM taxi \ WHERE pickup_date = DATE '2025-01-01'", - 0, ), ] { let (plan, eager_batches) = eager.query(query).await?; let (_, lazy_batches) = lazy.query(query).await?; - assert!( - plan.contains(&format!("planned_files={planned_files}")), - "eager scan not visible in plan: {plan}" - ); + insta::assert_snapshot!(name, format!("{plan}\n{eager_batches}")); assert_eq!(eager_batches, lazy_batches); } Ok(()) @@ -174,9 +171,8 @@ async fn filtered_and_empty_eager_scans_match_lazy_results() -> Result<()> { #[tokio::test] async fn eager_multi_partition_limit_is_global() -> Result<()> { let eager = IcebergTestHarness::with_config(eager_config()).await?; - let (plan, batches) = eager.query("SELECT pickup_date FROM taxi LIMIT 3").await?; - assert!(plan.contains("planned_files=7"), "unexpected plan: {plan}"); - assert_eq!(batches.matches("2024-01-").count(), 3, "{batches}"); + let (plan, batches) = eager.query("SELECT 1 AS n FROM taxi LIMIT 3").await?; + insta::assert_snapshot!(format!("{plan}\n{batches}")); Ok(()) } @@ -188,48 +184,60 @@ async fn eager_ordered_limit_matches_lazy_results() -> Result<()> { ORDER BY pickup_at, vendor_id, pickup_location_id LIMIT 3"; let (plan, eager_batches) = eager.query(sql).await?; let (_, lazy_batches) = lazy.query(sql).await?; - assert!(plan.contains("planned_files=7"), "unexpected plan: {plan}"); + insta::assert_snapshot!(format!("{plan}\n{eager_batches}")); assert_eq!(eager_batches, lazy_batches); Ok(()) } #[tokio::test] async fn eager_and_lazy_scans_apply_position_deletes() -> StdResult<(), Box> { - let fixture = taxi_with_position_delete().await?; - let eager = IcebergTestHarness::with_config(eager_config()).await?; + let fixture = table_with_position_delete().await?; + let eager = SessionContext::new_with_config(eager_config()); let table = fixture.table.clone(); let provider = IcebergStaticTableProvider::try_new_from_table(table).await?; - let lazy_state = IcebergTestHarness::with_config(SessionConfig::new()).await?; + let lazy = SessionContext::new(); - let state = eager.context().state(); + let state = eager.state(); let eager_plan = provider.scan(&state, None, &[], None).await?; + // Both unpartitioned files receive the delete task; the reader matches its path. assert_eq!( iceberg_source(eager_plan.as_ref()) .planned_files() .expect("eager plan has files") - .tasks() .iter() .filter(|task| !task.deletes().is_empty()) .count(), - 1 + 2 ); assert_eq!( iceberg_source(eager_plan.as_ref()) .planned_files() .expect("eager plan has files") - .tasks() .len(), - 7 + 2 ); - let lazy_plan = provider - .scan(&lazy_state.context().state(), None, &[], None) - .await?; + let lazy_plan = provider.scan(&lazy.state(), None, &[], None).await?; let eager_batches = collect_plan(eager_plan, state.task_ctx()).await?; - let lazy_batches = collect_plan(lazy_plan, lazy_state.context().task_ctx()).await?; - let eager_rows: usize = eager_batches.iter().map(|batch| batch.num_rows()).sum(); - let lazy_rows: usize = lazy_batches.iter().map(|batch| batch.num_rows()).sum(); - assert_eq!(lazy_rows, 174_999); - assert_eq!(eager_rows, lazy_rows); + let lazy_batches = collect_plan(lazy_plan, lazy.task_ctx()).await?; + let ids = |batches: &[datafusion::arrow::array::RecordBatch]| { + let mut ids: Vec = batches + .iter() + .flat_map(|batch| { + batch + .column(0) + .as_any() + .downcast_ref::() + .unwrap() + .values() + .iter() + .copied() + }) + .collect(); + ids.sort_unstable(); + ids + }; + assert_eq!(ids(&eager_batches), vec![1, 2, 3]); + assert_eq!(ids(&eager_batches), ids(&lazy_batches)); Ok(()) } @@ -260,11 +268,7 @@ fn eager_scan_uses_the_tables_custom_runtime() -> Result<()> { let state = harness.context().state(); let plan = provider.scan(&state, None, &[], None).await?; assert_eq!( - iceberg_source(plan.as_ref()) - .planned_files() - .unwrap() - .tasks() - .len(), + iceberg_source(plan.as_ref()).planned_files().unwrap().len(), 7 ); let mut rows = 0; @@ -293,7 +297,10 @@ async fn eager_tasks_respect_the_query_memory_pool() -> Result<()> { .await .expect_err("file metadata cannot fit in one byte"); assert!( - error.to_string().contains("Iceberg file planning"), + matches!( + error, + datafusion::error::DataFusionError::ResourcesExhausted(_) + ), "{error}" ); assert_eq!(runtime.memory_pool.reserved(), 0); @@ -329,7 +336,7 @@ async fn retained_tasks_follow_plan_and_stream_lifetimes() -> Result<()> { #[tokio::test] async fn public_builder_plans_projected_filtered_scans() -> Result<()> { - let harness = IcebergTestHarness::with_config(eager_config()).await?; + let harness = IcebergTestHarness::new().await?; let provider = harness.context().table_provider("taxi").await?; let provider = provider .downcast_ref::() @@ -345,11 +352,11 @@ async fn public_builder_plans_projected_filtered_scans() -> Result<()> { let source = IcebergDataSource::builder(table, schema) .with_projection(Some(&projection)) .with_filters(&[col("pickup_date").eq(lit(pickup_date))]) - .with_planned_files(harness.context().task_ctx()) - .await? - .build()?; + .with_eager_planning(true) + .build(harness.context().task_ctx()) + .await?; assert_eq!(source.projection(), &["pickup_date"]); - assert_eq!(source.planned_files().unwrap().tasks().len(), 1); + assert_eq!(source.planned_files().unwrap().len(), 1); assert_eq!(source.output_partitioning().partition_count(), 1); let batches = collect(source.open(0, harness.context().task_ctx())?).await?; assert!(!batches.is_empty()); diff --git a/crates/datafusion/tests/integration_datafusion_test.rs b/crates/datafusion/tests/integration_datafusion_test.rs index fd1ab67..f812d04 100644 --- a/crates/datafusion/tests/integration_datafusion_test.rs +++ b/crates/datafusion/tests/integration_datafusion_test.rs @@ -915,12 +915,16 @@ async fn run( } /// Rebuilds `scan` from its accessors alone, as a codec would. -fn rebuild_scan(scan: &IcebergDataSource) -> IcebergDataSource { +async fn rebuild_scan( + scan: &IcebergDataSource, + ctx: &SessionContext, +) -> IcebergDataSource { IcebergDataSource::builder(scan.table().clone(), scan.schema()) .with_snapshot_id(scan.snapshot_id()) .with_predicate(scan.predicate().cloned()) .with_fetch(scan.fetch()) - .build() + .build(ctx.task_ctx()) + .await .unwrap() } @@ -1039,7 +1043,7 @@ async fn test_plan_nodes_are_inspectable() -> Result<(), Box> { scan.predicate().map(ToString::to_string).as_deref(), Some("foo1 = 1") ); - let rebuilt = rebuild_scan(scan); + let rebuilt = rebuild_scan(scan, &ctx).await; assert_eq!(rebuilt.schema(), scan.schema()); assert_eq!(rebuilt.projection(), scan.projection()); let expected = run_source(scan, &ctx).await?; @@ -1062,7 +1066,7 @@ async fn test_plan_nodes_are_inspectable() -> Result<(), Box> { .downcast_ref::() .expect("an Iceberg data source"); assert_eq!(scan.projection(), ["foo1".to_string(), "foo2".to_string()]); - let rebuilt = rebuild_scan(scan); + let rebuilt = rebuild_scan(scan, &ctx).await; assert_eq!(rebuilt.fetch(), Some(1)); let expected = run_source(scan, &ctx).await?; insta::assert_snapshot!(&expected, @r#" @@ -1078,7 +1082,8 @@ async fn test_plan_nodes_are_inspectable() -> Result<(), Box> { let foo2_only = Arc::new(pinned.schema().project(&[1])?); let partial = IcebergDataSource::builder(pinned.table().clone(), foo2_only) .with_snapshot_id(Some(snapshot_id)) - .build()?; + .build(ctx.task_ctx()) + .await?; insta::assert_snapshot!(&run_source(&partial, &ctx).await?, @r#" +--------+ | foo2 | diff --git a/crates/datafusion/tests/snapshots/eager_file_planning__eager_multi_partition_limit_is_global.snap b/crates/datafusion/tests/snapshots/eager_file_planning__eager_multi_partition_limit_is_global.snap new file mode 100644 index 0000000..a4275c6 --- /dev/null +++ b/crates/datafusion/tests/snapshots/eager_file_planning__eager_multi_partition_limit_is_global.snap @@ -0,0 +1,15 @@ +--- +source: crates/datafusion/tests/eager_file_planning.rs +expression: "format!(\"{plan}\\n{batches}\")" +--- +ProjectionExec: expr=[1 as n] + CoalescePartitionsExec: fetch=3 + DataSourceExec: format=iceberg, projection=[], fetch=3, planned_files=7 + ++---+ +| n | ++---+ +| 1 | +| 1 | +| 1 | ++---+ diff --git a/crates/datafusion/tests/snapshots/eager_file_planning__eager_ordered_limit_matches_lazy_results.snap b/crates/datafusion/tests/snapshots/eager_file_planning__eager_ordered_limit_matches_lazy_results.snap new file mode 100644 index 0000000..24c752b --- /dev/null +++ b/crates/datafusion/tests/snapshots/eager_file_planning__eager_ordered_limit_matches_lazy_results.snap @@ -0,0 +1,16 @@ +--- +source: crates/datafusion/tests/eager_file_planning.rs +expression: "format!(\"{plan}\\n{eager_batches}\")" +--- +ProjectionExec: expr=[vendor_id@0 as vendor_id, pickup_at@1 as pickup_at] + SortPreservingMergeExec: [pickup_at@1 ASC NULLS LAST, vendor_id@0 ASC NULLS LAST, pickup_location_id@2 ASC NULLS LAST], fetch=3 + SortExec: TopK(fetch=3), expr=[pickup_at@1 ASC NULLS LAST, vendor_id@0 ASC NULLS LAST, pickup_location_id@2 ASC NULLS LAST], preserve_partitioning=[true] + DataSourceExec: format=iceberg, projection=[vendor_id, pickup_at, pickup_location_id], planned_files=7 + ++-----------+---------------------+ +| vendor_id | pickup_at | ++-----------+---------------------+ +| 2 | 2024-01-08T00:00:03 | +| 2 | 2024-01-08T00:00:07 | +| 2 | 2024-01-08T00:00:08 | ++-----------+---------------------+ diff --git a/crates/datafusion/tests/snapshots/eager_file_planning__empty_scan.snap b/crates/datafusion/tests/snapshots/eager_file_planning__empty_scan.snap new file mode 100644 index 0000000..cede0e6 --- /dev/null +++ b/crates/datafusion/tests/snapshots/eager_file_planning__empty_scan.snap @@ -0,0 +1,17 @@ +--- +source: crates/datafusion/tests/eager_file_planning.rs +expression: "format!(\"{plan}\\n{eager_batches}\")" +--- +ProjectionExec: expr=[count(Int64(1))@0 as trips] + AggregateExec: mode=Final, gby=[], aggr=[count(Int64(1))] + CoalescePartitionsExec + AggregateExec: mode=Partial, gby=[], aggr=[count(Int64(1))] + FilterExec: pickup_date@0 = 2025-01-01, projection=[] + RepartitionExec: partitioning=RoundRobinBatch(4), input_partitions=1 + DataSourceExec: format=iceberg, projection=[pickup_date], predicate=pickup_date = 2025-01-01, planned_files=0 + ++-------+ +| trips | ++-------+ +| 0 | ++-------+ diff --git a/crates/datafusion/tests/snapshots/eager_file_planning__partition_pruning.snap b/crates/datafusion/tests/snapshots/eager_file_planning__partition_pruning.snap new file mode 100644 index 0000000..cfcca93 --- /dev/null +++ b/crates/datafusion/tests/snapshots/eager_file_planning__partition_pruning.snap @@ -0,0 +1,18 @@ +--- +source: crates/datafusion/tests/eager_file_planning.rs +expression: "format!(\"{plan}\\n{eager_batches}\")" +--- +CoalescePartitionsExec + ProjectionExec: expr=[pickup_date@0 as pickup_date, count(Int64(1))@1 as trips] + AggregateExec: mode=FinalPartitioned, gby=[pickup_date@0 as pickup_date], aggr=[count(Int64(1))], ordering_mode=Sorted + RepartitionExec: partitioning=Hash([pickup_date@0], 4), input_partitions=4 + AggregateExec: mode=Partial, gby=[pickup_date@0 as pickup_date], aggr=[count(Int64(1))], ordering_mode=Sorted + FilterExec: pickup_date@0 = 2024-01-10 + RepartitionExec: partitioning=RoundRobinBatch(4), input_partitions=1 + DataSourceExec: format=iceberg, projection=[pickup_date], predicate=pickup_date = 2024-01-10, planned_files=1 + ++-------------+-------+ +| pickup_date | trips | ++-------------+-------+ +| 2024-01-10 | 25000 | ++-------------+-------+ diff --git a/crates/datafusion/tests/snapshots/eager_file_planning__residual_filter.snap b/crates/datafusion/tests/snapshots/eager_file_planning__residual_filter.snap new file mode 100644 index 0000000..f3653e7 --- /dev/null +++ b/crates/datafusion/tests/snapshots/eager_file_planning__residual_filter.snap @@ -0,0 +1,17 @@ +--- +source: crates/datafusion/tests/eager_file_planning.rs +expression: "format!(\"{plan}\\n{eager_batches}\")" +--- +ProjectionExec: expr=[count(Int64(1))@0 as trips] + AggregateExec: mode=Final, gby=[], aggr=[count(Int64(1))] + CoalescePartitionsExec + AggregateExec: mode=Partial, gby=[], aggr=[count(Int64(1))] + FilterExec: pickup_date@1 = 2024-01-08 AND trip_distance@0 > 50, projection=[] + RepartitionExec: partitioning=RoundRobinBatch(4), input_partitions=1 + DataSourceExec: format=iceberg, projection=[trip_distance, pickup_date], predicate=(pickup_date = 2024-01-08) AND (trip_distance > 50), planned_files=1 + ++-------+ +| trips | ++-------+ +| 2 | ++-------+