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::{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
121    pub async fn close(mut self) -> Result<()> {
122        self.write_header_once().await?;
123        self.write_footer().await?;
124        self.writer.close().await?;
125        Ok(())
126    }
127
128    async fn write(&mut self, bytes: Bytes) -> Result<()> {
129        let length = bytes.len();
130        self.writer.write(bytes).await?;
131        self.num_bytes_written += length as u64;
132        Ok(())
133    }
134
135    async fn write_header_once(&mut self) -> Result<()> {
136        if !self.is_header_written {
137            let bytes = Bytes::copy_from_slice(&FileMetadata::MAGIC);
138            self.write(bytes).await?;
139            self.is_header_written = true;
140        }
141        Ok(())
142    }
143
144    fn footer_payload_bytes(&self) -> Result<Vec<u8>> {
145        let file_metadata = FileMetadata {
146            blobs: self.written_blobs_metadata.clone(),
147            properties: self.properties.clone(),
148        };
149        let json = serde_json::to_string::<FileMetadata>(&file_metadata)?;
150        self.footer_compression_codec.compress(json.into_bytes())
151    }
152
153    fn flags_bytes(&self) -> [u8; FileMetadata::FOOTER_STRUCT_FLAGS_LENGTH as usize] {
154        let mut result = [0; FileMetadata::FOOTER_STRUCT_FLAGS_LENGTH as usize];
155        for flag in &self.flags {
156            let byte_idx: usize = flag.byte_idx().into();
157            result[byte_idx] |= 0x1 << flag.bit_idx();
158        }
159        result
160    }
161
162    async fn write_footer(&mut self) -> Result<()> {
163        let mut footer_payload_bytes = self.footer_payload_bytes()?;
164        let footer_payload_bytes_length = u32::to_le_bytes(footer_payload_bytes.len().try_into()?);
165
166        let mut footer_bytes = Vec::new();
167        footer_bytes.extend(&FileMetadata::MAGIC);
168        footer_bytes.append(&mut footer_payload_bytes);
169        footer_bytes.extend(footer_payload_bytes_length);
170        footer_bytes.extend(self.flags_bytes());
171        footer_bytes.extend(&FileMetadata::MAGIC);
172
173        self.write(footer_bytes.into()).await
174    }
175}
176
177#[cfg(test)]
178mod tests {
179    use std::collections::HashMap;
180
181    use tempfile::TempDir;
182
183    use crate::compression::CompressionCodec;
184    use crate::io::{FileIO, InputFile, OutputFile};
185    use crate::puffin::blob::Blob;
186    use crate::puffin::metadata::FileMetadata;
187    use crate::puffin::reader::PuffinReader;
188    use crate::puffin::test_utils::{
189        blob_0, blob_1, empty_footer_payload, empty_footer_payload_bytes, file_properties,
190        java_empty_uncompressed_input_file, java_uncompressed_metric_input_file,
191        java_zstd_compressed_metric_input_file, read_file_metadata,
192        uncompressed_metric_file_metadata, zstd_compressed_metric_file_metadata,
193    };
194    use crate::puffin::writer::PuffinWriter;
195    use crate::{ErrorKind, Result};
196
197    async fn write_puffin_file(
198        temp_dir: &TempDir,
199        blobs: Vec<(Blob, CompressionCodec)>,
200        properties: HashMap<String, String>,
201    ) -> Result<OutputFile> {
202        let file_io = FileIO::new_with_fs();
203
204        let path_buf = temp_dir.path().join("temp_puffin.bin");
205        let temp_path = path_buf.to_str().unwrap();
206        let output_file = file_io.new_output(temp_path)?;
207
208        let mut writer = PuffinWriter::new(&output_file, properties, false).await?;
209        for (blob, compression_codec) in blobs {
210            writer.add(blob, compression_codec).await?;
211        }
212        writer.close().await?;
213
214        Ok(output_file)
215    }
216
217    async fn read_all_blobs_from_puffin_file(input_file: InputFile) -> Vec<Blob> {
218        let puffin_reader = PuffinReader::new(input_file).await.unwrap();
219        let mut blobs = Vec::new();
220        let blobs_metadata = puffin_reader.file_metadata().await.unwrap().clone().blobs;
221        for blob_metadata in blobs_metadata {
222            blobs.push(puffin_reader.blob(&blob_metadata).await.unwrap());
223        }
224        blobs
225    }
226
227    #[tokio::test]
228    async fn test_write_uncompressed_empty_file() {
229        let temp_dir = TempDir::new().unwrap();
230
231        let input_file = write_puffin_file(&temp_dir, Vec::new(), HashMap::new())
232            .await
233            .unwrap()
234            .to_input_file();
235
236        assert_eq!(
237            read_file_metadata(&input_file).await.unwrap(),
238            empty_footer_payload()
239        );
240
241        assert_eq!(
242            input_file.read().await.unwrap().len(),
243            FileMetadata::MAGIC_LENGTH as usize
244                // no blobs since puffin file is empty
245                + FileMetadata::MAGIC_LENGTH as usize
246                + empty_footer_payload_bytes().len()
247                + FileMetadata::FOOTER_STRUCT_LENGTH as usize
248        )
249    }
250
251    fn blobs_with_compression(
252        blobs: Vec<Blob>,
253        compression_codec: CompressionCodec,
254    ) -> Vec<(Blob, CompressionCodec)> {
255        blobs
256            .into_iter()
257            .map(|blob| (blob, compression_codec))
258            .collect()
259    }
260
261    #[tokio::test]
262    async fn test_write_uncompressed_metric_data() {
263        let temp_dir = TempDir::new().unwrap();
264        let blobs = vec![blob_0(), blob_1()];
265        let blobs_with_compression = blobs_with_compression(blobs.clone(), CompressionCodec::None);
266
267        let input_file = write_puffin_file(&temp_dir, blobs_with_compression, file_properties())
268            .await
269            .unwrap()
270            .to_input_file();
271
272        assert_eq!(
273            read_file_metadata(&input_file).await.unwrap(),
274            uncompressed_metric_file_metadata()
275        );
276
277        assert_eq!(read_all_blobs_from_puffin_file(input_file).await, blobs)
278    }
279
280    #[tokio::test]
281    async fn test_write_zstd_compressed_metric_data() {
282        let temp_dir = TempDir::new().unwrap();
283        let blobs = vec![blob_0(), blob_1()];
284        let blobs_with_compression =
285            blobs_with_compression(blobs.clone(), CompressionCodec::zstd_default());
286
287        let input_file = write_puffin_file(&temp_dir, blobs_with_compression, file_properties())
288            .await
289            .unwrap()
290            .to_input_file();
291
292        assert_eq!(
293            read_file_metadata(&input_file).await.unwrap(),
294            zstd_compressed_metric_file_metadata()
295        );
296
297        assert_eq!(read_all_blobs_from_puffin_file(input_file).await, blobs)
298    }
299
300    #[tokio::test]
301    async fn test_write_lz4_compressed_metric_data() {
302        let temp_dir = TempDir::new().unwrap();
303        let blobs = vec![blob_0(), blob_1()];
304        let blobs_with_compression = blobs_with_compression(blobs.clone(), CompressionCodec::Lz4);
305
306        assert_eq!(
307            write_puffin_file(&temp_dir, blobs_with_compression, file_properties())
308                .await
309                .unwrap_err()
310                .to_string(),
311            "FeatureUnsupported => LZ4 compression is not supported currently"
312        );
313    }
314
315    async fn get_file_as_byte_vec(input_file: InputFile) -> Vec<u8> {
316        input_file.read().await.unwrap().to_vec()
317    }
318
319    async fn assert_files_are_bit_identical(actual: OutputFile, expected: InputFile) {
320        let actual_bytes = get_file_as_byte_vec(actual.to_input_file()).await;
321        let expected_bytes = get_file_as_byte_vec(expected).await;
322        assert_eq!(actual_bytes, expected_bytes);
323    }
324
325    #[tokio::test]
326    async fn test_uncompressed_empty_puffin_file_is_bit_identical_to_java_generated_file() {
327        let temp_dir = TempDir::new().unwrap();
328
329        assert_files_are_bit_identical(
330            write_puffin_file(&temp_dir, Vec::new(), HashMap::new())
331                .await
332                .unwrap(),
333            java_empty_uncompressed_input_file(),
334        )
335        .await
336    }
337
338    #[tokio::test]
339    async fn test_uncompressed_metric_data_is_bit_identical_to_java_generated_file() {
340        let temp_dir = TempDir::new().unwrap();
341        let blobs = vec![blob_0(), blob_1()];
342        let blobs_with_compression = blobs_with_compression(blobs, CompressionCodec::None);
343
344        assert_files_are_bit_identical(
345            write_puffin_file(&temp_dir, blobs_with_compression, file_properties())
346                .await
347                .unwrap(),
348            java_uncompressed_metric_input_file(),
349        )
350        .await
351    }
352
353    #[tokio::test]
354    async fn test_zstd_compressed_metric_data_is_bit_identical_to_java_generated_file() {
355        let temp_dir = TempDir::new().unwrap();
356        let blobs = vec![blob_0(), blob_1()];
357        let blobs_with_compression =
358            blobs_with_compression(blobs, CompressionCodec::zstd_default());
359
360        assert_files_are_bit_identical(
361            write_puffin_file(&temp_dir, blobs_with_compression, file_properties())
362                .await
363                .unwrap(),
364            java_zstd_compressed_metric_input_file(),
365        )
366        .await
367    }
368
369    #[tokio::test]
370    async fn test_gzip_compression_rejected() {
371        let temp_dir = TempDir::new().unwrap();
372        let blobs = vec![blob_0()];
373        let blobs_with_compression =
374            blobs_with_compression(blobs, CompressionCodec::gzip_default());
375
376        let result = write_puffin_file(&temp_dir, blobs_with_compression, file_properties()).await;
377
378        assert!(result.is_err());
379        let err = result.unwrap_err();
380        assert_eq!(err.kind(), ErrorKind::DataInvalid);
381        assert!(err.to_string().contains("gzip"));
382        assert!(
383            err.to_string()
384                .contains("is not supported for Puffin files")
385        );
386    }
387
388    #[tokio::test]
389    async fn test_encrypted_write_read_roundtrip() {
390        use crate::encryption::{EncryptedInputFile, EncryptedOutputFile, StandardKeyMetadata};
391
392        let key_metadata = || {
393            StandardKeyMetadata::try_new(b"0123456789abcdef")
394                .unwrap()
395                .with_aad_prefix(b"test-aad-prefix!")
396        };
397
398        let file_io = FileIO::new_with_memory();
399        let path = "memory:///test/encrypted.puffin";
400        let blobs = vec![blob_0(), blob_1()];
401
402        // Write through the encrypting writer.
403        let encrypted_output =
404            EncryptedOutputFile::new(file_io.new_output(path).unwrap(), key_metadata());
405        let mut writer =
406            PuffinWriter::new_from_encrypted(&encrypted_output, file_properties(), false)
407                .await
408                .unwrap();
409        for blob in blobs.clone() {
410            writer.add(blob, CompressionCodec::None).await.unwrap();
411        }
412        writer.close().await.unwrap();
413
414        // The ciphertext on disk must not equal a plaintext puffin file.
415        let raw = file_io.new_input(path).unwrap().read().await.unwrap();
416        assert_ne!(
417            &raw[..FileMetadata::MAGIC_LENGTH as usize],
418            FileMetadata::MAGIC
419        );
420
421        // Read back through the decrypting reader over plaintext offsets.
422        let encrypted_input =
423            EncryptedInputFile::new(file_io.new_input(path).unwrap(), key_metadata());
424        let reader = PuffinReader::new_from_encrypted(encrypted_input)
425            .await
426            .unwrap();
427
428        let file_metadata = reader.file_metadata().await.unwrap().clone();
429        assert_eq!(file_metadata, uncompressed_metric_file_metadata());
430
431        let mut read_blobs = Vec::new();
432        for blob_metadata in &file_metadata.blobs {
433            read_blobs.push(reader.blob(blob_metadata).await.unwrap());
434        }
435        assert_eq!(read_blobs, blobs);
436    }
437}
iceberg/puffin/writer.rs

iceberg/puffin/
writer.rs
