From b88e316f52756ec983588cf9806387a54d4a54f0 Mon Sep 17 00:00:00 2001 From: Gustavo Schneiter Date: Sun, 12 Jul 2026 19:31:09 -0300 Subject: [PATCH] fix(datafusion): honor projection in metadata table scan Closes #2819. --- crates/integrations/datafusion/public-api.txt | 4 +- .../src/physical_plan/metadata_scan.rs | 27 +++++-- .../datafusion/src/table/metadata_table.rs | 4 +- .../tests/integration_datafusion_test.rs | 80 +++++++++++++++++++ .../df_inspect_snapshots_manifests.toml | 23 ++++++ .../df_test/inspect_snapshots_manifests.slt | 61 ++++++++++++++ 6 files changed, 190 insertions(+), 9 deletions(-) create mode 100644 crates/sqllogictest/testdata/schedules/df_inspect_snapshots_manifests.toml create mode 100644 crates/sqllogictest/testdata/slts/df_test/inspect_snapshots_manifests.slt diff --git a/crates/integrations/datafusion/public-api.txt b/crates/integrations/datafusion/public-api.txt index e197c2057d..6494971f4d 100644 --- a/crates/integrations/datafusion/public-api.txt +++ b/crates/integrations/datafusion/public-api.txt @@ -8,7 +8,7 @@ pub fn iceberg_datafusion::metadata_table::IcebergMetadataTableProvider::clone(& impl core::fmt::Debug for iceberg_datafusion::metadata_table::IcebergMetadataTableProvider pub fn iceberg_datafusion::metadata_table::IcebergMetadataTableProvider::fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result impl datafusion_catalog::table::TableProvider for iceberg_datafusion::metadata_table::IcebergMetadataTableProvider -pub fn iceberg_datafusion::metadata_table::IcebergMetadataTableProvider::scan<'life0, 'life1, 'life2, 'life3, 'async_trait>(&'life0 self, _state: &'life1 dyn datafusion_session::session::Session, _projection: core::option::Option<&'life2 alloc::vec::Vec>, _filters: &'life3 [datafusion_expr::expr::Expr], _limit: core::option::Option) -> core::pin::Pin>> + core::marker::Send + 'async_trait)>> where Self: 'async_trait, 'life0: 'async_trait, 'life1: 'async_trait, 'life2: 'async_trait, 'life3: 'async_trait +pub fn iceberg_datafusion::metadata_table::IcebergMetadataTableProvider::scan<'life0, 'life1, 'life2, 'life3, 'async_trait>(&'life0 self, _state: &'life1 dyn datafusion_session::session::Session, projection: core::option::Option<&'life2 alloc::vec::Vec>, _filters: &'life3 [datafusion_expr::expr::Expr], _limit: core::option::Option) -> core::pin::Pin>> + core::marker::Send + 'async_trait)>> where Self: 'async_trait, 'life0: 'async_trait, 'life1: 'async_trait, 'life2: 'async_trait, 'life3: 'async_trait pub fn iceberg_datafusion::metadata_table::IcebergMetadataTableProvider::schema(&self) -> arrow_schema::schema::SchemaRef pub fn iceberg_datafusion::metadata_table::IcebergMetadataTableProvider::table_type(&self) -> datafusion_expr::table_source::TableType pub mod iceberg_datafusion::physical_plan @@ -41,7 +41,7 @@ pub fn iceberg_datafusion::metadata_table::IcebergMetadataTableProvider::clone(& impl core::fmt::Debug for iceberg_datafusion::metadata_table::IcebergMetadataTableProvider pub fn iceberg_datafusion::metadata_table::IcebergMetadataTableProvider::fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result impl datafusion_catalog::table::TableProvider for iceberg_datafusion::metadata_table::IcebergMetadataTableProvider -pub fn iceberg_datafusion::metadata_table::IcebergMetadataTableProvider::scan<'life0, 'life1, 'life2, 'life3, 'async_trait>(&'life0 self, _state: &'life1 dyn datafusion_session::session::Session, _projection: core::option::Option<&'life2 alloc::vec::Vec>, _filters: &'life3 [datafusion_expr::expr::Expr], _limit: core::option::Option) -> core::pin::Pin>> + core::marker::Send + 'async_trait)>> where Self: 'async_trait, 'life0: 'async_trait, 'life1: 'async_trait, 'life2: 'async_trait, 'life3: 'async_trait +pub fn iceberg_datafusion::metadata_table::IcebergMetadataTableProvider::scan<'life0, 'life1, 'life2, 'life3, 'async_trait>(&'life0 self, _state: &'life1 dyn datafusion_session::session::Session, projection: core::option::Option<&'life2 alloc::vec::Vec>, _filters: &'life3 [datafusion_expr::expr::Expr], _limit: core::option::Option) -> core::pin::Pin>> + core::marker::Send + 'async_trait)>> where Self: 'async_trait, 'life0: 'async_trait, 'life1: 'async_trait, 'life2: 'async_trait, 'life3: 'async_trait pub fn iceberg_datafusion::metadata_table::IcebergMetadataTableProvider::schema(&self) -> arrow_schema::schema::SchemaRef pub fn iceberg_datafusion::metadata_table::IcebergMetadataTableProvider::table_type(&self) -> datafusion_expr::table_source::TableType pub mod iceberg_datafusion::table::table_provider_factory diff --git a/crates/integrations/datafusion/src/physical_plan/metadata_scan.rs b/crates/integrations/datafusion/src/physical_plan/metadata_scan.rs index 340af9ba7e..80e8e996b6 100644 --- a/crates/integrations/datafusion/src/physical_plan/metadata_scan.rs +++ b/crates/integrations/datafusion/src/physical_plan/metadata_scan.rs @@ -22,7 +22,7 @@ use datafusion::physical_expr::EquivalenceProperties; use datafusion::physical_plan::execution_plan::{Boundedness, EmissionType}; use datafusion::physical_plan::stream::RecordBatchStreamAdapter; use datafusion::physical_plan::{DisplayAs, ExecutionPlan, Partitioning, PlanProperties}; -use futures::TryStreamExt; +use futures::{StreamExt, TryStreamExt}; use crate::metadata_table::IcebergMetadataTableProvider; @@ -30,12 +30,18 @@ use crate::metadata_table::IcebergMetadataTableProvider; pub struct IcebergMetadataScan { provider: IcebergMetadataTableProvider, properties: Arc, + /// Column indices to project, `None` means all columns. + projection: Option>, } impl IcebergMetadataScan { - pub fn new(provider: IcebergMetadataTableProvider) -> Self { + pub fn new(provider: IcebergMetadataTableProvider, projection: Option<&Vec>) -> Self { + let output_schema = match projection { + None => provider.schema(), + Some(projection) => Arc::new(provider.schema().project(projection).unwrap()), + }; let properties = Arc::new(PlanProperties::new( - EquivalenceProperties::new(provider.schema()), + EquivalenceProperties::new(output_schema), Partitioning::UnknownPartitioning(1), EmissionType::Incremental, Boundedness::Bounded, @@ -43,6 +49,7 @@ impl IcebergMetadataScan { Self { provider, properties, + projection: projection.cloned(), } } } @@ -82,9 +89,19 @@ impl ExecutionPlan for IcebergMetadataScan { _partition: usize, _context: std::sync::Arc, ) -> datafusion::error::Result { + let projection = self.projection.clone(); 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 stream = stream.map(move |batch| { + let batch = batch?; + match &projection { + Some(projection) => Ok(batch.project(projection)?), + None => Ok(batch), + } + }); + Ok(Box::pin(RecordBatchStreamAdapter::new( + self.schema(), + stream, + ))) } } diff --git a/crates/integrations/datafusion/src/table/metadata_table.rs b/crates/integrations/datafusion/src/table/metadata_table.rs index 43e069cf5a..05d6b4cf36 100644 --- a/crates/integrations/datafusion/src/table/metadata_table.rs +++ b/crates/integrations/datafusion/src/table/metadata_table.rs @@ -60,11 +60,11 @@ impl TableProvider for IcebergMetadataTableProvider { async fn scan( &self, _state: &dyn Session, - _projection: Option<&Vec>, + projection: Option<&Vec>, _filters: &[Expr], _limit: Option, ) -> DFResult> { - Ok(Arc::new(IcebergMetadataScan::new(self.clone()))) + Ok(Arc::new(IcebergMetadataScan::new(self.clone(), projection))) } } diff --git a/crates/integrations/datafusion/tests/integration_datafusion_test.rs b/crates/integrations/datafusion/tests/integration_datafusion_test.rs index cebac75dd9..c95fb6b54a 100644 --- a/crates/integrations/datafusion/tests/integration_datafusion_test.rs +++ b/crates/integrations/datafusion/tests/integration_datafusion_test.rs @@ -444,6 +444,86 @@ async fn test_metadata_table() -> Result<()> { Ok(()) } +#[tokio::test] +async fn test_metadata_table_projection() -> Result<()> { + // Regression test for issue #2819: metadata table scans must honor the + // projection argument. Before the fix, IcebergMetadataScan reported the + // full metadata-table schema and emitted unprojected batches, so `count(*)` + // failed with a physical/logical schema mismatch and single-column selects + // silently returned every column. Only deterministic columns are asserted. + let iceberg_catalog = get_iceberg_catalog().await; + let namespace = NamespaceIdent::new("test_metadata_projection".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); + + // Populate the table so its metadata tables have a snapshot. + ctx.sql("INSERT INTO catalog.test_metadata_projection.t1 VALUES (1, 'a')") + .await + .unwrap() + .collect() + .await + .unwrap(); + + // Empty projection (`count(*)`) must not fail with a schema mismatch. + let count = ctx + .sql("SELECT count(*) FROM catalog.test_metadata_projection.t1$snapshots") + .await + .unwrap() + .collect() + .await + .unwrap(); + check_record_batches( + count, + expect![[r#" + Field { "count(*)": Int64 }"#]], + expect![[r#" + count(*): PrimitiveArray + [ + 1, + ]"#]], + &[], + None, + ); + + // Single-column projection must return only the projected column. + let operation = ctx + .sql("SELECT operation FROM catalog.test_metadata_projection.t1$snapshots") + .await + .unwrap() + .collect() + .await + .unwrap(); + check_record_batches( + operation, + expect![[r#" + Field { "operation": nullable Utf8, metadata: {"PARQUET:field_id": "4"} }"#]], + expect![[r#" + operation: StringArray + [ + "append", + ]"#]], + &[], + None, + ); + + Ok(()) +} + #[tokio::test] async fn test_insert_into() -> Result<()> { let iceberg_catalog = get_iceberg_catalog().await; diff --git a/crates/sqllogictest/testdata/schedules/df_inspect_snapshots_manifests.toml b/crates/sqllogictest/testdata/schedules/df_inspect_snapshots_manifests.toml new file mode 100644 index 0000000000..4707d5557c --- /dev/null +++ b/crates/sqllogictest/testdata/schedules/df_inspect_snapshots_manifests.toml @@ -0,0 +1,23 @@ +# 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. + +[engines] +df = { type = "datafusion" } + +[[steps]] +engine = "df" +slt = "df_test/inspect_snapshots_manifests.slt" diff --git a/crates/sqllogictest/testdata/slts/df_test/inspect_snapshots_manifests.slt b/crates/sqllogictest/testdata/slts/df_test/inspect_snapshots_manifests.slt new file mode 100644 index 0000000000..e78b44c9b2 --- /dev/null +++ b/crates/sqllogictest/testdata/slts/df_test/inspect_snapshots_manifests.slt @@ -0,0 +1,61 @@ +# 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. + +# Regression coverage for projection on populated metadata tables (issue #2819). +# Only deterministic columns are asserted: count(*), operation, content, +# partition_spec_id. Non-deterministic columns (committed_at, snapshot_id, +# parent_id, manifest_list, summary, *_files_count, file paths) are never asserted. + +# Seed a single snapshot with one data file in one partition. +statement ok +INSERT INTO default.default.test_partitioned_table VALUES (1, 'electronics', 'laptop') + +# $manifests projection while exactly one snapshot exists: DATA content (0) in +# the seeded partition spec (id 0). Exercises projection on the manifests table. +query II +SELECT content, partition_spec_id FROM default.default.test_partitioned_table$manifests +---- +0 0 + +# count(*) on a populated $snapshots — the empty-projection shape that regressed +# with the DataFusion "(physical) 6 vs (logical) 0" internal error. +query I +SELECT count(*) FROM default.default.test_partitioned_table$snapshots +---- +1 + +# Single-column projection on $snapshots — the silent full-row regression: +# without honoring projection this returned all six snapshot columns. +query T +SELECT operation FROM default.default.test_partitioned_table$snapshots +---- +append + +# Add a second snapshot spanning multiple partitions. +statement ok +INSERT INTO default.default.test_partitioned_table VALUES (2, 'electronics', 'phone'), (3, 'books', 'novel') + +query I +SELECT count(*) FROM default.default.test_partitioned_table$snapshots +---- +2 + +query T rowsort +SELECT operation FROM default.default.test_partitioned_table$snapshots +---- +append +append