ZarrPythonTests.java

package dev.zarr.zarrjava;

import com.github.luben.zstd.Zstd;
import com.github.luben.zstd.ZstdCompressCtx;
import dev.zarr.zarrjava.core.Attributes;
import dev.zarr.zarrjava.store.FilesystemStore;
import dev.zarr.zarrjava.store.StoreHandle;
import dev.zarr.zarrjava.v2.Group;
import dev.zarr.zarrjava.v3.Array;
import dev.zarr.zarrjava.v3.ArrayMetadataBuilder;
import dev.zarr.zarrjava.v3.DataType;
import dev.zarr.zarrjava.v3.codec.CodecBuilder;
import org.junit.jupiter.api.Assertions;
import org.junit.jupiter.api.BeforeAll;
import org.junit.jupiter.api.Test;
import org.junit.jupiter.params.ParameterizedTest;
import org.junit.jupiter.params.provider.CsvSource;
import org.junit.jupiter.params.provider.MethodSource;

import java.io.BufferedReader;
import java.io.FileOutputStream;
import java.io.IOException;
import java.io.InputStreamReader;
import java.nio.ByteBuffer;
import java.nio.file.Path;
import java.nio.file.Paths;
import java.util.Arrays;
import java.util.stream.Stream;


public class ZarrPythonTests extends ZarrTest {

    final static Path PYTHON_TEST_PATH = Paths.get("src/test/python-scripts/");

    public static int runCommand(String... command) throws IOException, InterruptedException {
        ProcessBuilder pb = new ProcessBuilder();
        pb.command().addAll(Arrays.asList(command));
        Process process = pb.start();

        BufferedReader reader = new BufferedReader(new InputStreamReader(process.getInputStream()));
        String line;
        while ((line = reader.readLine()) != null) {
            System.out.println(line);
        }

        BufferedReader readerErr = new BufferedReader(new InputStreamReader(process.getErrorStream()));
        while ((line = readerErr.readLine()) != null) {
            System.err.println(line);
        }

        return process.waitFor();
    }

    @BeforeAll
    public static void setupUV() {
        try {
            int exitCode = runCommand("uv", "version");
            if (exitCode != 0) {
                //setup uv
                assert runCommand("uv", "venv") == 0;
                assert runCommand("uv", "init") == 0;
                assert runCommand("uv", "add", "zarr", "zstandard") == 0;
            }
        } catch (IOException | InterruptedException e) {
            throw new RuntimeException("uv not installed or not in PATH. See");
        }
    }

    static Stream<Object[]> compressorAndDataTypeProviderV3() {
        Stream<Object[]> datatypeTests = dataTypeProviderV3().flatMap(dt -> Stream.of(
                new Object[]{"sharding", "end", dt},
                new Object[]{"blosc", "blosclz_shuffle_3", dt}
        ));

        Stream<Object[]> codecsTests = Stream.of(
                new Object[]{"blosc", "blosclz_noshuffle_0", DataType.INT32},
                new Object[]{"blosc", "lz4_shuffle_6", DataType.INT32},
                new Object[]{"blosc", "lz4hc_bitshuffle_3", DataType.INT32},
                new Object[]{"blosc", "zlib_shuffle_5", DataType.INT32},
                new Object[]{"blosc", "zstd_bitshuffle_9", DataType.INT32},
                new Object[]{"gzip", "0", DataType.INT32},
                new Object[]{"gzip", "5", DataType.INT32},
                new Object[]{"zstd", "0_true", DataType.INT32},
                new Object[]{"zstd", "5_true", DataType.INT32},
                new Object[]{"zstd", "0_false", DataType.INT32},
                new Object[]{"zstd", "5_false", DataType.INT32},
                new Object[]{"bytes", "BIG", DataType.INT32},
                new Object[]{"bytes", "LITTLE", DataType.INT32},
                new Object[]{"transpose", "_", DataType.INT32},
                new Object[]{"sharding", "start", DataType.INT32},
                new Object[]{"sharding_nested", "_", DataType.INT32},
                new Object[]{"crc32c", "_", DataType.INT32}
        );

        return Stream.concat(datatypeTests, codecsTests);
    }

    static Stream<Object[]> compressorAndDataTypeProviderV2() {
        Stream<Object[]> datatypeTests = dataTypeProviderV2().flatMap(dt -> Stream.of(
                new Object[]{"zlib", "0", dt},
                new Object[]{"blosc", "blosclz_shuffle_3", dt}
        ));

        Stream<Object[]> bloscTests = Stream.of(
                new Object[]{"blosc", "blosclz_noshuffle_0", dev.zarr.zarrjava.v2.DataType.INT32},
                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[]{"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);
    }

    public void run_python_script(String scriptName, String... args) throws IOException, InterruptedException {
        int exitCode = runCommand(Stream.concat(Stream.of("uv", "run", PYTHON_TEST_PATH.resolve(scriptName)
                .toString()), Arrays.stream(args)).toArray(String[]::new));
        assert exitCode == 0;
    }

    @ParameterizedTest
    @MethodSource("compressorAndDataTypeProviderV3")
    public void testReadV3(String codec, String codecParam, DataType dataType) throws IOException, ZarrException, InterruptedException {
        StoreHandle storeHandle = new FilesystemStore(TESTOUTPUT).resolve("testReadV3", codec, codecParam, dataType.name());
        run_python_script("zarr_python_write.py", codec, codecParam, dataType.name().toLowerCase(), storeHandle.toPath().toString());
        Array array = Array.open(storeHandle);
        ucar.ma2.Array result = array.read();

        Assertions.assertArrayEquals(new int[]{16, 16, 16}, result.getShape());
        Assertions.assertEquals(dataType, array.metadata().dataType);
        Assertions.assertArrayEquals(new int[]{2, 4, 8}, array.metadata().chunkShape());
        Assertions.assertEquals(42, array.metadata().attributes.get("answer"));

        assertIsTestdata(result, dataType);
    }

    @ParameterizedTest
    @MethodSource("compressorAndDataTypeProviderV3")
    public void testWriteV3(String codec, String codecParam, DataType dataType) throws Exception {
        Attributes attributes = new Attributes();
        attributes.put("test_key", "test_value");
        StoreHandle storeHandle = new FilesystemStore(TESTOUTPUT).resolve("testWriteV3", codec, codecParam, dataType.name());

        ArrayMetadataBuilder builder = Array.metadataBuilder()
                .withShape(16, 16, 16)
                .withDataType(dataType)
                .withChunkShape(2, 4, 8)
                .withFillValue(0)
                .withAttributes(attributes);

        switch (codec) {
            case "blosc":
                String cname = codecParam.split("_")[0];
                String shuffle = codecParam.split("_")[1];
                int clevel_blosc = Integer.parseInt(codecParam.split("_")[2]);
                builder = builder.withCodecs(c -> c.withBlosc(cname, shuffle, clevel_blosc));
                break;
            case "gzip":
                builder = builder.withCodecs(c -> c.withGzip(Integer.parseInt(codecParam)));
                break;
            case "zstd":
                int level_zstd = Integer.parseInt(codecParam.split("_")[0]);
                boolean checksum = Boolean.parseBoolean(codecParam.split("_")[1]);
                builder = builder.withCodecs(c -> c.withZstd(level_zstd, checksum));
                break;
            case "bytes":
                builder = builder.withCodecs(c -> c.withBytes(codecParam));
                break;
            case "transpose":
                builder = builder.withCodecs(c -> c.withTranspose(new int[]{1, 0, 2}));
                break;
            case "sharding":
                builder = builder.withCodecs(c -> c.withSharding(new int[]{2, 2, 4}, c1 -> c1.withBytes("LITTLE"), codecParam));
                break;
            case "sharding_nested":
                builder = builder.withCodecs(c -> c.withSharding(new int[]{2, 2, 4}, c1 -> c1.withSharding(new int[]{2, 1, 2}, c2 -> c2.withBytes("LITTLE"))));
                break;
            case "crc32c":
                builder = builder.withCodecs(CodecBuilder::withCrc32c);
                break;
            default:
                throw new IllegalArgumentException("Invalid Codec: " + codec);
        }

        Array writeArray = Array.create(storeHandle, builder.build());
        writeArray.write(testdata(dataType));

        //read in zarr-java
        Array readArray = Array.open(storeHandle);
        ucar.ma2.Array result = readArray.read();

        Assertions.assertArrayEquals(new int[]{16, 16, 16}, result.getShape());
        Assertions.assertEquals(dataType, readArray.metadata().dataType);
        Assertions.assertArrayEquals(new int[]{2, 4, 8}, readArray.metadata().chunkShape());
        Assertions.assertEquals("test_value", readArray.metadata().attributes.get("test_key"));

        assertIsTestdata(result, dataType);

        //read in zarr_python
        run_python_script("zarr_python_read.py", codec, codecParam, dataType.name().toLowerCase(), storeHandle.toPath().toString());
    }

    @ParameterizedTest
    @MethodSource("compressorAndDataTypeProviderV2")
    public void testReadV2(String compressor, String compressorParam, dev.zarr.zarrjava.v2.DataType dt) throws IOException, ZarrException, InterruptedException {
        StoreHandle storeHandle = new FilesystemStore(TESTOUTPUT).resolve("testReadV2", compressor, compressorParam, dt.name());
        run_python_script("zarr_python_write_v2.py", compressor, compressorParam, dt.getValue(), storeHandle.toPath().toString());

        dev.zarr.zarrjava.v2.Array array = dev.zarr.zarrjava.v2.Array.open(storeHandle);
        ucar.ma2.Array result = array.read();

        Assertions.assertArrayEquals(new int[]{16, 16, 16}, result.getShape());
        Assertions.assertEquals(dt, array.metadata().dataType);
        Assertions.assertArrayEquals(new int[]{2, 4, 8}, array.metadata().chunkShape());
        Assertions.assertEquals(42, array.metadata().attributes().get("answer"));

        assertIsTestdata(result, dt);
    }


    @ParameterizedTest
    @MethodSource("compressorAndDataTypeProviderV2")
    public void testWriteV2(String compressor, String compressorParam, dev.zarr.zarrjava.v2.DataType dt) throws Exception {
        Attributes attributes = new Attributes();
        attributes.put("test_key", "test_value");
        StoreHandle storeHandle = new FilesystemStore(TESTOUTPUT).resolve("testCodecsWriteV2", compressor, compressorParam, dt.name());

        dev.zarr.zarrjava.v2.ArrayMetadataBuilder builder = dev.zarr.zarrjava.v2.Array.metadataBuilder()
                .withShape(16, 16, 16)
                .withDataType(dt)
                .withChunks(2, 4, 8)
                .withAttributes(attributes)
                .withFillValue(0);

        switch (compressor) {
            case "blosc":
                String cname = compressorParam.split("_")[0];
                String shuffle = compressorParam.split("_")[1];
                int clevel_blosc = Integer.parseInt(compressorParam.split("_")[2]);
                builder = builder.withBloscCompressor(cname, shuffle, clevel_blosc);
                break;
            case "zlib":
                builder = builder.withZlibCompressor(Integer.parseInt(compressorParam));
                break;
            case "zstd":
                builder = builder.withZstdCompressor(Integer.parseInt(compressorParam.split("_")[0]), Boolean.parseBoolean(compressorParam.split("_")[1]));
                break;
            default:
                throw new IllegalArgumentException("Invalid compressor: " + compressor);
        }

        dev.zarr.zarrjava.v2.Array writeArray = dev.zarr.zarrjava.v2.Array.create(storeHandle, builder.build());
        writeArray.write(testdata(dt));

        //read in zarr-java
        dev.zarr.zarrjava.v2.Array readArray = dev.zarr.zarrjava.v2.Array.open(storeHandle);
        ucar.ma2.Array result = readArray.read();

        Assertions.assertArrayEquals(new int[]{16, 16, 16}, result.getShape());
        Assertions.assertEquals(dt, readArray.metadata().dataType);
        Assertions.assertArrayEquals(new int[]{2, 4, 8}, readArray.metadata().chunkShape());
        Assertions.assertEquals("test_value", readArray.metadata().attributes().get("test_key"));
        assertIsTestdata(result, dt);

        //read in zarr_python
        run_python_script("zarr_python_read_v2.py", compressor, compressorParam, dt.getValue(), storeHandle.toPath().toString());
    }

    @CsvSource({"0,true", "0,false", "5, true", "10, false"})
    @ParameterizedTest
    public void testZstdLibrary(int level, boolean checksumFlag) throws IOException, InterruptedException {
        //compress using ZstdCompressCtx
        int number = 123456;
        byte[] src = ByteBuffer.allocate(4).putInt(number).array();
        byte[] compressed;
        try (ZstdCompressCtx ctx = new ZstdCompressCtx()) {
            ctx.setLevel(level);
            ctx.setChecksum(checksumFlag);
            compressed = ctx.compress(src);
        }
        //decompress with Zstd.decompress
        long originalSize = Zstd.getFrameContentSize(compressed);
        byte[] decompressed = Zstd.decompress(compressed, (int) originalSize);
        Assertions.assertEquals(number, ByteBuffer.wrap(decompressed).getInt());

        //write compressed to file
        String compressedDataPath = TESTOUTPUT.resolve("compressed" + level + checksumFlag + ".bin").toString();
        try (FileOutputStream fos = new FileOutputStream(compressedDataPath)) {
            fos.write(compressed);
        }

        //decompress in python
        int exitCode = ZarrPythonTests.runCommand(
                "uv",
                "run",
                PYTHON_TEST_PATH.resolve("zstd_decompress.py").toString(),
                compressedDataPath,
                Integer.toString(number)
        );
        assert exitCode == 0;
    }

    @Test
    public void testGroupReadWriteV2() throws Exception {
        StoreHandle storeHandle = new FilesystemStore(TESTOUTPUT).resolve("testGroupReadWriteV2", "write");
        StoreHandle storeHandle2 = new FilesystemStore(TESTOUTPUT).resolve("testGroupReadWriteV2", "read");
        Group group = Group.create(storeHandle, new Attributes(b -> b.set("attr", "value")));
        dev.zarr.zarrjava.v2.DataType dataType = dev.zarr.zarrjava.v2.DataType.INT32;
        dev.zarr.zarrjava.v2.Array array = group.createGroup("group").createArray("array", arrayMetadataBuilder -> arrayMetadataBuilder
                .withShape(16, 16, 16)
                .withDataType(dataType)
                .withChunks(2, 4, 8)
        );

        array.write(testdata(dataType));

        run_python_script("zarr_python_group.py", storeHandle.toPath().toString(), storeHandle2.toPath().toString(), "" + 2);

        Group group2 = Group.open(storeHandle2);
        Assertions.assertEquals("value", group2.metadata().attributes().get("attr"));
        Group subgroup = (Group) group2.get("group2");
        Assertions.assertNotNull(subgroup);
        dev.zarr.zarrjava.v2.Array array2 = (dev.zarr.zarrjava.v2.Array) subgroup.get("array2");
        Assertions.assertNotNull(array2);
        ucar.ma2.Array result = array2.read();
        Assertions.assertArrayEquals(new int[]{16, 16, 16}, result.getShape());
        assertIsTestdata(result, dataType);
    }

    @Test
    public void testGroupReadWriteV3() throws Exception {
        StoreHandle storeHandle = new FilesystemStore(TESTOUTPUT).resolve("testGroupReadWriteV3", "write");
        StoreHandle storeHandle2 = new FilesystemStore(TESTOUTPUT).resolve("testGroupReadWriteV3", "read");
        dev.zarr.zarrjava.v3.Group group = dev.zarr.zarrjava.v3.Group.create(storeHandle, new Attributes(b -> b.set("attr", "value")));
        dev.zarr.zarrjava.v3.DataType dataType = DataType.INT32;
        dev.zarr.zarrjava.v3.Array array = group.createGroup("group").createArray("array", arrayMetadataBuilder -> arrayMetadataBuilder
                .withShape(16, 16, 16)
                .withDataType(dataType)
                .withChunkShape(2, 4, 8)
        );

        array.write(testdata(dataType));

        run_python_script("zarr_python_group.py", storeHandle.toPath().toString(), storeHandle2.toPath().toString(), "" + 3);

        dev.zarr.zarrjava.v3.Group group2 = dev.zarr.zarrjava.v3.Group.open(storeHandle2);
        Assertions.assertEquals("value", group2.metadata().attributes().get("attr"));
        dev.zarr.zarrjava.v3.Group subgroup = (dev.zarr.zarrjava.v3.Group) group2.get("group2");
        Assertions.assertNotNull(subgroup);
        dev.zarr.zarrjava.v3.Array array2 = (dev.zarr.zarrjava.v3.Array) subgroup.get("array2");
        Assertions.assertNotNull(array2);
        ucar.ma2.Array result = array2.read();
        Assertions.assertArrayEquals(new int[]{16, 16, 16}, result.getShape());
        assertIsTestdata(result, dataType);
    }

    /**
     * The Zarr v2 {@code order} cases, as (shape, chunks, order).
     */
    static Stream<Object[]> orderProviderV2() {
        return Stream.of(
                new Object[]{"3,4", "3,4", "C"},
                new Object[]{"3,4", "3,4", "F"},
                new Object[]{"2,3,4", "2,3,4", "F"},
                // Chunks divide neither dimension evenly, so the trailing chunks are partial.
                new Object[]{"5,7", "2,3", "F"},
                new Object[]{"5,7,3", "2,3,2", "F"},
                // For rank 1 both orders describe the same layout.
                new Object[]{"10", "4", "F"}
        );
    }

    /**
     * Live counterpart to {@link ZarrV2OrderTest}, which covers the same ground offline against
     * committed fixtures. This runs against the installed zarr-python so that a change in the
     * reference implementation is noticed rather than silently diverging from the fixtures.
     */
    @ParameterizedTest
    @MethodSource("orderProviderV2")
    public void testReadOrderV2(String shape, String chunks, String order) throws Exception {
        StoreHandle storeHandle = new FilesystemStore(TESTOUTPUT)
                .resolve("testReadOrderV2", shape + "_" + chunks + "_" + order);
        run_python_script("zarr_python_order_v2.py", "write", storeHandle.toPath().toString(),
                shape, chunks, order);

        dev.zarr.zarrjava.v2.Array array = dev.zarr.zarrjava.v2.Array.open(storeHandle);
        ucar.ma2.Array result = array.read();

        Assertions.assertEquals(dev.zarr.zarrjava.v2.Order.valueOf(order), array.metadata().order);
        for (int i = 0; i < result.getSize(); i++) {
            Assertions.assertEquals(i, result.getInt(i),
                    "element " + i + " of a " + order + "-order array");
        }
    }

    /**
     * The write direction, verified by zarr-python. zarr-java used to write {@code "order": "F"}
     * into {@code .zarray} while laying the chunk out row-major, which its own reader then read back
     * correctly &mdash; so only an external reader can catch it.
     */
    @ParameterizedTest
    @MethodSource("orderProviderV2")
    public void testWriteOrderV2(String shape, String chunks, String order) throws Exception {
        StoreHandle storeHandle = new FilesystemStore(TESTOUTPUT)
                .resolve("testWriteOrderV2", shape + "_" + chunks + "_" + order);

        long[] shapeArray = Arrays.stream(shape.split(",")).mapToLong(Long::parseLong).toArray();
        int[] chunksArray = Arrays.stream(chunks.split(",")).mapToInt(Integer::parseInt).toArray();

        dev.zarr.zarrjava.v2.Array array = dev.zarr.zarrjava.v2.Array.create(
                storeHandle,
                dev.zarr.zarrjava.v2.Array.metadataBuilder()
                        .withShape(shapeArray)
                        .withChunks(chunksArray)
                        .withDataType(dev.zarr.zarrjava.v2.DataType.INT32)
                        .withOrder(dev.zarr.zarrjava.v2.Order.valueOf(order))
                        .withFillValue(0)
                        .build()
        );

        int[] intShape = new int[shapeArray.length];
        for (int i = 0; i < shapeArray.length; i++) {
            intShape[i] = (int) shapeArray[i];
        }
        ucar.ma2.Array data = ucar.ma2.Array.factory(ucar.ma2.DataType.INT, intShape);
        for (int i = 0; i < data.getSize(); i++) {
            data.setInt(i, i);
        }
        array.write(new long[shapeArray.length], data);

        run_python_script("zarr_python_order_v2.py", "read", storeHandle.toPath().toString(),
                shape, chunks, order);
    }
}