read and write continously

src/commands/merge.rs

@@ -1,9 +1,12 @@
 use crate::errors::PQRSError;
 use crate::errors::PQRSError::{FileExists, FileNotFound};
-use crate::utils::{check_path_present, get_row_batches, open_file, write_parquet};
+use crate::utils::{check_path_present, get_row_batches, open_file};
+use parquet::arrow::ArrowWriter;
 use clap::Parser;
+use arrow::datatypes::Schema;
 use log::debug;
-use std::ops::Add;
+use std::fs::File;
+use std::sync::Arc;
 use std::path::PathBuf;
 
 /// Merge file(s) into another parquet file
@@ -34,16 +37,40 @@ pub(crate) fn execute(opts: MergeCommandArgs) -> Result<(), PQRSError> {
         }
     }
 
-    let seed = open_file(&opts.input[0])?;
-    let mut combined = get_row_batches(seed)?;
+    let mut writer = {
+        let seed = open_file(&opts.input[0])?;
+        let data = get_row_batches(seed)?;
+
+        let file = File::create(&opts.output)?;
+        let fields = data.schema.fields().to_vec();
+
+        let schema_without_metadata = Schema::new(fields);
+
+        let mut writer = ArrowWriter::try_new(file, Arc::new(schema_without_metadata), None)?;
+
+        for record_batch in data.batches.iter() {
+            writer.write(record_batch)?;
+        }
+
+        writer
+    };
+
+
     for input in &opts.input[1..] {
         let current = open_file(input)?;
         let local = get_row_batches(current)?;
-        combined = combined.add(local);
+
+        // write record batches one at a time
+        // record batches are not combined
+        for record_batch in local.batches.iter() {
+            writer.write(record_batch)?;
+        }
     }
-    // debug!("The combined data looks like this: {:#?}", combined);
-    // debug!("This is the input schema: {:#?}", combined.schema);
-    write_parquet(combined, &opts.output)?;
+
+    // closing the writer writes out the FileMetaData
+    // if the writer is not closed properly, the metadata footer needed by the parquet
+    // format would be corrupt
+    writer.close()?;
 
     Ok(())
 }

src/utils.rs

@@ -2,7 +2,7 @@
 use crate::errors::PQRSError::CouldNotOpenFile;
 use arrow::{datatypes::Schema, record_batch::RecordBatch};
 use log::debug;
 use parquet::arrow::{arrow_reader::ArrowReaderBuilder, ArrowWriter};
 use parquet::file::reader::{FileReader, SerializedFileReader};
 use parquet::record::Row;
 use rand::seq::SliceRandom;
@@ -11,7 +11,7 @@
 use std::fs::File;
 use std::ops::Add;
 use std::path::Path;
 use std::sync::Arc;
 use walkdir::DirEntry;
 
 // calculate the sizes in bytes for one KiB, MiB, GiB, TiB, PiB
@@ -255,32 +255,6 @@
     })
 }
 
-/// Write a parquet file to the output location based on the given parquet input
-pub fn write_parquet<P: AsRef<Path>>(
-    data: ParquetData,
-    output: P,
-) -> Result<(), PQRSError> {
-    let file = File::create(output)?;
-    let fields = data.schema.fields().to_vec();
-    // the schema from the record batch might not contain the file specific metadata
-    // drop the schema to make sure that we don't fail in that case
-    let schema_without_metadata = Schema::new(fields);
-
-    let mut writer = ArrowWriter::try_new(file, Arc::new(schema_without_metadata), None)?;
-
-    // write record batches one at a time
-    // record batches are not combined
-    for record_batch in data.batches.iter() {
-        writer.write(record_batch)?;
-    }
-
-    // closing the writer writes out the FileMetaData
-    // if the writer is not closed properly, the metadata footer needed by the parquet
-    // format would be corrupt
-    writer.close()?;
-    Ok(())
-}
-
 /// Print the given parquet rows in json or json-like format
 fn print_row(row: &Row, format: Formats) {
     match format {