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}