-
Notifications
You must be signed in to change notification settings - Fork 329
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
feat(puffin): add file writer (#2776)
* feat(puffin): add file writer Signed-off-by: Zhenchi <[email protected]> * Update src/puffin/src/file_format/writer/file.rs Co-authored-by: dennis zhuang <[email protected]> * Update src/puffin/src/file_format/writer/file.rs Co-authored-by: dennis zhuang <[email protected]> * feat: footer bytes with capacity Signed-off-by: Zhenchi <[email protected]> * feat: footer bytes with capacity Signed-off-by: Zhenchi <[email protected]> * Update src/puffin/src/file_format/writer.rs Co-authored-by: Yingwen <[email protected]> * feat: add flush Signed-off-by: Zhenchi <[email protected]> * chore: specify default flags Signed-off-by: Zhenchi <[email protected]> * feat: close async writer Signed-off-by: Zhenchi <[email protected]> --------- Signed-off-by: Zhenchi <[email protected]> Co-authored-by: dennis zhuang <[email protected]> Co-authored-by: Yingwen <[email protected]>
- Loading branch information
3 people
authored
Nov 21, 2023
1 parent
efc5abf
commit a7bbd61
Showing
8 changed files
with
337 additions
and
11 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,62 @@ | ||
| // Copyright 2023 Greptime Team | ||
| // | ||
| // Licensed under the Apache License, Version 2.0 (the "License"); | ||
| // you may not use this file except in compliance with the License. | ||
| // You may obtain a copy of the License at | ||
| // | ||
| // http://www.apache.org/licenses/LICENSE-2.0 | ||
| // | ||
| // Unless required by applicable law or agreed to in writing, software | ||
| // distributed under the License is distributed on an "AS IS" BASIS, | ||
| // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. | ||
| // See the License for the specific language governing permissions and | ||
| // limitations under the License. | ||
|
|
||
| mod file; | ||
| mod footer; | ||
|
|
||
| use std::collections::HashMap; | ||
|
|
||
| use async_trait::async_trait; | ||
|
|
||
| use crate::error::Result; | ||
| pub use crate::file_format::writer::file::PuffinFileWriter; | ||
|
|
||
| /// Blob ready to be written | ||
| pub struct Blob<R> { | ||
| // TODO(zhongzc): ignore `input_fields`, `snapshot_id`, `sequence_number` | ||
| // and `compression_codec` for now to keep thing simple | ||
| /// The type of the blob | ||
| pub blob_type: String, | ||
|
|
||
| /// The data of the blob | ||
| pub data: R, | ||
|
|
||
| /// The properties of the blob | ||
| pub properties: HashMap<String, String>, | ||
| } | ||
|
|
||
| /// The trait for writing Puffin files synchronously | ||
| pub trait PuffinSyncWriter { | ||
| /// Set the properties of the Puffin file | ||
| fn set_properties(&mut self, properties: HashMap<String, String>); | ||
|
|
||
| /// Add a blob to the Puffin file | ||
| fn add_blob<R: std::io::Read>(&mut self, blob: Blob<R>) -> Result<()>; | ||
|
|
||
| /// Finish writing the Puffin file | ||
| fn finish(&mut self) -> Result<()>; | ||
| } | ||
|
|
||
| /// The trait for writing Puffin files asynchronously | ||
| #[async_trait] | ||
| pub trait PuffinAsyncWriter { | ||
| /// Set the properties of the Puffin file | ||
| fn set_properties(&mut self, properties: HashMap<String, String>); | ||
|
|
||
| /// Add a blob to the Puffin file | ||
| async fn add_blob<R: futures::AsyncRead + Send>(&mut self, blob: Blob<R>) -> Result<()>; | ||
|
|
||
| /// Finish writing the Puffin file | ||
| async fn finish(&mut self) -> Result<()>; | ||
| } |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,159 @@ | ||
| // Copyright 2023 Greptime Team | ||
| // | ||
| // Licensed under the Apache License, Version 2.0 (the "License"); | ||
| // you may not use this file except in compliance with the License. | ||
| // You may obtain a copy of the License at | ||
| // | ||
| // http://www.apache.org/licenses/LICENSE-2.0 | ||
| // | ||
| // Unless required by applicable law or agreed to in writing, software | ||
| // distributed under the License is distributed on an "AS IS" BASIS, | ||
| // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. | ||
| // See the License for the specific language governing permissions and | ||
| // limitations under the License. | ||
|
|
||
| use std::collections::HashMap; | ||
| use std::{io, mem}; | ||
|
|
||
| use async_trait::async_trait; | ||
| use futures::{AsyncRead, AsyncWrite, AsyncWriteExt}; | ||
| use snafu::ResultExt; | ||
|
|
||
| use crate::blob_metadata::{BlobMetadata, BlobMetadataBuilder}; | ||
| use crate::error::{CloseSnafu, FlushSnafu, Result, WriteSnafu}; | ||
| use crate::file_format::writer::footer::FooterWriter; | ||
| use crate::file_format::writer::{Blob, PuffinAsyncWriter, PuffinSyncWriter}; | ||
| use crate::file_format::MAGIC; | ||
|
|
||
| /// Puffin file writer, implements both [`PuffinSyncWriter`] and [`PuffinAsyncWriter`] | ||
| pub struct PuffinFileWriter<W> { | ||
| /// The writer to write to | ||
| writer: W, | ||
|
|
||
| /// The properties of the file | ||
| properties: HashMap<String, String>, | ||
|
|
||
| /// The metadata of the blobs | ||
| blob_metadata: Vec<BlobMetadata>, | ||
|
|
||
| /// The offset of the next blob | ||
| next_blob_offset: u64, | ||
| } | ||
|
|
||
| impl<W> PuffinFileWriter<W> { | ||
| pub fn new(writer: W) -> Self { | ||
| Self { | ||
| writer, | ||
| properties: HashMap::new(), | ||
| blob_metadata: Vec::new(), | ||
| next_blob_offset: 0, | ||
| } | ||
| } | ||
|
|
||
| fn create_blob_metadata( | ||
| &self, | ||
| typ: String, | ||
| properties: HashMap<String, String>, | ||
| size: u64, | ||
| ) -> BlobMetadata { | ||
| BlobMetadataBuilder::default() | ||
| .blob_type(typ) | ||
| .properties(properties) | ||
| .offset(self.next_blob_offset as _) | ||
| .length(size as _) | ||
| .build() | ||
| .expect("Required fields are not set") | ||
| } | ||
| } | ||
|
|
||
| impl<W: io::Write> PuffinSyncWriter for PuffinFileWriter<W> { | ||
| fn set_properties(&mut self, properties: HashMap<String, String>) { | ||
| self.properties = properties; | ||
| } | ||
|
|
||
| fn add_blob<R: io::Read>(&mut self, mut blob: Blob<R>) -> Result<()> { | ||
| self.write_header_if_needed_sync()?; | ||
|
|
||
| let size = io::copy(&mut blob.data, &mut self.writer).context(WriteSnafu)?; | ||
|
|
||
| let blob_metadata = self.create_blob_metadata(blob.blob_type, blob.properties, size); | ||
| self.blob_metadata.push(blob_metadata); | ||
|
|
||
| self.next_blob_offset += size; | ||
| Ok(()) | ||
| } | ||
|
|
||
| fn finish(&mut self) -> Result<()> { | ||
| self.write_header_if_needed_sync()?; | ||
| self.write_footer_sync()?; | ||
| self.writer.flush().context(FlushSnafu) | ||
| } | ||
| } | ||
|
|
||
| #[async_trait] | ||
| impl<W: AsyncWrite + Unpin + Send> PuffinAsyncWriter for PuffinFileWriter<W> { | ||
| fn set_properties(&mut self, properties: HashMap<String, String>) { | ||
| self.properties = properties; | ||
| } | ||
|
|
||
| async fn add_blob<R: AsyncRead + Send>(&mut self, blob: Blob<R>) -> Result<()> { | ||
| self.write_header_if_needed_async().await?; | ||
|
|
||
| let size = futures::io::copy(blob.data, &mut self.writer) | ||
| .await | ||
| .context(WriteSnafu)?; | ||
|
|
||
| let blob_metadata = self.create_blob_metadata(blob.blob_type, blob.properties, size); | ||
| self.blob_metadata.push(blob_metadata); | ||
|
|
||
| self.next_blob_offset += size; | ||
| Ok(()) | ||
| } | ||
|
|
||
| async fn finish(&mut self) -> Result<()> { | ||
| self.write_header_if_needed_async().await?; | ||
| self.write_footer_async().await?; | ||
| self.writer.flush().await.context(FlushSnafu)?; | ||
| self.writer.close().await.context(CloseSnafu) | ||
| } | ||
| } | ||
|
|
||
| impl<W: io::Write> PuffinFileWriter<W> { | ||
| fn write_header_if_needed_sync(&mut self) -> Result<()> { | ||
| if self.next_blob_offset == 0 { | ||
| self.writer.write_all(&MAGIC).context(WriteSnafu)?; | ||
| self.next_blob_offset += MAGIC.len() as u64; | ||
| } | ||
| Ok(()) | ||
| } | ||
|
|
||
| fn write_footer_sync(&mut self) -> Result<()> { | ||
| let bytes = FooterWriter::new( | ||
| mem::take(&mut self.blob_metadata), | ||
| mem::take(&mut self.properties), | ||
| ) | ||
| .into_footer_bytes()?; | ||
|
|
||
| self.writer.write_all(&bytes).context(WriteSnafu) | ||
| } | ||
| } | ||
|
|
||
| impl<W: AsyncWrite + Unpin> PuffinFileWriter<W> { | ||
| async fn write_header_if_needed_async(&mut self) -> Result<()> { | ||
| if self.next_blob_offset == 0 { | ||
| self.writer.write_all(&MAGIC).await.context(WriteSnafu)?; | ||
| self.next_blob_offset += MAGIC.len() as u64; | ||
| } | ||
| Ok(()) | ||
| } | ||
|
|
||
| async fn write_footer_async(&mut self) -> Result<()> { | ||
| let bytes = FooterWriter::new( | ||
| mem::take(&mut self.blob_metadata), | ||
| mem::take(&mut self.properties), | ||
| ) | ||
| .into_footer_bytes()?; | ||
|
|
||
| self.writer.write_all(&bytes).await.context(WriteSnafu) | ||
| } | ||
| } |
Oops, something went wrong.