Add Puffin reader and writer

Author: fqaiser94

Cargo.toml

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

crates/iceberg/src/writer/file_writer/mod.rs

@@ -26,7 +26,9 @@ use crate::Result;
 
 mod parquet_writer;
 pub use parquet_writer::{ParquetWriter, ParquetWriterBuilder};
+
 mod track_writer;
+pub use track_writer::TrackWriter;
 
 pub mod location_generator;
 

crates/iceberg/src/writer/file_writer/parquet_writer.rs

@@ -18,7 +18,7 @@
 //! The module contains the file writer for parquet file format.
 
 use std::collections::HashMap;
-use std::sync::atomic::AtomicI64;
+use std::sync::atomic::AtomicU64;
 use std::sync::Arc;
 
 use arrow_schema::SchemaRef as ArrowSchemaRef;
@@ -81,7 +81,7 @@ impl<T: LocationGenerator, F: FileNameGenerator> FileWriterBuilder for ParquetWr
 
     async fn build(self) -> crate::Result<Self::R> {
         let arrow_schema: ArrowSchemaRef = Arc::new(self.schema.as_ref().try_into()?);
-        let written_size = Arc::new(AtomicI64::new(0));
+        let written_size = Arc::new(AtomicU64::new(0));
         let out_file = self.file_io.new_output(
             self.location_generator
                 .generate_location(&self.file_name_generator.generate_file_name()),
@@ -214,7 +214,7 @@ pub struct ParquetWriter {
     schema: SchemaRef,
     out_file: OutputFile,
     writer: AsyncArrowWriter<AsyncFileWriter<TrackWriter>>,
-    written_size: Arc<AtomicI64>,
+    written_size: Arc<AtomicU64>,
     current_row_num: usize,
 }
 

crates/iceberg/src/writer/file_writer/track_writer.rs

@@ -15,7 +15,7 @@
 // specific language governing permissions and limitations
 // under the License.
 
-use std::sync::atomic::AtomicI64;
+use std::sync::atomic::AtomicU64;
 use std::sync::Arc;
 
 use bytes::Bytes;
@@ -24,18 +24,25 @@ use crate::io::FileWrite;
 use crate::Result;
 
 /// `TrackWriter` is used to track the written size.
-pub(crate) struct TrackWriter {
+pub struct TrackWriter {
     inner: Box<dyn FileWrite>,
-    written_size: Arc<AtomicI64>,
+    /// number of bytes written so far
+    written_size: Arc<AtomicU64>,
 }
 
 impl TrackWriter {
-    pub fn new(writer: Box<dyn FileWrite>, written_size: Arc<AtomicI64>) -> Self {
+    /// Create new writer
+    pub fn new(writer: Box<dyn FileWrite>, written_size: Arc<AtomicU64>) -> Self {
         Self {
             inner: writer,
             written_size,
         }
     }
+
+    /// Number of bytes written so far
+    pub fn bytes_written(&self) -> u64 {
+        self.written_size.load(std::sync::atomic::Ordering::SeqCst)
+    }
 }
 
 #[async_trait::async_trait]
@@ -44,7 +51,7 @@ impl FileWrite for TrackWriter {
         let size = bs.len();
         self.inner.write(bs).await.map(|v| {
             self.written_size
-                .fetch_add(size as i64, std::sync::atomic::Ordering::Relaxed);
+                .fetch_add(size as u64, std::sync::atomic::Ordering::Relaxed);
             v
         })
     }

crates/puffin/Cargo.toml

@@ -0,0 +1,42 @@
+# 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]
+bytes = { workspace = true }
+iceberg = { workspace = true }
+once_cell = { workspace = true }
+serde = { workspace = true }
+serde_derive = { workspace = true }
+serde_json = { workspace = true }
+zstd = { workspace = true }
+
+[dev-dependencies]
+tempfile = { workspace = true }
+tokio = { workspace = true }

crates/puffin/src/blob.rs

@@ -0,0 +1,126 @@
+// 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 iceberg::io::{FileRead, InputFile};
+use iceberg::Result;
+
+use crate::compression::CompressionCodec;
+use crate::metadata::BlobMetadata;
+
+/// A serialized form of a "compact" Theta sketch produced by the Apache DataSketches library.
+pub const APACHE_DATASKETCHES_THETA_V1: &str = "apache-datasketches-theta-v1";
+
+/// The blob
+#[derive(Debug)]
+pub struct Blob {
+    /// See Blob types: https://iceberg.apache.org/puffin-spec/#blob-types
+    pub r#type: String,
+    /// List of field IDs the blob was computed for; the order of items is used to compute sketches stored in the blob.
+    pub input_fields: Vec<i32>,
+    /// ID of the Iceberg table's snapshot the blob was computed from
+    pub snapshot_id: i64,
+    /// Sequence number of the Iceberg table's snapshot the blob was computed from
+    pub sequence_number: i64,
+    /// The actual blob data
+    pub data: Vec<u8>,
+    /// The compression codec to use to compress the blob data at write time
+    pub requested_compression_codec: Option<CompressionCodec>,
+    /// Arbitrary meta-information about the blob
+    pub properties: HashMap<String, String>,
+}
+
+impl Blob {
+    /// Returns blob from a Puffin file
+    pub async fn read(input_file: &InputFile, blob_metadata: &BlobMetadata) -> Result<Blob> {
+        let file_read = input_file.reader().await?;
+        let start = blob_metadata.offset;
+        let end = start + u64::try_from(blob_metadata.length)?;
+        let bytes = file_read.read(start..end).await?.to_vec();
+        let data = blob_metadata.compression_codec.decompress(bytes)?;
+        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 crate::test_utils::{
+        read_blob_as_utf8_string, read_test_file, METRIC_BLOB_0_DATA, METRIC_BLOB_1_DATA,
+    };
+
+    #[tokio::test]
+    async fn test_read_rust_generated_uncompressed_metric_data() {
+        let input_file = read_test_file("v1/rust-generated/sample-metric-data-uncompressed.bin");
+        assert_eq!(
+            read_blob_as_utf8_string(&input_file, 0).await,
+            METRIC_BLOB_0_DATA
+        );
+        assert_eq!(
+            read_blob_as_utf8_string(&input_file, 1).await,
+            METRIC_BLOB_1_DATA
+        );
+    }
+
+    #[tokio::test]
+    async fn test_read_rust_generated_ztd_compressed_metric_data() {
+        let input_file = read_test_file("v1/rust-generated/sample-metric-data-compressed-zstd.bin");
+        assert_eq!(
+            read_blob_as_utf8_string(&input_file, 0).await,
+            METRIC_BLOB_0_DATA
+        );
+        assert_eq!(
+            read_blob_as_utf8_string(&input_file, 1).await,
+            METRIC_BLOB_1_DATA
+        );
+    }
+
+    #[tokio::test]
+    async fn test_read_java_generated_uncompressed_metric_data() {
+        let input_file = read_test_file("v1/java-generated/sample-metric-data-uncompressed.bin");
+        assert_eq!(
+            read_blob_as_utf8_string(&input_file, 0).await,
+            METRIC_BLOB_0_DATA
+        );
+        assert_eq!(
+            read_blob_as_utf8_string(&input_file, 1).await,
+            METRIC_BLOB_1_DATA
+        );
+    }
+
+    #[tokio::test]
+    async fn test_read_java_generated_ztd_compressed_metric_data() {
+        let input_file = read_test_file("v1/java-generated/sample-metric-data-compressed-zstd.bin");
+        assert_eq!(
+            read_blob_as_utf8_string(&input_file, 0).await,
+            METRIC_BLOB_0_DATA
+        );
+        assert_eq!(
+            read_blob_as_utf8_string(&input_file, 1).await,
+            METRIC_BLOB_1_DATA
+        );
+    }
+}

crates/puffin/src/compression.rs

@@ -0,0 +1,123 @@
+// 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")]
+#[derive(Default)]
+/// Data compression formats
+pub enum CompressionCodec {
+    /// No compression
+    #[default]
+    None,
+    /// LZ4 single compression frame with content size present
+    Lz4,
+    /// Zstandard single compression frame with content size present
+    Zstd,
+}
+
+impl CompressionCodec {
+    pub(crate) fn decompress(&self, bytes: Vec<u8>) -> Result<Vec<u8>> {
+        match self {
+            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[..])?;
+                Ok(decompressed)
+            }
+        }
+    }
+
+    pub(crate) fn compress(&self, bytes: Vec<u8>) -> Result<Vec<u8>> {
+        match self {
+            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()?;
+                Ok(compressed)
+            }
+        }
+    }
+
+    pub(crate) fn is_none(&self) -> bool {
+        matches!(self, CompressionCodec::None)
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use crate::compression::CompressionCodec;
+
+    #[tokio::test]
+    async fn test_compression_codec_none() {
+        let compression_codec = CompressionCodec::None;
+        let bytes_vec = [0_u8; 100].to_vec();
+
+        let compressed = compression_codec.compress(bytes_vec.clone()).unwrap();
+        assert_eq!(bytes_vec, compressed);
+
+        let decompressed = compression_codec.decompress(compressed.clone()).unwrap();
+        assert_eq!(compressed, decompressed)
+    }
+
+    #[tokio::test]
+    async fn test_compression_codec_lz4() {
+        let compression_codec = CompressionCodec::Lz4;
+        let bytes_vec = [0_u8; 100].to_vec();
+
+        assert_eq!(
+            compression_codec
+                .compress(bytes_vec.clone())
+                .unwrap_err()
+                .to_string(),
+            "FeatureUnsupported => LZ4 compression is not supported currently",
+        );
+
+        assert_eq!(
+            compression_codec
+                .decompress(bytes_vec.clone())
+                .unwrap_err()
+                .to_string(),
+            "FeatureUnsupported => LZ4 decompression is not supported currently",
+        )
+    }
+
+    #[tokio::test]
+    async fn test_compression_codec_zstd() {
+        let compression_codec = CompressionCodec::Zstd;
+        let bytes_vec = [0_u8; 100].to_vec();
+
+        let compressed = compression_codec.compress(bytes_vec.clone()).unwrap();
+        assert!(compressed.len() < bytes_vec.len());
+
+        let decompressed = compression_codec.decompress(compressed.clone()).unwrap();
+        assert_eq!(decompressed, bytes_vec)
+    }
+}