**Refactor Parquet Writer Initialization and File Handling**

 - Updated `ParquetWriter` in `writer.rs` to handle `current_indexer` as an `Option`, allowing for more flexible initialization and management.
 - Introduced `finish_current_file` method to encapsulate logic for completing and transitioning between SST files, improving code clarity and maintainability.
 - Enhanced error handling and logging with `debug` statements for better traceability during file operations.

Author: v0y4g3r

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

@@ -15,11 +15,13 @@
 //! Parquet writer.
 
 use std::future::Future;
+use std::mem;
 use std::pin::Pin;
 use std::sync::atomic::{AtomicUsize, Ordering};
 use std::sync::Arc;
 use std::task::{Context, Poll};
 
+use common_telemetry::debug;
 use common_time::Timestamp;
 use datatypes::arrow::datatypes::SchemaRef;
 use object_store::{FuturesAsyncWriter, ObjectStore};
@@ -59,7 +61,7 @@ pub struct ParquetWriter<F: WriterFactory, I: IndexerBuilder, P: FilePathProvide
     /// Indexer build that can create indexer for multiple files.
     indexer_builder: I,
     /// Current active indexer.
-    current_indexer: Indexer,
+    current_indexer: Option<Indexer>,
     bytes_written: Arc<AtomicUsize>,
 }
 
@@ -121,8 +123,6 @@ where
         path_provider: P,
     ) -> ParquetWriter<F, I, P> {
         let init_file = FileId::random();
-        let index_file_path = path_provider.build_index_file_path(init_file);
-        let indexer = indexer_builder.build(init_file, index_file_path).await;
 
         ParquetWriter {
             path_provider,
@@ -131,22 +131,50 @@ where
             writer_factory: factory,
             metadata,
             indexer_builder,
-            current_indexer: indexer,
+            current_indexer: None,
             bytes_written: Arc::new(AtomicUsize::new(0)),
         }
     }
 
-    async fn roll_to_next_file(&mut self) {
-        self.current_file = FileId::random();
-        let index_file_path = self.path_provider.build_index_file_path(self.current_file);
-        let indexer = self
-            .indexer_builder
-            .build(self.current_file, index_file_path)
-            .await;
-        self.current_indexer = indexer;
-
+    /// Finishes current SST file and index file.
+    async fn finish_current_file(
+        &mut self,
+        ssts: &mut SstInfoArray,
+        stats: &mut SourceStats,
+    ) -> Result<()> {
         // maybe_init_writer will re-create a new file.
-        self.writer = None;
+        if let Some(mut current_writer) = mem::take(&mut self.writer) {
+            let stats = mem::take(stats);
+            // At least one row has been written.
+            assert!(stats.num_rows > 0);
+
+            // 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;
+            current_writer.flush().await.context(WriteParquetSnafu)?;
+
+            let file_meta = current_writer.close().await.context(WriteParquetSnafu)?;
+            let file_size = self.bytes_written.load(Ordering::Relaxed) as u64;
+
+            // Safety: num rows > 0 so we must have min/max.
+            let time_range = stats.time_range.unwrap();
+
+            // convert FileMetaData to ParquetMetaData
+            let parquet_metadata = parse_parquet_metadata(file_meta)?;
+            ssts.push(SstInfo {
+                file_id: self.current_file,
+                time_range,
+                file_size,
+                num_rows: stats.num_rows,
+                num_row_groups: parquet_metadata.num_row_groups() as u64,
+                file_metadata: Some(Arc::new(parquet_metadata)),
+                index_metadata: index_output,
+            });
+            self.current_file = FileId::random();
+            self.bytes_written.store(0, Ordering::Relaxed)
+        };
+
+        Ok(())
     }
 
     /// Iterates source and writes all rows to Parquet file.
@@ -158,6 +186,7 @@ where
         override_sequence: Option<SequenceNumber>, // override the `sequence` field from `Source`
         opts: &WriteOptions,
     ) -> Result<SstInfoArray> {
+        let mut results = smallvec![];
         let write_format =
             WriteFormat::new(self.metadata.clone()).with_override_sequence(override_sequence);
         let mut stats = SourceStats::default();
@@ -170,52 +199,31 @@ where
             match res {
                 Ok(batch) => {
                     stats.update(&batch);
-                    self.current_indexer.update(&batch).await;
+                    // 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 {
-                        self.roll_to_next_file().await;
+                        debug!(
+                            "Finishing current file {}, file size: {}, max file size: {}",
+                            self.current_file,
+                            self.bytes_written.load(Ordering::Relaxed),
+                            opts.max_file_size
+                        );
+                        self.finish_current_file(&mut results, &mut stats).await?;
                     }
                 }
                 Err(e) => {
-                    self.current_indexer.abort().await;
+                    if let Some(indexer) = &mut self.current_indexer {
+                        indexer.abort().await;
+                    }
                     return Err(e);
                 }
             }
         }
 
-        let index_output = self.current_indexer.finish().await;
-
-        if stats.num_rows == 0 {
-            return Ok(smallvec![]);
-        }
-
-        let Some(mut arrow_writer) = self.writer.take() else {
-            // No batch actually written.
-            return Ok(smallvec![]);
-        };
-
-        arrow_writer.flush().await.context(WriteParquetSnafu)?;
-
-        let file_meta = arrow_writer.close().await.context(WriteParquetSnafu)?;
-        let file_size = self.bytes_written.load(Ordering::Relaxed) as u64;
-
-        // Safety: num rows > 0 so we must have min/max.
-        let time_range = stats.time_range.unwrap();
-
-        // convert FileMetaData to ParquetMetaData
-        let parquet_metadata = parse_parquet_metadata(file_meta)?;
-
-        let file_id = self.current_file;
+        self.finish_current_file(&mut results, &mut stats).await?;
 
         // object_store.write will make sure all bytes are written or an error is raised.
-        Ok(smallvec![SstInfo {
-            file_id,
-            time_range,
-            file_size,
-            num_rows: stats.num_rows,
-            num_row_groups: parquet_metadata.num_row_groups() as u64,
-            file_metadata: Some(Arc::new(parquet_metadata)),
-            index_metadata: index_output,
-        }])
+        Ok(results)
     }
 
     /// Customizes per-column config according to schema and maybe column cardinality.
@@ -286,6 +294,14 @@ where
                 AsyncArrowWriter::try_new(writer, schema.clone(), Some(writer_props))
                     .context(WriteParquetSnafu)?;
             self.writer = Some(arrow_writer);
+
+            let index_file_path = self.path_provider.build_index_file_path(self.current_file);
+            let indexer = self
+                .indexer_builder
+                .build(self.current_file, index_file_path)
+                .await;
+            self.current_indexer = Some(indexer);
+
             // safety: self.writer is assigned above
             Ok(self.writer.as_mut().unwrap())
         }