feat: support projection pushdown for datafusion iceberg

Cargo.toml

@@ -64,6 +64,7 @@ iceberg = { version = "0.3.0", path = "./crates/iceberg" }
 iceberg-catalog-rest = { version = "0.3.0", path = "./crates/catalog/rest" }
 iceberg-catalog-hms = { version = "0.3.0", path = "./crates/catalog/hms" }
 iceberg-catalog-memory = { version = "0.3.0", path = "./crates/catalog/memory" }
+iceberg-datafusion = { version = "0.3.0", path = "./crates/integrations/datafusion" }
 itertools = "0.13"
 log = "0.4"
 mockito = "1"

crates/examples/Cargo.toml

@@ -27,6 +27,8 @@ rust-version = { workspace = true }
 [dependencies]
 iceberg = { workspace = true }
 iceberg-catalog-rest = { workspace = true }
+iceberg-datafusion = { workspace = true }
+datafusion = { version = "41.0.0" }
 tokio = { version = "1", features = ["full"] }
 
 [[example]]
@@ -36,3 +38,7 @@ path = "src/rest_catalog_namespace.rs"
 [[example]]
 name = "rest-catalog-table"
 path = "src/rest_catalog_table.rs"
+
+[[example]]
+name = "datafusion-read-data"
+path = "src/datafusion_read_data.rs"

crates/examples/src/datafusion_read_data.rs

@@ -0,0 +1,37 @@
+// 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::sync::Arc;
+
+use datafusion::prelude::SessionContext;
+use iceberg_datafusion::IcebergCatalogProvider;
+
+mod utils;
+
+#[tokio::main]
+async fn main() {
+    let iceberg_catalog = utils::get_rest_catalog();
+
+    let client = Arc::new(iceberg_catalog);
+    let catalog = Arc::new(IcebergCatalogProvider::try_new(client).await.unwrap());
+
+    let ctx = SessionContext::new();
+    ctx.register_catalog("catalog", catalog);
+    let df = ctx.sql("select * from catalog.ns.table1").await.unwrap();
+    let data = df.collect().await.unwrap();
+    println!("{:?}", data);
+}

crates/examples/src/utils.rs

@@ -0,0 +1,31 @@
+// 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::env;
+
+use iceberg_catalog_rest::{RestCatalog, RestCatalogConfig};
+
+fn get_catalog_uri_from_env() -> String {
+    env::var("CATALOG_URI").unwrap_or("http://localhost:8080".to_string())
+}
+
+pub fn get_rest_catalog() -> RestCatalog {
+    let config = RestCatalogConfig::builder()
+        .uri(get_catalog_uri_from_env())
+        .build();
+    RestCatalog::new(config)
+}

crates/integrations/datafusion/src/physical_plan/scan.rs

@@ -18,6 +18,7 @@
 use std::any::Any;
 use std::pin::Pin;
 use std::sync::Arc;
+use std::vec;
 
 use datafusion::arrow::array::RecordBatch;
 use datafusion::arrow::datatypes::SchemaRef as ArrowSchemaRef;
@@ -44,17 +45,25 @@ pub(crate) struct IcebergTableScan {
     /// Stores certain, often expensive to compute,
     /// plan properties used in query optimization.
     plan_properties: PlanProperties,
+    /// Projection column names
+    projection: Vec<String>,
 }
 
 impl IcebergTableScan {
     /// Creates a new [`IcebergTableScan`] object.
-    pub(crate) fn new(table: Table, schema: ArrowSchemaRef) -> Self {
+    pub(crate) fn new(
+        table: Table,
+        schema: ArrowSchemaRef,
+        projection: Option<&Vec<usize>>,
+    ) -> Self {
         let plan_properties = Self::compute_properties(schema.clone());
+        let projection = get_column_names(schema.clone(), projection);
 
         Self {
             table,
             schema,
             plan_properties,
+            projection,
         }
     }
 
@@ -100,7 +109,7 @@ impl ExecutionPlan for IcebergTableScan {
         _partition: usize,
         _context: Arc<TaskContext>,
     ) -> DFResult<SendableRecordBatchStream> {
-        let fut = get_batch_stream(self.table.clone());
+        let fut = get_batch_stream(self.table.clone(), self.projection.clone());
         let stream = futures::stream::once(fut).try_flatten();
 
         Ok(Box::pin(RecordBatchStreamAdapter::new(
@@ -116,7 +125,11 @@ impl DisplayAs for IcebergTableScan {
         _t: datafusion::physical_plan::DisplayFormatType,
         f: &mut std::fmt::Formatter,
     ) -> std::fmt::Result {
-        write!(f, "IcebergTableScan")
+        write!(
+            f,
+            "IcebergTableScan projection:[{}]",
+            self.projection.join(" ")
+        )
     }
 }
 
@@ -127,8 +140,13 @@ impl DisplayAs for IcebergTableScan {
 /// and then converts it into a stream of Arrow [`RecordBatch`]es.
 async fn get_batch_stream(
     table: Table,
+    column_names: Vec<String>,
 ) -> DFResult<Pin<Box<dyn Stream<Item = DFResult<RecordBatch>> + Send>>> {
-    let table_scan = table.scan().build().map_err(to_datafusion_error)?;
+    let table_scan = table
+        .scan()
+        .select(column_names)
+        .build()
+        .map_err(to_datafusion_error)?;
 
     let stream = table_scan
         .to_arrow()
@@ -138,3 +156,18 @@ async fn get_batch_stream(
 
     Ok(Box::pin(stream))
 }
+
+fn get_column_names(schema: ArrowSchemaRef, projection: Option<&Vec<usize>>) -> Vec<String> {
+    if let Some(projection) = projection {
+        projection
+            .iter()
+            .map(|p| schema.field(*p).name().clone())
+            .collect::<Vec<String>>()
+    } else {
+        schema
+            .fields()
+            .iter()
+            .map(|f| f.name().clone())
+            .collect::<Vec<String>>()
+    }
+}

crates/integrations/datafusion/src/table.rs

@@ -75,13 +75,14 @@ impl TableProvider for IcebergTableProvider {
     async fn scan(
         &self,
         _state: &dyn Session,
-        _projection: Option<&Vec<usize>>,
+        projection: Option<&Vec<usize>>,
         _filters: &[Expr],
         _limit: Option<usize>,
     ) -> DFResult<Arc<dyn ExecutionPlan>> {
         Ok(Arc::new(IcebergTableScan::new(
             self.table.clone(),
             self.schema.clone(),
+            projection,
         )))
     }
 }

crates/integrations/datafusion/tests/integration_datafusion_test.rs

@@ -19,11 +19,13 @@
 
 use std::collections::HashMap;
 use std::sync::Arc;
+use std::vec;
 
+use datafusion::arrow::array::{Array, StringArray};
 use datafusion::arrow::datatypes::DataType;
 use datafusion::execution::context::SessionContext;
 use iceberg::io::FileIOBuilder;
-use iceberg::spec::{NestedField, PrimitiveType, Schema, Type};
+use iceberg::spec::{NestedField, PrimitiveType, Schema, StructType, Type};
 use iceberg::{Catalog, NamespaceIdent, Result, TableCreation};
 use iceberg_catalog_memory::MemoryCatalog;
 use iceberg_datafusion::IcebergCatalogProvider;
@@ -39,6 +41,13 @@ fn get_iceberg_catalog() -> MemoryCatalog {
     MemoryCatalog::new(file_io, Some(temp_path()))
 }
 
+fn get_struct_type() -> StructType {
+    StructType::new(vec![
+        NestedField::required(13, "s_foo1", Type::Primitive(PrimitiveType::Int)).into(),
+        NestedField::required(14, "s_foo2", Type::Primitive(PrimitiveType::String)).into(),
+    ])
+}
+
 async fn set_test_namespace(catalog: &MemoryCatalog, namespace: &NamespaceIdent) -> Result<()> {
     let properties = HashMap::new();
 
@@ -47,14 +56,21 @@ async fn set_test_namespace(catalog: &MemoryCatalog, namespace: &NamespaceIdent)
     Ok(())
 }
 
-fn set_table_creation(location: impl ToString, name: impl ToString) -> Result<TableCreation> {
-    let schema = Schema::builder()
-        .with_schema_id(0)
-        .with_fields(vec![
-            NestedField::required(1, "foo", Type::Primitive(PrimitiveType::Int)).into(),
-            NestedField::required(2, "bar", Type::Primitive(PrimitiveType::String)).into(),
-        ])
-        .build()?;
+fn get_table_creation(
+    location: impl ToString,
+    name: impl ToString,
+    schema: Option<Schema>,
+) -> Result<TableCreation> {
+    let schema = match schema {
+        None => Schema::builder()
+            .with_schema_id(0)
+            .with_fields(vec![
+                NestedField::required(1, "foo1", Type::Primitive(PrimitiveType::Int)).into(),
+                NestedField::required(2, "foo2", Type::Primitive(PrimitiveType::String)).into(),
+            ])
+            .build()?,
+        Some(schema) => schema,
+    };
 
     let creation = TableCreation::builder()
         .location(location.to_string())
@@ -72,7 +88,7 @@ async fn test_provider_get_table_schema() -> Result<()> {
     let namespace = NamespaceIdent::new("test_provider_get_table_schema".to_string());
     set_test_namespace(&iceberg_catalog, &namespace).await?;
 
-    let creation = set_table_creation(temp_path(), "my_table")?;
+    let creation = get_table_creation(temp_path(), "my_table", None)?;
     iceberg_catalog.create_table(&namespace, creation).await?;
 
     let client = Arc::new(iceberg_catalog);
@@ -87,7 +103,7 @@ async fn test_provider_get_table_schema() -> Result<()> {
     let table = schema.table("my_table").await.unwrap().unwrap();
     let table_schema = table.schema();
 
-    let expected = [("foo", &DataType::Int32), ("bar", &DataType::Utf8)];
+    let expected = [("foo1", &DataType::Int32), ("foo2", &DataType::Utf8)];
 
     for (field, exp) in table_schema.fields().iter().zip(expected.iter()) {
         assert_eq!(field.name(), exp.0);
@@ -104,7 +120,7 @@ async fn test_provider_list_table_names() -> Result<()> {
     let namespace = NamespaceIdent::new("test_provider_list_table_names".to_string());
     set_test_namespace(&iceberg_catalog, &namespace).await?;
 
-    let creation = set_table_creation(temp_path(), "my_table")?;
+    let creation = get_table_creation(temp_path(), "my_table", None)?;
     iceberg_catalog.create_table(&namespace, creation).await?;
 
     let client = Arc::new(iceberg_catalog);
@@ -130,7 +146,6 @@ async fn test_provider_list_schema_names() -> Result<()> {
     let namespace = NamespaceIdent::new("test_provider_list_schema_names".to_string());
     set_test_namespace(&iceberg_catalog, &namespace).await?;
 
-    set_table_creation("test_provider_list_schema_names", "my_table")?;
     let client = Arc::new(iceberg_catalog);
     let catalog = Arc::new(IcebergCatalogProvider::try_new(client).await?);
 
@@ -147,3 +162,71 @@ async fn test_provider_list_schema_names() -> Result<()> {
         .all(|item| result.contains(&item.to_string())));
     Ok(())
 }
+
+#[tokio::test]
+async fn test_table_projection() -> Result<()> {
+    let iceberg_catalog = get_iceberg_catalog();
+    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, "foo1", Type::Primitive(PrimitiveType::Int)).into(),
+            NestedField::required(2, "foo2", Type::Primitive(PrimitiveType::String)).into(),
+            NestedField::optional(0, "foo3", Type::Struct(get_struct_type())).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);
+    let table_df = ctx.table("catalog.ns.t1").await.unwrap();
+
+    let records = table_df
+        .clone()
+        .explain(false, false)
+        .unwrap()
+        .collect()
+        .await
+        .unwrap();
+    assert_eq!(1, records.len());
+    let record = &records[0];
+    // the first column is plan_type, the second column plan string.
+    let s = record
+        .column(1)
+        .as_any()
+        .downcast_ref::<StringArray>()
+        .unwrap();
+    assert_eq!(2, s.len());
+    // the first row is logical_plan, the second row is physical_plan
+    assert_eq!(
+        "IcebergTableScan projection:[foo1 foo2 foo3]",
+        s.value(1).trim()
+    );
+
+    // datafusion doesn't support query foo3.s_foo1, use foo3 instead
+    let records = table_df
+        .select_columns(&["foo1", "foo3"])
+        .unwrap()
+        .explain(false, false)
+        .unwrap()
+        .collect()
+        .await
+        .unwrap();
+    assert_eq!(1, records.len());
+    let record = &records[0];
+    let s = record
+        .column(1)
+        .as_any()
+        .downcast_ref::<StringArray>()
+        .unwrap();
+    assert_eq!(2, s.len());
+    assert_eq!("IcebergTableScan projection:[foo1 foo3]", s.value(1).trim());
+
+    Ok(())
+}