Skip to content
This repository has been archived by the owner on Feb 18, 2024. It is now read-only.

Simplified IPC stream writer / remove finish on drop from stream writer #575

Merged
merged 1 commit into from
Nov 4, 2021
Merged
Changes from all commits
Commits
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
17 changes: 6 additions & 11 deletions src/io/ipc/write/stream.rs
Original file line number Diff line number Diff line change
Expand Up @@ -20,7 +20,7 @@
//! The `FileWriter` and `StreamWriter` have similar interfaces,
//! however the `FileWriter` expects a reader that supports `Seek`ing

use std::io::{BufWriter, Write};
use std::io::Write;

use super::common::{
encoded_batch, write_continuation, write_message, DictionaryTracker, EncodedData,
Expand All @@ -40,7 +40,7 @@ use crate::record_batch::RecordBatch;
/// For a usage walkthrough consult [this example](https://github.com/jorgecarleitao/arrow2/tree/main/examples/ipc_pyarrow).
pub struct StreamWriter<W: Write> {
/// The object to write to
writer: BufWriter<W>,
writer: W,
/// IPC write options
write_options: IpcWriteOptions,
/// Whether the writer footer has been written, and the writer is finished
Expand All @@ -57,11 +57,10 @@ impl<W: Write> StreamWriter<W> {
}

pub fn try_new_with_options(
writer: W,
mut writer: W,
schema: &Schema,
write_options: IpcWriteOptions,
) -> Result<Self> {
let mut writer = BufWriter::new(writer);
// write the schema, set the written bytes to the schema
let encoded_message = EncodedData {
ipc_message: schema_to_bytes(schema, *write_options.metadata_version()),
Expand Down Expand Up @@ -104,13 +103,9 @@ impl<W: Write> StreamWriter<W> {

Ok(())
}
}

/// Finish the stream if it is not 'finished' when it goes out of scope
impl<W: Write> Drop for StreamWriter<W> {
fn drop(&mut self) {
if !self.finished {
self.finish().unwrap();
}
/// Consumes itself, returning the inner writer.
pub fn into_inner(self) -> W {
self.writer
}
}