rolls the unit tests

Author: CTTY

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

@@ -149,3 +149,192 @@ impl<B: IcebergWriterBuilder> RollingFileWriter for RollingDataFileWriter<B> {
         self.written_size + input_size > self.target_size
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use std::collections::HashMap;
+    use std::sync::Arc;
+
+    use arrow_array::{Int32Array, StringArray};
+    use arrow_schema::{DataType, Field, Schema as ArrowSchema};
+    use parquet::arrow::PARQUET_FIELD_ID_META_KEY;
+    use parquet::file::properties::WriterProperties;
+    use tempfile::TempDir;
+
+    use super::*;
+    use crate::io::FileIOBuilder;
+    use crate::spec::{DataFileFormat, NestedField, PrimitiveType, Schema, Type};
+    use crate::writer::base_writer::data_file_writer::DataFileWriterBuilder;
+    use crate::writer::file_writer::ParquetWriterBuilder;
+    use crate::writer::file_writer::location_generator::DefaultFileNameGenerator;
+    use crate::writer::file_writer::location_generator::test::MockLocationGenerator;
+    use crate::writer::tests::check_parquet_data_file;
+    use crate::writer::{IcebergWriter, IcebergWriterBuilder, RecordBatch};
+
+    #[tokio::test]
+    async fn test_rolling_writer_basic() -> Result<()> {
+        let temp_dir = TempDir::new().unwrap();
+        let file_io = FileIOBuilder::new_fs_io().build().unwrap();
+        let location_gen =
+            MockLocationGenerator::new(temp_dir.path().to_str().unwrap().to_string());
+        let file_name_gen =
+            DefaultFileNameGenerator::new("test".to_string(), None, DataFileFormat::Parquet);
+
+        // Create schema
+        let schema = Schema::builder()
+            .with_schema_id(1)
+            .with_fields(vec![
+                NestedField::required(1, "id", Type::Primitive(PrimitiveType::Int)).into(),
+                NestedField::required(2, "name", Type::Primitive(PrimitiveType::String)).into(),
+            ])
+            .build()?;
+
+        // Create writer builders
+        let parquet_writer_builder = ParquetWriterBuilder::new(
+            WriterProperties::builder().build(),
+            Arc::new(schema),
+            file_io.clone(),
+            location_gen,
+            file_name_gen,
+        );
+        let data_file_writer_builder = DataFileWriterBuilder::new(parquet_writer_builder, None, 0);
+
+        // Set a large target size so no rolling occurs
+        let rolling_writer_builder = RollingDataFileWriterBuilder::new(
+            data_file_writer_builder,
+            1024 * 1024, // 1MB, large enough to not trigger rolling
+        );
+
+        // Create writer
+        let mut writer = rolling_writer_builder.build().await?;
+
+        // Create test data
+        let arrow_schema = ArrowSchema::new(vec![
+            Field::new("id", DataType::Int32, false).with_metadata(HashMap::from([(
+                PARQUET_FIELD_ID_META_KEY.to_string(),
+                1.to_string(),
+            )])),
+            Field::new("name", DataType::Utf8, false).with_metadata(HashMap::from([(
+                PARQUET_FIELD_ID_META_KEY.to_string(),
+                2.to_string(),
+            )])),
+        ]);
+
+        let batch = RecordBatch::try_new(Arc::new(arrow_schema), vec![
+            Arc::new(Int32Array::from(vec![1, 2, 3])),
+            Arc::new(StringArray::from(vec!["Alice", "Bob", "Charlie"])),
+        ])?;
+
+        // Write data
+        writer.write(batch.clone()).await?;
+
+        // Close writer and get data files
+        let data_files = writer.close().await?;
+
+        // Verify only one file was created
+        assert_eq!(
+            data_files.len(),
+            1,
+            "Expected only one data file to be created"
+        );
+
+        // Verify file content
+        check_parquet_data_file(&file_io, &data_files[0], &batch).await;
+
+        Ok(())
+    }
+
+    #[tokio::test]
+    async fn test_rolling_writer_with_rolling() -> Result<()> {
+        let temp_dir = TempDir::new().unwrap();
+        let file_io = FileIOBuilder::new_fs_io().build().unwrap();
+        let location_gen =
+            MockLocationGenerator::new(temp_dir.path().to_str().unwrap().to_string());
+        let file_name_gen =
+            DefaultFileNameGenerator::new("test".to_string(), None, DataFileFormat::Parquet);
+
+        // Create schema
+        let schema = Schema::builder()
+            .with_schema_id(1)
+            .with_fields(vec![
+                NestedField::required(1, "id", Type::Primitive(PrimitiveType::Int)).into(),
+                NestedField::required(2, "name", Type::Primitive(PrimitiveType::String)).into(),
+            ])
+            .build()?;
+
+        // Create writer builders
+        let parquet_writer_builder = ParquetWriterBuilder::new(
+            WriterProperties::builder().build(),
+            Arc::new(schema),
+            file_io.clone(),
+            location_gen,
+            file_name_gen,
+        );
+        let data_file_writer_builder = DataFileWriterBuilder::new(parquet_writer_builder, None, 0);
+
+        // Set a very small target size to trigger rolling
+        let rolling_writer_builder = RollingDataFileWriterBuilder::new(
+            data_file_writer_builder,
+            100, // Very small target size to ensure rolling
+        );
+
+        // Create writer
+        let mut writer = rolling_writer_builder.build().await?;
+
+        // Create test data
+        let arrow_schema = ArrowSchema::new(vec![
+            Field::new("id", DataType::Int32, false).with_metadata(HashMap::from([(
+                PARQUET_FIELD_ID_META_KEY.to_string(),
+                1.to_string(),
+            )])),
+            Field::new("name", DataType::Utf8, false).with_metadata(HashMap::from([(
+                PARQUET_FIELD_ID_META_KEY.to_string(),
+                2.to_string(),
+            )])),
+        ]);
+
+        // Create multiple batches to trigger rolling
+        let batch1 = RecordBatch::try_new(Arc::new(arrow_schema.clone()), vec![
+            Arc::new(Int32Array::from(vec![1, 2, 3])),
+            Arc::new(StringArray::from(vec!["Alice", "Bob", "Charlie"])),
+        ])?;
+
+        let batch2 = RecordBatch::try_new(Arc::new(arrow_schema.clone()), vec![
+            Arc::new(Int32Array::from(vec![4, 5, 6])),
+            Arc::new(StringArray::from(vec!["Dave", "Eve", "Frank"])),
+        ])?;
+
+        let batch3 = RecordBatch::try_new(Arc::new(arrow_schema), vec![
+            Arc::new(Int32Array::from(vec![7, 8, 9])),
+            Arc::new(StringArray::from(vec!["Grace", "Heidi", "Ivan"])),
+        ])?;
+
+        // Write data
+        writer.write(batch1.clone()).await?;
+        writer.write(batch2.clone()).await?;
+        writer.write(batch3.clone()).await?;
+
+        // Close writer and get data files
+        let data_files = writer.close().await?;
+
+        // Verify multiple files were created (at least 2)
+        assert!(
+            data_files.len() > 1,
+            "Expected multiple data files to be created, got {}",
+            data_files.len()
+        );
+
+        // Verify total record count across all files
+        let total_records: u64 = data_files.iter().map(|file| file.record_count).sum();
+        assert_eq!(
+            total_records, 9,
+            "Expected 9 total records across all files"
+        );
+
+        // Verify each file has the correct content
+        // Note: We can't easily verify which records went to which file without more complex logic,
+        // but we can verify the total count and that each file has valid content
+
+        Ok(())
+    }
+}