fix(iceberg): fix rewrite-files partition-spec-id (#54)

* fix(iceberg): fix rewrite-files partition-spec-id

* fix(docker): update docker file

* add test

* update minio

* Revert "update minio"

This reverts commit 4464d90.

Author: Li0k

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

@@ -274,6 +274,12 @@ impl DataFile {
     pub fn referenced_data_file(&self) -> Option<String> {
         self.referenced_data_file.clone()
     }
+
+    /// Get the partition spec id of the data file.
+    pub fn partition_spec_id(&self) -> i32 {
+        self.partition_spec_id
+    }
+
     /// Get the offset in the file where the blob content starts.
     /// Only meaningful for puffin blobs, and required for deletion vectors.
     pub fn content_offset(&self) -> Option<i64> {

crates/iceberg/src/transaction/mod.rs

@@ -183,7 +183,7 @@ impl Transaction {
     }
 
     /// Creates rewrite files action.
-    pub fn rewrite_files(self) -> RewriteFilesAction {
+    pub fn rewrite_files(&self) -> RewriteFilesAction {
         RewriteFilesAction::new()
     }
 

crates/iceberg/src/transaction/rewrite_files.rs

@@ -266,7 +266,7 @@ impl SnapshotProduceOperation for RewriteFilesOperation {
                 {
                     let mut manifest_writer = snapshot_produce.new_manifest_writer(
                         ManifestContentType::Data,
-                        table_metadata_ref.default_partition_spec_id(),
+                        manifest_file.partition_spec_id,
                     )?;
 
                     for entry in manifest.entries() {

crates/integration_tests/testdata/spark/Dockerfile

@@ -31,7 +31,7 @@ ENV SPARK_VERSION=3.5.7
 ENV ICEBERG_SPARK_RUNTIME_VERSION=3.5_2.12
 ENV ICEBERG_VERSION=1.10.0
 
-RUN curl --retry 5 -s -C - https://dlcdn.apache.org/spark/spark-${SPARK_VERSION}/spark-${SPARK_VERSION}-bin-hadoop3.tgz -o spark-${SPARK_VERSION}-bin-hadoop3.tgz \
+RUN curl --retry 5 -s -C - https://archive.apache.org/dist/spark/spark-${SPARK_VERSION}/spark-${SPARK_VERSION}-bin-hadoop3.tgz -o spark-${SPARK_VERSION}-bin-hadoop3.tgz \
  && tar xzf spark-${SPARK_VERSION}-bin-hadoop3.tgz --directory /opt/spark --strip-components 1 \
  && rm -rf spark-${SPARK_VERSION}-bin-hadoop3.tgz
 

crates/integration_tests/tests/shared_tests/rewrite_files_test.rs

@@ -21,6 +21,8 @@ use std::sync::Arc;
 
 use arrow_array::{ArrayRef, BooleanArray, Int32Array, RecordBatch, StringArray};
 use futures::TryStreamExt;
+use iceberg::spec::DataFile;
+use iceberg::table::Table;
 use iceberg::transaction::{ApplyTransactionAction, Transaction};
 use iceberg::writer::base_writer::data_file_writer::DataFileWriterBuilder;
 use iceberg::writer::file_writer::ParquetWriterBuilder;
@@ -453,3 +455,114 @@ async fn test_sequence_number_in_manifest_entry() {
         }
     }
 }
+
+#[tokio::test]
+async fn test_partition_spec_id_in_manifest() {
+    let fixture = get_shared_containers();
+    let rest_catalog = RestCatalogBuilder::default()
+        .load("rest", fixture.catalog_config.clone())
+        .await
+        .unwrap();
+    let ns = random_ns().await;
+    let schema = test_schema();
+
+    let table_creation = TableCreation::builder()
+        .name("t1".to_string())
+        .schema(schema.clone())
+        .build();
+
+    let mut table = rest_catalog
+        .create_table(ns.name(), table_creation)
+        .await
+        .unwrap();
+
+    // Create the writer and write the data
+    let schema: Arc<arrow_schema::Schema> = Arc::new(
+        table
+            .metadata()
+            .current_schema()
+            .as_ref()
+            .try_into()
+            .unwrap(),
+    );
+    let location_generator = DefaultLocationGenerator::new(table.metadata().clone()).unwrap();
+    let file_name_generator = DefaultFileNameGenerator::new(
+        "test".to_string(),
+        None,
+        iceberg::spec::DataFileFormat::Parquet,
+    );
+
+    // commit result
+    let mut data_files_vec = Vec::default();
+
+    async fn build_data_file_f(
+        schema: Arc<arrow_schema::Schema>,
+        table: &Table,
+        location_generator: DefaultLocationGenerator,
+        file_name_generator: DefaultFileNameGenerator,
+    ) -> DataFile {
+        let parquet_writer_builder = ParquetWriterBuilder::new(
+            WriterProperties::default(),
+            table.metadata().current_schema().clone(),
+        );
+        let rolling_writer_builder = RollingFileWriterBuilder::new_with_default_file_size(
+            parquet_writer_builder,
+            table.file_io().clone(),
+            location_generator.clone(),
+            file_name_generator.clone(),
+        );
+        let data_file_writer_builder = DataFileWriterBuilder::new(rolling_writer_builder);
+
+        let mut data_file_writer = data_file_writer_builder.build(None).await.unwrap();
+        let col1 = StringArray::from(vec![Some("foo"), Some("bar"), None, Some("baz")]);
+        let col2 = Int32Array::from(vec![Some(1), Some(2), Some(3), Some(4)]);
+        let col3 = BooleanArray::from(vec![Some(true), Some(false), None, Some(false)]);
+        let batch = RecordBatch::try_new(schema.clone(), vec![
+            Arc::new(col1) as ArrayRef,
+            Arc::new(col2) as ArrayRef,
+            Arc::new(col3) as ArrayRef,
+        ])
+        .unwrap();
+        data_file_writer.write(batch.clone()).await.unwrap();
+        data_file_writer.close().await.unwrap()[0].clone()
+    }
+
+    for _ in 0..10 {
+        let data_file = build_data_file_f(
+            schema.clone(),
+            &table,
+            location_generator.clone(),
+            file_name_generator.clone(),
+        )
+        .await;
+        data_files_vec.push(data_file.clone());
+        let tx = Transaction::new(&table);
+        let append_action = tx.fast_append().add_data_files(vec![data_file]);
+        let tx = append_action.apply(tx).unwrap();
+        table = tx.commit(&rest_catalog).await.unwrap();
+    }
+
+    let last_data_files = data_files_vec.last().unwrap();
+    let partition_spec_id = last_data_files.partition_spec_id();
+
+    // remove the data files by RewriteAction
+    for data_file in &data_files_vec {
+        let tx = Transaction::new(&table);
+        let rewrite_action = tx.rewrite_files().delete_files(vec![data_file.clone()]);
+        let tx = rewrite_action.apply(tx).unwrap();
+        table = tx.commit(&rest_catalog).await.unwrap();
+    }
+
+    // TODO: test update partition spec
+    // Verify that the partition spec ID is correctly set
+
+    let last_snapshot = table.metadata().current_snapshot().unwrap();
+    let manifest_list = last_snapshot
+        .load_manifest_list(table.file_io(), table.metadata())
+        .await
+        .unwrap();
+    assert_eq!(manifest_list.entries().len(), 1);
+    for manifest_file in manifest_list.entries() {
+        assert_eq!(manifest_file.partition_spec_id, partition_spec_id);
+    }
+}