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

@@ -236,6 +236,11 @@ impl DataFile {
         self.sort_order_id
     }
 
+    /// Get the partition spec id of the data file.
+    pub fn partition_spec_id(&self) -> i32 {
+        self.partition_spec_id
+    }
+
     pub(crate) fn rewrite_partition(&mut self, partition: Struct) {
         self.partition = partition;
     }

crates/iceberg/src/transaction/rewrite_files.rs

@@ -287,11 +287,7 @@ impl SnapshotProduceOperation for RewriteFilesOperation {
                 {
                     let mut manifest_writer = snapshot_produce.new_manifest_writer(
                         &ManifestContentType::Data,
-                        snapshot_produce
-                            .tx
-                            .current_table
-                            .metadata()
-                            .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.5
 ENV ICEBERG_SPARK_RUNTIME_VERSION=3.5_2.12
 ENV ICEBERG_VERSION=1.6.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::Transaction;
 use iceberg::writer::base_writer::data_file_writer::DataFileWriterBuilder;
 use iceberg::writer::file_writer::location_generator::{
@@ -458,3 +460,111 @@ 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 = RestCatalog::new(fixture.catalog_config.clone());
+    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(),
+            table.file_io().clone(),
+            location_generator,
+            file_name_generator,
+        );
+        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"), 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 mut append_action = tx.fast_append(None, None, vec![]).unwrap();
+        append_action.add_data_files(vec![data_file]).unwrap();
+        let tx = append_action.apply().await.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 mut rewrite_action = tx.rewrite_files(None, vec![]).unwrap();
+        rewrite_action = rewrite_action
+            .delete_files(vec![data_file.clone()])
+            .unwrap();
+        let tx = rewrite_action.apply().await.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);
+    }
+}