fix: project columns correctly when scanning iceberg table metadata

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,19 +28,33 @@ 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;
 
 /// Reads rows from an Iceberg metadata table through DataFusion's `DataSource` API.
 #[derive(Debug, Clone)]
 pub struct IcebergMetadataDataSource {
     provider: IcebergMetadataTableProvider,
+    schema: ArrowSchemaRef,
+    projection: Option<Vec<usize>>,
 }
 
 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<usize>>,
+    ) -> Result<Self> {
+        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<TaskContext>,
     ) -> Result<SendableRecordBatchStream> {
         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<usize>) -> Result<Arc<Statistics>> {
-        Ok(Arc::new(Statistics::new_unknown(&self.provider.schema())))
+        Ok(Arc::new(Statistics::new_unknown(&self.schema)))
     }
 
     fn with_fetch(&self, _fetch: Option<usize>) -> Option<Arc<dyn DataSource>> {

crates/datafusion/src/metadata/table.rs

@@ -62,12 +62,12 @@ impl TableProvider for IcebergMetadataTableProvider {
     async fn scan(
         &self,
         _state: &dyn Session,
-        _projection: Option<&Vec<usize>>,
+        projection: Option<&Vec<usize>>,
         _filters: &[Expr],
         _limit: Option<usize>,
     ) -> Result<Arc<dyn ExecutionPlan>> {
         Ok(DataSourceExec::from_data_source(
-            IcebergMetadataDataSource::new(self.clone()),
+            IcebergMetadataDataSource::try_new(self.clone(), projection)?,
         ))
     }
 }

crates/datafusion/tests/integration_datafusion_test.rs

@@ -693,8 +693,8 @@ async fn test_insert_into_nested() -> Result<(), Box<dyn Error>> {
     // 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<dyn Error>> {
             )
         ) 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<dyn Error>> {
     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<dyn Error>> {
     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<dyn Error>> {
         .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

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(())
+    }
+}