- **Refactor RegionFilePathFactory to RegionFilePathProvider:** Updated references and implementations in access_layer.rs, write_cache.rs, and related test files to use the new struct name.

 - **Add `max_file_size` support in compaction:** Introduced `max_file_size` option in `PickerOutput`, `SerializedPickerOutput`, and `WriteOptions` in `compactor.rs`, `picker.rs`, `twcs.rs`, and `window.rs`.
 - **Enhance Parquet writing logic:** Modified `parquet.rs` and `parquet/writer.rs` to support optional `max_file_size` and added a test case `test_write_multiple_files` to verify writing multiple files based on size constraints.

Author: v0y4g3r

src/mito2/src/access_layer.rs

@@ -135,7 +135,7 @@ impl AccessLayer {
                 .write_and_upload_sst(
                     request,
                     SstUploadRequest {
-                        dest_path_provider: RegionFilePathFactory {
+                        dest_path_provider: RegionFilePathProvider {
                             region_dir: self.region_dir.clone(),
                         },
                         remote_store: self.object_store.clone(),
@@ -161,7 +161,7 @@ impl AccessLayer {
                 self.object_store.clone(),
                 request.metadata,
                 indexer_builder,
-                RegionFilePathFactory {
+                RegionFilePathProvider {
                     region_dir: self.region_dir.clone(),
                 },
             )
@@ -266,11 +266,11 @@ impl FilePathProvider for WriteCachePathProvider {
 
 /// Path provider that builds paths in region storage path.
 #[derive(Clone, Debug)]
-pub(crate) struct RegionFilePathFactory {
+pub(crate) struct RegionFilePathProvider {
     pub(crate) region_dir: String,
 }
 
-impl FilePathProvider for RegionFilePathFactory {
+impl FilePathProvider for RegionFilePathProvider {
     fn build_index_file_path(&self, file_id: FileId) -> String {
         location::index_file_path(&self.region_dir, file_id)
     }

src/mito2/src/cache/write_cache.rs

@@ -24,7 +24,7 @@ use object_store::ObjectStore;
 use snafu::ResultExt;
 
 use crate::access_layer::{
-    new_fs_cache_store, FilePathProvider, RegionFilePathFactory, SstInfoArray, SstWriteRequest,
+    new_fs_cache_store, FilePathProvider, RegionFilePathProvider, SstInfoArray, SstWriteRequest,
     WriteCachePathProvider,
 };
 use crate::cache::file_cache::{FileCache, FileCacheRef, FileType, IndexKey, IndexValue};
@@ -328,7 +328,7 @@ impl WriteCache {
 /// Request to write and upload a SST.
 pub struct SstUploadRequest {
     /// Destination path provider of which SST files in write cache should be uploaded to.
-    pub dest_path_provider: RegionFilePathFactory,
+    pub dest_path_provider: RegionFilePathProvider,
     /// Remote object store to upload.
     pub remote_store: ObjectStore,
 }
@@ -355,7 +355,7 @@ mod tests {
         // and now just use local file system to mock.
         let mut env = TestEnv::new();
         let mock_store = env.init_object_store_manager();
-        let path_provider = RegionFilePathFactory {
+        let path_provider = RegionFilePathProvider {
             region_dir: "test".to_string(),
         };
 
@@ -488,7 +488,7 @@ mod tests {
             ..Default::default()
         };
         let upload_request = SstUploadRequest {
-            dest_path_provider: RegionFilePathFactory {
+            dest_path_provider: RegionFilePathProvider {
                 region_dir: data_home.clone(),
             },
             remote_store: mock_store.clone(),

src/mito2/src/compaction/compactor.rs

@@ -283,6 +283,7 @@ impl Compactor for DefaultCompactor {
 
             let write_opts = WriteOptions {
                 write_buffer_size: compaction_region.engine_config.sst_write_buffer_size,
+                max_file_size: picker_output.max_file_size,
                 ..Default::default()
             };
 

src/mito2/src/compaction/picker.rs

@@ -45,6 +45,8 @@ pub struct PickerOutput {
     pub outputs: Vec<CompactionOutput>,
     pub expired_ssts: Vec<FileHandle>,
     pub time_window_size: i64,
+    /// Max single output file size in bytes.
+    pub max_file_size: Option<usize>,
 }
 
 /// SerializedPickerOutput is a serialized version of PickerOutput by replacing [CompactionOutput] and [FileHandle] with [SerializedCompactionOutput] and [FileMeta].
@@ -53,6 +55,7 @@ pub struct SerializedPickerOutput {
     pub outputs: Vec<SerializedCompactionOutput>,
     pub expired_ssts: Vec<FileMeta>,
     pub time_window_size: i64,
+    pub max_file_size: Option<usize>,
 }
 
 impl From<&PickerOutput> for SerializedPickerOutput {
@@ -76,6 +79,7 @@ impl From<&PickerOutput> for SerializedPickerOutput {
             outputs,
             expired_ssts,
             time_window_size: input.time_window_size,
+            max_file_size: input.max_file_size,
         }
     }
 }
@@ -111,6 +115,7 @@ impl PickerOutput {
             outputs,
             expired_ssts,
             time_window_size: input.time_window_size,
+            max_file_size: input.max_file_size,
         }
     }
 }
@@ -179,6 +184,7 @@ mod tests {
             ],
             expired_ssts: expired_ssts_file_handle.clone(),
             time_window_size: 1000,
+            max_file_size: None,
         };
 
         let picker_output_str =

src/mito2/src/compaction/twcs.rs

@@ -196,10 +196,12 @@ impl Picker for TwcsPicker {
             return None;
         }
 
+        let max_file_size = self.max_output_file_size.map(|v| v as usize);
         Some(PickerOutput {
             outputs,
             expired_ssts,
             time_window_size,
+            max_file_size,
         })
     }
 }

src/mito2/src/compaction/window.rs

@@ -115,6 +115,7 @@ impl Picker for WindowedCompactionPicker {
             outputs,
             expired_ssts,
             time_window_size: time_window,
+            max_file_size: None, // todo (hl): we may need to support `max_file_size` parameter in manual compaction.
         })
     }
 }

src/mito2/src/sst/parquet.rs

@@ -50,15 +50,17 @@ pub struct WriteOptions {
     /// Row group size.
     pub row_group_size: usize,
     /// Max single output file size.
-    pub max_file_size: usize,
+    /// Note: This is not a hard limit as we can only observe the file size when
+    /// ArrowWrite writes to underlying writers.
+    pub max_file_size: Option<usize>,
 }
 
 impl Default for WriteOptions {
     fn default() -> Self {
         WriteOptions {
             write_buffer_size: DEFAULT_WRITE_BUFFER_SIZE,
             row_group_size: DEFAULT_ROW_GROUP_SIZE,
-            max_file_size: usize::MAX,
+            max_file_size: None,
         }
     }
 }
@@ -100,16 +102,18 @@ mod tests {
     use tokio_util::compat::FuturesAsyncWriteCompatExt;
 
     use super::*;
-    use crate::access_layer::FilePathProvider;
+    use crate::access_layer::{FilePathProvider, RegionFilePathProvider};
     use crate::cache::{CacheManager, CacheStrategy, PageKey};
+    use crate::read::BatchReader;
     use crate::sst::index::{Indexer, IndexerBuilder};
     use crate::sst::parquet::format::WriteFormat;
     use crate::sst::parquet::reader::ParquetReaderBuilder;
     use crate::sst::parquet::writer::ParquetWriter;
     use crate::sst::{location, DEFAULT_WRITE_CONCURRENCY};
     use crate::test_util::sst_util::{
         assert_parquet_metadata_eq, build_test_binary_test_region_metadata, new_batch_by_range,
-        new_batch_with_binary, new_source, sst_file_handle, sst_region_metadata,
+        new_batch_with_binary, new_source, sst_file_handle, sst_file_handle_with_file_id,
+        sst_region_metadata,
     };
     use crate::test_util::{check_reader_result, TestEnv};
 
@@ -539,4 +543,58 @@ mod tests {
         )
         .await;
     }
+
+    #[tokio::test]
+    async fn test_write_multiple_files() {
+        common_telemetry::init_default_ut_logging();
+        // create test env
+        let mut env = TestEnv::new();
+        let object_store = env.init_object_store_manager();
+        let metadata = Arc::new(sst_region_metadata());
+        let batches = &[
+            new_batch_by_range(&["a", "d"], 0, 1000),
+            new_batch_by_range(&["b", "f"], 0, 1000),
+            new_batch_by_range(&["b", "h"], 100, 200),
+            new_batch_by_range(&["b", "h"], 200, 300),
+            new_batch_by_range(&["b", "h"], 300, 1000),
+        ];
+        let total_rows: usize = batches.iter().map(|batch| batch.num_rows()).sum();
+
+        let source = new_source(batches);
+        let write_opts = WriteOptions {
+            row_group_size: 50,
+            max_file_size: Some(1024 * 16),
+            ..Default::default()
+        };
+
+        let path_provider = RegionFilePathProvider {
+            region_dir: "test".to_string(),
+        };
+        let mut writer = ParquetWriter::new_with_object_store(
+            object_store.clone(),
+            metadata.clone(),
+            NoopIndexBuilder,
+            path_provider,
+        )
+        .await;
+
+        let files = writer.write_all(source, None, &write_opts).await.unwrap();
+        assert_eq!(2, files.len());
+
+        let mut rows_read = 0;
+        for f in &files {
+            let file_handle = sst_file_handle_with_file_id(
+                f.file_id,
+                f.time_range.0.value(),
+                f.time_range.1.value(),
+            );
+            let builder =
+                ParquetReaderBuilder::new("test".to_string(), file_handle, object_store.clone());
+            let mut reader = builder.build().await.unwrap();
+            while let Some(batch) = reader.next_batch().await.unwrap() {
+                rows_read += batch.num_rows();
+            }
+        }
+        assert_eq!(total_rows, rows_read);
+    }
 }

src/mito2/src/sst/parquet/writer.rs

@@ -148,6 +148,13 @@ where
             // At least one row has been written.
             assert!(stats.num_rows > 0);
 
+            debug!(
+                "Finishing current file {}, file size: {}, num rows: {}",
+                self.current_file,
+                self.bytes_written.load(Ordering::Relaxed),
+                stats.num_rows
+            );
+
             // Finish indexer and writer.
             // safety: writer and index can only be both present or not.
             let index_output = self.current_indexer.as_mut().unwrap().finish().await;
@@ -201,13 +208,9 @@ where
                     stats.update(&batch);
                     // safety: self.current_indexer must be set when first batch has been written.
                     self.current_indexer.as_mut().unwrap().update(&batch).await;
-                    if self.bytes_written.load(Ordering::Relaxed) > opts.max_file_size {
-                        debug!(
-                            "Finishing current file {}, file size: {}, max file size: {}",
-                            self.current_file,
-                            self.bytes_written.load(Ordering::Relaxed),
-                            opts.max_file_size
-                        );
+                    if let Some(max_file_size) = opts.max_file_size
+                        && self.bytes_written.load(Ordering::Relaxed) > max_file_size
+                    {
                         self.finish_current_file(&mut results, &mut stats).await?;
                     }
                 }