Add ZSTD support to Rust library (google#84)

kralka · Jan 7, 2025 · 0261b66 · 0261b66
1 parent bef502e
commit 0261b66
Show file tree

Hide file tree

Showing 4 changed files with 52 additions and 0 deletions.
diff --git a/project-words.txt b/project-words.txt
@@ -45,3 +45,4 @@ tfrecord
 tfrecords
 tobytes
 uoffset
+zstd
diff --git a/rust/Cargo.lock b/rust/Cargo.lock
diff --git a/rust/Cargo.toml b/rust/Cargo.toml
@@ -28,3 +28,4 @@ rand = "0.8"
 strum = "0.26.3"
 strum_macros = "0.26.4"
 yoke = { version = "0.7", features = ["derive"] }
+zstd = "0.13.2"
diff --git a/rust/src/example_iteration.rs b/rust/src/example_iteration.rs
@@ -33,6 +33,7 @@ pub enum CompressionType {
     LZ4,
     Gzip,
     Zlib,
+    Zstd,
 }
 
 impl CompressionType {
@@ -48,6 +49,7 @@ impl std::fmt::Display for CompressionType {
             CompressionType::LZ4 => write!(f, "LZ4"),
             CompressionType::Gzip => write!(f, "GZIP"),
             CompressionType::Zlib => write!(f, "ZLIB"),
+            CompressionType::Zstd => write!(f, "ZSTD"),
         }
     }
 }
@@ -61,6 +63,7 @@ impl std::str::FromStr for CompressionType {
             "LZ4" => Ok(CompressionType::LZ4),
             "GZIP" => Ok(CompressionType::Gzip),
             "ZLIB" => Ok(CompressionType::Zlib),
+            "ZSTD" => Ok(CompressionType::Zstd),
             _ => Err(format!("{input} unimplemented")),
         }
     }
@@ -128,6 +131,7 @@ fn get_file_bytes(shard_info: &ShardInfo) -> Vec<u8> {
         CompressionType::Gzip | CompressionType::Zlib => {
             read_to_end(flate2::read::GzDecoder::new(open_file))
         }
+        CompressionType::Zstd => zstd::stream::decode_all(open_file).unwrap(),
     }
 }
-Original file line number
+Diff line change
@@ Expand Up / @@ -45,3 +45,4 @@ tfrecord @@
     tfrecords
     tobytes
     uoffset
+    zstd