parquet/arrow/async_writer/
mod.rs

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

parquet/arrow/async_writer/
mod.rs
