fix test compilation, rebase

Author: CTTY

crates/iceberg/src/spec/partition.rs

@@ -194,6 +194,15 @@ impl PartitionKey {
         Self { spec, schema, data }
     }
 
+    /// Creates a new partition key from another partition key, with a new data field.
+    pub fn copy_with_data(&self, data: Struct) -> Self {
+        Self {
+            spec: self.spec.clone(),
+            schema: self.schema.clone(),
+            data,
+        }
+    }
+
     /// Generates a partition path based on the partition values.
     pub fn to_path(&self) -> String {
         self.spec.partition_to_path(&self.data, self.schema.clone())

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

@@ -143,9 +143,8 @@ mod test {
     use crate::Result;
     use crate::io::FileIOBuilder;
     use crate::spec::{
-        DataContentType, DataFileFormat, Literal, NestedField,
-        PROPERTY_WRITE_TARGET_FILE_SIZE_BYTES_DEFAULT, PartitionKey, PartitionSpec, PrimitiveType,
-        Schema, Struct, Type,
+        DataContentType, DataFileFormat, Literal, NestedField, PartitionKey, PartitionSpec,
+        PrimitiveType, Schema, Struct, Type,
     };
     use crate::writer::base_writer::data_file_writer::DataFileWriterBuilder;
     use crate::writer::file_writer::ParquetWriterBuilder;
@@ -255,7 +254,7 @@ mod test {
         let partition_key = PartitionKey::new(
             PartitionSpec::builder(schema_ref.clone()).build()?,
             schema_ref.clone(),
-            partition_value,
+            partition_value.clone(),
         );
 
         let parquet_writer_builder =

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

@@ -425,8 +425,12 @@ mod test {
         // prepare writer
         let pb =
             ParquetWriterBuilder::new(WriterProperties::builder().build(), Arc::new(delete_schema));
-        let rolling_writer =
-            RollingFileWriter::new_with_default_file_size(pb, file_io, location_gen, file_name_gen);
+        let rolling_writer = RollingFileWriter::new_with_default_file_size(
+            pb,
+            file_io.clone(),
+            location_gen,
+            file_name_gen,
+        );
         let mut equality_delete_writer =
             EqualityDeleteFileWriterBuilder::new(rolling_writer, equality_config)
                 .build()

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

@@ -1624,35 +1624,31 @@ mod tests {
         };
         let col = Arc::new(Int64Array::from_iter_values(0..1024)) as ArrayRef;
         let to_write = RecordBatch::try_new(schema.clone(), vec![col]).unwrap();
-        let output_file = file_io.new_output(
-            location_gen.generate_location(None, &file_name_gen.generate_file_name()),
-        )?;
+        let file_path = location_gen.generate_location(None, &file_name_gen.generate_file_name());
+        let output_file = file_io.new_output(&file_path)?;
         let mut pw = ParquetWriterBuilder::new(
             WriterProperties::builder().build(),
             Arc::new(to_write.schema().as_ref().try_into().unwrap()),
         )
         .build(output_file)
         .await?;
         pw.write(&to_write).await?;
-        let file_path = output_file.location().to_string();
         pw.close().await.unwrap();
-        assert!(file_io.exists(file_path).await.unwrap());
+        assert!(file_io.exists(&file_path).await.unwrap());
 
         // Test that file will not create if no data to write
         let file_name_gen =
             DefaultFileNameGenerator::new("test_empty".to_string(), None, DataFileFormat::Parquet);
-        let output_file = file_io.new_output(
-            location_gen.generate_location(None, &file_name_gen.generate_file_name()),
-        )?;
+        let file_path = location_gen.generate_location(None, &file_name_gen.generate_file_name());
+        let output_file = file_io.new_output(&file_path)?;
         let pw = ParquetWriterBuilder::new(
             WriterProperties::builder().build(),
             Arc::new(to_write.schema().as_ref().try_into().unwrap()),
         )
         .build(output_file)
         .await?;
-        let file_path = output_file.location().to_string();
         pw.close().await.unwrap();
-        assert!(!file_io.exists(file_path).await.unwrap());
+        assert!(!file_io.exists(&file_path).await.unwrap());
 
         Ok(())
     }

crates/integration_tests/tests/shared_tests/append_data_file_test.rs

@@ -27,6 +27,7 @@ use iceberg::writer::file_writer::ParquetWriterBuilder;
 use iceberg::writer::file_writer::location_generator::{
     DefaultFileNameGenerator, DefaultLocationGenerator,
 };
+use iceberg::writer::file_writer::rolling_writer::RollingFileWriter;
 use iceberg::writer::{IcebergWriter, IcebergWriterBuilder};
 use iceberg::{Catalog, CatalogBuilder, TableCreation};
 use iceberg_catalog_rest::RestCatalogBuilder;
@@ -74,12 +75,14 @@ async fn test_append_data_file() {
     let parquet_writer_builder = ParquetWriterBuilder::new(
         WriterProperties::default(),
         table.metadata().current_schema().clone(),
-        None,
+    );
+    let rolling_file_writer = RollingFileWriter::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(parquet_writer_builder, None, 0);
+    let data_file_writer_builder = DataFileWriterBuilder::new(rolling_file_writer, None);
     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)]);

crates/integration_tests/tests/shared_tests/append_partition_data_file_test.rs

@@ -31,6 +31,7 @@ use iceberg::writer::file_writer::ParquetWriterBuilder;
 use iceberg::writer::file_writer::location_generator::{
     DefaultFileNameGenerator, DefaultLocationGenerator,
 };
+use iceberg::writer::file_writer::rolling_writer::RollingFileWriter;
 use iceberg::writer::{IcebergWriter, IcebergWriterBuilder};
 use iceberg::{Catalog, CatalogBuilder, TableCreation};
 use iceberg_catalog_rest::RestCatalogBuilder;
@@ -96,22 +97,20 @@ async fn test_append_partition_data_file() {
     let parquet_writer_builder = ParquetWriterBuilder::new(
         WriterProperties::default(),
         table.metadata().current_schema().clone(),
-        Some(partition_key),
+    );
+
+    let rolling_file_writer = RollingFileWriter::new_with_default_file_size(
+        parquet_writer_builder.clone(),
         table.file_io().clone(),
         location_generator.clone(),
         file_name_generator.clone(),
     );
 
-    let mut data_file_writer_valid = DataFileWriterBuilder::new(
-        parquet_writer_builder.clone(),
-        Some(Struct::from_iter([Some(Literal::Primitive(
-            PrimitiveLiteral::Int(first_partition_id_value),
-        ))])),
-        0,
-    )
-    .build()
-    .await
-    .unwrap();
+    let mut data_file_writer_valid =
+        DataFileWriterBuilder::new(rolling_file_writer.clone(), Some(partition_key.clone()))
+            .build()
+            .await
+            .unwrap();
 
     let col1 = StringArray::from(vec![Some("foo1"), Some("foo2")]);
     let col2 = Int32Array::from(vec![
@@ -148,44 +147,53 @@ async fn test_append_partition_data_file() {
     assert_eq!(batches.len(), 1);
     assert_eq!(batches[0], batch);
 
+    let partition_key = partition_key.copy_with_data(Struct::from_iter([Some(
+        Literal::Primitive(PrimitiveLiteral::Boolean(true)),
+    )]));
     test_schema_incompatible_partition_type(
-        parquet_writer_builder.clone(),
+        rolling_file_writer.clone(),
         batch.clone(),
+        partition_key.clone(),
         table.clone(),
         &rest_catalog,
     )
     .await;
 
+    let partition_key = partition_key.copy_with_data(Struct::from_iter([
+        Some(Literal::Primitive(PrimitiveLiteral::Int(
+            first_partition_id_value,
+        ))),
+        Some(Literal::Primitive(PrimitiveLiteral::Int(
+            first_partition_id_value,
+        ))),
+    ]));
     test_schema_incompatible_partition_fields(
-        parquet_writer_builder,
+        rolling_file_writer.clone(),
         batch,
+        partition_key,
         table,
         &rest_catalog,
-        first_partition_id_value,
     )
     .await;
 }
 
 async fn test_schema_incompatible_partition_type(
-    parquet_writer_builder: ParquetWriterBuilder<
+    rolling_file_writer: RollingFileWriter<
+        ParquetWriterBuilder,
         DefaultLocationGenerator,
         DefaultFileNameGenerator,
     >,
     batch: RecordBatch,
+    partition_key: PartitionKey,
     table: Table,
     catalog: &dyn Catalog,
 ) {
     // test writing different "type" of partition than mentioned in schema
-    let mut data_file_writer_invalid = DataFileWriterBuilder::new(
-        parquet_writer_builder.clone(),
-        Some(Struct::from_iter([Some(Literal::Primitive(
-            PrimitiveLiteral::Boolean(true),
-        ))])),
-        0,
-    )
-    .build()
-    .await
-    .unwrap();
+    let mut data_file_writer_invalid =
+        DataFileWriterBuilder::new(rolling_file_writer.clone(), Some(partition_key))
+            .build()
+            .await
+            .unwrap();
 
     data_file_writer_invalid.write(batch.clone()).await.unwrap();
     let data_file_invalid = data_file_writer_invalid.close().await.unwrap();
@@ -200,32 +208,22 @@ async fn test_schema_incompatible_partition_type(
 }
 
 async fn test_schema_incompatible_partition_fields(
-    parquet_writer_builder: ParquetWriterBuilder<
+    rolling_file_writer: RollingFileWriter<
+        ParquetWriterBuilder,
         DefaultLocationGenerator,
         DefaultFileNameGenerator,
     >,
     batch: RecordBatch,
+    partition_key: PartitionKey,
     table: Table,
     catalog: &dyn Catalog,
-    first_partition_id_value: i32,
 ) {
     // test writing different number of partition fields than mentioned in schema
-
-    let mut data_file_writer_invalid = DataFileWriterBuilder::new(
-        parquet_writer_builder,
-        Some(Struct::from_iter([
-            Some(Literal::Primitive(PrimitiveLiteral::Int(
-                first_partition_id_value,
-            ))),
-            Some(Literal::Primitive(PrimitiveLiteral::Int(
-                first_partition_id_value,
-            ))),
-        ])),
-        0,
-    )
-    .build()
-    .await
-    .unwrap();
+    let mut data_file_writer_invalid =
+        DataFileWriterBuilder::new(rolling_file_writer, Some(partition_key))
+            .build()
+            .await
+            .unwrap();
 
     data_file_writer_invalid.write(batch.clone()).await.unwrap();
     let data_file_invalid = data_file_writer_invalid.close().await.unwrap();

crates/integration_tests/tests/shared_tests/conflict_commit_test.rs

@@ -27,6 +27,7 @@ use iceberg::writer::file_writer::ParquetWriterBuilder;
 use iceberg::writer::file_writer::location_generator::{
     DefaultFileNameGenerator, DefaultLocationGenerator,
 };
+use iceberg::writer::file_writer::rolling_writer::RollingFileWriter;
 use iceberg::writer::{IcebergWriter, IcebergWriterBuilder};
 use iceberg::{Catalog, CatalogBuilder, TableCreation};
 use iceberg_catalog_rest::RestCatalogBuilder;
@@ -73,12 +74,14 @@ async fn test_append_data_file_conflict() {
     let parquet_writer_builder = ParquetWriterBuilder::new(
         WriterProperties::default(),
         table.metadata().current_schema().clone(),
-        None,
+    );
+    let rolling_file_writer = RollingFileWriter::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(parquet_writer_builder, None, 0);
+    let data_file_writer_builder = DataFileWriterBuilder::new(rolling_file_writer, None);
     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)]);

crates/integration_tests/tests/shared_tests/scan_all_type.rs

@@ -39,6 +39,7 @@ use iceberg::writer::file_writer::ParquetWriterBuilder;
 use iceberg::writer::file_writer::location_generator::{
     DefaultFileNameGenerator, DefaultLocationGenerator,
 };
+use iceberg::writer::file_writer::rolling_writer::RollingFileWriter;
 use iceberg::writer::{IcebergWriter, IcebergWriterBuilder};
 use iceberg::{Catalog, CatalogBuilder, TableCreation};
 use iceberg_catalog_rest::RestCatalogBuilder;
@@ -155,12 +156,14 @@ async fn test_scan_all_type() {
     let parquet_writer_builder = ParquetWriterBuilder::new(
         WriterProperties::default(),
         table.metadata().current_schema().clone(),
-        None,
+    );
+    let rolling_file_writer = RollingFileWriter::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(parquet_writer_builder, None, 0);
+    let data_file_writer_builder = DataFileWriterBuilder::new(rolling_file_writer, None);
     let mut data_file_writer = data_file_writer_builder.build().await.unwrap();
 
     // Prepare data