mito2/sst/parquet/
index_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 datatypes::arrow::datatypes::SchemaRef;
17use datatypes::arrow::record_batch::RecordBatch;
18use futures::future::BoxFuture;
19use object_store::{ObjectStore, Writer};
20use parquet::arrow::AsyncArrowWriter;
21use parquet::arrow::async_writer::AsyncFileWriter;
22use parquet::basic::{Compression, Encoding, ZstdLevel};
23use parquet::errors::ParquetError;
24use parquet::file::metadata::KeyValue;
25use parquet::file::properties::WriterProperties;
26use snafu::{OptionExt, ResultExt};
27
28use crate::access_layer::TempFileCleaner;
29use crate::error::{OpenDalSnafu, Result, UnexpectedSnafu, WriteParquetSnafu};
30use crate::sst::{DEFAULT_WRITE_BUFFER_SIZE, DEFAULT_WRITE_CONCURRENCY};
31
32type ArrowWriter = AsyncArrowWriter<AsyncWriter>;
33
34/// Bridges an OpenDAL [`Writer`] with Parquet's [`AsyncFileWriter`] and tracks
35/// the number of bytes successfully submitted to the object store.
36struct AsyncWriter {
37    inner: Writer,
38    output_bytes: u64,
39}
40
41impl AsyncWriter {
42    fn new(inner: Writer) -> Self {
43        Self {
44            inner,
45            output_bytes: 0,
46        }
47    }
48
49    fn output_bytes(&self) -> u64 {
50        self.output_bytes
51    }
52
53    fn into_inner(self) -> Writer {
54        self.inner
55    }
56}
57
58impl AsyncFileWriter for AsyncWriter {
59    fn write(&mut self, bytes: Bytes) -> BoxFuture<'_, parquet::errors::Result<()>> {
60        Box::pin(async move {
61            let len = bytes.len() as u64;
62            self.inner
63                .write(bytes)
64                .await
65                .map_err(|error| ParquetError::External(Box::new(error)))?;
66            self.output_bytes += len;
67            Ok(())
68        })
69    }
70
71    fn complete(&mut self) -> BoxFuture<'_, parquet::errors::Result<()>> {
72        Box::pin(async move {
73            self.inner
74                .close()
75                .await
76                .map(|_| ())
77                .map_err(|error| ParquetError::External(Box::new(error)))
78        })
79    }
80}
81
82/// Shared Parquet output and cleanup lifecycle for index writers.
83pub(crate) struct ParquetIndexWriter {
84    name: &'static str,
85    object_store: ObjectStore,
86    file_name: String,
87    writer: Option<ArrowWriter>,
88}
89
90impl ParquetIndexWriter {
91    /// Opens an index file with the common Parquet writer configuration.
92    pub(crate) async fn try_new(
93        name: &'static str,
94        object_store: ObjectStore,
95        path: &str,
96        schema: &SchemaRef,
97        row_group_size: usize,
98        key_value_metadata: Option<Vec<KeyValue>>,
99    ) -> Result<Self> {
100        let file_name = path.rsplit('/').next().unwrap_or(path).to_string();
101        let output = object_store
102            .writer_with(path)
103            .chunk(DEFAULT_WRITE_BUFFER_SIZE.as_bytes() as usize)
104            .concurrent(DEFAULT_WRITE_CONCURRENCY)
105            .await
106            .context(OpenDalSnafu)?;
107        let properties = WriterProperties::builder()
108            .set_compression(Compression::ZSTD(ZstdLevel::default()))
109            .set_encoding(Encoding::PLAIN)
110            .set_max_row_group_row_count(Some(row_group_size))
111            .set_column_index_truncate_length(None)
112            .set_statistics_truncate_length(None)
113            .set_key_value_metadata(key_value_metadata)
114            .build();
115        let writer =
116            AsyncArrowWriter::try_new(AsyncWriter::new(output), schema.clone(), Some(properties))
117                .context(WriteParquetSnafu)?;
118
119        Ok(Self {
120            name,
121            object_store,
122            file_name,
123            writer: Some(writer),
124        })
125    }
126
127    /// Writes one batch to the index file.
128    pub(crate) async fn write(&mut self, batch: &RecordBatch) -> Result<()> {
129        self.writer
130            .as_mut()
131            .with_context(|| UnexpectedSnafu {
132                reason: format!("{} Parquet writer is closed", self.name),
133            })?
134            .write(batch)
135            .await
136            .context(WriteParquetSnafu)
137    }
138
139    /// Finishes the index file and returns its committed size.
140    pub(crate) async fn finish(&mut self) -> Result<u64> {
141        self.writer
142            .as_mut()
143            .with_context(|| UnexpectedSnafu {
144                reason: format!("{} Parquet writer is closed", self.name),
145            })?
146            .finish()
147            .await
148            .context(WriteParquetSnafu)?;
149        let writer = self.writer.take().with_context(|| UnexpectedSnafu {
150            reason: format!("{} Parquet writer is closed", self.name),
151        })?;
152        Ok(writer.into_inner().output_bytes())
153    }
154
155    /// Aborts an incomplete output and removes its atomic-write temporary files.
156    pub(crate) async fn abort(&mut self) {
157        if let Some(writer) = self.writer.take() {
158            let mut writer = writer.into_inner().into_inner();
159            if let Err(error) = writer.abort().await {
160                common_telemetry::warn!(error; "Failed to abort {} writer", self.name);
161            }
162        }
163
164        TempFileCleaner::clean_atomic_dir_files(&self.object_store, &[&self.file_name]).await;
165    }
166}
mito2/sst/parquet/index_writer.rs

mito2/sst/parquet/
index_writer.rs
