feat: DiskWriter handles writing to arrow part file

Author: Devdutt Shenoi

src/parseable/staging/writer.rs

@@ -19,7 +19,8 @@
 
 use std::{
     collections::{HashMap, HashSet},
-    fs::File,
+    fs::{File, OpenOptions},
+    path::PathBuf,
     sync::Arc,
 };
 
@@ -28,13 +29,48 @@ use arrow_ipc::writer::StreamWriter;
 use arrow_schema::Schema;
 use arrow_select::concat::concat_batches;
 use itertools::Itertools;
+use tracing::error;
 
-use crate::utils::arrow::adapt_batch;
+use crate::{parseable::ARROW_FILE_EXTENSION, utils::arrow::adapt_batch};
+
+use super::StagingError;
 
 #[derive(Default)]
 pub struct Writer {
     pub mem: MemWriter<16384>,
-    pub disk: HashMap<String, StreamWriter<File>>,
+    pub disk: HashMap<String, DiskWriter>,
+}
+
+pub struct DiskWriter {
+    pub inner: StreamWriter<File>,
+    pub path: PathBuf,
+}
+
+impl DiskWriter {
+    pub fn new(path: PathBuf, schema: &Schema) -> Result<Self, StagingError> {
+        let file = OpenOptions::new().create(true).append(true).open(&path)?;
+
+        let inner = StreamWriter::try_new(file, schema)?;
+
+        Ok(Self { inner, path })
+    }
+
+    pub fn write(&mut self, rb: &RecordBatch) -> Result<(), StagingError> {
+        self.inner.write(rb).map_err(StagingError::Arrow)
+    }
+
+    pub fn finish(&mut self) {
+        if let Err(err) = self.inner.finish() {
+            error!("Couldn't finish arrow file {:?}, error = {err}", self.path);
+            return;
+        }
+
+        let mut arrow_path = self.path.to_owned();
+        arrow_path.set_extension(ARROW_FILE_EXTENSION);
+        if let Err(err) = std::fs::rename(&self.path, &arrow_path) {
+            error!("Couldn't rename file {:?}, error = {err}", self.path);
+        }
+    }
 }
 
 /// Structure to keep recordbatches in memory.

src/parseable/streams.rs

@@ -28,7 +28,6 @@ use std::{
 };
 
 use arrow_array::RecordBatch;
-use arrow_ipc::writer::StreamWriter;
 use arrow_schema::{Field, Fields, Schema};
 use chrono::{NaiveDateTime, Timelike};
 use derive_more::{Deref, DerefMut};
@@ -63,7 +62,7 @@ use crate::{
 use super::{
     staging::{
         reader::{MergedRecordReader, MergedReverseRecordReader},
-        writer::Writer,
+        writer::{DiskWriter, Writer},
         StagingError,
     },
     LogStream, ARROW_FILE_EXTENSION,
@@ -143,21 +142,14 @@ impl Stream {
                 }
                 None => {
                     // entry is not present thus we create it
-                    let file_path = self.path_by_current_time(
+                    let path = self.path_by_current_time(
                         schema_key,
                         parsed_timestamp,
                         custom_partition_values,
                     );
                     std::fs::create_dir_all(&self.data_path)?;
 
-                    let file = OpenOptions::new()
-                        .create(true)
-                        .append(true)
-                        .open(&file_path)?;
-
-                    let mut writer = StreamWriter::try_new(file, &record.schema())
-                        .expect("File and RecordBatch both are checked");
-
+                    let mut writer = DiskWriter::new(path, &record.schema())?;
                     writer.write(record)?;
                     guard.disk.insert(schema_key.to_owned(), writer);
                 }
@@ -180,7 +172,7 @@ impl Stream {
             hostname.push_str(id);
         }
         let filename = format!(
-            "{stream_hash}.date={}.hour={:02}.minute={}.{}{hostname}.data.{ARROW_FILE_EXTENSION}",
+            "{stream_hash}.date={}.hour={:02}.minute={}.{}{hostname}.data.part",
             parsed_timestamp.date(),
             parsed_timestamp.hour(),
             minute_to_slot(parsed_timestamp.minute(), OBJECT_STORE_DATA_GRANULARITY).unwrap(),
@@ -391,7 +383,7 @@ impl Stream {
 
         // Flush disk
         for writer in disk_writers.values_mut() {
-            _ = writer.finish();
+            writer.finish();
         }
     }