Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Add close method to RecordBatchWriter trait #4228

Merged
merged 4 commits into from
May 17, 2023
Merged

Add close method to RecordBatchWriter trait #4228

Show file tree
Hide file tree
Changes from 1 commit
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
5fd683f
Add finish method to RecordBatchWriter trait and implement it for CSV…
alexandreyc May 16, 2023
e6ac609
Simplify parquet::ArrowWriter::finish
alexandreyc May 16, 2023
71ae31f
Rename finish to close
alexandreyc May 16, 2023
4096bd8
Remove fully qualified method call
alexandreyc May 16, 2023
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
3 changes: 3 additions & 0 deletions arrow-array/src/record_batch.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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 finish(self) -> Result<(), ArrowError>;
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Suggested change
fn finish(self) -> Result<(), ArrowError>;
fn close(self) -> Result<(), ArrowError>;

Perhaps, this might avoid the name collisions?

}

/// A two-dimensional batch of column-oriented data with a defined
Expand Down
4 changes: 4 additions & 0 deletions arrow-csv/src/writer.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -197,6 +197,10 @@ impl<W: Write> RecordBatchWriter for Writer<W> {
fn write(&mut self, batch: &RecordBatch) -> Result<(), ArrowError> {
self.write(batch)
}

fn finish(self) -> Result<(), ArrowError> {
Ok(())
}
}

/// A CSV writer builder
Expand Down
5 changes: 0 additions & 5 deletions arrow-ipc/src/reader.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1510,8 +1510,6 @@ mod tests {
stream_writer.write(&batch).unwrap();
stream_writer.finish().unwrap();

drop(stream_writer);

file.rewind().unwrap();

// read stream back
Expand Down Expand Up @@ -1558,7 +1556,6 @@ mod tests {
crate::writer::FileWriter::try_new(&mut buf, &rb.schema()).unwrap();
writer.write(rb).unwrap();
writer.finish().unwrap();
drop(writer);

let mut reader =
crate::reader::FileReader::try_new(std::io::Cursor::new(buf), None).unwrap();
Expand All @@ -1571,7 +1568,6 @@ mod tests {
crate::writer::StreamWriter::try_new(&mut buf, &rb.schema()).unwrap();
writer.write(rb).unwrap();
writer.finish().unwrap();
drop(writer);

let mut reader =
crate::reader::StreamReader::try_new(std::io::Cursor::new(buf), None)
Expand All @@ -1591,7 +1587,6 @@ mod tests {
writer.write_metadata(k, v);
}
writer.finish().unwrap();
drop(writer);

let reader =
crate::reader::FileReader::try_new(std::io::Cursor::new(buf), None).unwrap();
Expand Down
18 changes: 13 additions & 5 deletions arrow-ipc/src/writer.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -851,7 +851,7 @@ impl<W: Write> FileWriter<W> {
/// writer.
pub fn into_inner(mut self) -> Result<W, ArrowError> {
if !self.finished {
self.finish()?;
FileWriter::finish(&mut self)?;
}
self.writer.into_inner().map_err(ArrowError::from)
}
Expand All @@ -861,6 +861,10 @@ impl<W: Write> RecordBatchWriter for FileWriter<W> {
fn write(&mut self, batch: &RecordBatch) -> Result<(), ArrowError> {
self.write(batch)
}

fn finish(mut self) -> Result<(), ArrowError> {
FileWriter::finish(&mut self)
}
}

pub struct StreamWriter<W: Write> {
Expand Down Expand Up @@ -991,7 +995,7 @@ impl<W: Write> StreamWriter<W> {
/// ```
pub fn into_inner(mut self) -> Result<W, ArrowError> {
if !self.finished {
self.finish()?;
StreamWriter::finish(&mut self)?;
}
self.writer.into_inner().map_err(ArrowError::from)
}
Expand All @@ -1001,6 +1005,10 @@ impl<W: Write> RecordBatchWriter for StreamWriter<W> {
fn write(&mut self, batch: &RecordBatch) -> Result<(), ArrowError> {
self.write(batch)
}

fn finish(mut self) -> Result<(), ArrowError> {
StreamWriter::finish(&mut self)
}
}

/// Stores the encoded data, which is an crate::Message, and optional Arrow data
Expand Down Expand Up @@ -1797,7 +1805,7 @@ mod tests {
let buffer: Vec<u8> = Vec::new();
let mut stream_writer = StreamWriter::try_new(buffer, &record.schema()).unwrap();
stream_writer.write(record).unwrap();
stream_writer.finish().unwrap();
StreamWriter::finish(&mut stream_writer).unwrap();
stream_writer.into_inner().unwrap()
}

Expand Down Expand Up @@ -2053,7 +2061,7 @@ mod tests {

let mut writer = StreamWriter::try_new(Vec::<u8>::new(), &schema).unwrap();
writer.write(&batch).unwrap();
writer.finish().unwrap();
StreamWriter::finish(&mut writer).unwrap();
let data = writer.into_inner().unwrap();

let mut reader = StreamReader::try_new(Cursor::new(data), None).unwrap();
Expand Down Expand Up @@ -2146,7 +2154,7 @@ mod tests {

let mut writer = FileWriter::try_new(Vec::<u8>::new(), &schema).unwrap();
writer.write(&batch).unwrap();
writer.finish().unwrap();
FileWriter::finish(&mut writer).unwrap();
let data = writer.into_inner().unwrap();

let mut reader = FileReader::try_new(Cursor::new(data), None).unwrap();
Expand Down
11 changes: 8 additions & 3 deletions arrow-json/src/writer.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -594,6 +594,10 @@ where
fn write(&mut self, batch: &RecordBatch) -> Result<(), ArrowError> {
self.write(batch)
}

fn finish(mut self) -> Result<(), ArrowError> {
Writer::finish(&mut self)
}
}

#[cfg(test)]
Expand Down Expand Up @@ -1262,15 +1266,16 @@ mod tests {
#[test]
fn json_writer_empty() {
let mut writer = ArrayWriter::new(vec![] as Vec<u8>);
writer.finish().unwrap();
ArrayWriter::finish(&mut writer).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>);
let v = json!({ "an": "object" });
writer.write_row(&v).unwrap();
writer.finish().unwrap();
ArrayWriter::finish(&mut writer).unwrap();
assert_eq!(
String::from_utf8(writer.into_inner()).unwrap(),
r#"[{"an":"object"}]"#
Expand All @@ -1284,7 +1289,7 @@ mod tests {
writer.write_row(&v).unwrap();
let v = json!({ "another": "object" });
writer.write_row(&v).unwrap();
writer.finish().unwrap();
ArrayWriter::finish(&mut writer).unwrap();
assert_eq!(
String::from_utf8(writer.into_inner()).unwrap(),
r#"[{"an":"object"},{"another":"object"}]"#
Expand Down
4 changes: 4 additions & 0 deletions parquet/src/arrow/arrow_writer/mod.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -250,6 +250,10 @@ impl<W: Write> RecordBatchWriter for ArrowWriter<W> {
fn write(&mut self, batch: &RecordBatch) -> Result<(), ArrowError> {
self.write(batch).map_err(|e| e.into())
}

fn finish(self) -> std::result::Result<(), ArrowError> {
self.close().map(|_| ()).map_err(ArrowError::from)
alexandreyc marked this conversation as resolved.
Show resolved Hide resolved
}
}

fn write_leaves<W: Write>(
Expand Down