feat: add process delete enrty in snapshot produce (#33)

* support spec id in data file

* support proccess delete entry

* fullfill partition spec id

* fix

* fix spelling mistake

Co-authored-by: Copilot <[email protected]>

---------

Co-authored-by: ZENOTME <[email protected]>
Co-authored-by: Copilot <[email protected]>
Signed-off-by: xxchan <[email protected]>

Author: 2 people

crates/iceberg/src/arrow/record_batch_partition_spliter.rs

@@ -27,7 +27,7 @@ use parquet::arrow::PARQUET_FIELD_ID_META_KEY;
 
 use super::record_batch_projector::RecordBatchProjector;
 use crate::arrow::{arrow_struct_to_literal, type_to_arrow_type};
-use crate::spec::{Literal, PartitionSpecRef, SchemaRef, Struct, StructType, Type};
+use crate::spec::{Literal, PartitionSpec, PartitionSpecRef, SchemaRef, Struct, StructType, Type};
 use crate::transform::{create_transform_function, BoxedTransformFunction};
 use crate::{Error, ErrorKind, Result};
 
@@ -186,6 +186,10 @@ impl RecordBatchPartitionSpliter {
         })
     }
 
+    pub(crate) fn partition_spec(&self) -> &PartitionSpec {
+        self.partition_spec.as_ref()
+    }
+
     /// Split the record batch into multiple record batches by the partition spec.
     pub(crate) fn split(&self, batch: &RecordBatch) -> Result<Vec<(OwnedRow, RecordBatch)>> {
         // get array using partition spec

crates/iceberg/src/spec/manifest/data_file.rs

@@ -239,6 +239,10 @@ impl DataFile {
     pub(crate) fn rewrite_partition(&mut self, partition: Struct) {
         self.partition = partition;
     }
+
+    pub(crate) fn rewrite_partition_id(&mut self, partition_spec_id: i32) {
+        self.partition_spec_id = partition_spec_id;
+    }
 }
 
 /// Convert data files to avro bytes and write to writer.

crates/iceberg/src/transaction/append.rs

@@ -310,6 +310,7 @@ mod tests {
 
         // check add data file with incompatible partition value
         let data_file = DataFileBuilder::default()
+            .partition_spec_id(0)
             .content(DataContentType::Data)
             .file_path("test/3.parquet".to_string())
             .file_format(DataFileFormat::Parquet)
@@ -322,6 +323,7 @@ mod tests {
         assert!(action.add_data_files(vec![data_file.clone()]).is_err());
 
         let data_file = DataFileBuilder::default()
+            .partition_spec_id(0)
             .content(DataContentType::Data)
             .file_path("test/3.parquet".to_string())
             .file_format(DataFileFormat::Parquet)

crates/iceberg/src/transaction/snapshot.rs

@@ -168,6 +168,7 @@ impl<'a> SnapshotProduceAction<'a> {
     fn new_manifest_writer(
         &mut self,
         content_type: &ManifestContentType,
+        partition_spec_id: i32,
     ) -> Result<ManifestWriter> {
         let new_manifest_path = format!(
             "{}/{}/{}-m{}.{}",
@@ -182,17 +183,26 @@ impl<'a> SnapshotProduceAction<'a> {
             .current_table
             .file_io()
             .new_output(new_manifest_path)?;
+        let partition_spec = self
+            .tx
+            .current_table
+            .metadata()
+            .partition_spec_by_id(partition_spec_id)
+            .ok_or_else(|| {
+                Error::new(
+                    ErrorKind::DataInvalid,
+                    "Invalid partition spec id for new manifest writer",
+                )
+                .with_context("partition spec id", partition_spec_id.to_string())
+            })?
+            .as_ref()
+            .clone();
         let builder = ManifestWriterBuilder::new(
             output,
             Some(self.snapshot_id),
             self.key_metadata.clone(),
             self.tx.current_table.metadata().current_schema().clone(),
-            self.tx
-                .current_table
-                .metadata()
-                .default_partition_spec()
-                .as_ref()
-                .clone(),
+            partition_spec,
         );
         if self.tx.current_table.metadata().format_version() == FormatVersion::V1 {
             Ok(builder.build_v1())
@@ -244,29 +254,95 @@ impl<'a> SnapshotProduceAction<'a> {
                 builder.build()
             }
         });
-        let mut writer = self.new_manifest_writer(&content_type)?;
+        let mut writer = self.new_manifest_writer(
+            &content_type,
+            self.tx.current_table.metadata().default_partition_spec_id(),
+        )?;
         for entry in manifest_entries {
             writer.add_entry(entry)?;
         }
         writer.write_manifest_file().await
     }
 
+    async fn write_delete_manifest(
+        &mut self,
+        deleted_entries: Vec<ManifestEntry>,
+    ) -> Result<Vec<ManifestFile>> {
+        if deleted_entries.is_empty() {
+            return Ok(vec![]);
+        }
+
+        // Group deleted entries by spec_id
+        let mut partition_groups = HashMap::new();
+        for entry in deleted_entries {
+            partition_groups
+                .entry(entry.data_file().partition_spec_id)
+                .or_insert_with(Vec::new)
+                .push(entry);
+        }
+
+        // Write a delete manifest per spec_id group
+        let mut deleted_manifests = Vec::new();
+        for (spec_id, entries) in partition_groups {
+            let mut data_file_writer: Option<ManifestWriter> = None;
+            let mut delete_file_writer: Option<ManifestWriter> = None;
+            for entry in entries {
+                match entry.content_type() {
+                    DataContentType::Data => {
+                        if data_file_writer.is_none() {
+                            data_file_writer = Some(
+                                self.new_manifest_writer(&ManifestContentType::Data, spec_id)?,
+                            );
+                        }
+                        data_file_writer.as_mut().unwrap().add_delete_entry(entry)?;
+                    }
+                    DataContentType::EqualityDeletes | DataContentType::PositionDeletes => {
+                        if delete_file_writer.is_none() {
+                            delete_file_writer = Some(
+                                self.new_manifest_writer(&ManifestContentType::Deletes, spec_id)?,
+                            );
+                        }
+                        delete_file_writer
+                            .as_mut()
+                            .unwrap()
+                            .add_delete_entry(entry)?;
+                    }
+                }
+            }
+            if let Some(writer) = data_file_writer {
+                deleted_manifests.push(writer.write_manifest_file().await?);
+            }
+            if let Some(writer) = delete_file_writer {
+                deleted_manifests.push(writer.write_manifest_file().await?);
+            }
+        }
+
+        Ok(deleted_manifests)
+    }
+
     async fn manifest_file<OP: SnapshotProduceOperation, MP: ManifestProcess>(
         &mut self,
         snapshot_produce_operation: &OP,
         manifest_process: &MP,
     ) -> Result<Vec<ManifestFile>> {
         let mut manifest_files = vec![];
         let data_files = std::mem::take(&mut self.added_data_files);
-        let delete_files = std::mem::take(&mut self.added_delete_files);
+        let added_delete_files = std::mem::take(&mut self.added_delete_files);
         if !data_files.is_empty() {
             let added_manifest = self.write_added_manifest(data_files).await?;
             manifest_files.push(added_manifest);
         }
-        if !delete_files.is_empty() {
-            let added_delete_manifest = self.write_added_manifest(delete_files).await?;
+
+        if !added_delete_files.is_empty() {
+            let added_delete_manifest = self.write_added_manifest(added_delete_files).await?;
             manifest_files.push(added_delete_manifest);
         }
+
+        let delete_manifests = self
+            .write_delete_manifest(snapshot_produce_operation.delete_entries(self).await?)
+            .await?;
+        manifest_files.extend(delete_manifests);
+
         let existing_manifests = snapshot_produce_operation.existing_manifest(self).await?;
 
         manifest_files.extend(existing_manifests);
@@ -505,7 +581,7 @@ impl MergeManifestManager {
                             Box<dyn Future<Output = Result<Vec<ManifestFile>>> + Send>,
                         >)
                 } else {
-                    let writer = snapshot_produce.new_manifest_writer(&self.content)?;
+                    let writer = snapshot_produce.new_manifest_writer(&self.content,snapshot_produce.tx.current_table.metadata().default_partition_spec_id())?;
                     let snapshot_id = snapshot_produce.snapshot_id;
                     let file_io = snapshot_produce.tx.current_table.file_io().clone();
                     Ok((Box::pin(async move {

crates/iceberg/src/writer/base_writer/sort_position_delete_writer.rs

@@ -35,15 +35,22 @@ pub struct SortPositionDeleteWriterBuilder<B: FileWriterBuilder> {
     inner: B,
     cache_num: usize,
     partition_value: Option<Struct>,
+    partition_spec_id: Option<i32>,
 }
 
 impl<B: FileWriterBuilder> SortPositionDeleteWriterBuilder<B> {
     /// Create a new `SortPositionDeleteWriterBuilder` using a `FileWriterBuilder`.
-    pub fn new(inner: B, cache_num: usize, partition_value: Option<Struct>) -> Self {
+    pub fn new(
+        inner: B,
+        cache_num: usize,
+        partition_value: Option<Struct>,
+        partition_spec_id: Option<i32>,
+    ) -> Self {
         Self {
             inner,
             cache_num,
             partition_value,
+            partition_spec_id,
         }
     }
 }
@@ -86,6 +93,7 @@ impl<B: FileWriterBuilder> IcebergWriterBuilder<PositionDeleteInput, Vec<DataFil
             cache: BTreeMap::new(),
             data_files: Vec::new(),
             partition_value: self.partition_value.unwrap_or(Struct::empty()),
+            partition_spec_id: self.partition_spec_id.unwrap_or(0),
         })
     }
 }
@@ -106,6 +114,7 @@ pub struct SortPositionDeleteWriter<B: FileWriterBuilder> {
     cache: BTreeMap<String, Vec<i64>>,
     data_files: Vec<DataFile>,
     partition_value: Struct,
+    partition_spec_id: i32,
 }
 
 impl<B: FileWriterBuilder> SortPositionDeleteWriter<B> {
@@ -140,6 +149,7 @@ impl<B: FileWriterBuilder> SortPositionDeleteWriter<B> {
             .extend(writer.close().await?.into_iter().map(|mut res| {
                 res.content(crate::spec::DataContentType::PositionDeletes);
                 res.partition(self.partition_value.clone());
+                res.partition_spec_id(self.partition_spec_id);
                 res.build().expect("Guaranteed to be valid")
             }));
         Ok(())
@@ -204,7 +214,7 @@ mod test {
             location_gen,
             file_name_gen,
         );
-        let mut position_delete_writer = SortPositionDeleteWriterBuilder::new(pw, 10, None)
+        let mut position_delete_writer = SortPositionDeleteWriterBuilder::new(pw, 10, None, None)
             .build()
             .await?;
 

crates/iceberg/src/writer/file_writer/parquet_writer.rs

@@ -438,6 +438,7 @@ impl ParquetWriter {
             .file_path(file_path)
             .file_format(DataFileFormat::Parquet)
             .partition(Struct::empty())
+            .partition_spec_id(0)
             .record_count(metadata.file_metadata().num_rows() as u64)
             .file_size_in_bytes(written_size as u64)
             .column_sizes(column_sizes)
@@ -1510,7 +1511,7 @@ mod tests {
         //     .next()
         //     .unwrap()
         //     .content(crate::spec::DataContentType::Data)
-        //     .partition(Struct::empty())
+        //     .partition(Struct::empty()).partition_spec_id(0)
         //     .build()
         //     .unwrap();
         // assert_eq!(

crates/iceberg/src/writer/function_writer/equality_delta_writer.rs

@@ -318,7 +318,7 @@ mod test {
                 location_gen.clone(),
                 file_name_gen.clone(),
             );
-            DataFileWriterBuilder::new(pw.clone(), None)
+            DataFileWriterBuilder::new(pw.clone(), None, None)
         };
         let position_delete_writer_builder = {
             let pw = ParquetWriterBuilder::new(
@@ -328,10 +328,10 @@ mod test {
                 location_gen.clone(),
                 file_name_gen.clone(),
             );
-            SortPositionDeleteWriterBuilder::new(pw.clone(), 100, None)
+            SortPositionDeleteWriterBuilder::new(pw.clone(), 100, None, None)
         };
         let equality_delete_writer_builder = {
-            let config = EqualityDeleteWriterConfig::new(vec![1, 2], schema, None)?;
+            let config = EqualityDeleteWriterConfig::new(vec![1, 2], schema, None, None)?;
             let pw = ParquetWriterBuilder::new(
                 WriterProperties::builder().build(),
                 arrow_schema_to_schema(config.projected_arrow_schema_ref())

crates/iceberg/src/writer/function_writer/fanout_partition_writer.rs

@@ -134,6 +134,7 @@ impl<B: IcebergWriterBuilder> IcebergWriter for FanoutPartitionWriter<B> {
             let mut data_files = writer.close().await?;
             for data_file in data_files.iter_mut() {
                 data_file.rewrite_partition(partition_value.clone());
+                data_file.rewrite_partition_id(self.partition_splitter.partition_spec().spec_id());
             }
             result.append(&mut data_files);
         }
@@ -211,7 +212,7 @@ mod test {
             location_gen,
             file_name_gen,
         );
-        let data_file_writer_builder = DataFileWriterBuilder::new(pw, None);
+        let data_file_writer_builder = DataFileWriterBuilder::new(pw, None, 0);
         let mut fanout_partition_writer = FanoutPartitionWriterBuilder::new(
             data_file_writer_builder,
             Arc::new(partition_spec),

crates/iceberg/src/writer/function_writer/precompute_partition_writer.rs

@@ -207,7 +207,7 @@ mod test {
             location_gen,
             file_name_gen,
         );
-        let data_file_writer_builder = DataFileWriterBuilder::new(pw, None);
+        let data_file_writer_builder = DataFileWriterBuilder::new(pw, None, 0);
         let mut precompute_partition_writer = PrecomputePartitionWriterBuilder::new(
             data_file_writer_builder,
             Arc::new(partition_spec),

crates/integration_tests/tests/shared_tests/merge_append_test.rs

@@ -63,7 +63,7 @@ async fn write_new_data_file(table: &Table) -> Vec<DataFile> {
         location_generator.clone(),
         file_name_generator.clone(),
     );
-    let data_file_writer_builder = DataFileWriterBuilder::new(parquet_writer_builder, None);
+    let data_file_writer_builder = DataFileWriterBuilder::new(parquet_writer_builder, None, 0);
     let mut data_file_writer = data_file_writer_builder.build().await.unwrap();
     let col1 = StringArray::from(vec![Some("foo"); 100]);
     let col2 = Int32Array::from(vec![Some(1); 100]);