feat(datafusion): Add opt-in eager file scan planning with output partitioning

crates/iceberg/public-api.txt

@@ -1331,6 +1331,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<iceberg::scan::FileScanTaskStream>
 pub fn iceberg::scan::TableScan::snapshot(&self) -> core::option::Option<&iceberg::spec::SnapshotRef>

crates/iceberg/src/scan/mod.rs

@@ -498,8 +498,8 @@ impl TableScan {
         Ok(file_scan_task_rx.boxed())
     }
 
-    /// Returns an [`ArrowRecordBatchStream`].
-    pub async fn to_arrow(&self) -> Result<ArrowRecordBatchStream> {
+    /// 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)
@@ -511,6 +511,11 @@ impl TableScan {
         }
 
         arrow_reader_builder
+    }
+
+    /// Returns an [`ArrowRecordBatchStream`].
+    pub async fn to_arrow(&self) -> Result<ArrowRecordBatchStream> {
+        self.arrow_reader_builder()
             .build()
             .read(self.plan_files().await?)
             .map(|result| result.stream())

crates/integrations/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::config::ConfigExtension;
+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.<option> = <value>` available.
+    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";
+}

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::*;
 

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;
 

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::FileScanTaskStream;
 use iceberg::table::Table;
 
-use super::expr_to_predicate::convert_filters_to_predicate;
+use super::scan_planning::{EagerScanPlan, IcebergScanConfig, build_table_scan};
 use crate::to_datafusion_error;
 
 /// Manages the scanning process of an Iceberg [`Table`], encapsulating the
@@ -41,44 +41,37 @@ use crate::to_datafusion_error;
 pub struct IcebergTableScan {
     /// A table in the catalog.
     table: Table,
-    /// Snapshot of the table to scan.
-    snapshot_id: Option<i64>,
+    /// 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<PlanProperties>,
-    /// Projection column names, None means all columns
-    projection: Option<Vec<String>>,
-    /// Filters to apply to the table scan
-    predicates: Option<Predicate>,
     /// Optional limit on the number of rows to return
     limit: Option<usize>,
+    /// Pre-planned file scan tasks and the scan that planned them. `None` keeps
+    /// planning lazy.
+    eager_plan: Option<EagerScanPlan>,
 }
 
 impl IcebergTableScan {
-    /// Creates a new [`IcebergTableScan`] object.
     pub(crate) fn new(
         table: Table,
-        snapshot_id: Option<i64>,
-        schema: ArrowSchemaRef,
-        projection: Option<&Vec<usize>>,
-        filters: &[Expr],
+        scan_config: IcebergScanConfig,
         limit: Option<usize>,
+        eager_plan: Option<EagerScanPlan>,
     ) -> 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 = eager_plan
+            .as_ref()
+            .map_or(1, EagerScanPlan::partition_count);
+        let plan_properties =
+            IcebergTableScan::compute_properties(scan_config.output_schema(), partition_count);
 
         Self {
             table,
-            snapshot_id,
+            scan_config,
             plan_properties,
-            projection,
-            predicates,
             limit,
+            eager_plan,
         }
     }
 
@@ -87,29 +80,26 @@ impl IcebergTableScan {
     }
 
     pub fn snapshot_id(&self) -> Option<i64> {
-        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<usize> {
         self.limit
     }
 
     /// Computes [`PlanProperties`] used in query optimization.
-    fn compute_properties(schema: ArrowSchemaRef) -> Arc<PlanProperties> {
-        // 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<PlanProperties> {
         Arc::new(PlanProperties::new(
             EquivalenceProperties::new(schema),
-            Partitioning::UnknownPartitioning(1),
+            Partitioning::UnknownPartitioning(partition_count),
             EmissionType::Incremental,
             Boundedness::Bounded,
         ))
@@ -127,8 +117,14 @@ impl ExecutionPlan for IcebergTableScan {
 
     fn with_new_children(
         self: Arc<Self>,
-        _children: Vec<Arc<dyn ExecutionPlan>>,
+        children: Vec<Arc<dyn ExecutionPlan>>,
     ) -> DFResult<Arc<dyn ExecutionPlan>> {
+        if !children.is_empty() {
+            return Err(datafusion::common::DataFusionError::Internal(
+                "IcebergTableScan is a leaf node and cannot have children".to_string(),
+            ));
+        }
+
         Ok(self)
     }
 
@@ -138,18 +134,61 @@ impl ExecutionPlan for IcebergTableScan {
 
     fn execute(
         &self,
-        _partition: usize,
+        partition: usize,
         _context: Arc<TaskContext>,
     ) -> DFResult<SendableRecordBatchStream> {
-        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();
-
-        // Apply limit if specified
+        let stream: Pin<Box<dyn Stream<Item = DFResult<RecordBatch>> + Send>> = match &self
+            .eager_plan
+        {
+            Some(eager_plan) => {
+                let Some(file_task_group) = eager_plan.task_group(partition) else {
+                    return Err(datafusion::common::DataFusionError::Internal(format!(
+                        "IcebergTableScan partition {partition} does not exist; scan has {} partitions",
+                        eager_plan.partition_count()
+                    )));
+                };
+
+                let tasks: FileScanTaskStream = Box::pin(futures::stream::iter(
+                    (0..file_task_group.len()).map(move |idx| Ok(file_task_group[idx].clone())),
+                ));
+                let stream = eager_plan
+                    .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
+                    // 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.
+                    // Tracked in https://github.com/apache/iceberg-rust/issues/2964.
+                    .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 a scan-partition bound if specified. In eager planning this is only
+        // a per-partition bound. The optimized DataFusion plan enforces the final
+        // global bound above the scan, currently through CoalescePartitionsExec::fetch.
         let limited_stream: Pin<Box<dyn Stream<Item = DFResult<RecordBatch>> + Send>> =
             if let Some(limit) = self.limit {
                 let mut remaining = limit;
@@ -185,60 +224,21 @@ 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(eager_plan) = &self.eager_plan {
+            write!(
+                f,
+                " task_groups:[{}] tasks:[{}]",
+                eager_plan.partition_count(),
+                eager_plan.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<i64>,
-    column_names: Option<Vec<String>>,
-    predicates: Option<Predicate>,
-) -> DFResult<Pin<Box<dyn Stream<Item = DFResult<RecordBatch>> + 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<usize>>,
-) -> Option<Vec<String>> {
-    projection.map(|v| {
-        v.iter()
-            .map(|p| schema.field(*p).name().clone())
-            .collect::<Vec<String>>()
-    })
-}