common_datasource/
parquet_writer.rs

1// Copyright 2023 Greptime Team
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//     http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15use bytes::Bytes;
16use futures::future::BoxFuture;
17use object_store::Writer;
18use parquet::arrow::async_writer::AsyncFileWriter;
19use parquet::errors::ParquetError;
20
21/// Bridges opendal [Writer] with parquet [AsyncFileWriter].
22pub struct AsyncWriter {
23    inner: Writer,
24}
25
26impl AsyncWriter {
27    /// Create a [`AsyncWriter`] by given [`Writer`].
28    pub fn new(writer: Writer) -> Self {
29        Self { inner: writer }
30    }
31}
32
33impl AsyncFileWriter for AsyncWriter {
34    fn write(&mut self, bs: Bytes) -> BoxFuture<'_, parquet::errors::Result<()>> {
35        Box::pin(async move {
36            self.inner
37                .write(bs)
38                .await
39                .map_err(|err| ParquetError::External(Box::new(err)))
40        })
41    }
42
43    fn complete(&mut self) -> BoxFuture<'_, parquet::errors::Result<()>> {
44        Box::pin(async move {
45            self.inner
46                .close()
47                .await
48                .map(|_| ())
49                .map_err(|err| ParquetError::External(Box::new(err)))
50        })
51    }
52}