Add Puffin reader and writer

Author: fqaiser94

Cargo.toml

@@ -23,6 +23,7 @@ members = [
   "crates/iceberg",
   "crates/integrations/*",
   "crates/test_utils",
+  "crates/puffin",
 ]
 exclude = ["bindings/python"]
 
@@ -97,3 +98,4 @@ uuid = { version = "1.6.1", features = ["v7"] }
 volo-thrift = "0.10"
 hive_metastore = "0.1"
 tera = "1"
+zstd = "0.13.2"

crates/puffin/Cargo.toml

@@ -0,0 +1,41 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you 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.
+
+[package]
+name = "iceberg-puffin"
+version = { workspace = true }
+edition = { workspace = true }
+homepage = { workspace = true }
+rust-version = { workspace = true }
+
+categories = ["database"]
+description = "Apache Iceberg Puffin"
+repository = { workspace = true }
+license = { workspace = true }
+keywords = ["iceberg", "puffin"]
+
+[dependencies]
+iceberg = { workspace = true }
+serde = { workspace = true }
+serde_derive = { workspace = true }
+serde_json = { workspace = true }
+once_cell = { workspace = true }
+zstd = { workspace = true }
+
+[dev-dependencies]
+tempfile = { workspace = true }
+tokio = { workspace = true }

crates/puffin/src/blob.rs

@@ -0,0 +1,96 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements.  See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership.  The ASF licenses this file
+// to you 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::fs::File;
+use std::io::{Read, Seek, SeekFrom};
+
+use iceberg::Result;
+
+use crate::compression::CompressionCodec;
+use crate::metadata::BlobMetadata;
+
+pub const APACHE_DATASKETCHES_THETA_V1: &str = "apache-datasketches-theta-v1";
+
+#[derive(Debug)]
+pub struct Blob {
+    pub(crate) r#type: String,
+    pub(crate) input_fields: Vec<i32>,
+    pub(crate) snapshot_id: i64,
+    pub(crate) sequence_number: i64,
+    pub(crate) data: Vec<u8>,
+    pub(crate) requested_compression_codec: Option<CompressionCodec>,
+    pub(crate) properties: HashMap<String, String>,
+}
+
+impl Blob {
+    /// Reads blob from a Puffin file
+    pub fn read(file: &mut File, blob_metadata: &BlobMetadata) -> Result<Blob> {
+        let mut bytes = vec![0; blob_metadata.length];
+        file.seek(SeekFrom::Start(blob_metadata.offset))?;
+        file.read(&mut bytes)?;
+        let data = CompressionCodec::decompress(blob_metadata.compression_codec, bytes)?;
+        return Ok(Blob {
+            r#type: blob_metadata.r#type.clone(),
+            input_fields: blob_metadata.input_fields.clone(),
+            snapshot_id: blob_metadata.snapshot_id,
+            sequence_number: blob_metadata.sequence_number,
+            data,
+            requested_compression_codec: Some(blob_metadata.compression_codec),
+            properties: blob_metadata.properties.clone(),
+        });
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use std::fs::File;
+
+    use crate::blob::Blob;
+    use crate::metadata::FileMetadata;
+
+    fn test_data_path() -> String {
+        env!("CARGO_MANIFEST_DIR").to_owned() + "/testdata/"
+    }
+
+    const EXPECTED_BLOB_0: &str = "abcdefghi";
+    const EXPECTED_BLOB_1: &str =
+        "some blob \u{0000} binary data 🤯 that is not very very very very very very long, is it?";
+
+    fn read_blob_as_utf8_string(file_path: &String, idx: usize) -> String {
+        let mut file = File::open(file_path).unwrap();
+        let file_metadata = FileMetadata::from_file(&mut file).unwrap();
+        let blob_metadata = file_metadata.blobs.get(idx).unwrap();
+        let blob = Blob::read(&mut file, blob_metadata).unwrap();
+        return String::from_utf8(blob.data).unwrap();
+    }
+
+    #[tokio::test]
+    async fn test_read_java_generated_uncompressed_metric_data() {
+        let file_path = test_data_path() + "v1/java-generated/sample-metric-data-uncompressed.bin";
+        assert_eq!(read_blob_as_utf8_string(&file_path, 0), EXPECTED_BLOB_0);
+        assert_eq!(read_blob_as_utf8_string(&file_path, 1), EXPECTED_BLOB_1);
+    }
+
+    #[tokio::test]
+    async fn test_read_java_generated_ztd_compressed_metric_data() {
+        let file_path =
+            test_data_path() + "v1/java-generated/sample-metric-data-compressed-zstd.bin";
+        assert_eq!(read_blob_as_utf8_string(&file_path, 0), EXPECTED_BLOB_0);
+        assert_eq!(read_blob_as_utf8_string(&file_path, 1), EXPECTED_BLOB_1);
+    }
+}

crates/puffin/src/compression.rs

@@ -0,0 +1,81 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements.  See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership.  The ASF licenses this file
+// to you 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 iceberg::{Error, ErrorKind, Result};
+use serde_derive::{Deserialize, Serialize};
+
+#[derive(Debug, PartialEq, Eq, Serialize, Deserialize, Clone, Copy)]
+#[serde(rename_all = "lowercase")]
+pub enum CompressionCodec {
+    /** No compression */
+    NONE,
+    /** LZ4 single compression frame with content size present */
+    LZ4,
+    /** Zstandard single compression frame with content size present */
+    ZSTD,
+}
+
+impl Default for CompressionCodec {
+    fn default() -> Self {
+        CompressionCodec::NONE
+    }
+}
+
+impl CompressionCodec {
+    pub(crate) fn decompress(
+        compression_codec: CompressionCodec,
+        bytes: Vec<u8>,
+    ) -> Result<Vec<u8>> {
+        match compression_codec {
+            CompressionCodec::NONE => Ok(bytes),
+            CompressionCodec::LZ4 => Err(Error::new(
+                ErrorKind::FeatureUnsupported,
+                "LZ4 decompression is not supported currently",
+            )),
+            CompressionCodec::ZSTD => {
+                let decompressed = zstd::stream::decode_all(&bytes[..])?;
+                return Ok(decompressed);
+            }
+        }
+    }
+
+    pub(crate) fn compress(compression_codec: CompressionCodec, bytes: Vec<u8>) -> Result<Vec<u8>> {
+        match compression_codec {
+            CompressionCodec::NONE => Ok(bytes),
+            CompressionCodec::LZ4 => Err(Error::new(
+                ErrorKind::FeatureUnsupported,
+                "LZ4 compression is not supported currently",
+            )),
+            CompressionCodec::ZSTD => {
+                let writer = Vec::<u8>::new();
+                let mut encoder = zstd::stream::Encoder::new(writer, 3)?;
+                encoder.include_checksum(true)?;
+                encoder.set_pledged_src_size(Some(bytes.len().try_into()?))?;
+                std::io::copy(&mut &bytes[..], &mut encoder)?;
+                let compressed = encoder.finish()?;
+                return Ok(compressed);
+            }
+        }
+    }
+
+    pub(crate) fn is_none(&self) -> bool {
+        return match self {
+            CompressionCodec::NONE => true,
+            _ => false,
+        };
+    }
+}

crates/puffin/src/lib.rs

@@ -0,0 +1,25 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements.  See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership.  The ASF licenses this file
+// to you 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.
+
+//! Iceberg Puffin implementation.
+
+#![deny(missing_docs)]
+
+mod blob;
+mod compression;
+mod metadata;
+mod writer;