Add ZstdCompressor (#180)

* Add ZstdCompressor

* fix typo

* Use decode!

Author: nhz2

Project.toml

@@ -6,6 +6,8 @@ version = "0.9.4"
 [deps]
 AWSS3 = "1c724243-ef5b-51ab-93f4-b0a88ac62a95"
 Blosc = "a74b3585-a348-5f62-a45c-50e91977d574"
+ChunkCodecCore = "0b6fb165-00bc-4d37-ab8b-79f91016dbe1"
+ChunkCodecLibZstd = "55437552-ac27-4d47-9aa3-63184e8fd398"
 CodecZlib = "944b1d66-785c-5afd-91f1-9de20f533193"
 DataStructures = "864edb3b-99cc-5e75-8d2d-829cb0a9cfe8"
 DateTimes64 = "b342263e-b350-472a-b1a9-8dfd21b51589"
@@ -22,6 +24,8 @@ ZipArchives = "49080126-0e18-4c2a-b176-c102e4b3760c"
 [compat]
 AWSS3 = "0.10, 0.11"
 Blosc = "0.5, 0.6, 0.7"
+ChunkCodecCore = "0.4.2"
+ChunkCodecLibZstd = "0.1.2"
 CodecZlib = "0.6, 0.7"
 DataStructures = "0.17, 0.18"
 DateTimes64 = "1"
@@ -32,4 +36,4 @@ OffsetArrays = "0.11, 1.0"
 OpenSSL = "1"
 URIs = "1"
 ZipArchives = "2"
-julia = "1.2"
+julia = "1.10"

docs/src/reference.md

@@ -18,5 +18,5 @@ Pages = ["ZGroup.jl"]
 
 ```@autodocs
 Modules = [Zarr]
-Pages = ["Compressors/Compressors.jl", "Compressors/blosc.jl", "Compressors/zlib.jl"]
+Pages = ["Compressors/Compressors.jl", "Compressors/blosc.jl", "Compressors/zlib.jl", "Compressors/zstd.jl"]
 ```

src/Compressors/Compressors.jl

@@ -48,6 +48,7 @@ const compressortypes = Dict{Union{String,Nothing}, Type{<: Compressor}}()
 # Include the compressor implementations
 include("blosc.jl")
 include("zlib.jl")
+include("zstd.jl")
 
 # ## Fallback definitions for the compressor interface
 # Define fallbacks and generic methods for the compressor interface

src/Compressors/zstd.jl

@@ -0,0 +1,45 @@
+#=
+# Zstd compression
+
+This file implements a Zstd compressor via ChunkCodecLibZstd.jl.
+
+=#
+using ChunkCodecLibZstd: ZstdEncodeOptions
+using ChunkCodecCore: encode, decode, decode!
+
+"""
+    ZstdCompressor(;level=0, checksum=false)
+Returns a `ZstdCompressor` struct that can serve as a Zarr array compressor. Keyword arguments are:
+* `level=0`: the compression level, regular levels are 1 to 22, 0 is a special value for default, there are also even faster negative levels.
+* `checksum=false`: flag to enable saving checksums.
+"""
+struct ZstdCompressor <: Compressor
+    config::ZstdEncodeOptions
+end
+
+ZstdCompressor(;level=0, checksum::Bool=false) = ZstdCompressor(ZstdEncodeOptions(;compressionLevel=level, checksum))
+
+function getCompressor(::Type{ZstdCompressor}, d::Dict)
+    ZstdCompressor(;
+        level=get(Returns(0), d, "level"),
+        checksum=Bool(get(Returns(false), d, "checksum")),
+    )
+end
+
+function zuncompress(a, z::ZstdCompressor, T)
+    result = decode(z.config.codec, a)
+    _reinterpret(Base.nonmissingtype(T),result)
+end
+
+function zuncompress!(data::DenseArray, compressed, z::ZstdCompressor)
+    decode!(z.config.codec, reinterpret(UInt8, vec(data)), compressed)
+    data
+end
+
+function zcompress(a, z::ZstdCompressor)
+    encode(z.config, reinterpret(UInt8, vec(a)))
+end
+
+JSON.lower(z::ZstdCompressor) = Dict("id"=>"zstd", "level" => z.config.compressionLevel, "checksum" => z.config.checksum)
+
+Zarr.compressortypes["zstd"] = ZstdCompressor

test/python.jl

@@ -21,7 +21,7 @@ groupattrs = Dict("String attribute"=>"One", "Int attribute"=>5, "Float attribut
 g = zgroup(pjulia,attrs=groupattrs)
 
 # Test all supported data types and compressors
-import Zarr: NoCompressor, BloscCompressor, ZlibCompressor, MaxLengthString, 
+import Zarr: NoCompressor, BloscCompressor, ZlibCompressor, ZstdCompressor, MaxLengthString, 
        Fletcher32Filter, FixedScaleOffsetFilter, ShuffleFilter, QuantizeFilter, DeltaFilter
 using Random: randstring
 numeric_dtypes = (UInt8, UInt16, UInt32, UInt64,
@@ -38,7 +38,9 @@ compressors = (
     "blosc_autoshuffle"=>BloscCompressor(cname="zstd",shuffle=-1),
     "blosc_noshuffle"=>BloscCompressor(cname="zstd",shuffle=0),
     "blosc_bitshuffle"=>BloscCompressor(cname="zstd",shuffle=2),
-    "zlib"=>ZlibCompressor())
+    "zlib"=>ZlibCompressor(),
+    "zstd"=>ZstdCompressor(),
+)
 filters = (
     "fletcher32"=>Fletcher32Filter(),
     "scale_offset"=>FixedScaleOffsetFilter(offset=1000, scale=10^6, T=Float64, Tenc=Int32),