From 15f1987170568cf4d49a9cd7805ccda757fa407c Mon Sep 17 00:00:00 2001 From: Perseus Date: Sat, 10 Oct 2026 13:36:30 +0200 Subject: [PATCH 1/3] fix: project columns correctly when scanning iceberg table metadata --- crates/datafusion/src/metadata/scan.rs | 43 +++++- crates/datafusion/src/metadata/table.rs | 4 +- .../tests/integration_datafusion_test.rs | 137 ++++++++++++++++-- 3 files changed, 161 insertions(+), 23 deletions(-) diff --git a/crates/datafusion/src/metadata/scan.rs b/crates/datafusion/src/metadata/scan.rs index 3d20cb8..eed13c4 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,21 @@ 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(); + 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 +100,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..f071b24 100644 --- a/crates/datafusion/tests/integration_datafusion_test.rs +++ b/crates/datafusion/tests/integration_datafusion_test.rs @@ -518,6 +518,114 @@ async fn test_metadata_table() -> Result<(), Box> { Ok(()) } +#[tokio::test] +async fn test_metadata_table_projections() -> Result<(), Box> { + let iceberg_catalog = get_iceberg_catalog().await; + let namespace = NamespaceIdent::new("ns".to_string()); + set_test_namespace(&iceberg_catalog, &namespace).await?; + + let schema = Schema::builder() + .with_schema_id(0) + .with_fields(vec![ + NestedField::required(1, "foo", Type::Primitive(PrimitiveType::Int)).into(), + NestedField::optional(2, "bar", Type::Primitive(PrimitiveType::String)) + .into(), + ]) + .build()?; + let creation = get_table_creation(temp_path(), "t1", Some(schema))?; + iceberg_catalog.create_table(&namespace, creation).await?; + + let client = Arc::new(iceberg_catalog); + let catalog = Arc::new(IcebergCatalogProvider::try_new(client).await?); + + let ctx = SessionContext::new(); + ctx.register_catalog("catalog", catalog); + + // commit twice to create two snapshots + ctx.sql("INSERT INTO catalog.ns.t1 VALUES (1, 'foo')") + .await + .unwrap() + .collect() + .await + .unwrap(); + + ctx.sql("INSERT INTO catalog.ns.t1 VALUES (2, 'bar')") + .await + .unwrap() + .collect() + .await + .unwrap(); + + let single_projected_column_result = ctx + .sql("SELECT operation FROM catalog.ns.t1$snapshots") + .await + .unwrap() + .collect() + .await + .unwrap(); + + check_record_batches( + single_projected_column_result, + expect![[ + r#"Field { "operation": nullable Utf8, metadata: {"PARQUET:field_id": "4"} }"# + ]], + expect![[r#" + operation: StringArray + [ + "append", + "append", + ]"#]], + &[], + None, + ); + + let multiple_projected_columns_result = ctx + .sql("SELECT operation, snapshot_id FROM catalog.ns.t1$snapshots") + .await + .unwrap() + .collect() + .await + .unwrap(); + + check_record_batches( + multiple_projected_columns_result, + expect![[r#" + Field { "operation": nullable Utf8, metadata: {"PARQUET:field_id": "4"} }, + Field { "snapshot_id": Int64, metadata: {"PARQUET:field_id": "2"} }"#]], + expect![[r#" + operation: StringArray + [ + "append", + "append", + ], + snapshot_id: (skipped)"#]], + &["snapshot_id"], + None, + ); + + let aggregated_columns_result = ctx + .sql("SELECT count(*) FROM catalog.ns.t1$snapshots") + .await + .unwrap() + .collect() + .await + .unwrap(); + + check_record_batches( + aggregated_columns_result, + expect![[r#"Field { "count(*)": Int64 }"#]], + expect![[r#" + count(*): PrimitiveArray + [ + 2, + ]"#]], + &[], + None, + ); + + Ok(()) +} + #[tokio::test] async fn test_insert_into() -> Result<(), Box> { let iceberg_catalog = get_iceberg_catalog().await; @@ -693,8 +801,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 +816,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 +939,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 +1089,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 +1452,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 From f040471f246d0ee6edd6a316a307293300dfdd41 Mon Sep 17 00:00:00 2001 From: Perseus Date: Sat, 10 Oct 2026 13:40:17 +0200 Subject: [PATCH 2/3] chore: add todo for projection pushdown --- crates/datafusion/src/metadata/scan.rs | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/crates/datafusion/src/metadata/scan.rs b/crates/datafusion/src/metadata/scan.rs index eed13c4..37eb689 100644 --- a/crates/datafusion/src/metadata/scan.rs +++ b/crates/datafusion/src/metadata/scan.rs @@ -71,6 +71,10 @@ impl DataSource for IcebergMetadataDataSource { ) -> Result { let fut = self.provider.clone().scan(); 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() From f1383a9d0440a4e90b9492fd28b6418ad8eecfc7 Mon Sep 17 00:00:00 2001 From: Perseus Date: Mon, 12 Oct 2026 00:01:32 +0200 Subject: [PATCH 3/3] tests: use the new test harness for metadata projection tests --- .../tests/integration_datafusion_test.rs | 108 ------------------ .../datafusion/tests/metadata_projection.rs | 81 +++++++++++++ 2 files changed, 81 insertions(+), 108 deletions(-) create mode 100644 crates/datafusion/tests/metadata_projection.rs diff --git a/crates/datafusion/tests/integration_datafusion_test.rs b/crates/datafusion/tests/integration_datafusion_test.rs index f071b24..67dfcfa 100644 --- a/crates/datafusion/tests/integration_datafusion_test.rs +++ b/crates/datafusion/tests/integration_datafusion_test.rs @@ -518,114 +518,6 @@ async fn test_metadata_table() -> Result<(), Box> { Ok(()) } -#[tokio::test] -async fn test_metadata_table_projections() -> Result<(), Box> { - let iceberg_catalog = get_iceberg_catalog().await; - let namespace = NamespaceIdent::new("ns".to_string()); - set_test_namespace(&iceberg_catalog, &namespace).await?; - - let schema = Schema::builder() - .with_schema_id(0) - .with_fields(vec![ - NestedField::required(1, "foo", Type::Primitive(PrimitiveType::Int)).into(), - NestedField::optional(2, "bar", Type::Primitive(PrimitiveType::String)) - .into(), - ]) - .build()?; - let creation = get_table_creation(temp_path(), "t1", Some(schema))?; - iceberg_catalog.create_table(&namespace, creation).await?; - - let client = Arc::new(iceberg_catalog); - let catalog = Arc::new(IcebergCatalogProvider::try_new(client).await?); - - let ctx = SessionContext::new(); - ctx.register_catalog("catalog", catalog); - - // commit twice to create two snapshots - ctx.sql("INSERT INTO catalog.ns.t1 VALUES (1, 'foo')") - .await - .unwrap() - .collect() - .await - .unwrap(); - - ctx.sql("INSERT INTO catalog.ns.t1 VALUES (2, 'bar')") - .await - .unwrap() - .collect() - .await - .unwrap(); - - let single_projected_column_result = ctx - .sql("SELECT operation FROM catalog.ns.t1$snapshots") - .await - .unwrap() - .collect() - .await - .unwrap(); - - check_record_batches( - single_projected_column_result, - expect![[ - r#"Field { "operation": nullable Utf8, metadata: {"PARQUET:field_id": "4"} }"# - ]], - expect![[r#" - operation: StringArray - [ - "append", - "append", - ]"#]], - &[], - None, - ); - - let multiple_projected_columns_result = ctx - .sql("SELECT operation, snapshot_id FROM catalog.ns.t1$snapshots") - .await - .unwrap() - .collect() - .await - .unwrap(); - - check_record_batches( - multiple_projected_columns_result, - expect![[r#" - Field { "operation": nullable Utf8, metadata: {"PARQUET:field_id": "4"} }, - Field { "snapshot_id": Int64, metadata: {"PARQUET:field_id": "2"} }"#]], - expect![[r#" - operation: StringArray - [ - "append", - "append", - ], - snapshot_id: (skipped)"#]], - &["snapshot_id"], - None, - ); - - let aggregated_columns_result = ctx - .sql("SELECT count(*) FROM catalog.ns.t1$snapshots") - .await - .unwrap() - .collect() - .await - .unwrap(); - - check_record_batches( - aggregated_columns_result, - expect![[r#"Field { "count(*)": Int64 }"#]], - expect![[r#" - count(*): PrimitiveArray - [ - 2, - ]"#]], - &[], - None, - ); - - Ok(()) -} - #[tokio::test] async fn test_insert_into() -> Result<(), Box> { let iceberg_catalog = get_iceberg_catalog().await; 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(()) + } +}