Merge remote-tracking branch 'upstream/main'

Author: gruuya

datafusion_iceberg/src/error.rs

@@ -56,7 +56,7 @@ pub enum Error {
     ParseInt(#[from] std::num::ParseIntError),
     /// parse int error
     #[error(transparent)]
-    ConfigBuilder(#[from] crate::table::DataFusionTableConfigBuilderError),
+    DeriveBuilder(#[from] derive_builder::UninitializedFieldError),
 }
 
 impl From<Error> for DataFusionError {

datafusion_iceberg/src/table.rs

@@ -107,6 +107,7 @@ impl From<MaterializedView> for DataFusionTable {
 }
 
 #[derive(Clone, Debug, Builder)]
+#[builder(build_fn(error = "Error"))]
 pub struct DataFusionTableConfig {
     /// With this option, an additional "__data_file_path" column is added to the output of the
     /// TableProvider that contains the path of the data-file the row originates from.
@@ -407,30 +408,22 @@ async fn table_scan(
 
     let file_schema: SchemaRef = Arc::new((schema.fields()).try_into().unwrap());
 
-    let projection = projection.cloned().or_else(|| {
-        Some(
-            arrow_schema
-                .fields()
-                .iter()
-                .enumerate()
-                .map(|(i, _)| i)
-                .collect(),
-        )
-    });
+    // If no projection was specified default to projecting all the fields
+    let projection = projection
+        .cloned()
+        .unwrap_or((0..arrow_schema.fields().len()).collect_vec());
 
-    let projection_expr: Option<Vec<_>> = projection.as_ref().map(|projection| {
-        projection
-            .iter()
-            .enumerate()
-            .map(|(i, id)| {
-                let name = arrow_schema.fields[*id].name();
-                (
-                    Arc::new(Column::new(name, i)) as Arc<dyn PhysicalExpr>,
-                    name.to_owned(),
-                )
-            })
-            .collect()
-    });
+    let projection_expr: Vec<_> = projection
+        .iter()
+        .enumerate()
+        .map(|(i, id)| {
+            let name = arrow_schema.fields[*id].name();
+            (
+                Arc::new(Column::new(name, i)) as Arc<dyn PhysicalExpr>,
+                name.to_owned(),
+            )
+        })
+        .collect();
 
     if enable_data_file_path_column {
         table_partition_cols.push(Field::new(DATA_FILE_PATH_COLUMN, DataType::Utf8, false));
@@ -615,6 +608,31 @@ async fn table_scan(
 
                 let mut data_file_iter = data_files.into_iter().peekable();
 
+                // Gather the complete equality projection up-front, since in general the requested
+                // projection may differ from the equality delete columns. Moreover, in principle
+                // each equality delete file may have different deletion columns.
+                // And since we need to reconcile them all with data files using joins and unions,
+                // we need to make sure their schemas are fully compatible in all intermediate nodes.
+                let mut equality_projection = projection.clone();
+                delete_files
+                    .iter()
+                    .flat_map(|delete_manifest| delete_manifest.1.data_file().equality_ids())
+                    .flatten()
+                    .unique()
+                    .for_each(|eq_id| {
+                        // Look up the zero-based index of the column based on its equality id
+                        if let Some((id, _)) = schema
+                            .fields()
+                            .iter()
+                            .enumerate()
+                            .find(|(_, f)| f.id == *eq_id)
+                        {
+                            if !equality_projection.contains(&id) {
+                                equality_projection.push(id);
+                            }
+                        }
+                    });
+
                 let mut plan = stream::iter(delete_files.iter())
                     .map(Ok::<_, DataFusionError>)
                     .try_fold(None, |acc, delete_manifest| {
@@ -626,6 +644,8 @@ async fn table_scan(
                         let file_schema: Arc<ArrowSchema> = file_schema.clone();
                         let file_source = file_source.clone();
                         let mut data_files = Vec::new();
+                        let equality_projection = equality_projection.clone();
+
                         while let Some(data_manifest) = data_file_iter.next_if(|x| {
                             x.1.sequence_number().unwrap()
                                 < delete_manifest.1.sequence_number().unwrap()
@@ -657,26 +677,6 @@ async fn table_scan(
                             );
                             let delete_file_schema: SchemaRef =
                                 Arc::new((delete_schema.fields()).try_into().unwrap());
-                            let equality_projection: Option<Vec<usize>> =
-                                match (&projection, delete_manifest.1.data_file().equality_ids()) {
-                                    (Some(projection), Some(equality_ids)) => {
-                                        let collect: Vec<usize> = schema
-                                            .iter()
-                                            .enumerate()
-                                            .filter_map(|(id, x)| {
-                                                if equality_ids.contains(&x.id)
-                                                    && !projection.contains(&id)
-                                                {
-                                                    Some(id)
-                                                } else {
-                                                    None
-                                                }
-                                            })
-                                            .collect();
-                                        Some([projection.as_slice(), &collect].concat())
-                                    }
-                                    _ => None,
-                                };
 
                             let last_updated_ms = table.metadata().last_updated_ms;
                             let manifest_path = if enable_manifest_file_path_column {
@@ -724,7 +724,7 @@ async fn table_scan(
                             )
                             .with_file_group(FileGroup::new(data_files))
                             .with_statistics(statistics)
-                            .with_projection(equality_projection)
+                            .with_projection(Some(equality_projection))
                             .with_limit(limit)
                             .with_table_partition_cols(table_partition_cols)
                             .build();
@@ -804,7 +804,7 @@ async fn table_scan(
                     )
                     .with_file_group(FileGroup::new(additional_data_files))
                     .with_statistics(statistics)
-                    .with_projection(projection.as_ref().cloned())
+                    .with_projection(Some(equality_projection))
                     .with_limit(limit)
                     .with_table_partition_cols(table_partition_cols)
                     .build();
@@ -816,14 +816,8 @@ async fn table_scan(
                     plan = Arc::new(UnionExec::new(vec![plan, data_files_scan]));
                 }
 
-                if let Some(projection_expr) = projection_expr {
-                    Ok::<_, DataFusionError>(Arc::new(ProjectionExec::try_new(
-                        projection_expr,
-                        plan,
-                    )?) as Arc<dyn ExecutionPlan>)
-                } else {
-                    Ok(plan)
-                }
+                Ok::<_, DataFusionError>(Arc::new(ProjectionExec::try_new(projection_expr, plan)?)
+                    as Arc<dyn ExecutionPlan>)
             }
         })
         .try_collect::<Vec<_>>()
@@ -859,7 +853,7 @@ async fn table_scan(
             FileScanConfigBuilder::new(object_store_url, file_schema, file_source)
                 .with_file_groups(file_groups)
                 .with_statistics(statistics)
-                .with_projection(projection.clone())
+                .with_projection(Some(projection.clone()))
                 .with_limit(limit)
                 .with_table_partition_cols(table_partition_cols)
                 .build();
@@ -873,10 +867,7 @@ async fn table_scan(
 
     match plans.len() {
         0 => {
-            let projected_schema = projection
-                .map(|p| arrow_schema.project(&p))
-                .transpose()?
-                .unwrap_or(arrow_schema.as_ref().clone());
+            let projected_schema = arrow_schema.project(&projection)?;
             Ok(Arc::new(EmptyExec::new(Arc::new(projected_schema))))
         }
         1 => Ok(plans.remove(0)),

datafusion_iceberg/tests/equality_delete.rs

@@ -215,4 +215,34 @@ pub async fn test_equality_delete() {
         .unwrap()
         .collect();
     assert_batches_eq!(expected, &duckdb_batches);
+
+    // Test that projecting a column that is not included in equality deletes works
+    ctx.sql(
+        "INSERT INTO warehouse.test.orders (id, customer_id, product_id, date, amount) VALUES
+                (7, 3, 2, '2020-01-01', 2),
+                (8, 2, 1, '2020-02-02', 3),
+                (9, 1, 3, '2020-01-01', 1);",
+    )
+    .await
+    .expect("Failed to create query plan for insert")
+    .collect()
+    .await
+    .expect("Failed to insert values into table");
+
+    let batches = ctx
+        .sql("select sum(amount) from warehouse.test.orders")
+        .await
+        .expect("Failed to create plan for select")
+        .collect()
+        .await
+        .expect("Failed to execute select query");
+
+    let expected = [
+        "+-----------------------------------+",
+        "| sum(warehouse.test.orders.amount) |",
+        "+-----------------------------------+",
+        "| 13                                |",
+        "+-----------------------------------+",
+    ];
+    assert_batches_eq!(expected, &batches);
 }

iceberg-rust-spec/src/error.rs

@@ -55,30 +55,9 @@ pub enum Error {
     /// parse int error
     #[error(transparent)]
     ParseInt(#[from] std::num::ParseIntError),
-    /// table metadata builder
+    /// derive builder
     #[error(transparent)]
-    TableMetadataBuilder(#[from] crate::spec::table_metadata::TableMetadataBuilderError),
-    /// view metadata builder
-    #[error(transparent)]
-    ViewMetadataBuilder(#[from] crate::spec::view_metadata::GeneralViewMetadataBuilderError),
-    /// version builder
-    #[error(transparent)]
-    VersionBuilder(#[from] crate::spec::view_metadata::VersionBuilderError),
-    /// manifest builder
-    #[error(transparent)]
-    ManifestEntryBuilder(#[from] crate::spec::manifest::ManifestEntryBuilderError),
-    /// datafile builder
-    #[error(transparent)]
-    DatafileBuilder(#[from] crate::spec::manifest::DataFileBuilderError),
-    /// snapshot builder
-    #[error(transparent)]
-    SnapshotBuilder(#[from] crate::spec::snapshot::SnapshotBuilderError),
-    /// structype builder
-    #[error(transparent)]
-    StructTypeBuilder(#[from] crate::spec::types::StructTypeBuilderError),
-    /// partition spec builder
-    #[error(transparent)]
-    PartitionSpec(#[from] crate::spec::partition::PartitionSpecBuilderError),
+    DeriveBuilder(#[from] derive_builder::UninitializedFieldError),
 }
 
 impl From<apache_avro::Error> for Error {

iceberg-rust-spec/src/spec/manifest.rs

@@ -34,7 +34,7 @@ use super::{
 /// Entry in manifest with the iceberg spec version 2.
 #[derive(Debug, Serialize, PartialEq, Clone, Getters, Builder)]
 #[serde(into = "ManifestEntryEnum")]
-#[builder(setter(prefix = "with"))]
+#[builder(build_fn(error = "Error"), setter(prefix = "with"))]
 pub struct ManifestEntry {
     /// Table format version
     format_version: FormatVersion,
@@ -516,7 +516,7 @@ impl From<HashMap<i32, Value>> for AvroMap<ByteBuf> {
 }
 
 #[derive(Debug, PartialEq, Clone, Getters, Builder)]
-#[builder(setter(prefix = "with"))]
+#[builder(build_fn(error = "Error"), setter(prefix = "with"))]
 /// DataFile found in Manifest.
 pub struct DataFile {
     ///Type of content in data file.

iceberg-rust-spec/src/spec/materialized_view_metadata.rs

@@ -29,7 +29,7 @@ pub type MaterializedViewMetadata = GeneralViewMetadata<FullIdentifier>;
 pub type MaterializedViewMetadataBuilder = GeneralViewMetadataBuilder<FullIdentifier>;
 
 impl MaterializedViewMetadata {
-    pub fn as_ref(&self) -> TabularMetadataRef {
+    pub fn as_ref(&self) -> TabularMetadataRef<'_> {
         TabularMetadataRef::MaterializedView(self)
     }
 }

iceberg-rust-spec/src/spec/partition.rs

@@ -169,7 +169,7 @@ impl PartitionField {
 
 #[derive(Debug, Serialize, Deserialize, PartialEq, Eq, Clone, Default, Builder, Getters)]
 #[serde(rename_all = "kebab-case")]
-#[builder(setter(prefix = "with"))]
+#[builder(build_fn(error = "Error"), setter(prefix = "with"))]
 ///  Partition spec that defines how to produce a tuple of partition values from a record.
 pub struct PartitionSpec {
     /// Identifier for PartitionSpec

iceberg-rust-spec/src/spec/snapshot.rs

@@ -29,7 +29,7 @@ use _serde::SnapshotEnum;
 
 #[derive(Debug, PartialEq, Eq, Clone, Serialize, Deserialize, Builder, Getters)]
 #[serde(from = "SnapshotEnum", into = "SnapshotEnum")]
-#[builder(setter(prefix = "with"))]
+#[builder(build_fn(error = "Error"), setter(prefix = "with"))]
 /// A snapshot represents the state of a table at some time and is used to access the complete set of data files in the table.
 pub struct Snapshot {
     /// A unique long ID

iceberg-rust-spec/src/spec/table_metadata.rs

@@ -204,7 +204,7 @@ impl TableMetadata {
     pub fn current_partition_fields(
         &self,
         branch: Option<&str>,
-    ) -> Result<Vec<BoundPartitionField>, Error> {
+    ) -> Result<Vec<BoundPartitionField<'_>>, Error> {
         let schema = self.current_schema(branch)?;
         let partition_spec = self.default_partition_spec()?;
         partition_fields(partition_spec, schema)
@@ -218,7 +218,10 @@ impl TableMetadata {
     /// # Returns
     /// * `Result<Vec<BoundPartitionField>, Error>` - Vector of partition fields bound to their source schema fields,
     ///   or an error if the schema or partition spec cannot be found
-    pub fn partition_fields(&self, snapshot_id: i64) -> Result<Vec<BoundPartitionField>, Error> {
+    pub fn partition_fields(
+        &self,
+        snapshot_id: i64,
+    ) -> Result<Vec<BoundPartitionField<'_>>, Error> {
         let schema = self.schema(snapshot_id)?;
         self.default_partition_spec()?
             .fields()
@@ -325,7 +328,7 @@ impl TableMetadata {
             .map(|x| *x.sequence_number())
     }
 
-    pub fn as_ref(&self) -> TabularMetadataRef {
+    pub fn as_ref(&self) -> TabularMetadataRef<'_> {
         TabularMetadataRef::Table(self)
     }
 }

iceberg-rust-spec/src/spec/types.rs

@@ -196,6 +196,7 @@ impl fmt::Display for PrimitiveType {
 /// DataType for a specific struct
 #[derive(Debug, Serialize, PartialEq, Eq, Clone, Builder)]
 #[serde(rename = "struct", tag = "type")]
+#[builder(build_fn(error = "Error"))]
 pub struct StructType {
     /// Struct fields
     #[builder(setter(each(name = "with_struct_field")))]