fix: don't flush current data

Author: Devdutt Shenoi

src/parseable/staging/reader.rs

@@ -335,11 +335,16 @@ mod tests {
         write_message, DictionaryTracker, IpcDataGenerator, IpcWriteOptions, StreamWriter,
     };
     use arrow_schema::{DataType, Field, Schema};
+    use chrono::Utc;
     use temp_dir::TempDir;
 
-    use crate::parseable::staging::{
-        reader::{MergedReverseRecordReader, OffsetReader},
-        writer::DiskWriter,
+    use crate::{
+        parseable::staging::{
+            reader::{MergedReverseRecordReader, OffsetReader},
+            writer::DiskWriter,
+        },
+        utils::time::TimeRange,
+        OBJECT_STORE_DATA_GRANULARITY,
     };
 
     use super::get_reverse_reader;
@@ -484,7 +489,9 @@ mod tests {
         schema: &Arc<Schema>,
         batches: &[RecordBatch],
     ) -> io::Result<()> {
-        let mut writer = DiskWriter::try_new(path, schema).expect("Failed to create StreamWriter");
+        let range = TimeRange::granularity_range(Utc::now(), OBJECT_STORE_DATA_GRANULARITY);
+        let mut writer =
+            DiskWriter::try_new(path, schema, range).expect("Failed to create StreamWriter");
 
         for batch in batches {
             writer.write(batch).expect("Failed to write batch");

src/parseable/staging/writer.rs

@@ -29,12 +29,13 @@ use arrow_array::RecordBatch;
 use arrow_ipc::writer::StreamWriter;
 use arrow_schema::Schema;
 use arrow_select::concat::concat_batches;
+use chrono::Utc;
 use itertools::Itertools;
 use tracing::{error, warn};
 
 use crate::{
     parseable::{ARROW_FILE_EXTENSION, PART_FILE_EXTENSION},
-    utils::arrow::adapt_batch,
+    utils::{arrow::adapt_batch, time::TimeRange},
 };
 
 use super::StagingError;
@@ -48,11 +49,16 @@ pub struct Writer {
 pub struct DiskWriter {
     inner: StreamWriter<BufWriter<File>>,
     path: PathBuf,
+    range: TimeRange,
 }
 
 impl DiskWriter {
     /// Try to create a file to stream arrows into
-    pub fn try_new(path: impl Into<PathBuf>, schema: &Schema) -> Result<Self, StagingError> {
+    pub fn try_new(
+        path: impl Into<PathBuf>,
+        schema: &Schema,
+        range: TimeRange,
+    ) -> Result<Self, StagingError> {
         let mut path = path.into();
         path.set_extension(PART_FILE_EXTENSION);
         let file = OpenOptions::new()
@@ -62,7 +68,11 @@ impl DiskWriter {
             .open(&path)?;
         let inner = StreamWriter::try_new_buffered(file, schema)?;
 
-        Ok(Self { inner, path })
+        Ok(Self { inner, path, range })
+    }
+
+    pub fn is_current(&self) -> bool {
+        self.range.contains(Utc::now())
     }
 
     /// Write a single recordbatch into file

src/parseable/streams.rs

@@ -29,7 +29,7 @@ use std::{
 
 use arrow_array::RecordBatch;
 use arrow_schema::{Field, Fields, Schema};
-use chrono::{NaiveDateTime, Timelike};
+use chrono::{NaiveDateTime, Timelike, Utc};
 use derive_more::{Deref, DerefMut};
 use itertools::Itertools;
 use parquet::{
@@ -51,7 +51,7 @@ use crate::{
     metrics,
     option::Mode,
     storage::{object_storage::to_bytes, retention::Retention, StreamType},
-    utils::time::Minute,
+    utils::time::{Minute, TimeRange},
     LOCK_EXPECT, OBJECT_STORE_DATA_GRANULARITY,
 };
 
@@ -132,7 +132,11 @@ impl Stream {
                     );
                     std::fs::create_dir_all(&self.data_path)?;
 
-                    let mut writer = DiskWriter::try_new(file_path, &record.schema())
+                    let range = TimeRange::granularity_range(
+                        parsed_timestamp.and_local_timezone(Utc).unwrap(),
+                        OBJECT_STORE_DATA_GRANULARITY,
+                    );
+                    let mut writer = DiskWriter::try_new(file_path, &record.schema(), range)
                         .expect("File and RecordBatch both are checked");
 
                     writer.write(record)?;
@@ -357,12 +361,12 @@ impl Stream {
         self.writer.lock().unwrap().mem.clear();
     }
 
-    pub fn flush(&self) {
+    pub fn flush(&self, forced: bool) {
         let mut writer = self.writer.lock().unwrap();
         // Flush memory
         writer.mem.clear();
         // Drop schema -> disk writer mapping, triggers flush to disk
-        writer.disk.drain();
+        writer.disk.retain(|_, w| !forced && w.is_current());
     }
 
     fn parquet_writer_props(
@@ -662,7 +666,7 @@ impl Stream {
 
     /// First flushes arrows onto disk and then converts the arrow into parquet
     pub fn flush_and_convert(&self, shutdown_signal: bool) -> Result<(), StagingError> {
-        self.flush();
+        self.flush(shutdown_signal);
 
         self.prepare_parquet(shutdown_signal)
     }
@@ -974,7 +978,7 @@ mod tests {
                 StreamType::UserDefined,
             )
             .unwrap();
-        staging.flush();
+        staging.flush(true);
     }
 
     #[test]