add zstd codec for v2

Author: brokkoli71

src/main/java/dev/zarr/zarrjava/core/codec/core/ZstdCodec.java

@@ -0,0 +1,35 @@
+package dev.zarr.zarrjava.core.codec.core;
+
+import com.github.luben.zstd.Zstd;
+import com.github.luben.zstd.ZstdCompressCtx;
+import dev.zarr.zarrjava.ZarrException;
+import dev.zarr.zarrjava.core.codec.BytesBytesCodec;
+import dev.zarr.zarrjava.utils.Utils;
+
+import java.nio.ByteBuffer;
+
+public abstract class ZstdCodec extends BytesBytesCodec {
+
+    @Override
+    public ByteBuffer decode(ByteBuffer compressedBytes) throws ZarrException {
+        byte[] compressedArray = Utils.toArray(compressedBytes);
+        long originalSize = Zstd.getFrameContentSize(compressedArray);
+        if (originalSize < 0) {
+            throw new ZarrException("Failed to get decompressed zstd size.");
+        }
+        byte[] decompressed = Zstd.decompress(compressedArray, (int) originalSize);
+        return ByteBuffer.wrap(decompressed);
+    }
+
+    protected ByteBuffer encodeInternal(int level, boolean checksum, ByteBuffer chunkBytes)
+            throws ZarrException {
+        byte[] arr = Utils.toArray(chunkBytes);
+        byte[] compressed;
+        try (ZstdCompressCtx ctx = new ZstdCompressCtx()) {
+            ctx.setLevel(level);
+            ctx.setChecksum(checksum);
+            compressed = ctx.compress(arr);
+        }
+        return ByteBuffer.wrap(compressed);
+    }
+}

src/main/java/dev/zarr/zarrjava/v2/codec/CodecRegistry.java

@@ -3,6 +3,7 @@
 import com.fasterxml.jackson.databind.jsontype.NamedType;
 import dev.zarr.zarrjava.v2.codec.core.BloscCodec;
 import dev.zarr.zarrjava.v2.codec.core.ZlibCodec;
+import dev.zarr.zarrjava.v2.codec.core.ZstdCodec;
 
 import java.util.HashMap;
 import java.util.Map;
@@ -14,6 +15,7 @@ public class CodecRegistry {
     static {
         addType("blosc", BloscCodec.class);
         addType("zlib", ZlibCodec.class);
+        addType("zstd", ZstdCodec.class);
     }
 
     public static void addType(String name, Class<? extends Codec> codecClass) {

src/main/java/dev/zarr/zarrjava/v2/codec/core/ZstdCodec.java

@@ -0,0 +1,39 @@
+package dev.zarr.zarrjava.v2.codec.core;
+
+import com.fasterxml.jackson.annotation.JsonCreator;
+import com.fasterxml.jackson.annotation.JsonIgnore;
+import com.fasterxml.jackson.annotation.JsonProperty;
+import dev.zarr.zarrjava.ZarrException;
+import dev.zarr.zarrjava.core.ArrayMetadata;
+import dev.zarr.zarrjava.v2.codec.Codec;
+
+import java.nio.ByteBuffer;
+
+public class ZstdCodec extends dev.zarr.zarrjava.core.codec.core.ZstdCodec implements Codec {
+
+    @JsonIgnore
+    public final String id = "zstd";
+    public final int level;
+    public final boolean checksum;
+
+    @JsonCreator(mode = JsonCreator.Mode.PROPERTIES)
+    public ZstdCodec(
+            @JsonProperty(value = "level", defaultValue = "0") int level,
+            @JsonProperty(value = "checksum", defaultValue = "false") boolean checksum) throws ZarrException {
+        if (level < -131072 || level > 22) {
+            throw new ZarrException("'level' needs to be between -131072 and 22.");
+        }
+        this.level = level;
+        this.checksum = checksum;
+    }
+
+    @Override
+    public ByteBuffer encode(ByteBuffer chunkBytes) throws ZarrException {
+        return encodeInternal(this.level, this.checksum, chunkBytes);
+    }
+
+    @Override
+    public Codec evolveFromCoreArrayMetadata(ArrayMetadata.CoreArrayMetadata arrayMetadata) {
+        return this;
+    }
+}

src/main/java/dev/zarr/zarrjava/v3/codec/core/ZstdCodec.java

@@ -3,18 +3,14 @@
 import com.fasterxml.jackson.annotation.JsonCreator;
 import com.fasterxml.jackson.annotation.JsonIgnore;
 import com.fasterxml.jackson.annotation.JsonProperty;
-import com.github.luben.zstd.Zstd;
-import com.github.luben.zstd.ZstdCompressCtx;
 import dev.zarr.zarrjava.ZarrException;
-import dev.zarr.zarrjava.core.codec.BytesBytesCodec;
-import dev.zarr.zarrjava.utils.Utils;
 import dev.zarr.zarrjava.v3.ArrayMetadata;
 import dev.zarr.zarrjava.v3.codec.Codec;
 
 import javax.annotation.Nonnull;
 import java.nio.ByteBuffer;
 
-public class ZstdCodec extends BytesBytesCodec implements Codec {
+public class ZstdCodec extends dev.zarr.zarrjava.core.codec.core.ZstdCodec implements Codec {
 
     @JsonIgnore
     public final String name = "zstd";
@@ -27,29 +23,9 @@ public ZstdCodec(
         this.configuration = configuration;
     }
 
-    @Override
-    public ByteBuffer decode(ByteBuffer compressedBytes) throws ZarrException {
-        byte[] compressedArray = Utils.toArray(compressedBytes);
-
-        long originalSize = Zstd.getFrameContentSize(compressedArray);
-        if (originalSize == 0) {
-            throw new ZarrException("Failed to get decompressed size");
-        }
-
-        byte[] decompressed = Zstd.decompress(compressedArray, (int) originalSize);
-        return ByteBuffer.wrap(decompressed);
-    }
-
     @Override
     public ByteBuffer encode(ByteBuffer chunkBytes) throws ZarrException {
-        byte[] arr = Utils.toArray(chunkBytes);
-        byte[] compressed;
-        try (ZstdCompressCtx ctx = new ZstdCompressCtx()) {
-            ctx.setLevel(configuration.level);
-            ctx.setChecksum(configuration.checksum);
-            compressed = ctx.compress(arr);
-        }
-        return ByteBuffer.wrap(compressed);
+        return encodeInternal(configuration.level, configuration.checksum, chunkBytes);
     }
 
     @Override
@@ -75,5 +51,3 @@ public Configuration(@JsonProperty(value = "level", defaultValue = "5") int leve
         }
     }
 }
-
-

src/test/java/dev/zarr/zarrjava/ZarrPythonTests.java

@@ -106,7 +106,9 @@ static Stream<Object[]> compressorAndDataTypeProviderV2() {
                 new Object[]{"blosc", "lz4_shuffle_6", dev.zarr.zarrjava.v2.DataType.INT32},
                 new Object[]{"blosc", "lz4hc_bitshuffle_3", dev.zarr.zarrjava.v2.DataType.INT32},
                 new Object[]{"blosc", "zlib_shuffle_5", dev.zarr.zarrjava.v2.DataType.INT32},
-                new Object[]{"blosc", "zstd_bitshuffle_9", dev.zarr.zarrjava.v2.DataType.INT32}
+                new Object[]{"blosc", "zstd_bitshuffle_9", dev.zarr.zarrjava.v2.DataType.INT32},
+                new Object[]{"zstd", "0_true", dev.zarr.zarrjava.v2.DataType.INT32},
+                new Object[]{"zstd", "5_false", dev.zarr.zarrjava.v2.DataType.INT32}
         );
 
         return Stream.concat(datatypeTests, bloscTests);

src/test/python-scripts/parse_codecs.py

@@ -48,6 +48,9 @@ def parse_codecs_zarr_python(codec_string: str, param_string: str, zarr_version:
                                                                                 codecs=[BytesCodec(endian="little")]),))
     elif codec_string == "crc32c" and zarr_version == 3:
         compressor = Crc32cCodec()
+    elif codec_string == "zstd" and zarr_version == 2:
+        level, checksum = param_string.split("_")
+        compressor = numcodecs.Zstd(level=int(level), checksum=checksum == 'true')
     else:
         raise ValueError(f"Invalid codec: {codec_string}, zarr_version: {zarr_version}")