From 049917c63de8ffef3070878742f3c5019c332a21 Mon Sep 17 00:00:00 2001 From: "alex.bianchi" Date: Sun, 11 Oct 2026 23:38:11 -0400 Subject: [PATCH] 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 6dbb79d..6744d18 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(()) +}