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

parquet/arrow/async_writer/
store.rs
