Skip to content

Commit 3984a9e

Browse files
authored
fix: zstd with sharding (#56)
`zstd` would try end decode the entire shard instead of a subchunk. The added test fails without the patch with: `ZarrV3Test.testShardingWithZstdCodecReadWrite:225 » Zstd Destination buffer is too small`
1 parent 32f785f commit 3984a9e

2 files changed

Lines changed: 24 additions & 2 deletions

File tree

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

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -7,6 +7,7 @@
77
import com.github.luben.zstd.ZstdCompressCtx;
88
import dev.zarr.zarrjava.ZarrException;
99
import dev.zarr.zarrjava.core.codec.BytesBytesCodec;
10+
import dev.zarr.zarrjava.utils.Utils;
1011
import dev.zarr.zarrjava.v3.ArrayMetadata;
1112
import dev.zarr.zarrjava.v3.codec.Codec;
1213

@@ -28,7 +29,7 @@ public ZstdCodec(
2829

2930
@Override
3031
public ByteBuffer decode(ByteBuffer compressedBytes) throws ZarrException {
31-
byte[] compressedArray = compressedBytes.array();
32+
byte[] compressedArray = Utils.toArray(compressedBytes);
3233

3334
long originalSize = Zstd.getFrameContentSize(compressedArray);
3435
if (originalSize == 0) {
@@ -41,7 +42,7 @@ public ByteBuffer decode(ByteBuffer compressedBytes) throws ZarrException {
4142

4243
@Override
4344
public ByteBuffer encode(ByteBuffer chunkBytes) throws ZarrException {
44-
byte[] arr = chunkBytes.array();
45+
byte[] arr = Utils.toArray(chunkBytes);
4546
byte[] compressed;
4647
try (ZstdCompressCtx ctx = new ZstdCompressCtx()) {
4748
ctx.setLevel(configuration.level);

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

Lines changed: 21 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -206,6 +206,27 @@ public void testZstdCodecReadWrite(int clevel, boolean checksum) throws ZarrExce
206206
Assertions.assertArrayEquals(testData, (int[]) result.get1DJavaArray(ucar.ma2.DataType.UINT));
207207
}
208208

209+
@Test
210+
public void testShardingWithZstdCodecReadWrite() throws ZarrException, IOException {
211+
int[] testData = new int[16 * 16 * 16];
212+
Arrays.setAll(testData, p -> p);
213+
214+
StoreHandle storeHandle = new FilesystemStore(TESTOUTPUT).resolve("testShardingWithZstdCodecReadWrite");
215+
ArrayMetadataBuilder builder = Array.metadataBuilder()
216+
.withShape(16, 16, 16)
217+
.withDataType(DataType.UINT32)
218+
.withChunkShape(8, 8, 8)
219+
.withFillValue(0)
220+
.withCodecs(c -> c.withSharding(new int[]{2, 4, 8}, c1 -> c1.withZstd()));
221+
Array writeArray = Array.create(storeHandle, builder.build());
222+
writeArray.write(ucar.ma2.Array.factory(ucar.ma2.DataType.UINT, new int[]{16, 16, 16}, testData));
223+
224+
Array readArray = Array.open(storeHandle);
225+
ucar.ma2.Array result = readArray.read();
226+
227+
Assertions.assertArrayEquals(testData, (int[]) result.get1DJavaArray(ucar.ma2.DataType.UINT));
228+
}
229+
209230
@Test
210231
public void testTransposeCodec() throws ZarrException {
211232
ucar.ma2.Array testData = ucar.ma2.Array.factory(ucar.ma2.DataType.UINT, new int[]{2, 3, 3}, new int[]{

0 commit comments

Comments
 (0)