diff --git a/crates/datafusion/src/metadata/scan.rs b/crates/datafusion/src/metadata/scan.rs index 3d20cb8..37eb689 100644 --- a/crates/datafusion/src/metadata/scan.rs +++ b/crates/datafusion/src/metadata/scan.rs @@ -17,6 +17,7 @@ use std::sync::Arc; +use datafusion::arrow::datatypes::SchemaRef as ArrowSchemaRef; use datafusion::common::Statistics; use datafusion::common::tree_node::TreeNodeRecursion; use datafusion::datasource::TableProvider; @@ -27,7 +28,7 @@ use datafusion::physical_expr::projection::ProjectionExprs; use datafusion::physical_expr::{EquivalenceProperties, PhysicalExpr}; use datafusion::physical_plan::stream::RecordBatchStreamAdapter; use datafusion::physical_plan::{DisplayFormatType, Partitioning}; -use futures::TryStreamExt; +use futures::{StreamExt, TryStreamExt}; use super::table::IcebergMetadataTableProvider; @@ -35,11 +36,25 @@ use super::table::IcebergMetadataTableProvider; #[derive(Debug, Clone)] pub struct IcebergMetadataDataSource { provider: IcebergMetadataTableProvider, + schema: ArrowSchemaRef, + projection: Option>, } impl IcebergMetadataDataSource { - pub fn new(provider: IcebergMetadataTableProvider) -> Self { - Self { provider } + /// Creates a source with the requested output columns. + pub fn try_new( + provider: IcebergMetadataTableProvider, + projection: Option<&Vec>, + ) -> Result { + let schema = match projection { + Some(indices) => Arc::new(provider.schema().project(indices)?), + None => provider.schema(), + }; + Ok(Self { + provider, + schema, + projection: projection.cloned(), + }) } /// The provider this source scans. @@ -55,9 +70,25 @@ impl DataSource for IcebergMetadataDataSource { _context: Arc, ) -> Result { let fut = self.provider.clone().scan(); - let stream = futures::stream::once(fut).try_flatten(); - let schema = self.provider.schema(); - Ok(Box::pin(RecordBatchStreamAdapter::new(schema, stream))) + let projection = self.projection.clone(); + + // TODO: Push these projections down into the scan layer instead of manually iterating over the result set + // and applying them. + // This will be possible once this issue is addressed in iceberg-rust: https://github.com/apache/iceberg-rust/issues/3391 + let stream = + futures::stream::once(fut) + .try_flatten() + .map(move |result| -> Result<_> { + let batch = result?; + match &projection { + Some(indices) => Ok(batch.project(indices)?), + None => Ok(batch), + } + }); + Ok(Box::pin(RecordBatchStreamAdapter::new( + self.schema.clone(), + stream, + ))) } fn fmt_as( @@ -73,11 +104,11 @@ impl DataSource for IcebergMetadataDataSource { } fn eq_properties(&self) -> EquivalenceProperties { - EquivalenceProperties::new(self.provider.schema()) + EquivalenceProperties::new(self.schema.clone()) } fn partition_statistics(&self, _partition: Option) -> Result> { - Ok(Arc::new(Statistics::new_unknown(&self.provider.schema()))) + Ok(Arc::new(Statistics::new_unknown(&self.schema))) } fn with_fetch(&self, _fetch: Option) -> Option> { diff --git a/crates/datafusion/src/metadata/table.rs b/crates/datafusion/src/metadata/table.rs index 6bd9448..3cf351d 100644 --- a/crates/datafusion/src/metadata/table.rs +++ b/crates/datafusion/src/metadata/table.rs @@ -62,12 +62,12 @@ impl TableProvider for IcebergMetadataTableProvider { async fn scan( &self, _state: &dyn Session, - _projection: Option<&Vec>, + projection: Option<&Vec>, _filters: &[Expr], _limit: Option, ) -> Result> { Ok(DataSourceExec::from_data_source( - IcebergMetadataDataSource::new(self.clone()), + IcebergMetadataDataSource::try_new(self.clone(), projection)?, )) } } diff --git a/crates/datafusion/tests/integration_datafusion_test.rs b/crates/datafusion/tests/integration_datafusion_test.rs index 187ae2f..67dfcfa 100644 --- a/crates/datafusion/tests/integration_datafusion_test.rs +++ b/crates/datafusion/tests/integration_datafusion_test.rs @@ -693,8 +693,8 @@ async fn test_insert_into_nested() -> Result<(), Box> { // Insert data with nested structs let insert_sql = r#" INSERT INTO catalog.test_insert_nested.nested_table - SELECT - 1 as id, + SELECT + 1 as id, 'Alice' as name, named_struct( 'address', named_struct( @@ -708,8 +708,8 @@ async fn test_insert_into_nested() -> Result<(), Box> { ) ) as profile UNION ALL - SELECT - 2 as id, + SELECT + 2 as id, 'Bob' as name, named_struct( 'address', named_struct( @@ -831,15 +831,15 @@ async fn test_insert_into_nested() -> Result<(), Box> { let df = ctx .sql( r#" - SELECT - id, + SELECT + id, name, profile.address.street, profile.address.city, profile.address.zip, profile.contact.email, profile.contact.phone - FROM catalog.test_insert_nested.nested_table + FROM catalog.test_insert_nested.nested_table ORDER BY id "#, ) @@ -981,8 +981,8 @@ async fn test_insert_into_partitioned() -> Result<(), Box> { let df = ctx .sql( r#" - INSERT INTO catalog.test_partitioned_write.partitioned_table - VALUES + INSERT INTO catalog.test_partitioned_write.partitioned_table + VALUES (1, 'electronics', 'laptop', CAST(123.456 AS REAL), CAST(123.456 AS DOUBLE)), (2, 'electronics', 'phone', CAST(123.456 AS REAL), CAST(123.456 AS DOUBLE)), (3, 'books', 'novel', CAST('NaN' AS REAL), CAST('NaN' AS DOUBLE)), @@ -1344,10 +1344,13 @@ async fn test_plan_nodes_are_inspectable() -> Result<(), Box> { .expect("an Iceberg metadata source"); let provider = metadata_scan.provider(); assert_eq!(provider.table().identifier(), &ident); - let rebuilt = IcebergMetadataDataSource::new(IcebergMetadataTableProvider::new( - provider.table().clone(), - provider.metadata_type().clone(), - )); + let rebuilt = IcebergMetadataDataSource::try_new( + IcebergMetadataTableProvider::new( + provider.table().clone(), + provider.metadata_type().clone(), + ), + None, + )?; let batches = run_batches(metadata_exec, &ctx).await?; // Snapshots come back in no set order, so put the first, which has no // parent, first. Their ids, times and paths differ on every run, so each diff --git a/crates/datafusion/tests/metadata_projection.rs b/crates/datafusion/tests/metadata_projection.rs new file mode 100644 index 0000000..8d77899 --- /dev/null +++ b/crates/datafusion/tests/metadata_projection.rs @@ -0,0 +1,81 @@ +// 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. + +#[cfg(test)] +mod tests { + use datafusion::error::Result; + use iceberg_datafusion::test_utils::IcebergTestHarness; + + #[tokio::test] + async fn applies_metadata_table_projections() -> Result<()> { + let harness = IcebergTestHarness::new().await?; + + let single_projected_column_result = harness + .query("SELECT operation FROM datafusion.public.taxi$snapshots") + .await?; + insta::assert_snapshot!(single_projected_column_result.0, @" + CooperativeExec + DataSourceExec: format=iceberg_metadata + "); + insta::assert_snapshot!(single_projected_column_result.1, @" + +-----------+ + | operation | + +-----------+ + | append | + +-----------+ + "); + + let multiple_projected_columns_result = harness + .query("SELECT operation, snapshot_id FROM datafusion.public.taxi$snapshots") + .await?; + + insta::assert_snapshot!(multiple_projected_columns_result.0, @" + ProjectionExec: expr=[operation@1 as operation, snapshot_id@0 as snapshot_id] + CooperativeExec + DataSourceExec: format=iceberg_metadata + "); + insta::assert_snapshot!(multiple_projected_columns_result.1, @" + +-----------+---------------------+ + | operation | snapshot_id | + +-----------+---------------------+ + | append | 3167948105555765929 | + +-----------+---------------------+ + "); + + let aggregated_columns_result = harness + .query("SELECT COUNT(*) FROM datafusion.public.taxi$snapshots") + .await?; + + insta::assert_snapshot!(aggregated_columns_result.0, @" + ProjectionExec: expr=[count(Int64(1))@0 as count(*)] + AggregateExec: mode=Final, gby=[], aggr=[count(Int64(1))] + CoalescePartitionsExec + AggregateExec: mode=Partial, gby=[], aggr=[count(Int64(1))] + RepartitionExec: partitioning=RoundRobinBatch(16), input_partitions=1 + DataSourceExec: format=iceberg_metadata + "); + insta::assert_snapshot!(aggregated_columns_result.1, @" + +----------+ + | count(*) | + +----------+ + | 1 | + +----------+ + "); + + Ok(()) + } +}