parquet/arrow/async_writer/
mod.rs

1// Licensed to the Apache Software Foundation (ASF) under one
2// or more contributor license agreements.  See the NOTICE file
3// distributed with this work for additional information
4// regarding copyright ownership.  The ASF licenses this file
5// to you under the Apache License, Version 2.0 (the
6// "License"); you may not use this file except in compliance
7// with the License.  You may obtain a copy of the License at
8//
9//   http://www.apache.org/licenses/LICENSE-2.0
10//
11// Unless required by applicable law or agreed to in writing,
12// software distributed under the License is distributed on an
13// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
14// KIND, either express or implied.  See the License for the
15// specific language governing permissions and limitations
16// under the License.
17
18//! `async` API for writing [`RecordBatch`]es to Parquet files
19//!
20//! See the [crate-level documentation](crate) for more details.
21//!
22//! The `async` API for writing [`RecordBatch`]es is
23//! similar to the [`sync` API](ArrowWriter), so please
24//! read the documentation there before using this API.
25//!
26//! Here is an example for using [`AsyncArrowWriter`]:
27//!
28//! ```
29//! # #[tokio::main(flavor="current_thread")]
30//! # async fn main() {
31//! #
32//! # use std::sync::Arc;
33//! # use arrow_array::{ArrayRef, Int64Array, RecordBatch, RecordBatchReader};
34//! # use bytes::Bytes;
35//! # use parquet::arrow::{AsyncArrowWriter, arrow_reader::ParquetRecordBatchReaderBuilder};
36//! #
37//! let col = Arc::new(Int64Array::from_iter_values([1, 2, 3])) as ArrayRef;
38//! let to_write = RecordBatch::try_from_iter([("col", col)]).unwrap();
39//!
40//! let mut buffer = Vec::new();
41//! let mut writer = AsyncArrowWriter::try_new(&mut buffer, to_write.schema(), None).unwrap();
42//! writer.write(&to_write).await.unwrap();
43//! writer.close().await.unwrap();
44//!
45//! let buffer = Bytes::from(buffer);
46//! let mut reader = ParquetRecordBatchReaderBuilder::try_new(buffer.clone())
47//!     .unwrap()
48//!     .build()
49//!     .unwrap();
50//! let read = reader.next().unwrap().unwrap();
51//!
52//! assert_eq!(to_write, read);
53//! # }
54//! ```
55//!
56//! There is a blanket implementation of [`AsyncFileWriter`] for all types that
57//! implement [`AsyncWrite`], so writers such as `tokio::fs::File`, or
58//! `object_store::buffered::BufWriter` for writing to object storage, can be
59//! passed to [`AsyncArrowWriter`] directly.
60
61#[cfg(feature = "object_store")]
62mod store;
63#[allow(deprecated)]
64#[cfg(feature = "object_store")]
65pub use store::*;
66
67use crate::{
68    arrow::ArrowWriter,
69    arrow::arrow_writer::ArrowWriterOptions,
70    errors::{ParquetError, Result},
71    file::{
72        metadata::{KeyValue, ParquetMetaData, RowGroupMetaData},
73        properties::WriterProperties,
74    },
75};
76use arrow_array::RecordBatch;
77use arrow_schema::SchemaRef;
78use bytes::Bytes;
79use futures::FutureExt;
80use futures::future::BoxFuture;
81use std::mem;
82use tokio::io::{AsyncWrite, AsyncWriteExt};
83
84/// The asynchronous interface used by [`AsyncArrowWriter`] to write parquet files.
85pub trait AsyncFileWriter: Send {
86    /// Write the provided bytes to the underlying writer
87    ///
88    /// The underlying writer CAN decide to buffer the data or write it immediately.
89    /// This design allows the writer implementer to control the buffering and I/O scheduling.
90    ///
91    /// The underlying writer MAY implement retry logic to prevent breaking users write process.
92    fn write(&mut self, bs: Bytes) -> BoxFuture<'_, Result<()>>;
93
94    /// Flush any buffered data to the underlying writer and finish writing process.
95    ///
96    /// After `complete` returns `Ok(())`, caller SHOULD not call write again.
97    fn complete(&mut self) -> BoxFuture<'_, Result<()>>;
98}
99
100impl AsyncFileWriter for Box<dyn AsyncFileWriter + '_> {
101    fn write(&mut self, bs: Bytes) -> BoxFuture<'_, Result<()>> {
102        self.as_mut().write(bs)
103    }
104
105    fn complete(&mut self) -> BoxFuture<'_, Result<()>> {
106        self.as_mut().complete()
107    }
108}
109
110impl<T: AsyncWrite + Unpin + Send> AsyncFileWriter for T {
111    fn write(&mut self, bs: Bytes) -> BoxFuture<'_, Result<()>> {
112        async move {
113            self.write_all(&bs).await?;
114            Ok(())
115        }
116        .boxed()
117    }
118
119    fn complete(&mut self) -> BoxFuture<'_, Result<()>> {
120        async move {
121            self.flush().await?;
122            self.shutdown().await?;
123            Ok(())
124        }
125        .boxed()
126    }
127}
128
129/// Encodes [`RecordBatch`] to parquet, outputting to an [`AsyncFileWriter`]
130///
131/// ## Memory Usage
132///
133/// This writer eagerly writes data as soon as possible to the underlying [`AsyncFileWriter`],
134/// permitting fine-grained control over buffering and I/O scheduling. However, the columnar
135/// nature of parquet forces data for an entire row group to be buffered in memory, before
136/// it can be flushed. Depending on the data and the configured row group size, this buffering
137/// may be substantial.
138///
139/// Memory usage can be limited by calling [`Self::flush`] to flush the in progress row group,
140/// although this will likely increase overall file size and reduce query performance.
141/// See [ArrowWriter] for more information.
142///
143/// ```no_run
144/// # use tokio::fs::File;
145/// # use arrow_array::RecordBatch;
146/// # use parquet::arrow::AsyncArrowWriter;
147/// # async fn test() {
148/// let mut writer: AsyncArrowWriter<File> = todo!();
149/// let batch: RecordBatch = todo!();
150/// writer.write(&batch).await.unwrap();
151/// // Trigger an early flush if buffered size exceeds 1_000_000
152/// if writer.in_progress_size() > 1_000_000 {
153///     writer.flush().await.unwrap()
154/// }
155/// # }
156/// ```
157pub struct AsyncArrowWriter<W> {
158    /// Underlying sync writer
159    sync_writer: ArrowWriter<Vec<u8>>,
160
161    /// Async writer provided by caller
162    async_writer: W,
163}
164
165impl<W: AsyncFileWriter> AsyncArrowWriter<W> {
166    /// Try to create a new Async Arrow Writer
167    pub fn try_new(
168        writer: W,
169        arrow_schema: SchemaRef,
170        props: Option<WriterProperties>,
171    ) -> Result<Self> {
172        let options = ArrowWriterOptions::new().with_properties(props.unwrap_or_default());
173        Self::try_new_with_options(writer, arrow_schema, options)
174    }
175
176    /// Try to create a new Async Arrow Writer with [`ArrowWriterOptions`]
177    pub fn try_new_with_options(
178        writer: W,
179        arrow_schema: SchemaRef,
180        options: ArrowWriterOptions,
181    ) -> Result<Self> {
182        let sync_writer = ArrowWriter::try_new_with_options(Vec::new(), arrow_schema, options)?;
183
184        Ok(Self {
185            sync_writer,
186            async_writer: writer,
187        })
188    }
189
190    /// Returns metadata for any flushed row groups
191    pub fn flushed_row_groups(&self) -> &[RowGroupMetaData] {
192        self.sync_writer.flushed_row_groups()
193    }
194
195    /// Estimated memory usage, in bytes, of this `ArrowWriter`
196    ///
197    /// See [ArrowWriter::memory_size] for more information.
198    pub fn memory_size(&self) -> usize {
199        self.sync_writer.memory_size()
200    }
201
202    /// Anticipated encoded size of the in progress row group.
203    ///
204    /// See [ArrowWriter::memory_size] for more information.
205    pub fn in_progress_size(&self) -> usize {
206        self.sync_writer.in_progress_size()
207    }
208
209    /// Returns the number of rows buffered in the in progress row group
210    pub fn in_progress_rows(&self) -> usize {
211        self.sync_writer.in_progress_rows()
212    }
213
214    /// Returns the number of bytes written by this instance
215    pub fn bytes_written(&self) -> usize {
216        self.sync_writer.bytes_written()
217    }
218
219    /// Enqueues the provided `RecordBatch` to be written
220    ///
221    /// After every sync write by the inner [ArrowWriter], the inner buffer will be
222    /// checked and flush if at least half full
223    pub async fn write(&mut self, batch: &RecordBatch) -> Result<()> {
224        let before = self.sync_writer.flushed_row_groups().len();
225        self.sync_writer.write(batch)?;
226        if before != self.sync_writer.flushed_row_groups().len() {
227            self.do_write().await?;
228        }
229        Ok(())
230    }
231
232    /// Flushes all buffered rows into a new row group
233    pub async fn flush(&mut self) -> Result<()> {
234        self.sync_writer.flush()?;
235        self.do_write().await?;
236
237        Ok(())
238    }
239
240    /// Append [`KeyValue`] metadata in addition to those in [`WriterProperties`]
241    ///
242    /// This method allows to append metadata after [`RecordBatch`]es are written.
243    pub fn append_key_value_metadata(&mut self, kv_metadata: KeyValue) {
244        self.sync_writer.append_key_value_metadata(kv_metadata);
245    }
246
247    /// Close and finalize the writer.
248    ///
249    /// All the data in the inner buffer will be force flushed.
250    ///
251    /// Unlike [`Self::close`] this does not consume self
252    ///
253    /// Attempting to write after calling finish will result in an error
254    pub async fn finish(&mut self) -> Result<ParquetMetaData> {
255        let metadata = self.sync_writer.finish()?;
256
257        // Force to flush the remaining data.
258        self.do_write().await?;
259        self.async_writer.complete().await?;
260
261        Ok(metadata)
262    }
263
264    /// Close and finalize the writer.
265    ///
266    /// All the data in the inner buffer will be force flushed.
267    pub async fn close(mut self) -> Result<ParquetMetaData> {
268        self.finish().await
269    }
270
271    /// Consumes the [`AsyncArrowWriter`] and returns the underlying [`AsyncFileWriter`]
272    ///
273    /// # Notes
274    ///
275    /// This method does **not** flush or finalize the writer, so buffered data
276    /// will be lost if you have not called [`Self::finish`].
277    pub fn into_inner(self) -> W {
278        self.async_writer
279    }
280
281    /// Flush the data written by `sync_writer` into the `async_writer`
282    ///
283    /// # Notes
284    ///
285    /// This method will take the inner buffer from the `sync_writer` and write it into the
286    /// async writer. After the write, the inner buffer will be empty.
287    async fn do_write(&mut self) -> Result<()> {
288        let buffer = mem::take(self.sync_writer.inner_mut());
289
290        self.async_writer
291            .write(Bytes::from(buffer))
292            .await
293            .map_err(|e| ParquetError::External(Box::new(e)))?;
294
295        Ok(())
296    }
297}
298
299#[cfg(test)]
300mod tests {
301    use crate::arrow::arrow_reader::{ParquetRecordBatchReader, ParquetRecordBatchReaderBuilder};
302    use arrow::datatypes::{DataType, Field, Schema};
303    use arrow_array::{ArrayRef, BinaryArray, Int32Array, Int64Array, RecordBatchReader};
304    use bytes::Bytes;
305    use std::sync::Arc;
306
307    use super::*;
308
309    fn get_test_reader() -> ParquetRecordBatchReader {
310        let testdata = arrow::util::test_util::parquet_test_data();
311        // This test file is large enough to generate multiple row groups.
312        let path = format!("{testdata}/alltypes_tiny_pages_plain.parquet");
313        let original_data = Bytes::from(std::fs::read(path).unwrap());
314        ParquetRecordBatchReaderBuilder::try_new(original_data)
315            .unwrap()
316            .build()
317            .unwrap()
318    }
319
320    #[tokio::test]
321    async fn test_async_writer() {
322        let col = Arc::new(Int64Array::from_iter_values([1, 2, 3])) as ArrayRef;
323        let to_write = RecordBatch::try_from_iter([("col", col)]).unwrap();
324
325        let mut buffer = Vec::new();
326        let mut writer = AsyncArrowWriter::try_new(&mut buffer, to_write.schema(), None).unwrap();
327        writer.write(&to_write).await.unwrap();
328        writer.close().await.unwrap();
329
330        let buffer = Bytes::from(buffer);
331        let mut reader = ParquetRecordBatchReaderBuilder::try_new(buffer)
332            .unwrap()
333            .build()
334            .unwrap();
335        let read = reader.next().unwrap().unwrap();
336
337        assert_eq!(to_write, read);
338    }
339
340    // Read the data from the test file and write it by the async writer and sync writer.
341    // And then compares the results of the two writers.
342    #[tokio::test]
343    async fn test_async_writer_with_sync_writer() {
344        let reader = get_test_reader();
345
346        let write_props = WriterProperties::builder()
347            .set_max_row_group_row_count(Some(64))
348            .build();
349
350        let mut async_buffer = Vec::new();
351        let mut async_writer = AsyncArrowWriter::try_new(
352            &mut async_buffer,
353            reader.schema(),
354            Some(write_props.clone()),
355        )
356        .unwrap();
357
358        let mut sync_buffer = Vec::new();
359        let mut sync_writer =
360            ArrowWriter::try_new(&mut sync_buffer, reader.schema(), Some(write_props)).unwrap();
361        for record_batch in reader {
362            let record_batch = record_batch.unwrap();
363            async_writer.write(&record_batch).await.unwrap();
364            sync_writer.write(&record_batch).unwrap();
365        }
366        sync_writer.close().unwrap();
367        async_writer.close().await.unwrap();
368
369        assert_eq!(sync_buffer, async_buffer);
370    }
371
372    #[tokio::test]
373    async fn test_async_writer_bytes_written() {
374        let col = Arc::new(Int64Array::from_iter_values([1, 2, 3])) as ArrayRef;
375        let to_write = RecordBatch::try_from_iter([("col", col)]).unwrap();
376
377        let temp = tempfile::tempfile().unwrap();
378
379        let file = tokio::fs::File::from_std(temp.try_clone().unwrap());
380        let mut writer =
381            AsyncArrowWriter::try_new(file.try_clone().await.unwrap(), to_write.schema(), None)
382                .unwrap();
383        writer.write(&to_write).await.unwrap();
384        let _metadata = writer.finish().await.unwrap();
385        // After `finish` this should include the metadata and footer
386        let reported = writer.bytes_written();
387
388        // Get actual size from file metadata
389        let actual = file.metadata().await.unwrap().len() as usize;
390
391        assert_eq!(reported, actual);
392    }
393
394    #[tokio::test]
395    async fn test_async_writer_file() {
396        let col = Arc::new(Int64Array::from_iter_values([1, 2, 3])) as ArrayRef;
397        let col2 = Arc::new(BinaryArray::from_iter_values(vec![
398            vec![0; 500000],
399            vec![0; 500000],
400            vec![0; 500000],
401        ])) as ArrayRef;
402        let to_write = RecordBatch::try_from_iter([("col", col), ("col2", col2)]).unwrap();
403
404        let temp = tempfile::tempfile().unwrap();
405
406        let file = tokio::fs::File::from_std(temp.try_clone().unwrap());
407        let mut writer = AsyncArrowWriter::try_new(file, to_write.schema(), None).unwrap();
408        writer.write(&to_write).await.unwrap();
409        writer.close().await.unwrap();
410
411        let mut reader = ParquetRecordBatchReaderBuilder::try_new(temp)
412            .unwrap()
413            .build()
414            .unwrap();
415        let read = reader.next().unwrap().unwrap();
416
417        assert_eq!(to_write, read);
418    }
419
420    #[tokio::test]
421    async fn in_progress_accounting() {
422        // define schema
423        let schema = Schema::new(vec![Field::new("a", DataType::Int32, false)]);
424
425        // create some data
426        let a = Int32Array::from_value(0_i32, 512);
427
428        // build a record batch
429        let batch = RecordBatch::try_new(Arc::new(schema), vec![Arc::new(a)]).unwrap();
430
431        let temp = tempfile::tempfile().unwrap();
432        let file = tokio::fs::File::from_std(temp.try_clone().unwrap());
433        let mut writer = AsyncArrowWriter::try_new(file, batch.schema(), None).unwrap();
434
435        // starts empty
436        assert_eq!(writer.in_progress_size(), 0);
437        assert_eq!(writer.in_progress_rows(), 0);
438        assert_eq!(writer.bytes_written(), 4); // Initial Parquet header
439        writer.write(&batch).await.unwrap();
440
441        // updated on write
442        let initial_size = writer.in_progress_size();
443        assert!(initial_size > 0);
444        assert_eq!(writer.in_progress_rows(), batch.num_rows());
445        let initial_memory = writer.memory_size();
446        // memory estimate is larger than estimated encoded size
447        assert!(
448            initial_size <= initial_memory,
449            "{initial_size} <= {initial_memory}"
450        );
451
452        // updated on second write
453        writer.write(&batch).await.unwrap();
454        assert!(writer.in_progress_size() > initial_size);
455        assert_eq!(writer.in_progress_rows(), batch.num_rows() * 2);
456        assert!(writer.memory_size() > initial_memory);
457        assert!(
458            writer.in_progress_size() <= writer.memory_size(),
459            "in_progress_size {} <= memory_size {}",
460            writer.in_progress_size(),
461            writer.memory_size()
462        );
463
464        // in progress tracking is cleared, but the overall data written is updated
465        let pre_flush_bytes_written = writer.bytes_written();
466        writer.flush().await.unwrap();
467        assert_eq!(writer.in_progress_size(), 0);
468        assert_eq!(writer.memory_size(), 0);
469        assert_eq!(writer.in_progress_rows(), 0);
470        assert!(writer.bytes_written() > pre_flush_bytes_written);
471
472        writer.close().await.unwrap();
473    }
474}
parquet/arrow/async_writer/mod.rs

parquet/arrow/async_writer/
mod.rs
