iceberg/puffin/
writer.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 std::collections::{HashMap, HashSet};
19
20use bytes::Bytes;
21
22use super::validate_puffin_compression;
23use crate::Result;
24use crate::compression::CompressionCodec;
25use crate::encryption::EncryptedOutputFile;
26use crate::io::{FileMetadata as IoFileMetadata, FileWrite, OutputFile};
27use crate::puffin::blob::Blob;
28use crate::puffin::metadata::{BlobMetadata, FileMetadata, Flag};
29
30/// Puffin writer
31pub struct PuffinWriter {
32    writer: Box<dyn FileWrite>,
33    is_header_written: bool,
34    num_bytes_written: u64,
35    written_blobs_metadata: Vec<BlobMetadata>,
36    properties: HashMap<String, String>,
37    footer_compression_codec: CompressionCodec,
38    flags: HashSet<Flag>,
39}
40
41impl PuffinWriter {
42    /// Returns a new Puffin writer for an unencrypted file.
43    pub async fn new(
44        output_file: &OutputFile,
45        properties: HashMap<String, String>,
46        compress_footer: bool,
47    ) -> Result<Self> {
48        Ok(Self::from_writer(
49            output_file.writer().await?,
50            properties,
51            compress_footer,
52        ))
53    }
54
55    /// Returns a new Puffin writer from an [`EncryptedOutputFile`].
56    ///
57    /// Use this when writing Puffin files with transparent encryption. Blob
58    /// and footer offsets are recorded as plaintext positions, matching what
59    /// an unencrypted writer would produce.
60    pub async fn new_from_encrypted(
61        encrypted_output: &EncryptedOutputFile,
62        properties: HashMap<String, String>,
63        compress_footer: bool,
64    ) -> Result<Self> {
65        Ok(Self::from_writer(
66            encrypted_output.writer().await?,
67            properties,
68            compress_footer,
69        ))
70    }
71
72    fn from_writer(
73        writer: Box<dyn FileWrite>,
74        properties: HashMap<String, String>,
75        compress_footer: bool,
76    ) -> Self {
77        let mut flags = HashSet::<Flag>::new();
78        let footer_compression_codec = if compress_footer {
79            flags.insert(Flag::FooterPayloadCompressed);
80            CompressionCodec::Lz4
81        } else {
82            CompressionCodec::None
83        };
84
85        Self {
86            writer,
87            is_header_written: false,
88            num_bytes_written: 0,
89            written_blobs_metadata: Vec::new(),
90            properties,
91            footer_compression_codec,
92            flags,
93        }
94    }
95
96    /// Adds blob to Puffin file
97    pub async fn add(&mut self, blob: Blob, compression_codec: CompressionCodec) -> Result<()> {
98        validate_puffin_compression(compression_codec)?;
99
100        self.write_header_once().await?;
101
102        let offset = self.num_bytes_written;
103        let compressed_bytes: Bytes = compression_codec.compress(blob.data)?.into();
104        let length = compressed_bytes.len().try_into()?;
105        self.write(compressed_bytes).await?;
106        self.written_blobs_metadata.push(BlobMetadata {
107            r#type: blob.r#type,
108            fields: blob.fields,
109            snapshot_id: blob.snapshot_id,
110            sequence_number: blob.sequence_number,
111            offset,
112            length,
113            compression_codec,
114            properties: blob.properties,
115        });
116
117        Ok(())
118    }
119
120    /// Finalizes the Puffin file and returns its stored size.
121    pub async fn close(mut self) -> Result<IoFileMetadata> {
122        self.write_header_once().await?;
123        self.write_footer().await?;
124        self.writer.close().await
125    }
126
127    async fn write(&mut self, bytes: Bytes) -> Result<()> {
128        let length = bytes.len();
129        self.writer.write(bytes).await?;
130        self.num_bytes_written += length as u64;
131        Ok(())
132    }
133
134    async fn write_header_once(&mut self) -> Result<()> {
135        if !self.is_header_written {
136            let bytes = Bytes::copy_from_slice(&FileMetadata::MAGIC);
137            self.write(bytes).await?;
138            self.is_header_written = true;
139        }
140        Ok(())
141    }
142
143    fn footer_payload_bytes(&self) -> Result<Vec<u8>> {
144        let file_metadata = FileMetadata {
145            blobs: self.written_blobs_metadata.clone(),
146            properties: self.properties.clone(),
147        };
148        let json = serde_json::to_string::<FileMetadata>(&file_metadata)?;
149        self.footer_compression_codec.compress(json.into_bytes())
150    }
151
152    fn flags_bytes(&self) -> [u8; FileMetadata::FOOTER_STRUCT_FLAGS_LENGTH as usize] {
153        let mut result = [0; FileMetadata::FOOTER_STRUCT_FLAGS_LENGTH as usize];
154        for flag in &self.flags {
155            let byte_idx: usize = flag.byte_idx().into();
156            result[byte_idx] |= 0x1 << flag.bit_idx();
157        }
158        result
159    }
160
161    async fn write_footer(&mut self) -> Result<()> {
162        let mut footer_payload_bytes = self.footer_payload_bytes()?;
163        let footer_payload_bytes_length = u32::to_le_bytes(footer_payload_bytes.len().try_into()?);
164
165        let mut footer_bytes = Vec::new();
166        footer_bytes.extend(&FileMetadata::MAGIC);
167        footer_bytes.append(&mut footer_payload_bytes);
168        footer_bytes.extend(footer_payload_bytes_length);
169        footer_bytes.extend(self.flags_bytes());
170        footer_bytes.extend(&FileMetadata::MAGIC);
171
172        self.write(footer_bytes.into()).await
173    }
174}
175
176#[cfg(test)]
177mod tests {
178    use std::collections::HashMap;
179
180    use tempfile::TempDir;
181
182    use crate::compression::CompressionCodec;
183    use crate::io::{FileIO, InputFile, OutputFile};
184    use crate::puffin::blob::Blob;
185    use crate::puffin::metadata::FileMetadata;
186    use crate::puffin::reader::PuffinReader;
187    use crate::puffin::test_utils::{
188        blob_0, blob_1, empty_footer_payload, empty_footer_payload_bytes, file_properties,
189        java_empty_uncompressed_input_file, java_uncompressed_metric_input_file,
190        java_zstd_compressed_metric_input_file, read_file_metadata,
191        uncompressed_metric_file_metadata, zstd_compressed_metric_file_metadata,
192    };
193    use crate::puffin::writer::PuffinWriter;
194    use crate::{ErrorKind, Result};
195
196    async fn write_puffin_file(
197        temp_dir: &TempDir,
198        blobs: Vec<(Blob, CompressionCodec)>,
199        properties: HashMap<String, String>,
200    ) -> Result<OutputFile> {
201        let file_io = FileIO::new_with_fs();
202
203        let path_buf = temp_dir.path().join("temp_puffin.bin");
204        let temp_path = path_buf.to_str().unwrap();
205        let output_file = file_io.new_output(temp_path)?;
206
207        let mut writer = PuffinWriter::new(&output_file, properties, false).await?;
208        for (blob, compression_codec) in blobs {
209            writer.add(blob, compression_codec).await?;
210        }
211        writer.close().await?;
212
213        Ok(output_file)
214    }
215
216    async fn read_all_blobs_from_puffin_file(input_file: InputFile) -> Vec<Blob> {
217        let puffin_reader = PuffinReader::new(input_file).await.unwrap();
218        let mut blobs = Vec::new();
219        let blobs_metadata = puffin_reader.file_metadata().await.unwrap().clone().blobs;
220        for blob_metadata in blobs_metadata {
221            blobs.push(puffin_reader.blob(&blob_metadata).await.unwrap());
222        }
223        blobs
224    }
225
226    #[tokio::test]
227    async fn test_write_uncompressed_empty_file() {
228        let temp_dir = TempDir::new().unwrap();
229
230        let input_file = write_puffin_file(&temp_dir, Vec::new(), HashMap::new())
231            .await
232            .unwrap()
233            .to_input_file();
234
235        assert_eq!(
236            read_file_metadata(&input_file).await.unwrap(),
237            empty_footer_payload()
238        );
239
240        assert_eq!(
241            input_file.read().await.unwrap().len(),
242            FileMetadata::MAGIC_LENGTH as usize
243                // no blobs since puffin file is empty
244                + FileMetadata::MAGIC_LENGTH as usize
245                + empty_footer_payload_bytes().len()
246                + FileMetadata::FOOTER_STRUCT_LENGTH as usize
247        )
248    }
249
250    fn blobs_with_compression(
251        blobs: Vec<Blob>,
252        compression_codec: CompressionCodec,
253    ) -> Vec<(Blob, CompressionCodec)> {
254        blobs
255            .into_iter()
256            .map(|blob| (blob, compression_codec))
257            .collect()
258    }
259
260    #[tokio::test]
261    async fn test_write_uncompressed_metric_data() {
262        let temp_dir = TempDir::new().unwrap();
263        let blobs = vec![blob_0(), blob_1()];
264        let blobs_with_compression = blobs_with_compression(blobs.clone(), CompressionCodec::None);
265
266        let input_file = write_puffin_file(&temp_dir, blobs_with_compression, file_properties())
267            .await
268            .unwrap()
269            .to_input_file();
270
271        assert_eq!(
272            read_file_metadata(&input_file).await.unwrap(),
273            uncompressed_metric_file_metadata()
274        );
275
276        assert_eq!(read_all_blobs_from_puffin_file(input_file).await, blobs)
277    }
278
279    #[tokio::test]
280    async fn test_write_zstd_compressed_metric_data() {
281        let temp_dir = TempDir::new().unwrap();
282        let blobs = vec![blob_0(), blob_1()];
283        let blobs_with_compression =
284            blobs_with_compression(blobs.clone(), CompressionCodec::zstd_default());
285
286        let input_file = write_puffin_file(&temp_dir, blobs_with_compression, file_properties())
287            .await
288            .unwrap()
289            .to_input_file();
290
291        assert_eq!(
292            read_file_metadata(&input_file).await.unwrap(),
293            zstd_compressed_metric_file_metadata()
294        );
295
296        assert_eq!(read_all_blobs_from_puffin_file(input_file).await, blobs)
297    }
298
299    #[tokio::test]
300    async fn test_write_lz4_compressed_metric_data() {
301        let temp_dir = TempDir::new().unwrap();
302        let blobs = vec![blob_0(), blob_1()];
303        let blobs_with_compression = blobs_with_compression(blobs.clone(), CompressionCodec::Lz4);
304
305        assert_eq!(
306            write_puffin_file(&temp_dir, blobs_with_compression, file_properties())
307                .await
308                .unwrap_err()
309                .to_string(),
310            "FeatureUnsupported => LZ4 compression is not supported currently"
311        );
312    }
313
314    async fn get_file_as_byte_vec(input_file: InputFile) -> Vec<u8> {
315        input_file.read().await.unwrap().to_vec()
316    }
317
318    async fn assert_files_are_bit_identical(actual: OutputFile, expected: InputFile) {
319        let actual_bytes = get_file_as_byte_vec(actual.to_input_file()).await;
320        let expected_bytes = get_file_as_byte_vec(expected).await;
321        assert_eq!(actual_bytes, expected_bytes);
322    }
323
324    #[tokio::test]
325    async fn test_uncompressed_empty_puffin_file_is_bit_identical_to_java_generated_file() {
326        let temp_dir = TempDir::new().unwrap();
327
328        assert_files_are_bit_identical(
329            write_puffin_file(&temp_dir, Vec::new(), HashMap::new())
330                .await
331                .unwrap(),
332            java_empty_uncompressed_input_file(),
333        )
334        .await
335    }
336
337    #[tokio::test]
338    async fn test_uncompressed_metric_data_is_bit_identical_to_java_generated_file() {
339        let temp_dir = TempDir::new().unwrap();
340        let blobs = vec![blob_0(), blob_1()];
341        let blobs_with_compression = blobs_with_compression(blobs, CompressionCodec::None);
342
343        assert_files_are_bit_identical(
344            write_puffin_file(&temp_dir, blobs_with_compression, file_properties())
345                .await
346                .unwrap(),
347            java_uncompressed_metric_input_file(),
348        )
349        .await
350    }
351
352    #[tokio::test]
353    async fn test_zstd_compressed_metric_data_is_bit_identical_to_java_generated_file() {
354        let temp_dir = TempDir::new().unwrap();
355        let blobs = vec![blob_0(), blob_1()];
356        let blobs_with_compression =
357            blobs_with_compression(blobs, CompressionCodec::zstd_default());
358
359        assert_files_are_bit_identical(
360            write_puffin_file(&temp_dir, blobs_with_compression, file_properties())
361                .await
362                .unwrap(),
363            java_zstd_compressed_metric_input_file(),
364        )
365        .await
366    }
367
368    #[tokio::test]
369    async fn test_gzip_compression_rejected() {
370        let temp_dir = TempDir::new().unwrap();
371        let blobs = vec![blob_0()];
372        let blobs_with_compression =
373            blobs_with_compression(blobs, CompressionCodec::gzip_default());
374
375        let result = write_puffin_file(&temp_dir, blobs_with_compression, file_properties()).await;
376
377        assert!(result.is_err());
378        let err = result.unwrap_err();
379        assert_eq!(err.kind(), ErrorKind::DataInvalid);
380        assert!(err.to_string().contains("gzip"));
381        assert!(
382            err.to_string()
383                .contains("is not supported for Puffin files")
384        );
385    }
386
387    #[tokio::test]
388    async fn test_encrypted_write_read_roundtrip() {
389        use crate::encryption::{EncryptedInputFile, EncryptedOutputFile, StandardKeyMetadata};
390
391        let key_metadata = || {
392            StandardKeyMetadata::try_new(b"0123456789abcdef")
393                .unwrap()
394                .with_aad_prefix(b"test-aad-prefix!")
395        };
396
397        let file_io = FileIO::new_with_memory();
398        let path = "memory:///test/encrypted.puffin";
399        let blobs = vec![blob_0(), blob_1()];
400
401        // Write through the encrypting writer.
402        let encrypted_output =
403            EncryptedOutputFile::new(file_io.new_output(path).unwrap(), key_metadata());
404        let mut writer =
405            PuffinWriter::new_from_encrypted(&encrypted_output, file_properties(), false)
406                .await
407                .unwrap();
408        for blob in blobs.clone() {
409            writer.add(blob, CompressionCodec::None).await.unwrap();
410        }
411        let file_metadata = writer.close().await.unwrap();
412
413        // The ciphertext on disk must not equal a plaintext puffin file.
414        let raw = file_io.new_input(path).unwrap().read().await.unwrap();
415        assert_ne!(
416            &raw[..FileMetadata::MAGIC_LENGTH as usize],
417            FileMetadata::MAGIC
418        );
419
420        // Read back through the decrypting reader over plaintext offsets.
421        let key_metadata = encrypted_output.key_metadata_with_saved_file_metadata(&file_metadata);
422        assert_eq!(key_metadata.file_length(), Some(raw.len() as u64));
423        let encrypted_input =
424            EncryptedInputFile::new(file_io.new_input(path).unwrap(), key_metadata);
425        let reader = PuffinReader::new_from_encrypted(encrypted_input)
426            .await
427            .unwrap();
428
429        let file_metadata = reader.file_metadata().await.unwrap().clone();
430        assert_eq!(file_metadata, uncompressed_metric_file_metadata());
431
432        let mut read_blobs = Vec::new();
433        for blob_metadata in &file_metadata.blobs {
434            read_blobs.push(reader.blob(blob_metadata).await.unwrap());
435        }
436        assert_eq!(read_blobs, blobs);
437    }
438}
iceberg/puffin/writer.rs

iceberg/puffin/
writer.rs
