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

Commit

Permalink
Simplified IPC stream writer (#575)
Browse files Browse the repository at this point in the history
  • Loading branch information
jorgecarleitao authored Nov 4, 2021
1 parent 8f2c9ea commit 2fb2995
Showing 1 changed file with 6 additions and 11 deletions.
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
}
}

0 comments on commit 2fb2995

Please sign in to comment.