This is an automated email from the ASF dual-hosted git repository.
tustvold pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/arrow-rs.git
The following commit(s) were added to refs/heads/master by this push:
new bccbf2354 Add close method to RecordBatchWriter trait (#4228)
bccbf2354 is described below
commit bccbf2354aec165e77592014bd0c487ca24d002f
Author: Alexandre Crayssac <[email protected]>
AuthorDate: Wed May 17 13:12:00 2023 +0200
Add close method to RecordBatchWriter trait (#4228)
* Add finish method to RecordBatchWriter trait and implement it for CSV,
JSON, IPC and Parquet
* Simplify parquet::ArrowWriter::finish
Co-authored-by: Raphael Taylor-Davies
<[email protected]>
* Rename finish to close
* Remove fully qualified method call
---------
Co-authored-by: Raphael Taylor-Davies
<[email protected]>
---
arrow-array/src/record_batch.rs | 3 +++
arrow-csv/src/writer.rs | 4 ++++
arrow-ipc/src/writer.rs | 8 ++++++++
arrow-json/src/writer.rs | 5 +++++
parquet/src/arrow/arrow_writer/mod.rs | 5 +++++
5 files changed, 25 insertions(+)
diff --git a/arrow-array/src/record_batch.rs b/arrow-array/src/record_batch.rs
index aea49c047..d2e36780a 100644
--- a/arrow-array/src/record_batch.rs
+++ b/arrow-array/src/record_batch.rs
@@ -47,6 +47,9 @@ pub trait RecordBatchReader: Iterator<Item =
Result<RecordBatch, ArrowError>> {
pub trait RecordBatchWriter {
/// Write a single batch to the writer.
fn write(&mut self, batch: &RecordBatch) -> Result<(), ArrowError>;
+
+ /// Write footer or termination data, then mark the writer as done.
+ fn close(self) -> Result<(), ArrowError>;
}
/// A two-dimensional batch of column-oriented data with a defined
diff --git a/arrow-csv/src/writer.rs b/arrow-csv/src/writer.rs
index ba2123a09..840e8e8a9 100644
--- a/arrow-csv/src/writer.rs
+++ b/arrow-csv/src/writer.rs
@@ -197,6 +197,10 @@ impl<W: Write> RecordBatchWriter for Writer<W> {
fn write(&mut self, batch: &RecordBatch) -> Result<(), ArrowError> {
self.write(batch)
}
+
+ fn close(self) -> Result<(), ArrowError> {
+ Ok(())
+ }
}
/// A CSV writer builder
diff --git a/arrow-ipc/src/writer.rs b/arrow-ipc/src/writer.rs
index fcfd4d97a..59657bc4b 100644
--- a/arrow-ipc/src/writer.rs
+++ b/arrow-ipc/src/writer.rs
@@ -861,6 +861,10 @@ impl<W: Write> RecordBatchWriter for FileWriter<W> {
fn write(&mut self, batch: &RecordBatch) -> Result<(), ArrowError> {
self.write(batch)
}
+
+ fn close(mut self) -> Result<(), ArrowError> {
+ self.finish()
+ }
}
pub struct StreamWriter<W: Write> {
@@ -1001,6 +1005,10 @@ impl<W: Write> RecordBatchWriter for StreamWriter<W> {
fn write(&mut self, batch: &RecordBatch) -> Result<(), ArrowError> {
self.write(batch)
}
+
+ fn close(mut self) -> Result<(), ArrowError> {
+ self.finish()
+ }
}
/// Stores the encoded data, which is an crate::Message, and optional Arrow
data
diff --git a/arrow-json/src/writer.rs b/arrow-json/src/writer.rs
index 6f241be40..e6c960aef 100644
--- a/arrow-json/src/writer.rs
+++ b/arrow-json/src/writer.rs
@@ -594,6 +594,10 @@ where
fn write(&mut self, batch: &RecordBatch) -> Result<(), ArrowError> {
self.write(batch)
}
+
+ fn close(mut self) -> Result<(), ArrowError> {
+ self.finish()
+ }
}
#[cfg(test)]
@@ -1265,6 +1269,7 @@ mod tests {
writer.finish().unwrap();
assert_eq!(String::from_utf8(writer.into_inner()).unwrap(), "");
}
+
#[test]
fn json_writer_one_row() {
let mut writer = ArrayWriter::new(vec![] as Vec<u8>);
diff --git a/parquet/src/arrow/arrow_writer/mod.rs
b/parquet/src/arrow/arrow_writer/mod.rs
index 075ecc034..af8202182 100644
--- a/parquet/src/arrow/arrow_writer/mod.rs
+++ b/parquet/src/arrow/arrow_writer/mod.rs
@@ -250,6 +250,11 @@ impl<W: Write> RecordBatchWriter for ArrowWriter<W> {
fn write(&mut self, batch: &RecordBatch) -> Result<(), ArrowError> {
self.write(batch).map_err(|e| e.into())
}
+
+ fn close(self) -> std::result::Result<(), ArrowError> {
+ self.close()?;
+ Ok(())
+ }
}
fn write_leaves<W: Write>(
