parquet/arrow/async_writer/
store.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
5use bytes::Bytes;
6use futures::future::BoxFuture;
7use std::sync::Arc;
8
9use crate::arrow::async_writer::AsyncFileWriter;
10use crate::errors::{ParquetError, Result};
11use object_store::ObjectStore;
12use object_store::buffered::BufWriter;
13use object_store::path::Path;
14use tokio::io::AsyncWriteExt;
15
16/// [`ParquetObjectWriter`] for writing to parquet to [`ObjectStore`]
17///
18/// This type is deprecated: [`BufWriter`] implements [`AsyncWrite`] and can
19/// therefore be passed to [`AsyncArrowWriter`] directly via the blanket
20/// [`AsyncFileWriter`] implementation for [`AsyncWrite`] types:
21///
22/// ```
23/// # use arrow_array::{ArrayRef, Int64Array, RecordBatch};
24/// # use object_store::buffered::BufWriter;
25/// # use object_store::memory::InMemory;
26/// # use object_store::path::Path;
27/// # use object_store::{ObjectStore, ObjectStoreExt};
28/// # use std::sync::Arc;
29///
30/// # use parquet::arrow::arrow_reader::ParquetRecordBatchReaderBuilder;
31/// # use parquet::arrow::AsyncArrowWriter;
32///
33/// # #[tokio::main(flavor="current_thread")]
34/// # async fn main() {
35///     let store = Arc::new(InMemory::new());
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 object_store_writer = BufWriter::new(store.clone(), Path::from("test"));
41///     let mut writer =
42///         AsyncArrowWriter::try_new(object_store_writer, to_write.schema(), None).unwrap();
43///     writer.write(&to_write).await.unwrap();
44///     writer.close().await.unwrap();
45///
46///     let buffer = store
47///         .get(&Path::from("test"))
48///         .await
49///         .unwrap()
50///         .bytes()
51///         .await
52///         .unwrap();
53///     let mut reader = ParquetRecordBatchReaderBuilder::try_new(buffer)
54///         .unwrap()
55///         .build()
56///         .unwrap();
57///     let read = reader.next().unwrap().unwrap();
58///
59///     assert_eq!(to_write, read);
60/// # }
61/// ```
62///
63/// [`AsyncWrite`]: tokio::io::AsyncWrite
64/// [`AsyncArrowWriter`]: crate::arrow::async_writer::AsyncArrowWriter
65#[deprecated(
66    since = "59.2.0",
67    note = "Pass an `object_store::buffered::BufWriter` to `AsyncArrowWriter` directly instead; see https://github.com/apache/arrow-rs/issues/10308 and `parquet/examples/object_store.rs`."
68)]
69#[derive(Debug)]
70pub struct ParquetObjectWriter {
71    w: BufWriter,
72}
73
74#[expect(deprecated)]
75impl ParquetObjectWriter {
76    /// Create a new [`ParquetObjectWriter`] that writes to the specified path in the given store.
77    ///
78    /// To configure the writer behavior, please build [`BufWriter`] and then use [`Self::from_buf_writer`]
79    pub fn new(store: Arc<dyn ObjectStore>, path: Path) -> Self {
80        Self::from_buf_writer(BufWriter::new(store, path))
81    }
82
83    /// Construct a new ParquetObjectWriter via a existing BufWriter.
84    pub fn from_buf_writer(w: BufWriter) -> Self {
85        Self { w }
86    }
87
88    /// Consume the writer and return the underlying BufWriter.
89    pub fn into_inner(self) -> BufWriter {
90        self.w
91    }
92}
93
94#[expect(deprecated)]
95impl AsyncFileWriter for ParquetObjectWriter {
96    fn write(&mut self, bs: Bytes) -> BoxFuture<'_, Result<()>> {
97        Box::pin(async {
98            self.w
99                .put(bs)
100                .await
101                .map_err(|err| ParquetError::External(Box::new(err)))
102        })
103    }
104
105    fn complete(&mut self) -> BoxFuture<'_, Result<()>> {
106        Box::pin(async {
107            self.w
108                .shutdown()
109                .await
110                .map_err(|err| ParquetError::External(Box::new(err)))
111        })
112    }
113}
114#[expect(deprecated)]
115impl From<BufWriter> for ParquetObjectWriter {
116    fn from(w: BufWriter) -> Self {
117        Self::from_buf_writer(w)
118    }
119}
120#[cfg(test)]
121#[expect(deprecated)]
122mod tests {
123    use arrow_array::{ArrayRef, Int64Array, RecordBatch};
124    use object_store::memory::InMemory;
125    use std::sync::Arc;
126
127    use super::*;
128    use crate::arrow::AsyncArrowWriter;
129    use crate::arrow::arrow_reader::ParquetRecordBatchReaderBuilder;
130    use object_store::ObjectStoreExt;
131
132    #[tokio::test]
133    async fn test_async_writer() {
134        let store = Arc::new(InMemory::new());
135
136        let col = Arc::new(Int64Array::from_iter_values([1, 2, 3])) as ArrayRef;
137        let to_write = RecordBatch::try_from_iter([("col", col)]).unwrap();
138
139        let object_store_writer = ParquetObjectWriter::new(store.clone(), Path::from("test"));
140        let mut writer =
141            AsyncArrowWriter::try_new(object_store_writer, to_write.schema(), None).unwrap();
142        writer.write(&to_write).await.unwrap();
143        writer.close().await.unwrap();
144
145        let buffer = store
146            .get(&Path::from("test"))
147            .await
148            .unwrap()
149            .bytes()
150            .await
151            .unwrap();
152        let mut reader = ParquetRecordBatchReaderBuilder::try_new(buffer)
153            .unwrap()
154            .build()
155            .unwrap();
156        let read = reader.next().unwrap().unwrap();
157
158        assert_eq!(to_write, read);
159    }
160}
parquet/arrow/async_writer/store.rs

parquet/arrow/async_writer/
store.rs
