gabotechs commented on code in PR #67: URL: https://github.com/apache/datafusion-iceberg/pull/67#discussion_r4561293658
########## README.md: ########## @@ -93,3 +93,31 @@ cargo test --workspace --locked This project is licensed under the Apache License, Version 2.0. See [`LICENSE`](LICENSE) and [`NOTICE`](NOTICE). + +### Eager file planning Review Comment: I'd try to keep the entry in the README.md as lean as possible, you can expect people reading this file to be 99% lurkers that don't know anything about the project and might not even know what "planning" means. I think we are fine without an entry in the README.md here as long as it's properly documented in the doc comments. ########## crates/datafusion/src/read/eager.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::arrow::datatypes::SchemaRef; +use datafusion::common::tree_node::TreeNodeRecursion; +use datafusion::common::{Statistics, exec_err}; +use datafusion::datasource::source::DataSource; +use datafusion::error::Result; +use datafusion::execution::{SendableRecordBatchStream, TaskContext}; +use datafusion::physical_expr::projection::ProjectionExprs; +use datafusion::physical_expr::{EquivalenceProperties, Partitioning, PhysicalExpr}; +use datafusion::physical_plan::DisplayFormatType; +use datafusion::physical_plan::metrics::ExecutionPlanMetricsSet; +use futures::StreamExt; +use iceberg::arrow::ArrowReader; +use iceberg::scan::FileScanTask; + +use super::file_planning::PlannedFileScan; +use super::reader::ScanReader; + +/// An Iceberg scan whose complete file assignment is available before execution. +/// Reads already-selected file groups without retaining a table or discovery request. +/// Consumers can inspect and regroup these tasks, or reconstruct a source on another +/// process using a locally configured reader. +#[derive(Clone)] +pub struct EagerIcebergDataSource { + schema: SchemaRef, + fetch: Option<usize>, + reader: ScanReader, + files: Arc<PlannedFileScan>, + partitions: Arc<[Vec<usize>]>, +} +impl std::fmt::Debug for EagerIcebergDataSource { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("EagerIcebergDataSource") + .field("schema", &self.schema) + .field("fetch", &self.fetch) + .field("partitions", &self.partitions) + .finish_non_exhaustive() + } +} +impl EagerIcebergDataSource { + pub(crate) fn from_planned_files( + schema: SchemaRef, + files: PlannedFileScan, + groups: Vec<Vec<usize>>, + reader: ArrowReader, + fetch: Option<usize>, + ) -> Self { + Self { + schema, + fetch, + files: Arc::new(files), + partitions: groups.into(), + reader: ScanReader::new(reader), + } + } + + /// Construct a scan from already selected files grouped by local output partition. + /// Preserves group order without discovering, pruning, or assigning files. + /// The caller supplies a process-local reader configured for these files. Reader + /// clones share its delete cache across partitions; regrouping preserves its + /// settings, including file concurrency. The schema must match the tasks' projection. + /// Descriptors are charged to `context`'s memory pool; statistics are unknown. + /// + /// # Errors + /// Returns an error if there are no partition groups or memory reservation fails. + /// Empty groups are valid, including for an empty scan. + pub fn from_file_groups( + schema: SchemaRef, + groups: Vec<Vec<FileScanTask>>, + reader: ArrowReader, + fetch: Option<usize>, + context: &TaskContext, + ) -> Result<Self> { + if groups.is_empty() { + return exec_err!("Iceberg scan requires at least one file group"); + } + let mut tasks = Vec::new(); + let mut partitions = Vec::with_capacity(groups.len()); + for group in groups { + let start = tasks.len(); + tasks.extend(group); + partitions.push((start..tasks.len()).collect()); + } + Ok(Self { + schema, + fetch, + reader: ScanReader::new(reader), + files: Arc::new(PlannedFileScan::from_tasks(tasks, context)?), + partitions: partitions.into(), + }) + } + + /// Output schema, including the scan's projection. + pub fn schema(&self) -> &SchemaRef { + &self.schema + } + + /// Borrow this scan's files in retained-file order, independent of partition groups. + /// Positions in this iterator are the inputs to [`Self::with_file_groups`]. + /// This sorts file indices, without cloning descriptors or performing I/O. + /// File iteration order does not promise row ordering. + pub fn planned_files(&self) -> impl ExactSizeIterator<Item = &FileScanTask> { + self.file_indices() + .into_iter() + .map(|index| &self.files.tasks()[index]) + } + + fn file_indices(&self) -> Vec<usize> { + let mut indices: Vec<_> = self.partitions.iter().flatten().copied().collect(); + indices.sort_unstable(); + indices + } + + /// Borrow the files that execution will read for one local partition. + /// + /// # Errors + /// Returns an error for an out-of-bounds partition. + pub fn planned_partition_tasks( Review Comment: I'd be consistent with the naming of this public method: `planned_files` and `planned_partition_files`, or `planned_tasks` and `planned_partition_tasks`. ########## crates/datafusion/src/read/eager.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::arrow::datatypes::SchemaRef; +use datafusion::common::tree_node::TreeNodeRecursion; +use datafusion::common::{Statistics, exec_err}; +use datafusion::datasource::source::DataSource; +use datafusion::error::Result; +use datafusion::execution::{SendableRecordBatchStream, TaskContext}; +use datafusion::physical_expr::projection::ProjectionExprs; +use datafusion::physical_expr::{EquivalenceProperties, Partitioning, PhysicalExpr}; +use datafusion::physical_plan::DisplayFormatType; +use datafusion::physical_plan::metrics::ExecutionPlanMetricsSet; +use futures::StreamExt; +use iceberg::arrow::ArrowReader; +use iceberg::scan::FileScanTask; + +use super::file_planning::PlannedFileScan; +use super::reader::ScanReader; + +/// An Iceberg scan whose complete file assignment is available before execution. +/// Reads already-selected file groups without retaining a table or discovery request. +/// Consumers can inspect and regroup these tasks, or reconstruct a source on another +/// process using a locally configured reader. +#[derive(Clone)] +pub struct EagerIcebergDataSource { + schema: SchemaRef, + fetch: Option<usize>, + reader: ScanReader, + files: Arc<PlannedFileScan>, + partitions: Arc<[Vec<usize>]>, +} +impl std::fmt::Debug for EagerIcebergDataSource { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("EagerIcebergDataSource") + .field("schema", &self.schema) + .field("fetch", &self.fetch) + .field("partitions", &self.partitions) + .finish_non_exhaustive() + } +} +impl EagerIcebergDataSource { + pub(crate) fn from_planned_files( + schema: SchemaRef, + files: PlannedFileScan, + groups: Vec<Vec<usize>>, + reader: ArrowReader, + fetch: Option<usize>, + ) -> Self { + Self { + schema, + fetch, + files: Arc::new(files), + partitions: groups.into(), + reader: ScanReader::new(reader), + } + } + + /// Construct a scan from already selected files grouped by local output partition. + /// Preserves group order without discovering, pruning, or assigning files. + /// The caller supplies a process-local reader configured for these files. Reader + /// clones share its delete cache across partitions; regrouping preserves its + /// settings, including file concurrency. The schema must match the tasks' projection. + /// Descriptors are charged to `context`'s memory pool; statistics are unknown. + /// + /// # Errors + /// Returns an error if there are no partition groups or memory reservation fails. + /// Empty groups are valid, including for an empty scan. + pub fn from_file_groups( Review Comment: We should not expose methods other than the existing builder for constructing the data sources, as these are subject to breaking API changes. I see this method only exists for the tests, so I'd: - Either move it inside the `#[cfg(test)]` block with the tests that use it. - Just remove it, and test the `EagerIcebergDataSource` through the public API I have a strong preference for the second option, as otherwise the test is subject to change as implementations detail change, getting in the way rather than helping. But if you don't see it being possible, that's fine. ########## crates/datafusion/src/read/eager.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::arrow::datatypes::SchemaRef; +use datafusion::common::tree_node::TreeNodeRecursion; +use datafusion::common::{Statistics, exec_err}; +use datafusion::datasource::source::DataSource; +use datafusion::error::Result; +use datafusion::execution::{SendableRecordBatchStream, TaskContext}; +use datafusion::physical_expr::projection::ProjectionExprs; +use datafusion::physical_expr::{EquivalenceProperties, Partitioning, PhysicalExpr}; +use datafusion::physical_plan::DisplayFormatType; +use datafusion::physical_plan::metrics::ExecutionPlanMetricsSet; +use futures::StreamExt; +use iceberg::arrow::ArrowReader; +use iceberg::scan::FileScanTask; + +use super::file_planning::PlannedFileScan; +use super::reader::ScanReader; + +/// An Iceberg scan whose complete file assignment is available before execution. +/// Reads already-selected file groups without retaining a table or discovery request. +/// Consumers can inspect and regroup these tasks, or reconstruct a source on another +/// process using a locally configured reader. +#[derive(Clone)] +pub struct EagerIcebergDataSource { + schema: SchemaRef, + fetch: Option<usize>, + reader: ScanReader, + files: Arc<PlannedFileScan>, + partitions: Arc<[Vec<usize>]>, +} +impl std::fmt::Debug for EagerIcebergDataSource { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + f.debug_struct("EagerIcebergDataSource") + .field("schema", &self.schema) + .field("fetch", &self.fetch) + .field("partitions", &self.partitions) + .finish_non_exhaustive() + } +} +impl EagerIcebergDataSource { + pub(crate) fn from_planned_files( + schema: SchemaRef, + files: PlannedFileScan, + groups: Vec<Vec<usize>>, + reader: ArrowReader, + fetch: Option<usize>, + ) -> Self { + Self { + schema, + fetch, + files: Arc::new(files), + partitions: groups.into(), + reader: ScanReader::new(reader), + } + } + + /// Construct a scan from already selected files grouped by local output partition. + /// Preserves group order without discovering, pruning, or assigning files. + /// The caller supplies a process-local reader configured for these files. Reader + /// clones share its delete cache across partitions; regrouping preserves its + /// settings, including file concurrency. The schema must match the tasks' projection. + /// Descriptors are charged to `context`'s memory pool; statistics are unknown. + /// + /// # Errors + /// Returns an error if there are no partition groups or memory reservation fails. + /// Empty groups are valid, including for an empty scan. + pub fn from_file_groups( + schema: SchemaRef, + groups: Vec<Vec<FileScanTask>>, + reader: ArrowReader, + fetch: Option<usize>, + context: &TaskContext, + ) -> Result<Self> { + if groups.is_empty() { + return exec_err!("Iceberg scan requires at least one file group"); + } + let mut tasks = Vec::new(); + let mut partitions = Vec::with_capacity(groups.len()); + for group in groups { + let start = tasks.len(); + tasks.extend(group); + partitions.push((start..tasks.len()).collect()); + } + Ok(Self { + schema, + fetch, + reader: ScanReader::new(reader), + files: Arc::new(PlannedFileScan::from_tasks(tasks, context)?), + partitions: partitions.into(), + }) + } + + /// Output schema, including the scan's projection. + pub fn schema(&self) -> &SchemaRef { + &self.schema + } + + /// Borrow this scan's files in retained-file order, independent of partition groups. + /// Positions in this iterator are the inputs to [`Self::with_file_groups`]. + /// This sorts file indices, without cloning descriptors or performing I/O. + /// File iteration order does not promise row ordering. + pub fn planned_files(&self) -> impl ExactSizeIterator<Item = &FileScanTask> { + self.file_indices() + .into_iter() + .map(|index| &self.files.tasks()[index]) + } + + fn file_indices(&self) -> Vec<usize> { + let mut indices: Vec<_> = self.partitions.iter().flatten().copied().collect(); + indices.sort_unstable(); + indices + } + + /// Borrow the files that execution will read for one local partition. + /// + /// # Errors + /// Returns an error for an out-of-bounds partition. + pub fn planned_partition_tasks( + &self, + partition: usize, + ) -> Result<impl ExactSizeIterator<Item = &FileScanTask>> { + let Some(indices) = self.partitions.get(partition) else { + return exec_err!("Invalid Iceberg scan partition {partition}"); + }; + Ok(indices.iter().map(|&index| &self.files.tasks()[index])) + } + + /// Derive a scan using positions in this source's [`Self::planned_files`] iterator. + /// Each group is one local output partition; supplied order is preserved. + /// Files may be omitted, but each position may appear at most once. + /// + /// The result shares descriptors and their memory reservation with this source. + /// It has fresh metrics and unknown statistics because its file set may differ. + /// No file discovery, assignment policy, or file I/O is performed. + /// + /// # Errors + /// Returns an error for zero groups, out-of-bounds positions, or repeated positions. + pub fn with_file_groups(&self, mut groups: Vec<Vec<usize>>) -> Result<Self> { Review Comment: This is a weird method to have in the public API. In the same way we have getter methods for the planned files, why not just have setter methods consistent in API with the getters? Something like: ```rust pub fn with_planned_files(&self, mut groups: Vec<Vec<FileScanTask>>) -> Result<Self> { ``` -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected] --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
