toutane commented on code in PR #2298: URL: https://github.com/apache/iceberg-rust/pull/2298#discussion_r3111525487
########## crates/integrations/datafusion/src/table/partitioned.rs: ########## @@ -0,0 +1,315 @@ +// 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::any::Any; +use std::sync::Arc; + +use async_trait::async_trait; +use datafusion::arrow::datatypes::SchemaRef as ArrowSchemaRef; +use datafusion::catalog::Session; +use datafusion::common::DataFusionError; +use datafusion::datasource::{TableProvider, TableType}; +use datafusion::error::Result as DFResult; +use datafusion::logical_expr::{Expr, TableProviderFilterPushDown}; +use datafusion::physical_plan::ExecutionPlan; +use futures::TryStreamExt; +use iceberg::arrow::schema_to_arrow_schema; +use iceberg::{Catalog, Error, ErrorKind, NamespaceIdent, Result, TableIdent}; + +use crate::error::to_datafusion_error; +use crate::physical_plan::expr_to_predicate::convert_filters_to_predicate; +use crate::physical_plan::partitioned_scan::IcebergPartitionedScan; + +/// Catalog-backed table provider that scans each data file in a separate DataFusion partition. +/// +/// This provider reloads table metadata from the catalog on every [`scan`][Self::scan] call +/// to guarantee freshness, then issues one DataFusion partition per data file so that +/// DataFusion's scheduler can execute file reads in parallel. +/// +/// Write operations are not supported. Use [`IcebergTableProvider`] for write access. +/// +/// For consistent read-only access to a fixed snapshot without per-scan catalog overhead, +/// use [`IcebergStaticTableProvider`] instead. +#[derive(Debug, Clone)] +pub struct IcebergPartitionedTableProvider { + catalog: Arc<dyn Catalog>, + table_ident: TableIdent, + schema: ArrowSchemaRef, +} + +impl IcebergPartitionedTableProvider { + pub async fn try_new( + catalog: Arc<dyn Catalog>, + namespace: NamespaceIdent, + name: impl Into<String>, + ) -> Result<Self> { + let table_ident = TableIdent::new(namespace, name.into()); + // First load: used only to snapshot the Arrow schema for DataFusion planning. + // A second load_table is issued at scan time to guarantee the freshest snapshot. + let table = catalog.load_table(&table_ident).await?; + let schema = Arc::new(schema_to_arrow_schema(table.metadata().current_schema())?); + Ok(Self { + catalog, + table_ident, + schema, + }) + } +} + +#[async_trait] +impl TableProvider for IcebergPartitionedTableProvider { + fn as_any(&self) -> &dyn Any { + self + } + + fn schema(&self) -> ArrowSchemaRef { + self.schema.clone() + } + + fn table_type(&self) -> TableType { + TableType::Base + } + + async fn scan( + &self, + _state: &dyn Session, + projection: Option<&Vec<usize>>, + filters: &[Expr], + _limit: Option<usize>, + ) -> DFResult<Arc<dyn ExecutionPlan>> { + // Per-partition row limits are not yet implemented for IcebergPartitionedScan. + // DataFusion will apply a GlobalLimitExec on top of this node when needed. + + // Second load: fetch the latest snapshot so scans always reflect current table state. + let table = self + .catalog + .load_table(&self.table_ident) + .await + .map_err(to_datafusion_error)?; + + // Projection indices are resolved against self.schema (captured at try_new time), + // same as IcebergTableProvider / IcebergTableScan. + let col_names = projection.map(|indices| { + indices + .iter() + .map(|&i| self.schema.field(i).name().clone()) + .collect::<Vec<_>>() + }); + + let predicate = convert_filters_to_predicate(filters); + + let mut builder = table.scan(); + builder = match col_names { + Some(names) => builder.select(names), + None => builder.select_all(), + }; + if let Some(pred) = predicate { + builder = builder.with_filter(pred); + } + + let tasks = builder + .build() + .map_err(to_datafusion_error)? + .plan_files() + .await + .map_err(to_datafusion_error)? + .try_collect::<Vec<_>>() + .await + .map_err(to_datafusion_error)?; Review Comment: Thanks for raising this, please push back if any of the below is off. For context, the long-term direction for this is tracked in the EPIC #1604 (row-group-based parallel scan with a `GroupPruner` that can split/merge `FileScanTask` below the file grain). What I was hoping to land with this PR is a more immediate, scoped optimization that stays within the current file-grain contract, so we don't preempt the design choices in #1604. The file-grouping step you're pointing at is essentially what #2220 describes as the intermediate improvement on the path toward #1604. If you think it's appropriate, I'd be happy to pick up a short-term follow-up along these lines: 1. Switch `IcebergPartitionedScan` from tasks: `Vec<FileScanTask>` to file_groups: `Vec<Vec<FileScanTask>>`, to follow the convention used by DataFusion's own `FileScanConfig`, each group = one DataFusion partition that streams its files sequentially through `ArrowReaderBuilder::read`. 2. In `IcebergPartitionedTableProvider::scan`, read `state.config().target_partitions()` and group tasks into `min(n_files, target_partitions)` buckets. 3. When `n_files` < `target_partitions`, parallelism is still capped at `n_files`. I think that's inherent to the file grain, but let me know if I'm missing something. I'm happy to open the follow-up issue/PR myself, or defer to you if you'd rather frame it, whatever works best. -- 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]
