Skip to content

Commit

Permalink
feat(parquet): Add try_new_with_options to async writer
Browse files Browse the repository at this point in the history
  • Loading branch information
evenyag committed Jan 13, 2024
1 parent fb46773 commit 0ea5554
Showing 1 changed file with 22 additions and 1 deletion.
23 changes: 22 additions & 1 deletion parquet/src/arrow/async_writer/mod.rs
Original file line number Diff line number Diff line change
Expand Up @@ -54,6 +54,7 @@
use std::{io::Write, sync::Arc};

use crate::{
arrow::arrow_writer::ArrowWriterOptions,
arrow::ArrowWriter,
errors::{ParquetError, Result},
file::properties::WriterProperties,
Expand Down Expand Up @@ -97,9 +98,29 @@ impl<W: AsyncWrite + Unpin + Send> AsyncArrowWriter<W> {
arrow_schema: SchemaRef,
buffer_size: usize,
props: Option<WriterProperties>,
) -> Result<Self> {
Self::try_new_with_options(writer, arrow_schema, buffer_size, props, Default::default())
}

/// Try to create a new Async Arrow Writer with [`ArrowWriterOptions`].
///
/// `buffer_size` determines the number of bytes to buffer before flushing
/// to the underlying [`AsyncWrite`]
///
/// The intermediate buffer will automatically be resized if necessary
///
/// [`Self::write`] will flush this intermediate buffer if it is at least
/// half full
pub fn try_new_with_options(
writer: W,
arrow_schema: SchemaRef,
buffer_size: usize,
props: Option<WriterProperties>,
options: ArrowWriterOptions,
) -> Result<Self> {
let shared_buffer = SharedBuffer::new(buffer_size);
let sync_writer = ArrowWriter::try_new(shared_buffer.clone(), arrow_schema, props)?;
let sync_writer =
ArrowWriter::try_new_with_options(shared_buffer.clone(), arrow_schema, props, options)?;

Ok(Self {
sync_writer,
Expand Down

0 comments on commit 0ea5554

Please sign in to comment.