feat: plan Iceberg file tasks before execution using public APIs

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.

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";
+}

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

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<Arc<FileScanTask>>,
+    _reservation: MemoryReservation,
+}
+
+impl PlannedFileScan {
+    /// Tasks in Iceberg's planning order. Each must be assigned exactly once by a consumer.
+    pub fn tasks(&self) -> &[Arc<FileScanTask>] {
+        &self.tasks
+    }
+
+    pub(crate) async fn collect(
+        scan: TableScan,
+        context: &TaskContext,
+        max_files: usize,
+    ) -> Result<Self> {
+        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::<FileScanTask>()
+        + size_of::<Arc<FileScanTask>>()
+        + 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
+}

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};