test: use DiskWriter

Author: Devdutt Shenoi

src/parseable/mod.rs

@@ -61,6 +61,9 @@ mod streams;
 /// File extension for arrow files in staging
 const ARROW_FILE_EXTENSION: &str = "arrows";
 
+/// File extension for incomplete arrow files
+const PART_FILE_EXTENSION: &str = "part";
+
 /// Name of a Stream
 /// NOTE: this used to be a struct, flattened out for simplicity
 pub type LogStream = String;

src/parseable/staging/reader.rs

@@ -322,7 +322,6 @@ fn find_limit_and_type(
 #[cfg(test)]
 mod tests {
     use std::{
-        fs::File,
         io::{self, Cursor, Read},
         path::Path,
         sync::Arc,
@@ -338,7 +337,10 @@ mod tests {
     use arrow_schema::{DataType, Field, Schema};
     use temp_dir::TempDir;
 
-    use crate::parseable::staging::reader::{MergedReverseRecordReader, OffsetReader};
+    use crate::parseable::staging::{
+        reader::{MergedReverseRecordReader, OffsetReader},
+        writer::DiskWriter,
+    };
 
     use super::get_reverse_reader;
 
@@ -482,15 +484,12 @@ mod tests {
         schema: &Arc<Schema>,
         batches: &[RecordBatch],
     ) -> io::Result<()> {
-        let file = File::create(path)?;
-        let mut writer =
-            StreamWriter::try_new(file, schema).expect("Failed to create StreamWriter");
+        let mut writer = DiskWriter::try_new(path, schema).expect("Failed to create StreamWriter");
 
         for batch in batches {
             writer.write(batch).expect("Failed to write batch");
         }
 
-        writer.finish().expect("Failed to finalize writer");
         Ok(())
     }
 
@@ -524,7 +523,7 @@ mod tests {
     #[test]
     fn test_merged_reverse_record_reader() -> io::Result<()> {
         let dir = TempDir::new().unwrap();
-        let file_path = dir.path().join("test.arrow");
+        let file_path = dir.path().join("test.data.arrows");
 
         // Create a schema
         let schema = Arc::new(Schema::new(vec![
@@ -627,7 +626,7 @@ mod tests {
     #[test]
     fn test_get_reverse_reader_single_message() -> io::Result<()> {
         let dir = TempDir::new().unwrap();
-        let file_path = dir.path().join("test_single.arrow");
+        let file_path = dir.path().join("test_single.data.arrows");
 
         // Create a schema
         let schema = Arc::new(Schema::new(vec![Field::new("id", DataType::Int32, false)]));

src/parseable/staging/writer.rs

@@ -32,7 +32,10 @@ use arrow_select::concat::concat_batches;
 use itertools::Itertools;
 use tracing::error;
 
-use crate::{parseable::ARROW_FILE_EXTENSION, utils::arrow::adapt_batch};
+use crate::{
+    parseable::{ARROW_FILE_EXTENSION, PART_FILE_EXTENSION},
+    utils::arrow::adapt_batch,
+};
 
 use super::StagingError;
 
@@ -50,8 +53,9 @@ pub struct DiskWriter {
 impl DiskWriter {
     /// Try to create a file to stream arrows into
     pub fn try_new(path: impl Into<PathBuf>, schema: &Schema) -> Result<Self, StagingError> {
-        let path = path.into();
-        let file = OpenOptions::new().create(true).append(true).open(&path)?;
+        let mut path = path.into();
+        path.set_extension(PART_FILE_EXTENSION);
+        let file = OpenOptions::new().write(true).create(true).open(&path)?;
         let inner = StreamWriter::try_new_buffered(file, schema)?;
 
         Ok(Self { inner, path })