Refactor record-batch writing through a shared sink harness

Made-with: Cursor

Author: aisrael

src/pipeline.rs

@@ -1,6 +1,7 @@
 //! The `pipeline` module is the core of the datu crate.
 
 pub mod avro;
+pub mod batch_write;
 pub mod csv;
 pub mod dataframe;
 pub mod datasource;

src/pipeline/avro.rs

@@ -13,6 +13,8 @@ use crate::pipeline::RecordBatchReaderSource;
 use crate::pipeline::Source;
 use crate::pipeline::Step;
 use crate::pipeline::WriteArgs;
+use crate::pipeline::batch_write::BatchWriteSink;
+use crate::pipeline::batch_write::write_record_batches_with_sink;
 
 /// Pipeline step that reads an Avro file and produces a record batch reader.
 pub struct ReadAvroStep {
@@ -144,15 +146,30 @@ pub struct WriteAvroResult {}
 
 /// Write record batches from a reader to an Avro file.
 pub fn write_record_batches(path: &str, reader: &mut dyn RecordBatchReader) -> Result<()> {
-    let file = std::fs::File::create(path).map_err(Error::IoError)?;
-    let schema = reader.schema();
-    let mut writer = AvroWriter::new(file, (*schema).clone()).map_err(Error::ArrowError)?;
-    for batch in reader {
-        let batch = batch.map_err(Error::ArrowError)?;
-        writer.write(&batch).map_err(Error::ArrowError)?;
+    write_record_batches_with_sink(path, reader, AvroSink::new)
+}
+
+struct AvroSink {
+    writer: AvroWriter<std::fs::File>,
+}
+
+impl AvroSink {
+    fn new(path: &str, schema: arrow::datatypes::SchemaRef) -> Result<Self> {
+        let file = std::fs::File::create(path).map_err(Error::IoError)?;
+        let writer = AvroWriter::new(file, (*schema).clone()).map_err(Error::ArrowError)?;
+        Ok(Self { writer })
+    }
+}
+
+impl BatchWriteSink for AvroSink {
+    fn write_batch(&mut self, batch: &RecordBatch) -> Result<()> {
+        self.writer.write(batch).map_err(Error::ArrowError)
+    }
+
+    fn finish(mut self) -> Result<()> {
+        self.writer.finish().map_err(Error::ArrowError)?;
+        Ok(())
     }
-    writer.finish().map_err(Error::ArrowError)?;
-    Ok(())
 }
 
 #[async_trait(?Send)]

src/pipeline/batch_write.rs

@@ -0,0 +1,33 @@
+use arrow::array::RecordBatchReader;
+use arrow::datatypes::SchemaRef;
+use arrow::record_batch::RecordBatch;
+
+use crate::Error;
+use crate::Result;
+
+/// Per-format sink adapter for writing record batches.
+pub trait BatchWriteSink {
+    fn write_batch(&mut self, batch: &RecordBatch) -> Result<()>;
+    fn finish(self) -> Result<()>;
+}
+
+/// Shared harness for batch-oriented file writers.
+pub fn write_record_batches_with_sink<S, BuildSink>(
+    path: &str,
+    reader: &mut dyn RecordBatchReader,
+    build_sink: BuildSink,
+) -> Result<()>
+where
+    S: BatchWriteSink,
+    BuildSink: FnOnce(&str, SchemaRef) -> Result<S>,
+{
+    let schema = reader.schema();
+    let mut sink = build_sink(path, schema)?;
+
+    for batch in reader {
+        let batch = batch.map_err(Error::ArrowError)?;
+        sink.write_batch(&batch)?;
+    }
+
+    sink.finish()
+}

src/pipeline/csv.rs

@@ -10,6 +10,8 @@ use crate::pipeline::Source;
 use crate::pipeline::Step;
 use crate::pipeline::VecRecordBatchReader;
 use crate::pipeline::WriteArgs;
+use crate::pipeline::batch_write::BatchWriteSink;
+use crate::pipeline::batch_write::write_record_batches_with_sink;
 
 /// Pipeline step that reads a CSV file and produces a record batch reader.
 /// Uses DataFusion for schema inference and type detection.
@@ -47,21 +49,43 @@ pub struct WriteCsvStep {
 /// Result of successfully writing a CSV file.
 pub struct WriteCsvResult {}
 
+/// Write record batches from a reader to a CSV file.
+pub fn write_record_batches(path: &str, reader: &mut dyn RecordBatchReader) -> Result<()> {
+    write_record_batches_with_sink(path, reader, CsvSink::new)
+}
+
+struct CsvSink {
+    writer: arrow::csv::Writer<std::fs::File>,
+}
+
+impl CsvSink {
+    fn new(path: &str, _schema: arrow::datatypes::SchemaRef) -> Result<Self> {
+        let file = std::fs::File::create(path).map_err(Error::IoError)?;
+        Ok(Self {
+            writer: arrow::csv::Writer::new(file),
+        })
+    }
+}
+
+impl BatchWriteSink for CsvSink {
+    fn write_batch(&mut self, batch: &arrow::record_batch::RecordBatch) -> Result<()> {
+        self.writer.write(batch).map_err(Error::ArrowError)
+    }
+
+    fn finish(self) -> Result<()> {
+        Ok(())
+    }
+}
+
 #[async_trait(?Send)]
 impl Step for WriteCsvStep {
     type Input = ();
     type Output = WriteCsvResult;
 
     async fn execute(self, _input: Self::Input) -> Result<Self::Output> {
-        let path = self.args.path.as_str();
-        let file = std::fs::File::create(path).map_err(Error::IoError)?;
-        let mut writer = arrow::csv::Writer::new(file);
         let mut source = self.source;
-        let reader = source.get()?;
-        for batch in reader {
-            let batch = batch.map_err(Error::ArrowError)?;
-            writer.write(&batch).map_err(Error::ArrowError)?;
-        }
+        let mut reader = source.get()?;
+        write_record_batches(self.args.path.as_str(), &mut *reader)?;
         Ok(WriteCsvResult {})
     }
 }

src/pipeline/dataframe.rs

@@ -86,12 +86,7 @@ impl Step for DataFrameWriter {
                 parquet::write_record_batches(output_path, &mut reader)?;
             }
             FileType::Csv => {
-                let file = std::fs::File::create(output_path).map_err(Error::IoError)?;
-                let mut writer = arrow::csv::Writer::new(file);
-                for batch in &mut reader {
-                    let batch = batch.map_err(Error::ArrowError)?;
-                    writer.write(&batch).map_err(Error::ArrowError)?;
-                }
+                crate::pipeline::csv::write_record_batches(output_path, &mut reader)?;
             }
             FileType::Json => {
                 let file = std::fs::File::create(output_path).map_err(Error::IoError)?;

src/pipeline/orc.rs

@@ -11,6 +11,8 @@ use crate::pipeline::RecordBatchReaderSource;
 use crate::pipeline::Source;
 use crate::pipeline::Step;
 use crate::pipeline::WriteArgs;
+use crate::pipeline::batch_write::BatchWriteSink;
+use crate::pipeline::batch_write::write_record_batches_with_sink;
 
 /// Pipeline step that reads an ORC file and produces a record batch reader.
 pub struct ReadOrcStep {
@@ -52,17 +54,32 @@ pub struct WriteOrcResult {}
 
 /// Write record batches from a reader to an ORC file.
 pub fn write_record_batches(path: &str, reader: &mut dyn RecordBatchReader) -> Result<()> {
-    let file = std::fs::File::create(path).map_err(Error::IoError)?;
-    let schema = reader.schema();
-    let mut writer = ArrowWriterBuilder::new(file, schema)
-        .try_build()
-        .map_err(Error::OrcError)?;
-    for batch in reader {
-        let batch = batch.map_err(Error::ArrowError)?;
-        writer.write(&batch).map_err(Error::OrcError)?;
+    write_record_batches_with_sink(path, reader, OrcSink::new)
+}
+
+struct OrcSink {
+    writer: orc_rust::arrow_writer::ArrowWriter<std::fs::File>,
+}
+
+impl OrcSink {
+    fn new(path: &str, schema: arrow::datatypes::SchemaRef) -> Result<Self> {
+        let file = std::fs::File::create(path).map_err(Error::IoError)?;
+        let writer = ArrowWriterBuilder::new(file, schema)
+            .try_build()
+            .map_err(Error::OrcError)?;
+        Ok(Self { writer })
+    }
+}
+
+impl BatchWriteSink for OrcSink {
+    fn write_batch(&mut self, batch: &arrow::record_batch::RecordBatch) -> Result<()> {
+        self.writer.write(batch).map_err(Error::OrcError)
+    }
+
+    fn finish(self) -> Result<()> {
+        self.writer.close().map_err(Error::OrcError)?;
+        Ok(())
     }
-    writer.close().map_err(Error::OrcError)?;
-    Ok(())
 }
 
 #[async_trait(?Send)]

src/pipeline/parquet.rs

@@ -11,6 +11,8 @@ use crate::pipeline::RecordBatchReaderSource;
 use crate::pipeline::Source;
 use crate::pipeline::Step;
 use crate::pipeline::WriteArgs;
+use crate::pipeline::batch_write::BatchWriteSink;
+use crate::pipeline::batch_write::write_record_batches_with_sink;
 
 /// Pipeline step that reads a Parquet file and produces a record batch reader.
 pub struct ReadParquetStep {
@@ -62,15 +64,30 @@ impl Step for WriteParquetStep {
 
 /// Write record batches from a reader to a Parquet file.
 pub fn write_record_batches(path: &str, reader: &mut dyn RecordBatchReader) -> Result<()> {
-    let file = std::fs::File::create(path).map_err(Error::IoError)?;
-    let schema = reader.schema();
-    let mut writer = ArrowWriter::try_new(file, schema, None).map_err(Error::ParquetError)?;
-    for batch in reader {
-        let batch = batch.map_err(Error::ArrowError)?;
-        writer.write(&batch).map_err(Error::ParquetError)?;
+    write_record_batches_with_sink(path, reader, ParquetSink::new)
+}
+
+struct ParquetSink {
+    writer: ArrowWriter<std::fs::File>,
+}
+
+impl ParquetSink {
+    fn new(path: &str, schema: arrow::datatypes::SchemaRef) -> Result<Self> {
+        let file = std::fs::File::create(path).map_err(Error::IoError)?;
+        let writer = ArrowWriter::try_new(file, schema, None).map_err(Error::ParquetError)?;
+        Ok(Self { writer })
+    }
+}
+
+impl BatchWriteSink for ParquetSink {
+    fn write_batch(&mut self, batch: &arrow::record_batch::RecordBatch) -> Result<()> {
+        self.writer.write(batch).map_err(Error::ParquetError)
+    }
+
+    fn finish(self) -> Result<()> {
+        self.writer.close().map_err(Error::ParquetError)?;
+        Ok(())
     }
-    writer.close().map_err(Error::ParquetError)?;
-    Ok(())
 }
 
 #[cfg(test)]