From 0ea5554d5f2e6aac354351ddbdffe4e79bd3e701 Mon Sep 17 00:00:00 2001
From: evenyag <realevenyag@gmail.com>
Date: Sat, 13 Jan 2024 18:34:11 +0800
Subject: [PATCH] feat(parquet): Add try_new_with_options to async writer

---
 parquet/src/arrow/async_writer/mod.rs | 23 ++++++++++++++++++++++-
 1 file changed, 22 insertions(+), 1 deletion(-)

diff --git a/parquet/src/arrow/async_writer/mod.rs b/parquet/src/arrow/async_writer/mod.rs
index 30080c579e8f..d9710219c1e7 100644
--- a/parquet/src/arrow/async_writer/mod.rs
+++ b/parquet/src/arrow/async_writer/mod.rs
@@ -54,6 +54,7 @@
 use std::{io::Write, sync::Arc};
 
 use crate::{
+    arrow::arrow_writer::ArrowWriterOptions,
     arrow::ArrowWriter,
     errors::{ParquetError, Result},
     file::properties::WriterProperties,
@@ -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,