From 0785d948963effd05bb899aca6454c77c6206363 Mon Sep 17 00:00:00 2001 From: Charles-Antoine Leger Date: Thu, 18 Jun 2026 17:43:26 +0200 Subject: [PATCH 01/10] feat(datafusion): add opt-in eager scan planning --- crates/integrations/datafusion/public-api.txt | 19 ++ crates/integrations/datafusion/src/config.rs | 31 ++++ crates/integrations/datafusion/src/lib.rs | 3 + .../datafusion/src/physical_plan/mod.rs | 1 + .../datafusion/src/physical_plan/scan.rs | 173 +++++++++--------- .../src/physical_plan/scan_planning.rs | 150 +++++++++++++++ .../integrations/datafusion/src/table/mod.rs | 56 ++++-- 7 files changed, 331 insertions(+), 102 deletions(-) create mode 100644 crates/integrations/datafusion/src/config.rs create mode 100644 crates/integrations/datafusion/src/physical_plan/scan_planning.rs diff --git a/crates/integrations/datafusion/public-api.txt b/crates/integrations/datafusion/public-api.txt index e197c2057d..78c7b08ed2 100644 --- a/crates/integrations/datafusion/public-api.txt +++ b/crates/integrations/datafusion/public-api.txt @@ -99,6 +99,25 @@ pub fn iceberg_datafusion::IcebergCatalogProvider::fmt(&self, f: &mut core::fmt: impl datafusion_catalog::catalog::CatalogProvider for iceberg_datafusion::IcebergCatalogProvider pub fn iceberg_datafusion::IcebergCatalogProvider::schema(&self, name: &str) -> core::option::Option> pub fn iceberg_datafusion::IcebergCatalogProvider::schema_names(&self) -> alloc::vec::Vec +#[non_exhaustive] pub struct iceberg_datafusion::IcebergDataFusionConfig +pub iceberg_datafusion::IcebergDataFusionConfig::enable_eager_scan_planning: bool +impl core::clone::Clone for iceberg_datafusion::IcebergDataFusionConfig +pub fn iceberg_datafusion::IcebergDataFusionConfig::clone(&self) -> iceberg_datafusion::IcebergDataFusionConfig +impl core::default::Default for iceberg_datafusion::IcebergDataFusionConfig +pub fn iceberg_datafusion::IcebergDataFusionConfig::default() -> Self +impl core::fmt::Debug for iceberg_datafusion::IcebergDataFusionConfig +pub fn iceberg_datafusion::IcebergDataFusionConfig::fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result +impl datafusion_common::config::ConfigExtension for iceberg_datafusion::IcebergDataFusionConfig +pub const iceberg_datafusion::IcebergDataFusionConfig::PREFIX: &'static str +impl datafusion_common::config::ConfigField for iceberg_datafusion::IcebergDataFusionConfig +pub fn iceberg_datafusion::IcebergDataFusionConfig::set(&mut self, key: &str, value: &str) -> datafusion_common::error::Result<()> +pub fn iceberg_datafusion::IcebergDataFusionConfig::visit(&self, v: &mut V, _key_prefix: &str, _description: &'static str) +impl datafusion_common::config::ExtensionOptions for iceberg_datafusion::IcebergDataFusionConfig +pub fn iceberg_datafusion::IcebergDataFusionConfig::as_any(&self) -> &dyn core::any::Any +pub fn iceberg_datafusion::IcebergDataFusionConfig::as_any_mut(&mut self) -> &mut dyn core::any::Any +pub fn iceberg_datafusion::IcebergDataFusionConfig::cloned(&self) -> alloc::boxed::Box +pub fn iceberg_datafusion::IcebergDataFusionConfig::entries(&self) -> alloc::vec::Vec +pub fn iceberg_datafusion::IcebergDataFusionConfig::set(&mut self, key: &str, value: &str) -> datafusion_common::error::Result<()> pub struct iceberg_datafusion::IcebergStaticTableProvider impl iceberg_datafusion::IcebergStaticTableProvider pub async fn iceberg_datafusion::IcebergStaticTableProvider::try_new_from_table(table: iceberg::table::Table) -> iceberg::error::Result diff --git a/crates/integrations/datafusion/src/config.rs b/crates/integrations/datafusion/src/config.rs new file mode 100644 index 0000000000..39d5f99119 --- /dev/null +++ b/crates/integrations/datafusion/src/config.rs @@ -0,0 +1,31 @@ +// 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::config::ConfigExtension; +use datafusion::common::extensions_options; + +extensions_options! { + /// Configuration options for Iceberg's DataFusion integration. + pub struct IcebergDataFusionConfig { + /// Plan Iceberg file scan tasks during TableProvider::scan(). + pub enable_eager_scan_planning: bool, default = false + } +} + +impl ConfigExtension for IcebergDataFusionConfig { + const PREFIX: &'static str = "iceberg"; +} diff --git a/crates/integrations/datafusion/src/lib.rs b/crates/integrations/datafusion/src/lib.rs index 4b0ea8606d..3e415336d3 100644 --- a/crates/integrations/datafusion/src/lib.rs +++ b/crates/integrations/datafusion/src/lib.rs @@ -18,6 +18,9 @@ mod catalog; pub use catalog::*; +mod config; +pub use config::IcebergDataFusionConfig; + mod error; pub use error::*; diff --git a/crates/integrations/datafusion/src/physical_plan/mod.rs b/crates/integrations/datafusion/src/physical_plan/mod.rs index aeac30de32..d83b6b9c4f 100644 --- a/crates/integrations/datafusion/src/physical_plan/mod.rs +++ b/crates/integrations/datafusion/src/physical_plan/mod.rs @@ -21,6 +21,7 @@ pub(crate) mod metadata_scan; pub(crate) mod project; pub(crate) mod repartition; pub(crate) mod scan; +pub(crate) mod scan_planning; pub(crate) mod sort; pub(crate) mod write; diff --git a/crates/integrations/datafusion/src/physical_plan/scan.rs b/crates/integrations/datafusion/src/physical_plan/scan.rs index 1e53c20cec..5359a32d4a 100644 --- a/crates/integrations/datafusion/src/physical_plan/scan.rs +++ b/crates/integrations/datafusion/src/physical_plan/scan.rs @@ -27,12 +27,12 @@ use datafusion::physical_expr::EquivalenceProperties; use datafusion::physical_plan::execution_plan::{Boundedness, EmissionType}; use datafusion::physical_plan::stream::RecordBatchStreamAdapter; use datafusion::physical_plan::{DisplayAs, ExecutionPlan, Partitioning, PlanProperties}; -use datafusion::prelude::Expr; use futures::{Stream, TryStreamExt}; use iceberg::expr::Predicate; +use iceberg::scan::{FileScanTask, FileScanTaskStream}; use iceberg::table::Table; -use super::expr_to_predicate::convert_filters_to_predicate; +use super::scan_planning::{IcebergScanConfig, build_table_scan}; use crate::to_datafusion_error; /// Manages the scanning process of an Iceberg [`Table`], encapsulating the @@ -41,44 +41,40 @@ use crate::to_datafusion_error; pub struct IcebergTableScan { /// A table in the catalog. table: Table, - /// Snapshot of the table to scan. - snapshot_id: Option, + /// Snapshot, projection, output schema, and pushed predicates for this scan. + scan_config: IcebergScanConfig, /// Stores certain, often expensive to compute, /// plan properties used in query optimization. plan_properties: Arc, - /// Projection column names, None means all columns - projection: Option>, - /// Filters to apply to the table scan - predicates: Option, /// Optional limit on the number of rows to return limit: Option, + /// Pre-planned file scan tasks, grouped by partition. `None` keeps planning lazy. + file_task_groups: Option>>, } impl IcebergTableScan { - /// Creates a new [`IcebergTableScan`] object. pub(crate) fn new( table: Table, - snapshot_id: Option, - schema: ArrowSchemaRef, - projection: Option<&Vec>, - filters: &[Expr], + scan_config: IcebergScanConfig, limit: Option, + file_task_groups: Option>>, ) -> Self { - let output_schema = match projection { - None => schema.clone(), - Some(projection) => Arc::new(schema.project(projection).unwrap()), - }; - let plan_properties = Self::compute_properties(output_schema.clone()); - let projection = get_column_names(schema.clone(), projection); - let predicates = convert_filters_to_predicate(filters); + let partition_count = file_task_groups.as_ref().map_or(1, |groups| groups.len()); + let plan_properties = + IcebergTableScan::compute_properties(scan_config.output_schema(), partition_count); + let file_task_groups = file_task_groups.map(|groups| { + groups + .into_iter() + .map(Arc::<[FileScanTask]>::from) + .collect() + }); Self { table, - snapshot_id, + scan_config, plan_properties, - projection, - predicates, limit, + file_task_groups, } } @@ -87,15 +83,15 @@ impl IcebergTableScan { } pub fn snapshot_id(&self) -> Option { - self.snapshot_id + self.scan_config.snapshot_id() } pub fn projection(&self) -> Option<&[String]> { - self.projection.as_deref() + self.scan_config.column_names() } pub fn predicates(&self) -> Option<&Predicate> { - self.predicates.as_ref() + self.scan_config.predicates() } pub fn limit(&self) -> Option { @@ -103,13 +99,10 @@ impl IcebergTableScan { } /// Computes [`PlanProperties`] used in query optimization. - fn compute_properties(schema: ArrowSchemaRef) -> Arc { - // TODO: - // This is more or less a placeholder, to be replaced - // once we support output-partitioning + fn compute_properties(schema: ArrowSchemaRef, partition_count: usize) -> Arc { Arc::new(PlanProperties::new( EquivalenceProperties::new(schema), - Partitioning::UnknownPartitioning(1), + Partitioning::UnknownPartitioning(partition_count), EmissionType::Incremental, Boundedness::Bounded, )) @@ -141,13 +134,55 @@ impl ExecutionPlan for IcebergTableScan { _partition: usize, _context: Arc, ) -> DFResult { - let fut = get_batch_stream( - self.table.clone(), - self.snapshot_id, - self.projection.clone(), - self.predicates.clone(), - ); - let stream = futures::stream::once(fut).try_flatten(); + let stream: Pin> + Send>> = match &self + .file_task_groups + { + Some(file_task_groups) => { + let Some(file_task_group) = file_task_groups.get(_partition).cloned() else { + return Err(datafusion::common::DataFusionError::Internal(format!( + "IcebergTableScan partition {_partition} does not exist; scan has {} partitions", + file_task_groups.len() + ))); + }; + + let task_count = file_task_group.len(); + let tasks: FileScanTaskStream = Box::pin(futures::stream::iter( + (0..task_count).map(move |idx| Ok(file_task_group[idx].clone())), + )); + let stream = self + .table + .reader_builder() + // Eager planning lets DataFusion drive scan concurrency via output + // partitions. Match DataFusion's FileStream model, where each + // output partition owns one ScanState; keep one data file in + // flight per output partition here. + // https://github.com/apache/datafusion/blob/ad8e7b7f2babe3fcddc3a4f9b5cd1ac0d1b16ad9/datafusion/datasource/src/file_stream/scan_state.rs#L42-L43 + .with_data_file_concurrency_limit(1) + .build() + // TODO: Avoid cloning FileScanTasks here once ArrowReader can accept shared tasks. + .read(tasks) + .map_err(to_datafusion_error)? + .stream() + .map_err(to_datafusion_error); + + Box::pin(stream) + } + None => { + let table = self.table.clone(); + let scan_config = self.scan_config.clone(); + let fut = async move { + let table_scan = build_table_scan(&table, &scan_config)?; + let stream = table_scan + .to_arrow() + .await + .map_err(to_datafusion_error)? + .map_err(to_datafusion_error); + Ok::<_, datafusion::common::DataFusionError>(stream) + }; + + Box::pin(futures::stream::once(fut).try_flatten()) + } + }; // Apply limit if specified let limited_stream: Pin> + Send>> = @@ -185,60 +220,22 @@ impl DisplayAs for IcebergTableScan { write!( f, "IcebergTableScan projection:[{}] predicate:[{}]", - self.projection - .clone() - .map_or(String::new(), |v| v.join(",")), - self.predicates - .clone() - .map_or(String::from(""), |p| format!("{p}")) + self.projection().map_or(String::new(), |v| v.join(",")), + self.predicates() + .map_or(String::from(""), |p| format!("{p}")), )?; + if let Some(file_task_groups) = &self.file_task_groups { + let task_count: usize = file_task_groups.iter().map(|group| group.len()).sum(); + write!( + f, + " task_groups:[{}] tasks:[{}]", + file_task_groups.len(), + task_count, + )?; + } if let Some(limit) = self.limit { write!(f, " limit:[{limit}]")?; } Ok(()) } } - -/// 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>, - predicates: Option, -) -> DFResult> + 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(pred) = predicates { - scan_builder = scan_builder.with_filter(pred); - } - let table_scan = scan_builder.build().map_err(to_datafusion_error)?; - - let stream = table_scan - .to_arrow() - .await - .map_err(to_datafusion_error)? - .map_err(to_datafusion_error); - Ok(Box::pin(stream)) -} - -fn get_column_names( - schema: ArrowSchemaRef, - projection: Option<&Vec>, -) -> Option> { - projection.map(|v| { - v.iter() - .map(|p| schema.field(*p).name().clone()) - .collect::>() - }) -} diff --git a/crates/integrations/datafusion/src/physical_plan/scan_planning.rs b/crates/integrations/datafusion/src/physical_plan/scan_planning.rs new file mode 100644 index 0000000000..181e95fbf2 --- /dev/null +++ b/crates/integrations/datafusion/src/physical_plan/scan_planning.rs @@ -0,0 +1,150 @@ +// 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::arrow::datatypes::SchemaRef as ArrowSchemaRef; +use datafusion::error::Result as DFResult; +use datafusion::prelude::Expr; +use futures::TryStreamExt; +use iceberg::expr::Predicate; +use iceberg::scan::{FileScanTask, TableScan}; +use iceberg::table::Table; + +use super::expr_to_predicate::convert_filters_to_predicate; +use crate::to_datafusion_error; + +#[derive(Debug, Clone)] +pub(crate) struct IcebergScanConfig { + /// Snapshot of the table to scan. + snapshot_id: Option, + /// Output schema after projection. + output_schema: ArrowSchemaRef, + /// Projection column names, None means all columns. + column_names: Option>, + /// Filters to apply to the table scan. + predicates: Option, +} + +impl IcebergScanConfig { + pub(crate) fn new( + schema: ArrowSchemaRef, + snapshot_id: Option, + projection: Option<&Vec>, + filters: &[Expr], + ) -> Self { + let output_schema = match projection { + None => schema.clone(), + Some(projection) => Arc::new(schema.project(projection).unwrap()), + }; + + Self { + snapshot_id, + output_schema, + column_names: get_column_names(schema, projection), + predicates: convert_filters_to_predicate(filters), + } + } + + pub(crate) fn snapshot_id(&self) -> Option { + self.snapshot_id + } + + pub(crate) fn output_schema(&self) -> ArrowSchemaRef { + self.output_schema.clone() + } + + pub(crate) fn column_names(&self) -> Option<&[String]> { + self.column_names.as_deref() + } + + pub(crate) fn predicates(&self) -> Option<&Predicate> { + self.predicates.as_ref() + } +} + +pub(crate) async fn plan_file_task_groups( + table: &Table, + scan_config: &IcebergScanConfig, + target_partitions: usize, +) -> DFResult>> { + // Do not cache planned FileScanTasks in the provider in v1. They are query-specific + // because projection, predicate binding, snapshot schema, and delete planning can differ + // between scans. Catalog-backed providers also need fresh metadata on each scan. + // TODO: Revisit provider-level caching for static tables with a precise cache key. + let tasks: Vec = build_table_scan(table, scan_config)? + .plan_files() + .await + .map_err(to_datafusion_error)? + .try_collect::>() + .await + .map_err(to_datafusion_error)?; + + Ok(group_file_scan_tasks_round_robin(tasks, target_partitions)) +} + +fn get_column_names( + schema: ArrowSchemaRef, + projection: Option<&Vec>, +) -> Option> { + projection.map(|v| { + v.iter() + .map(|p| schema.field(*p).name().clone()) + .collect::>() + }) +} + +/// Groups file scan tasks into `target_partitions` groups using a naive +/// round-robin assignment. `target_partitions` is clamped to a minimum of 1. +// TODO: Replace this naive round-robin grouping with size-based grouping once the +// first parallel scan path is stable. Keep this v1 simple and deterministic. +fn group_file_scan_tasks_round_robin( + tasks: Vec, + target_partitions: usize, +) -> Vec> { + if tasks.is_empty() { + return vec![vec![]]; + } + + let target_partitions = target_partitions.max(1); + + let mut groups: Vec> = vec![Vec::new(); target_partitions]; + for (i, task) in tasks.into_iter().enumerate() { + groups[i % target_partitions].push(task); + } + + groups.retain(|group| !group.is_empty()); + groups +} + +pub(crate) fn build_table_scan( + table: &Table, + scan_config: &IcebergScanConfig, +) -> DFResult { + let builder = match scan_config.snapshot_id { + Some(id) => table.scan().snapshot_id(id), + None => table.scan(), + }; + let mut builder = match scan_config.column_names.clone() { + Some(names) => builder.select(names), + None => builder.select_all(), + }; + if let Some(pred) = scan_config.predicates.clone() { + builder = builder.with_filter(pred); + } + builder.build().map_err(to_datafusion_error) +} diff --git a/crates/integrations/datafusion/src/table/mod.rs b/crates/integrations/datafusion/src/table/mod.rs index 2fd958dff4..91308324e9 100644 --- a/crates/integrations/datafusion/src/table/mod.rs +++ b/crates/integrations/datafusion/src/table/mod.rs @@ -48,14 +48,45 @@ use iceberg::table::Table; use iceberg::{Catalog, Error, ErrorKind, NamespaceIdent, Result, TableIdent}; use metadata_table::IcebergMetadataTableProvider; +use crate::IcebergDataFusionConfig; use crate::error::to_datafusion_error; use crate::physical_plan::commit::IcebergCommitExec; use crate::physical_plan::project::project_with_partition; use crate::physical_plan::repartition::repartition; use crate::physical_plan::scan::IcebergTableScan; +use crate::physical_plan::scan_planning::{IcebergScanConfig, plan_file_task_groups}; use crate::physical_plan::sort::sort_by_partition; use crate::physical_plan::write::IcebergWriteExec; +fn enable_eager_scan_planning(state: &dyn Session) -> bool { + state + .config() + .options() + .extensions + .get::() + .is_some_and(|config| config.enable_eager_scan_planning) +} + +async fn create_scan_plan( + state: &dyn Session, + table: Table, + scan_config: IcebergScanConfig, + limit: Option, +) -> DFResult> { + let file_task_groups = if enable_eager_scan_planning(state) { + Some(plan_file_task_groups(&table, &scan_config, state.config().target_partitions()).await?) + } else { + None + }; + + Ok(Arc::new(IcebergTableScan::new( + table, + scan_config, + limit, + file_task_groups, + ))) +} + /// Catalog-backed table provider with automatic metadata refresh. /// /// This provider loads fresh table metadata from the catalog on every scan and write @@ -131,15 +162,15 @@ impl TableProvider for IcebergTableProvider { .await .map_err(to_datafusion_error)?; - // Create scan with fresh metadata (always use current snapshot) - Ok(Arc::new(IcebergTableScan::new( - table, - None, // Always use current snapshot for catalog-backed provider + let scan_config = IcebergScanConfig::new( self.schema.clone(), + None, // Always use current snapshot for catalog-backed provider. projection, filters, - limit, - ))) + ); + + // Create scan with fresh metadata (always use current snapshot) + create_scan_plan(_state, table, scan_config, limit).await } fn supports_filters_pushdown( @@ -311,15 +342,11 @@ impl TableProvider for IcebergStaticTableProvider { filters: &[Expr], limit: Option, ) -> DFResult> { + let scan_config = + IcebergScanConfig::new(self.schema.clone(), self.snapshot_id, projection, filters); + // Use cached table (no refresh) - Ok(Arc::new(IcebergTableScan::new( - self.table.clone(), - self.snapshot_id, - self.schema.clone(), - projection, - filters, - limit, - ))) + create_scan_plan(_state, self.table.clone(), scan_config, limit).await } fn supports_filters_pushdown( @@ -361,6 +388,7 @@ mod tests { use tempfile::TempDir; use super::*; + use crate::physical_plan::scan::IcebergTableScan; async fn get_test_table_from_metadata_file() -> Table { let metadata_file_name = "TableMetadataV2Valid.json"; From 29acf0ff8a58b8b3d4a49aa4d00f1e8c418a6862 Mon Sep 17 00:00:00 2001 From: Charles-Antoine Leger Date: Thu, 18 Jun 2026 17:43:35 +0200 Subject: [PATCH 02/10] test(datafusion): cover lazy and eager scan planning --- .../tests/integration_datafusion_test.rs | 239 +++++++++++++++++- 1 file changed, 237 insertions(+), 2 deletions(-) diff --git a/crates/integrations/datafusion/tests/integration_datafusion_test.rs b/crates/integrations/datafusion/tests/integration_datafusion_test.rs index 9aa6a1a501..2b16fcd463 100644 --- a/crates/integrations/datafusion/tests/integration_datafusion_test.rs +++ b/crates/integrations/datafusion/tests/integration_datafusion_test.rs @@ -21,10 +21,13 @@ use std::collections::HashMap; use std::sync::Arc; use std::vec; -use datafusion::arrow::array::{Array, StringArray, UInt64Array}; +use datafusion::arrow::array::{Array, ArrayRef, Int32Array, StringArray, UInt64Array}; use datafusion::arrow::datatypes::{DataType, Field, Schema as ArrowSchema}; +use datafusion::arrow::record_batch::RecordBatch; +use datafusion::datasource::MemTable; use datafusion::execution::context::SessionContext; use datafusion::parquet::arrow::PARQUET_FIELD_ID_META_KEY; +use datafusion::prelude::SessionConfig; use expect_test::expect; use iceberg::io::LocalFsStorageFactory; use iceberg::memory::{MEMORY_CATALOG_WAREHOUSE, MemoryCatalogBuilder}; @@ -35,7 +38,8 @@ use iceberg::test_utils::check_record_batches; use iceberg::{ Catalog, CatalogBuilder, MemoryCatalog, NamespaceIdent, Result, TableCreation, TableIdent, }; -use iceberg_datafusion::IcebergCatalogProvider; +use iceberg_datafusion::physical_plan::IcebergTableScan; +use iceberg_datafusion::{IcebergCatalogProvider, IcebergDataFusionConfig}; use tempfile::TempDir; fn temp_path() -> String { @@ -95,6 +99,237 @@ fn get_table_creation( Ok(creation) } +async fn get_multi_file_table_context( + namespace_name: &str, + table_name: &str, + data_file_count: usize, +) -> Result<(Arc, NamespaceIdent, String)> { + let iceberg_catalog = Arc::new(get_iceberg_catalog().await); + let namespace = NamespaceIdent::new(namespace_name.to_string()); + set_test_namespace(&iceberg_catalog, &namespace).await?; + + let creation = get_table_creation(temp_path(), table_name, None)?; + iceberg_catalog.create_table(&namespace, creation).await?; + + let write_ctx = SessionContext::new_with_config( + SessionConfig::new().with_target_partitions(data_file_count), + ); + let arrow_schema = Arc::new(ArrowSchema::new(vec![ + Field::new("foo1", DataType::Int32, false), + Field::new("foo2", DataType::Utf8, false), + ])); + + let batches: Vec = (1..=data_file_count as i32) + .map(|idx| { + RecordBatch::try_new(arrow_schema.clone(), vec![ + Arc::new(Int32Array::from(vec![idx])) as ArrayRef, + Arc::new(StringArray::from(vec![format!("row-{idx}")])) as ArrayRef, + ]) + }) + .collect::>()?; + + let partitions = batches.into_iter().map(|batch| vec![batch]).collect(); + let source_table = Arc::new(MemTable::try_new(arrow_schema, partitions).unwrap()); + write_ctx + .register_table("source_table", source_table) + .unwrap(); + + let catalog = Arc::new(IcebergCatalogProvider::try_new(iceberg_catalog.clone()).await?); + write_ctx.register_catalog("catalog", catalog); + + let insert_sql = + format!("INSERT INTO catalog.{namespace_name}.{table_name} SELECT * FROM source_table"); + let batches = write_ctx + .sql(&insert_sql) + .await + .unwrap() + .collect() + .await + .unwrap(); + assert_eq!(batches.len(), 1); + + let rows_inserted = batches[0] + .column(0) + .as_any() + .downcast_ref::() + .unwrap(); + assert_eq!(rows_inserted.value(0), data_file_count as u64); + + Ok((iceberg_catalog, namespace, table_name.to_string())) +} + +async fn get_read_context( + catalog: Arc, + target_partitions: usize, + enable_eager_scan_planning: Option, +) -> Result { + let ctx = SessionContext::new_with_config(read_session_config( + target_partitions, + enable_eager_scan_planning, + )); + let catalog = Arc::new(IcebergCatalogProvider::try_new(catalog).await?); + ctx.register_catalog("catalog", catalog); + Ok(ctx) +} + +fn read_session_config( + target_partitions: usize, + enable_eager_scan_planning: Option, +) -> SessionConfig { + let config = SessionConfig::new().with_target_partitions(target_partitions); + + match enable_eager_scan_planning { + Some(enabled) => { + let mut iceberg_config = IcebergDataFusionConfig::default(); + iceberg_config.enable_eager_scan_planning = enabled; + config.with_option_extension(iceberg_config) + } + None => config, + } +} + +async fn scan_partition_count( + ctx: &SessionContext, + namespace: &NamespaceIdent, + table_name: &str, +) -> usize { + let provider = ctx.catalog("catalog").unwrap(); + let namespace_name = &namespace[0]; + let schema = provider.schema(namespace_name).unwrap(); + let table = schema.table(table_name).await.unwrap().unwrap(); + + let state = ctx.state(); + let plan = table.scan(&state, None, &[], None).await.unwrap(); + plan.downcast_ref::() + .expect("Expected IcebergTableScan"); + plan.properties().output_partitioning().partition_count() +} + +#[tokio::test] +async fn test_multi_file_scan_produces_multiple_partitions() -> Result<()> { + let data_file_count = 3; + // Ask for more partitions than files to verify scan planning does not expose empty partitions. + let target_partitions = data_file_count + 1; + let (iceberg_catalog, namespace, table_name) = get_multi_file_table_context( + "test_multi_file_scan_partitions", + "my_table", + data_file_count, + ) + .await?; + let ctx = get_read_context(iceberg_catalog, target_partitions, Some(true)).await?; + let actual_partition_count = scan_partition_count(&ctx, &namespace, &table_name).await; + + assert_eq!(actual_partition_count, data_file_count); + + Ok(()) +} + +#[tokio::test] +async fn test_multi_file_scan_defaults_to_single_lazy_partition() -> Result<()> { + let data_file_count = 3; + let target_partitions = data_file_count + 1; + let (iceberg_catalog, namespace, table_name) = get_multi_file_table_context( + "test_multi_file_scan_default_lazy", + "my_table", + data_file_count, + ) + .await?; + let ctx = get_read_context(iceberg_catalog, target_partitions, None).await?; + + let actual_partition_count = scan_partition_count(&ctx, &namespace, &table_name).await; + + assert_eq!(actual_partition_count, 1); + + Ok(()) +} + +#[tokio::test] +async fn test_set_enable_eager_scan_planning() -> Result<()> { + let data_file_count = 3; + let target_partitions = data_file_count + 1; + let (iceberg_catalog, namespace, table_name) = + get_multi_file_table_context("test_set_eager_scan_planning", "my_table", data_file_count) + .await?; + let ctx = get_read_context(iceberg_catalog, target_partitions, Some(false)).await?; + + ctx.sql("SET iceberg.enable_eager_scan_planning = true") + .await + .unwrap() + .collect() + .await + .unwrap(); + + let actual_partition_count = scan_partition_count(&ctx, &namespace, &table_name).await; + + assert_eq!(actual_partition_count, data_file_count); + + Ok(()) +} + +#[tokio::test] +async fn test_multi_partition_scan_matches_single_partition_results() -> Result<()> { + let data_file_count = 3; + let target_partitions = data_file_count + 1; + let (iceberg_catalog, namespace, table_name) = get_multi_file_table_context( + "test_multi_partition_scan_results", + "my_table", + data_file_count, + ) + .await?; + let namespace_name = &namespace[0]; + + let single_partition_ctx = get_read_context(iceberg_catalog.clone(), 1, Some(true)).await?; + let multi_partition_ctx = + get_read_context(iceberg_catalog, target_partitions, Some(true)).await?; + + let query = + format!("SELECT foo1, foo2 FROM catalog.{namespace_name}.{table_name} ORDER BY foo1"); + + let single_partition_batches = single_partition_ctx + .sql(&query) + .await + .unwrap() + .collect() + .await + .unwrap(); + let multi_partition_batches = multi_partition_ctx + .sql(&query) + .await + .unwrap() + .collect() + .await + .unwrap(); + + let assert_expected_batches = |batches| { + check_record_batches( + batches, + expect![[r#" + Field { "foo1": Int32, metadata: {"PARQUET:field_id": "1"} }, + Field { "foo2": Utf8, metadata: {"PARQUET:field_id": "2"} }"#]], + expect![[r#" + foo1: PrimitiveArray + [ + 1, + 2, + 3, + ], + foo2: StringArray + [ + "row-1", + "row-2", + "row-3", + ]"#]], + &[], + None, + ); + }; + + assert_expected_batches(single_partition_batches); + assert_expected_batches(multi_partition_batches); + + Ok(()) +} + #[tokio::test] async fn test_provider_plan_stream_schema() -> Result<()> { let iceberg_catalog = get_iceberg_catalog().await; From c6f7a1f9968bee7bb8bd635a961b36683d489671 Mon Sep 17 00:00:00 2001 From: toutane Date: Fri, 10 Jul 2026 14:51:34 +0200 Subject: [PATCH 03/10] fix(datafusion): address eager scan review feedback --- crates/integrations/datafusion/public-api.txt | 10 +++++----- .../integrations/datafusion/src/physical_plan/scan.rs | 9 ++++----- .../datafusion/src/physical_plan/scan_planning.rs | 5 ++--- crates/integrations/datafusion/src/table/mod.rs | 8 ++++---- 4 files changed, 15 insertions(+), 17 deletions(-) diff --git a/crates/integrations/datafusion/public-api.txt b/crates/integrations/datafusion/public-api.txt index 78c7b08ed2..821525dea7 100644 --- a/crates/integrations/datafusion/public-api.txt +++ b/crates/integrations/datafusion/public-api.txt @@ -25,7 +25,7 @@ impl datafusion_physical_plan::display::DisplayAs for iceberg_datafusion::physic pub fn iceberg_datafusion::physical_plan::IcebergTableScan::fmt_as(&self, _t: datafusion_physical_plan::display::DisplayFormatType, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result impl datafusion_physical_plan::execution_plan::ExecutionPlan for iceberg_datafusion::physical_plan::IcebergTableScan pub fn iceberg_datafusion::physical_plan::IcebergTableScan::children(&self) -> alloc::vec::Vec<&alloc::sync::Arc<(dyn datafusion_physical_plan::execution_plan::ExecutionPlan + 'static)>> -pub fn iceberg_datafusion::physical_plan::IcebergTableScan::execute(&self, _partition: usize, _context: alloc::sync::Arc) -> datafusion_common::error::Result +pub fn iceberg_datafusion::physical_plan::IcebergTableScan::execute(&self, partition: usize, _context: alloc::sync::Arc) -> datafusion_common::error::Result pub fn iceberg_datafusion::physical_plan::IcebergTableScan::name(&self) -> &str pub fn iceberg_datafusion::physical_plan::IcebergTableScan::properties(&self) -> &alloc::sync::Arc pub fn iceberg_datafusion::physical_plan::IcebergTableScan::with_new_children(self: alloc::sync::Arc, _children: alloc::vec::Vec>) -> datafusion_common::error::Result> @@ -65,7 +65,7 @@ impl core::fmt::Debug for iceberg_datafusion::IcebergStaticTableProvider pub fn iceberg_datafusion::IcebergStaticTableProvider::fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result impl datafusion_catalog::table::TableProvider for iceberg_datafusion::IcebergStaticTableProvider pub fn iceberg_datafusion::IcebergStaticTableProvider::insert_into<'life0, 'life1, 'async_trait>(&'life0 self, _state: &'life1 dyn datafusion_session::session::Session, _input: alloc::sync::Arc, _insert_op: datafusion_expr::logical_plan::dml::InsertOp) -> core::pin::Pin>> + core::marker::Send + 'async_trait)>> where Self: 'async_trait, 'life0: 'async_trait, 'life1: 'async_trait -pub fn iceberg_datafusion::IcebergStaticTableProvider::scan<'life0, 'life1, 'life2, 'life3, 'async_trait>(&'life0 self, _state: &'life1 dyn datafusion_session::session::Session, projection: core::option::Option<&'life2 alloc::vec::Vec>, filters: &'life3 [datafusion_expr::expr::Expr], limit: core::option::Option) -> core::pin::Pin>> + core::marker::Send + 'async_trait)>> where Self: 'async_trait, 'life0: 'async_trait, 'life1: 'async_trait, 'life2: 'async_trait, 'life3: 'async_trait +pub fn iceberg_datafusion::IcebergStaticTableProvider::scan<'life0, 'life1, 'life2, 'life3, 'async_trait>(&'life0 self, state: &'life1 dyn datafusion_session::session::Session, projection: core::option::Option<&'life2 alloc::vec::Vec>, filters: &'life3 [datafusion_expr::expr::Expr], limit: core::option::Option) -> core::pin::Pin>> + core::marker::Send + 'async_trait)>> where Self: 'async_trait, 'life0: 'async_trait, 'life1: 'async_trait, 'life2: 'async_trait, 'life3: 'async_trait pub fn iceberg_datafusion::IcebergStaticTableProvider::schema(&self) -> arrow_schema::schema::SchemaRef pub fn iceberg_datafusion::IcebergStaticTableProvider::supports_filters_pushdown(&self, filters: &[&datafusion_expr::expr::Expr]) -> datafusion_common::error::Result> pub fn iceberg_datafusion::IcebergStaticTableProvider::table_type(&self) -> datafusion_expr::table_source::TableType @@ -76,7 +76,7 @@ impl core::fmt::Debug for iceberg_datafusion::IcebergTableProvider pub fn iceberg_datafusion::IcebergTableProvider::fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result impl datafusion_catalog::table::TableProvider for iceberg_datafusion::IcebergTableProvider pub fn iceberg_datafusion::IcebergTableProvider::insert_into<'life0, 'life1, 'async_trait>(&'life0 self, state: &'life1 dyn datafusion_session::session::Session, input: alloc::sync::Arc, _insert_op: datafusion_expr::logical_plan::dml::InsertOp) -> core::pin::Pin>> + core::marker::Send + 'async_trait)>> where Self: 'async_trait, 'life0: 'async_trait, 'life1: 'async_trait -pub fn iceberg_datafusion::IcebergTableProvider::scan<'life0, 'life1, 'life2, 'life3, 'async_trait>(&'life0 self, _state: &'life1 dyn datafusion_session::session::Session, projection: core::option::Option<&'life2 alloc::vec::Vec>, filters: &'life3 [datafusion_expr::expr::Expr], limit: core::option::Option) -> core::pin::Pin>> + core::marker::Send + 'async_trait)>> where Self: 'async_trait, 'life0: 'async_trait, 'life1: 'async_trait, 'life2: 'async_trait, 'life3: 'async_trait +pub fn iceberg_datafusion::IcebergTableProvider::scan<'life0, 'life1, 'life2, 'life3, 'async_trait>(&'life0 self, state: &'life1 dyn datafusion_session::session::Session, projection: core::option::Option<&'life2 alloc::vec::Vec>, filters: &'life3 [datafusion_expr::expr::Expr], limit: core::option::Option) -> core::pin::Pin>> + core::marker::Send + 'async_trait)>> where Self: 'async_trait, 'life0: 'async_trait, 'life1: 'async_trait, 'life2: 'async_trait, 'life3: 'async_trait pub fn iceberg_datafusion::IcebergTableProvider::schema(&self) -> arrow_schema::schema::SchemaRef pub fn iceberg_datafusion::IcebergTableProvider::supports_filters_pushdown(&self, filters: &[&datafusion_expr::expr::Expr]) -> datafusion_common::error::Result> pub fn iceberg_datafusion::IcebergTableProvider::table_type(&self) -> datafusion_expr::table_source::TableType @@ -128,7 +128,7 @@ impl core::fmt::Debug for iceberg_datafusion::IcebergStaticTableProvider pub fn iceberg_datafusion::IcebergStaticTableProvider::fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result impl datafusion_catalog::table::TableProvider for iceberg_datafusion::IcebergStaticTableProvider pub fn iceberg_datafusion::IcebergStaticTableProvider::insert_into<'life0, 'life1, 'async_trait>(&'life0 self, _state: &'life1 dyn datafusion_session::session::Session, _input: alloc::sync::Arc, _insert_op: datafusion_expr::logical_plan::dml::InsertOp) -> core::pin::Pin>> + core::marker::Send + 'async_trait)>> where Self: 'async_trait, 'life0: 'async_trait, 'life1: 'async_trait -pub fn iceberg_datafusion::IcebergStaticTableProvider::scan<'life0, 'life1, 'life2, 'life3, 'async_trait>(&'life0 self, _state: &'life1 dyn datafusion_session::session::Session, projection: core::option::Option<&'life2 alloc::vec::Vec>, filters: &'life3 [datafusion_expr::expr::Expr], limit: core::option::Option) -> core::pin::Pin>> + core::marker::Send + 'async_trait)>> where Self: 'async_trait, 'life0: 'async_trait, 'life1: 'async_trait, 'life2: 'async_trait, 'life3: 'async_trait +pub fn iceberg_datafusion::IcebergStaticTableProvider::scan<'life0, 'life1, 'life2, 'life3, 'async_trait>(&'life0 self, state: &'life1 dyn datafusion_session::session::Session, projection: core::option::Option<&'life2 alloc::vec::Vec>, filters: &'life3 [datafusion_expr::expr::Expr], limit: core::option::Option) -> core::pin::Pin>> + core::marker::Send + 'async_trait)>> where Self: 'async_trait, 'life0: 'async_trait, 'life1: 'async_trait, 'life2: 'async_trait, 'life3: 'async_trait pub fn iceberg_datafusion::IcebergStaticTableProvider::schema(&self) -> arrow_schema::schema::SchemaRef pub fn iceberg_datafusion::IcebergStaticTableProvider::supports_filters_pushdown(&self, filters: &[&datafusion_expr::expr::Expr]) -> datafusion_common::error::Result> pub fn iceberg_datafusion::IcebergStaticTableProvider::table_type(&self) -> datafusion_expr::table_source::TableType @@ -139,7 +139,7 @@ impl core::fmt::Debug for iceberg_datafusion::IcebergTableProvider pub fn iceberg_datafusion::IcebergTableProvider::fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result impl datafusion_catalog::table::TableProvider for iceberg_datafusion::IcebergTableProvider pub fn iceberg_datafusion::IcebergTableProvider::insert_into<'life0, 'life1, 'async_trait>(&'life0 self, state: &'life1 dyn datafusion_session::session::Session, input: alloc::sync::Arc, _insert_op: datafusion_expr::logical_plan::dml::InsertOp) -> core::pin::Pin>> + core::marker::Send + 'async_trait)>> where Self: 'async_trait, 'life0: 'async_trait, 'life1: 'async_trait -pub fn iceberg_datafusion::IcebergTableProvider::scan<'life0, 'life1, 'life2, 'life3, 'async_trait>(&'life0 self, _state: &'life1 dyn datafusion_session::session::Session, projection: core::option::Option<&'life2 alloc::vec::Vec>, filters: &'life3 [datafusion_expr::expr::Expr], limit: core::option::Option) -> core::pin::Pin>> + core::marker::Send + 'async_trait)>> where Self: 'async_trait, 'life0: 'async_trait, 'life1: 'async_trait, 'life2: 'async_trait, 'life3: 'async_trait +pub fn iceberg_datafusion::IcebergTableProvider::scan<'life0, 'life1, 'life2, 'life3, 'async_trait>(&'life0 self, state: &'life1 dyn datafusion_session::session::Session, projection: core::option::Option<&'life2 alloc::vec::Vec>, filters: &'life3 [datafusion_expr::expr::Expr], limit: core::option::Option) -> core::pin::Pin>> + core::marker::Send + 'async_trait)>> where Self: 'async_trait, 'life0: 'async_trait, 'life1: 'async_trait, 'life2: 'async_trait, 'life3: 'async_trait pub fn iceberg_datafusion::IcebergTableProvider::schema(&self) -> arrow_schema::schema::SchemaRef pub fn iceberg_datafusion::IcebergTableProvider::supports_filters_pushdown(&self, filters: &[&datafusion_expr::expr::Expr]) -> datafusion_common::error::Result> pub fn iceberg_datafusion::IcebergTableProvider::table_type(&self) -> datafusion_expr::table_source::TableType diff --git a/crates/integrations/datafusion/src/physical_plan/scan.rs b/crates/integrations/datafusion/src/physical_plan/scan.rs index 5359a32d4a..d9433f6cdf 100644 --- a/crates/integrations/datafusion/src/physical_plan/scan.rs +++ b/crates/integrations/datafusion/src/physical_plan/scan.rs @@ -131,23 +131,22 @@ impl ExecutionPlan for IcebergTableScan { fn execute( &self, - _partition: usize, + partition: usize, _context: Arc, ) -> DFResult { let stream: Pin> + Send>> = match &self .file_task_groups { Some(file_task_groups) => { - let Some(file_task_group) = file_task_groups.get(_partition).cloned() else { + let Some(file_task_group) = file_task_groups.get(partition).cloned() else { return Err(datafusion::common::DataFusionError::Internal(format!( - "IcebergTableScan partition {_partition} does not exist; scan has {} partitions", + "IcebergTableScan partition {partition} does not exist; scan has {} partitions", file_task_groups.len() ))); }; - let task_count = file_task_group.len(); let tasks: FileScanTaskStream = Box::pin(futures::stream::iter( - (0..task_count).map(move |idx| Ok(file_task_group[idx].clone())), + (0..file_task_group.len()).map(move |idx| Ok(file_task_group[idx].clone())), )); let stream = self .table diff --git a/crates/integrations/datafusion/src/physical_plan/scan_planning.rs b/crates/integrations/datafusion/src/physical_plan/scan_planning.rs index 181e95fbf2..678984d94b 100644 --- a/crates/integrations/datafusion/src/physical_plan/scan_planning.rs +++ b/crates/integrations/datafusion/src/physical_plan/scan_planning.rs @@ -109,7 +109,7 @@ fn get_column_names( } /// Groups file scan tasks into `target_partitions` groups using a naive -/// round-robin assignment. `target_partitions` is clamped to a minimum of 1. +/// round-robin assignment. Non-empty groups are bounded by `tasks.len()`. // TODO: Replace this naive round-robin grouping with size-based grouping once the // first parallel scan path is stable. Keep this v1 simple and deterministic. fn group_file_scan_tasks_round_robin( @@ -120,14 +120,13 @@ fn group_file_scan_tasks_round_robin( return vec![vec![]]; } - let target_partitions = target_partitions.max(1); + let target_partitions = target_partitions.max(1).min(tasks.len()); let mut groups: Vec> = vec![Vec::new(); target_partitions]; for (i, task) in tasks.into_iter().enumerate() { groups[i % target_partitions].push(task); } - groups.retain(|group| !group.is_empty()); groups } diff --git a/crates/integrations/datafusion/src/table/mod.rs b/crates/integrations/datafusion/src/table/mod.rs index 91308324e9..e669802630 100644 --- a/crates/integrations/datafusion/src/table/mod.rs +++ b/crates/integrations/datafusion/src/table/mod.rs @@ -150,7 +150,7 @@ impl TableProvider for IcebergTableProvider { async fn scan( &self, - _state: &dyn Session, + state: &dyn Session, projection: Option<&Vec>, filters: &[Expr], limit: Option, @@ -170,7 +170,7 @@ impl TableProvider for IcebergTableProvider { ); // Create scan with fresh metadata (always use current snapshot) - create_scan_plan(_state, table, scan_config, limit).await + create_scan_plan(state, table, scan_config, limit).await } fn supports_filters_pushdown( @@ -337,7 +337,7 @@ impl TableProvider for IcebergStaticTableProvider { async fn scan( &self, - _state: &dyn Session, + state: &dyn Session, projection: Option<&Vec>, filters: &[Expr], limit: Option, @@ -346,7 +346,7 @@ impl TableProvider for IcebergStaticTableProvider { IcebergScanConfig::new(self.schema.clone(), self.snapshot_id, projection, filters); // Use cached table (no refresh) - create_scan_plan(_state, self.table.clone(), scan_config, limit).await + create_scan_plan(state, self.table.clone(), scan_config, limit).await } fn supports_filters_pushdown( From 1b0504d93faa3b585fc80f1a9b3977e4ffe05e26 Mon Sep 17 00:00:00 2001 From: toutane Date: Mon, 27 Jul 2026 16:30:25 +0200 Subject: [PATCH 04/10] fix(datafusion): address eager scan review feedback Second review round for apache/iceberg-rust#2671. --- crates/iceberg/public-api.txt | 1 + crates/iceberg/src/scan/mod.rs | 9 +- crates/integrations/datafusion/public-api.txt | 2 +- .../datafusion/src/physical_plan/scan.rs | 17 +- .../tests/integration_datafusion_test.rs | 198 ++++++++++++++++-- 5 files changed, 199 insertions(+), 28 deletions(-) diff --git a/crates/iceberg/public-api.txt b/crates/iceberg/public-api.txt index 610675c83d..713cde84ad 100644 --- a/crates/iceberg/public-api.txt +++ b/crates/iceberg/public-api.txt @@ -1326,6 +1326,7 @@ pub fn iceberg::scan::ScanResult::metrics(&self) -> &iceberg::scan::ScanMetrics pub fn iceberg::scan::ScanResult::stream(self) -> iceberg::scan::ArrowRecordBatchStream pub struct iceberg::scan::TableScan impl iceberg::scan::TableScan +pub fn iceberg::scan::TableScan::arrow_reader_builder(&self) -> iceberg::arrow::ArrowReaderBuilder pub fn iceberg::scan::TableScan::column_names(&self) -> core::option::Option<&[alloc::string::String]> pub async fn iceberg::scan::TableScan::plan_files(&self) -> iceberg::Result pub fn iceberg::scan::TableScan::snapshot(&self) -> core::option::Option<&iceberg::spec::SnapshotRef> diff --git a/crates/iceberg/src/scan/mod.rs b/crates/iceberg/src/scan/mod.rs index 19f46c2cb1..8a0e1ce7dc 100644 --- a/crates/iceberg/src/scan/mod.rs +++ b/crates/iceberg/src/scan/mod.rs @@ -480,8 +480,8 @@ impl TableScan { Ok(file_scan_task_rx.boxed()) } - /// Returns an [`ArrowRecordBatchStream`]. - pub async fn to_arrow(&self) -> Result { + /// Returns an [`ArrowReaderBuilder`] configured for this table scan. + pub fn arrow_reader_builder(&self) -> ArrowReaderBuilder { let mut arrow_reader_builder = ArrowReaderBuilder::new(self.file_io.clone(), self.runtime.clone()) .with_data_file_concurrency_limit(self.concurrency_limit_data_files) @@ -493,6 +493,11 @@ impl TableScan { } arrow_reader_builder + } + + /// Returns an [`ArrowRecordBatchStream`]. + pub async fn to_arrow(&self) -> Result { + self.arrow_reader_builder() .build() .read(self.plan_files().await?) .map(|result| result.stream()) diff --git a/crates/integrations/datafusion/public-api.txt b/crates/integrations/datafusion/public-api.txt index 821525dea7..58071e109a 100644 --- a/crates/integrations/datafusion/public-api.txt +++ b/crates/integrations/datafusion/public-api.txt @@ -28,7 +28,7 @@ pub fn iceberg_datafusion::physical_plan::IcebergTableScan::children(&self) -> a pub fn iceberg_datafusion::physical_plan::IcebergTableScan::execute(&self, partition: usize, _context: alloc::sync::Arc) -> datafusion_common::error::Result pub fn iceberg_datafusion::physical_plan::IcebergTableScan::name(&self) -> &str pub fn iceberg_datafusion::physical_plan::IcebergTableScan::properties(&self) -> &alloc::sync::Arc -pub fn iceberg_datafusion::physical_plan::IcebergTableScan::with_new_children(self: alloc::sync::Arc, _children: alloc::vec::Vec>) -> datafusion_common::error::Result> +pub fn iceberg_datafusion::physical_plan::IcebergTableScan::with_new_children(self: alloc::sync::Arc, children: alloc::vec::Vec>) -> datafusion_common::error::Result> pub fn iceberg_datafusion::physical_plan::convert_filters_to_predicate(filters: &[datafusion_expr::expr::Expr]) -> core::option::Option pub fn iceberg_datafusion::physical_plan::project_with_partition(input: alloc::sync::Arc, table: &iceberg::table::Table) -> datafusion_common::error::Result> pub mod iceberg_datafusion::table diff --git a/crates/integrations/datafusion/src/physical_plan/scan.rs b/crates/integrations/datafusion/src/physical_plan/scan.rs index d9433f6cdf..54bcef85b9 100644 --- a/crates/integrations/datafusion/src/physical_plan/scan.rs +++ b/crates/integrations/datafusion/src/physical_plan/scan.rs @@ -120,8 +120,14 @@ impl ExecutionPlan for IcebergTableScan { fn with_new_children( self: Arc, - _children: Vec>, + children: Vec>, ) -> DFResult> { + if !children.is_empty() { + return Err(datafusion::common::DataFusionError::Internal( + "IcebergTableScan is a leaf node and cannot have children".to_string(), + )); + } + Ok(self) } @@ -148,9 +154,8 @@ impl ExecutionPlan for IcebergTableScan { let tasks: FileScanTaskStream = Box::pin(futures::stream::iter( (0..file_task_group.len()).map(move |idx| Ok(file_task_group[idx].clone())), )); - let stream = self - .table - .reader_builder() + let stream = build_table_scan(&self.table, &self.scan_config)? + .arrow_reader_builder() // Eager planning lets DataFusion drive scan concurrency via output // partitions. Match DataFusion's FileStream model, where each // output partition owns one ScanState; keep one data file in @@ -183,7 +188,9 @@ impl ExecutionPlan for IcebergTableScan { } }; - // Apply limit if specified + // Apply a scan-partition bound if specified. In eager planning this is only + // a per-partition bound; DataFusion's GlobalLimitExec remains responsible + // for enforcing the final global limit. let limited_stream: Pin> + Send>> = if let Some(limit) = self.limit { let mut remaining = limit; diff --git a/crates/integrations/datafusion/tests/integration_datafusion_test.rs b/crates/integrations/datafusion/tests/integration_datafusion_test.rs index 2b16fcd463..3e9dbe829b 100644 --- a/crates/integrations/datafusion/tests/integration_datafusion_test.rs +++ b/crates/integrations/datafusion/tests/integration_datafusion_test.rs @@ -18,6 +18,7 @@ //! Integration tests for Iceberg Datafusion with Hive Metastore. use std::collections::HashMap; +use std::fs::File; use std::sync::Arc; use std::vec; @@ -27,6 +28,7 @@ use datafusion::arrow::record_batch::RecordBatch; use datafusion::datasource::MemTable; use datafusion::execution::context::SessionContext; use datafusion::parquet::arrow::PARQUET_FIELD_ID_META_KEY; +use datafusion::physical_plan::ExecutionPlan; use datafusion::prelude::SessionConfig; use expect_test::expect; use iceberg::io::LocalFsStorageFactory; @@ -35,11 +37,21 @@ use iceberg::spec::{ NestedField, PrimitiveType, Schema, StructType, Transform, Type, UnboundPartitionSpec, }; use iceberg::test_utils::check_record_batches; +use iceberg::transaction::{ApplyTransactionAction, Transaction}; +use iceberg::writer::base_writer::data_file_writer::DataFileWriterBuilder; +use iceberg::writer::file_writer::ParquetWriterBuilder; +use iceberg::writer::file_writer::location_generator::{ + DefaultFileNameGenerator, DefaultLocationGenerator, +}; +use iceberg::writer::file_writer::rolling_writer::RollingFileWriterBuilder; +use iceberg::writer::{IcebergWriter, IcebergWriterBuilder}; use iceberg::{ Catalog, CatalogBuilder, MemoryCatalog, NamespaceIdent, Result, TableCreation, TableIdent, }; use iceberg_datafusion::physical_plan::IcebergTableScan; use iceberg_datafusion::{IcebergCatalogProvider, IcebergDataFusionConfig}; +use parquet::file::properties::WriterProperties; +use parquet::file::reader::{FileReader, SerializedFileReader}; use tempfile::TempDir; fn temp_path() -> String { @@ -158,6 +170,106 @@ async fn get_multi_file_table_context( Ok((iceberg_catalog, namespace, table_name.to_string())) } +async fn get_multi_row_group_table_context( + namespace_name: &str, + table_name: &str, +) -> Result<(Arc, NamespaceIdent, String)> { + let iceberg_catalog = Arc::new(get_iceberg_catalog().await); + let namespace = NamespaceIdent::new(namespace_name.to_string()); + set_test_namespace(&iceberg_catalog, &namespace).await?; + + let creation = get_table_creation(temp_path(), table_name, None)?; + let table = iceberg_catalog.create_table(&namespace, creation).await?; + let arrow_schema: Arc = Arc::new( + table + .metadata() + .current_schema() + .as_ref() + .try_into() + .unwrap(), + ); + + let file_rows = [ + vec![ + (1, "row-1"), + (2, "row-2"), + (100, "row-100"), + (101, "row-101"), + ], + vec![ + (50, "row-50"), + (150, "row-150"), + (151, "row-151"), + (152, "row-152"), + ], + vec![ + (99, "row-99"), + (1000, "row-1000"), + (1001, "row-1001"), + (1002, "row-1002"), + ], + ]; + + let location_generator = DefaultLocationGenerator::new(table.metadata()).unwrap(); + let writer_properties = WriterProperties::builder() + .set_max_row_group_row_count(Some(2)) + .build(); + let mut data_files = Vec::with_capacity(file_rows.len()); + + for (file_idx, rows) in file_rows.into_iter().enumerate() { + let parquet_writer_builder = ParquetWriterBuilder::new( + writer_properties.clone(), + table.metadata().current_schema().clone(), + ); + let file_name_generator = DefaultFileNameGenerator::new( + format!("multi-row-group-{file_idx}"), + None, + iceberg::spec::DataFileFormat::Parquet, + ); + let rolling_file_writer_builder = RollingFileWriterBuilder::new_with_default_file_size( + parquet_writer_builder, + table.file_io().clone(), + location_generator.clone(), + file_name_generator, + ); + let data_file_writer_builder = DataFileWriterBuilder::new(rolling_file_writer_builder); + let mut data_file_writer = data_file_writer_builder.build(None).await?; + let batch = RecordBatch::try_new(arrow_schema.clone(), vec![ + Arc::new(Int32Array::from( + rows.iter().map(|(foo1, _)| *foo1).collect::>(), + )) as ArrayRef, + Arc::new(StringArray::from( + rows.iter().map(|(_, foo2)| *foo2).collect::>(), + )) as ArrayRef, + ])?; + + data_file_writer.write(batch).await?; + let file_data_files = data_file_writer.close().await?; + assert_eq!(file_data_files.len(), 1); + data_files.extend(file_data_files); + } + + for data_file in &data_files { + let file_path = data_file + .file_path() + .strip_prefix("file://") + .unwrap_or(data_file.file_path()); + let reader = SerializedFileReader::new(File::open(file_path)?)?; + assert!( + reader.metadata().num_row_groups() > 1, + "expected multiple row groups for {}", + data_file.file_path() + ); + } + + let tx = Transaction::new(&table); + let action = tx.fast_append().add_data_files(data_files); + let tx = action.apply(tx)?; + tx.commit(iceberg_catalog.as_ref()).await?; + + Ok((iceberg_catalog, namespace, table_name.to_string())) +} + async fn get_read_context( catalog: Arc, target_partitions: usize, @@ -270,29 +382,29 @@ async fn test_set_enable_eager_scan_planning() -> Result<()> { async fn test_multi_partition_scan_matches_single_partition_results() -> Result<()> { let data_file_count = 3; let target_partitions = data_file_count + 1; - let (iceberg_catalog, namespace, table_name) = get_multi_file_table_context( - "test_multi_partition_scan_results", - "my_table", - data_file_count, - ) - .await?; + let (iceberg_catalog, namespace, table_name) = + get_multi_row_group_table_context("test_multi_partition_scan_results", "my_table").await?; let namespace_name = &namespace[0]; - let single_partition_ctx = get_read_context(iceberg_catalog.clone(), 1, Some(true)).await?; - let multi_partition_ctx = + let lazy_ctx = get_read_context(iceberg_catalog.clone(), target_partitions, None).await?; + let eager_single_partition_ctx = + get_read_context(iceberg_catalog.clone(), 1, Some(true)).await?; + let eager_multi_partition_ctx = get_read_context(iceberg_catalog, target_partitions, Some(true)).await?; - let query = - format!("SELECT foo1, foo2 FROM catalog.{namespace_name}.{table_name} ORDER BY foo1"); + let query = format!( + "SELECT foo1, foo2 FROM catalog.{namespace_name}.{table_name} WHERE foo1 >= 100 ORDER BY foo1" + ); - let single_partition_batches = single_partition_ctx + let lazy_batches = lazy_ctx.sql(&query).await.unwrap().collect().await.unwrap(); + let eager_single_partition_batches = eager_single_partition_ctx .sql(&query) .await .unwrap() .collect() .await .unwrap(); - let multi_partition_batches = multi_partition_ctx + let eager_multi_partition_batches = eager_multi_partition_ctx .sql(&query) .await .unwrap() @@ -309,23 +421,69 @@ async fn test_multi_partition_scan_matches_single_partition_results() -> Result< expect![[r#" foo1: PrimitiveArray [ - 1, - 2, - 3, + 100, + 101, + 150, + 151, + 152, + 1000, + 1001, + 1002, ], foo2: StringArray [ - "row-1", - "row-2", - "row-3", + "row-100", + "row-101", + "row-150", + "row-151", + "row-152", + "row-1000", + "row-1001", + "row-1002", ]"#]], &[], None, ); }; - assert_expected_batches(single_partition_batches); - assert_expected_batches(multi_partition_batches); + assert_expected_batches(lazy_batches); + assert_expected_batches(eager_single_partition_batches); + assert_expected_batches(eager_multi_partition_batches); + + Ok(()) +} + +#[tokio::test] +async fn test_iceberg_table_scan_rejects_non_empty_children() -> Result<()> { + use datafusion::common::DataFusionError; + use datafusion::physical_plan::empty::EmptyExec; + + let (iceberg_catalog, namespace, table_name) = + get_multi_file_table_context("test_iceberg_table_scan_rejects_children", "my_table", 1) + .await?; + let ctx = get_read_context(iceberg_catalog, 1, Some(true)).await?; + let provider = ctx.catalog("catalog").unwrap(); + let schema = provider.schema(&namespace[0]).unwrap(); + let table = schema.table(&table_name).await.unwrap().unwrap(); + let state = ctx.state(); + let scan_plan = table.scan(&state, None, &[], None).await.unwrap(); + scan_plan + .downcast_ref::() + .expect("Expected IcebergTableScan"); + + let child = Arc::new(EmptyExec::new(scan_plan.schema())) as Arc; + let error = scan_plan + .with_new_children(vec![child]) + .expect_err("IcebergTableScan should reject children"); + + assert!( + matches!( + error, + DataFusionError::Internal(ref message) + if message == "IcebergTableScan is a leaf node and cannot have children" + ), + "unexpected error: {error}" + ); Ok(()) } From 5ae5544b7f12c1399bad9a855085349ca8c99b90 Mon Sep 17 00:00:00 2001 From: toutane Date: Wed, 29 Jul 2026 17:39:11 +0200 Subject: [PATCH 05/10] test(datafusion): cover eager scan global limits --- crates/integrations/datafusion/src/config.rs | 3 + .../datafusion/src/physical_plan/scan.rs | 4 +- .../tests/integration_datafusion_test.rs | 75 +++++++++++++++++++ 3 files changed, 80 insertions(+), 2 deletions(-) diff --git a/crates/integrations/datafusion/src/config.rs b/crates/integrations/datafusion/src/config.rs index 39d5f99119..1154c3d266 100644 --- a/crates/integrations/datafusion/src/config.rs +++ b/crates/integrations/datafusion/src/config.rs @@ -20,6 +20,9 @@ use datafusion::common::extensions_options; extensions_options! { /// Configuration options for Iceberg's DataFusion integration. + /// + /// Register this extension with [`SessionConfig::with_option_extension`](datafusion::prelude::SessionConfig::with_option_extension) + /// before creating a session to make `SET iceberg.