From 52aab0108f27cb12bcb64325b48e3762a961979f Mon Sep 17 00:00:00 2001
From: Brandon Neth <brandon.neth@hpe.com>
Date: Tue, 5 Mar 2024 15:32:15 -0700
Subject: [PATCH 1/3] Implementation and testing for zarr reader and writer ---
 Signed-off-by: Brandon Neth <brandon.neth@hpe.com>

---
 test/library/draft/Zarr/COMPOPTS            |   1 +
 test/library/draft/Zarr/EXECENV             |   1 +
 test/library/draft/Zarr/PRECOMP             |   9 +
 test/library/draft/Zarr/SKIPIF              |   2 +
 test/library/draft/Zarr/Zarr.chpl           | 403 ++++++++++++++++++++
 test/library/draft/Zarr/Zarr.notest         |   0
 test/library/draft/Zarr/ZarrPerf.chpl       |  64 ++++
 test/library/draft/Zarr/ZarrPerf.notest     |   0
 test/library/draft/Zarr/ZarrTest.chpl       |  82 ++++
 test/library/draft/Zarr/ZarrTest.cleanfiles |   6 +
 test/library/draft/Zarr/ZarrTest.good       |   5 +
 test/library/draft/Zarr/ZarrTest.numlocales |   1 +
 12 files changed, 574 insertions(+)
 create mode 100644 test/library/draft/Zarr/COMPOPTS
 create mode 100644 test/library/draft/Zarr/EXECENV
 create mode 100755 test/library/draft/Zarr/PRECOMP
 create mode 100644 test/library/draft/Zarr/SKIPIF
 create mode 100644 test/library/draft/Zarr/Zarr.chpl
 create mode 100644 test/library/draft/Zarr/Zarr.notest
 create mode 100644 test/library/draft/Zarr/ZarrPerf.chpl
 create mode 100644 test/library/draft/Zarr/ZarrPerf.notest
 create mode 100644 test/library/draft/Zarr/ZarrTest.chpl
 create mode 100644 test/library/draft/Zarr/ZarrTest.cleanfiles
 create mode 100644 test/library/draft/Zarr/ZarrTest.good
 create mode 100644 test/library/draft/Zarr/ZarrTest.numlocales

diff --git a/test/library/draft/Zarr/COMPOPTS b/test/library/draft/Zarr/COMPOPTS
new file mode 100644
index 000000000000..e41f09c7b90b
--- /dev/null
+++ b/test/library/draft/Zarr/COMPOPTS
@@ -0,0 +1 @@
+-I./c-blosc/blosc -L./c-blosc/build/blosc  
diff --git a/test/library/draft/Zarr/EXECENV b/test/library/draft/Zarr/EXECENV
new file mode 100644
index 000000000000..531c310e2040
--- /dev/null
+++ b/test/library/draft/Zarr/EXECENV
@@ -0,0 +1 @@
+LD_LIBRARY_PATH=.
diff --git a/test/library/draft/Zarr/PRECOMP b/test/library/draft/Zarr/PRECOMP
new file mode 100755
index 000000000000..e3b57207a1b0
--- /dev/null
+++ b/test/library/draft/Zarr/PRECOMP
@@ -0,0 +1,9 @@
+#!/bin/bash
+git clone https://github.com/Blosc/c-blosc.git
+cd c-blosc
+mkdir build
+cd build
+cmake ..
+make -j4
+cd ../..
+cp c-blosc/build/blosc/*.so* .
diff --git a/test/library/draft/Zarr/SKIPIF b/test/library/draft/Zarr/SKIPIF
new file mode 100644
index 000000000000..3281d31f1011
--- /dev/null
+++ b/test/library/draft/Zarr/SKIPIF
@@ -0,0 +1,2 @@
+CHPL_LIB_PIC != PIC
+CHPL_TARGET_PLATFORM == darwin
diff --git a/test/library/draft/Zarr/Zarr.chpl b/test/library/draft/Zarr/Zarr.chpl
new file mode 100644
index 000000000000..1b77e8b9290b
--- /dev/null
+++ b/test/library/draft/Zarr/Zarr.chpl
@@ -0,0 +1,403 @@
+use IO;
+use FileSystem;
+use JSON;
+use Map;
+use List;
+use Path;
+use CTypes;
+use BlockDist;
+use Time;
+
+require "blosc.h";
+require "-lblosc";
+extern proc blosc_init();
+extern proc blosc_compress(clevel: c_int, doshuffle: c_int, typesize: c_size_t,
+                           nbytes: c_size_t, src: c_ptrConst(void), 
+                           dest: c_ptr(void), destsize: c_size_t): int;
+extern proc blosc_decompress(src: c_ptrConst(void), dest: c_ptr(void), destsize: c_size_t): int;
+extern proc blosc_destroy();
+extern proc blosc_set_nthreads(nthreads_new: c_int) : c_int;
+extern proc blosc_get_nthreads() : c_int;
+
+// checks values based on how the data is written
+proc verifyCorrectness(ref A: [?D]) {
+  param dimCount = D.rank;
+  if dimCount == 1 then
+    forall i in A.domain do assert(A[i] == i);
+
+  if dimCount == 2 {
+    forall (i,j) in A.domain {
+      if (A[i,j] != i*j) {
+        writeln("Failure for indices %i %i".format(i,j));
+        writeln("Expected: %i\nReceived: %s".format(i*j, A[i,j]:string));
+        
+        assert(A[i,j] == i*j);
+      }
+      
+    }
+  }
+    
+  if dimCount == 3 then
+    forall (i,j,k) in A.domain do
+      assert(A[i,j,k] == k + i*j);
+}
+
+record zarrMetadataV2 {
+  var zarr_format: int;
+  var chunks: list(int);
+  var dtype: string;
+  var shape: list(int);
+};
+
+record zarrMetadataV3 {
+  var zarr_format: int;
+  var node_type: string;
+  var shape: list(int);
+  var data_type: string;
+  var dimension_names: list(string);
+};
+
+proc dtypeString(type dtype) throws {
+  select dtype {
+    when real(32) do return "f4";
+    when real(64) do return "f8";
+    when int(32) do return "i4";
+    when int(64) do return "i8";
+  }
+  throw Error("Unexpected data type, only real and int types are supported.");
+}
+
+proc getMetadata(directoryPath: string) {
+  var metadataPath = joinPath(directoryPath, ".zarray");
+  var r = openReader(metadataPath, deserializer = new jsonDeserializer(), locking=true);
+  var md: zarrMetadataV2;
+  r.readf("%?", md);
+  return md;
+}
+
+proc validateMetadata(metadata: zarrMetadataV2, type dtype, param dimCount) throws {
+  //dimensionality matches
+  if dimCount != metadata.shape.size then
+    throw new Error("Expected metadata shape field to have %i dimensions: %?".format(dimCount, metadata.shape));
+  if dimCount != metadata.chunks.size then
+    throw new Error("Expected metadata chunks field to have %i dimensions: %?".format(dimCount, metadata.chunks));
+  //positive, integer sizes
+  for i in 0..<dimCount {
+    if metadata.shape[i] <= 0 then
+      throw new Error("Metadata shape field must have positive side lengths: %?".format(metadata.shape));
+    if metadata.chunks[i] <= 0 then
+      throw new Error("Metadata chunks field must have positive side lengths: %?".format(metadata.chunks));
+  }
+
+  var chplType: string;
+  select metadata.dtype {
+    when "i4", "<i4" do chplType = "int(32)";
+    when "i8", "<i8" do chplType = "int(64)";
+    when "f4", "<f4" do chplType = "real(32)";
+    when "f8", "<f8" do chplType = "real(64)";
+    otherwise {
+      throw new Error("Only integer and floating point data types currently supported: %s".format(metadata.dtype));
+    }
+  }
+
+  if chplType != dtype:string then
+    throw new Error("Expected entries of type %s. Found %s".format(dtype:string, chplType));
+}
+
+
+
+proc buildChunkPath(directoryPath: string, delimiter: string, const chunkIndices: ?dimCount * int) {
+  var indexStrings: dimCount*string;
+  for i in 0..<dimCount do indexStrings[i] = chunkIndices[i] : string;
+  return joinPath(directoryPath, delimiter.join(indexStrings));
+}
+proc buildChunkPath(directoryPath: string, delimiter: string, chunkIndex: int) {
+  return joinPath(directoryPath, chunkIndex:string);
+}
+
+proc getLocalChunks(D: domain(?), localD: domain(?), chunkShape: ?dimCount*int): domain(dimCount) {
+  
+  const totalShape = D.shape;
+  var chunkCounts: dimCount*int;
+  for i in 0..<dimCount {
+    chunkCounts[i] = ceil(totalShape[i]:real / chunkShape[i]: real) : int;
+  }
+
+  var localChunks: dimCount*range(int);
+  for i in 0..<dimCount {
+    var l = if dimCount != 1 then localD.low[i] else localD.low;
+    var h = if dimCount != 1 then localD.high[i] else localD.high;
+    var low = floor(l:real / chunkShape[i]:real):int;
+    var high = ceil(h / chunkShape[i]:real):int;
+    localChunks[i] = max(low,0)..<min(high,chunkCounts[i]); 
+  }
+  const localChunkDomain: domain(dimCount) = localChunks;
+  return localChunkDomain;
+}
+
+
+/* Returns the domain of the `chunkIndices`-th chunk for chunks of size `chunkShape` */
+proc getChunkDomain(chunkShape: ?dimCount*int, chunkIndices: dimCount*int) {
+  var thisChunkRange: dimCount*range(int);
+  for i in 0..<dimCount {
+    const start = chunkIndices[i] * chunkShape[i];
+    thisChunkRange[i] = start..<start+chunkShape[i];
+  }
+  const thisChunkDomain: domain(dimCount) = thisChunkRange;
+  return thisChunkDomain;
+}
+proc getChunkDomain(chunkShape: ?dimCount*int, chunkIndices: int) {
+  return getChunkDomain(chunkShape, (chunkIndices,));
+}
+
+
+/* 
+   Reads a chunk from storage and fills `arraySlice` with its corresponding
+   values.
+   
+   :arg dimCount: Dimensionality of the array being read.
+
+   :arg chunkPath: Relative or absolute path to the chunk being read.
+
+   :arg chunkDomain: Array subdomain the chunk contains.
+
+   :arg arraySlice: Reference to the portion of the array the calling locale stores.
+
+   :throws Error: If the decompression fails
+*/
+proc readChunk(param dimCount: int, chunkPath: string, chunkDomain: domain(dimCount), ref arraySlice: [] ?t) throws {
+  const f: file;
+  // if the file does not exist, the chunk is empty
+  try {
+    f = open(chunkPath, ioMode.r);
+  } catch {
+    arraySlice[arraySlice.domain] = 0;
+    return;
+  }
+
+  const r = f.reader(deserializer = new binaryDeserializer(), locking=true);
+  var compressedChunk = r.readBytes(f.size);
+  var readBytes = compressedChunk.size;
+
+  var copyIn: [chunkDomain] t;
+  var numRead = blosc_decompress(compressedChunk.c_str(), c_ptrTo(copyIn), copyIn.size*c_sizeof(t));
+  if numRead <= 0 {
+    throw new Error("Failed to decompress data from %?. Blosc error code: %?".format(chunkPath, numRead));
+  }
+
+  arraySlice[arraySlice.domain] = copyIn[arraySlice.domain];
+}
+
+/* 
+   Updates a chunk in storage with a locale's contribution to that chunk.
+   The calling function is expected to manage synchronization among locales.
+   If the locale contributes the entire chunk, it will immediately compress
+   and write the chunk's data. If the contribution is partial, it decompresses
+   the chunk, updates the necessary values, then compresses and writes the
+   chunk to storage.
+
+   :arg dimCount: Dimensionality of the array being written.
+
+   :arg chunkPath: Relative or absolute path to the chunk being written.
+
+   :arg chunkDomain: Array subdomain that the chunk contains.
+
+   :arg arraySlice: The portion of the array that the calling locale 
+    contributes to this chunk.
+
+   :arg bloscLevel: Compression level to use. 0 indicates no compression,
+    9 (default) indicates maximum compression. Values outside of this range
+    will be clipped to a value between 0 and 9.
+
+   :throws Error: If the compression fails
+*/
+proc writeChunk(param dimCount, chunkPath: string, chunkDomain: domain(dimCount), ref arraySlice: [] ?t, bloscLevel: int(32) = 9) throws {
+  //bloscLevel must be between 0 and 9
+  var _bloscLevel = min(9,max(0,bloscLevel));
+
+  // If this chunk is entirely contained in the array slice, we can write 
+  // it out immediately. Otherwise, we need to read in the chunk and update
+  // it with the partial data before writing
+  var copyOut: [chunkDomain] t;
+  if (chunkDomain != arraySlice.domain) {
+    readChunk(dimCount, chunkPath, chunkDomain, copyOut);
+  }
+  copyOut[arraySlice.domain] = arraySlice[arraySlice.domain];
+
+  // Create buffer for compressed bytes 
+  var compressedBuffer = allocate(t, copyOut.size + 16);
+
+  // Compress the chunk's data
+  var bytesCompressed = blosc_compress(_bloscLevel, 0, c_sizeof(t), copyOut.size*c_sizeof(t), c_ptrTo(copyOut), compressedBuffer, (copyOut.size + 16) * c_sizeof(t));
+  if bytesCompressed == 0 then
+    throw new Error("Failed to compress bytes");
+
+  // Write it to storage
+  const f = open(chunkPath, ioMode.cw);
+  const w = f.writer(serializer = new binarySerializer(),locking=true);
+  w.writeBinary(compressedBuffer: c_ptr(void),bytesCompressed);
+}
+
+/* 
+   Reads a zarr store from storage, returning a block distributed array. Each
+   locale reads and decompresses the chunks with elements in their subdomain.
+   
+   :arg directoryPath: Relative or absolute path to the root of the zarr store. 
+    The store is expected to contain a '.zarray' metadata file
+
+   :arg dtype: Chapel type of the store's data
+
+   :arg dimCount: Dimensionality of the zarr array
+
+   :arg bloscThreads: The number of threads to use during decompression (default=1)
+*/
+proc readZarrArray(directoryPath: string, type dtype, param dimCount: int, bloscThreads: int(32) = 1) {
+  var md = getMetadata(directoryPath);
+  validateMetadata(md, dtype, dimCount);
+  // Size and shape tuples
+  var totalShape, chunkShape : dimCount*int;
+  var chunkCounts: dimCount*int;
+  var totalRanges,chunkRanges: dimCount*range(int);
+  for i in 0..<dimCount {
+    totalShape[i] = md.shape[i];
+    chunkShape[i] = md.chunks[i];
+    chunkCounts[i] = ceil(totalShape[i]:real / chunkShape[i]:real) : int;
+    totalRanges[i] = 0..<totalShape[i];
+    chunkRanges[i] = 0..<chunkCounts[i];
+  }
+  const fullChunkDomain: domain(dimCount) = chunkRanges;
+
+  // Initialize the distributed domain and array
+  const undistD : domain(dimCount) = totalRanges;
+  const Dist = new blockDist(boundingBox=undistD);
+  const D = Dist.createDomain(undistD);
+  var A: [D] dtype;
+
+
+  coforall loc in Locales do on loc {
+    blosc_init();
+    blosc_set_nthreads(bloscThreads);
+    const hereD = A.localSubdomain();
+    ref hereA = A[hereD];
+
+    const localChunks = getLocalChunks(D, hereD, chunkShape);
+    for chunkIndices in localChunks do {
+      
+      const chunkPath = buildChunkPath(directoryPath, ".", chunkIndices);
+      
+      const thisChunkDomain = getChunkDomain(chunkShape, chunkIndices);
+      const thisChunkHere = hereD[thisChunkDomain];
+      
+      ref thisChunkSlice = hereA.localSlice(thisChunkHere);
+      readChunk(dimCount, chunkPath, thisChunkDomain, thisChunkSlice);
+    }
+    blosc_destroy();
+  }
+  return A;
+}
+
+/*
+   Writes an array to storage as a v2.0 zarr store. The array metadata and
+   chunks will be stored within the `directoryPath` directory, which is created
+   if it does not yet exist. The chunks will have the dimensions given in the
+   `chunkShape` argument. This function writes chunks in parallel, and supports
+   distributed execution.
+
+   :arg directoryPath: Relative or absolute path to the root of the zarr store. 
+    The directory and all necessary parent directories will be created if it 
+    does not exist.
+
+   :arg A: The array to write to storage.
+
+   :arg chunkShape: The dimension extents to use when breaking A into chunks.
+
+   :arg bloscThreads: The number of threads to use during compression (default=1)
+
+   :arg bloscLevel: Compression level to use. 0 indicates no compression,
+    9 (default) indicates maximum compression. 
+*/
+proc writeZarrArray(directoryPath: string, ref A: [?domainType] ?dtype, chunkShape: ?dimCount*int, bloscThreads: int(32) = 1, bloscLevel: int(32) = 9) {
+
+  // Create the metadata record that is written before the chunks
+  var shape, chunks: list(int);
+  for size in A.shape do shape.pushBack(size);
+  for size in chunkShape do chunks.pushBack(size);
+  const md: zarrMetadataV2 = new zarrMetadataV2(2, chunks, dtypeString(dtype), shape);
+
+  // Clear the directory before writing
+  if exists(directoryPath) then rmTree(directoryPath);
+  mkdir(directoryPath, parents=true);
+
+  // Write the metadata
+  const metadataPath = joinPath(directoryPath, ".zarray");
+  const w = openWriter(metadataPath, serializer = new jsonSerializer(),locking=true);
+  w.writef("%?\n", md);
+
+  // Normalize the array's domain to be zero-indexed
+  var normalizedRanges: dimCount*range(int);
+  for i in 0..<dimCount do
+    normalizedRanges[i] = 0..<shape[i];
+  const undistD: domain(dimCount) = normalizedRanges;
+  const dist = new blockDist(boundingBox=undistD);
+  const D = dist.createDomain(undistD);
+  ref normA = A.reindex(D);
+
+  // Locks to synchronize locales writing to the same chunks
+  const allChunks = getLocalChunks(normA.domain, normA.domain, chunkShape);
+  var locks: [allChunks] sync bool;
+
+  // Write the chunks
+  coforall loc in Locales do on loc {
+    // Initialize blosc on each locale
+    blosc_init();
+    blosc_set_nthreads(bloscThreads);
+
+    // Get the part of the array that belongs to this locale
+    const hereD = normA.localSubdomain();
+    ref hereA = normA[hereD];
+
+    // Identify the range of chunks this locale will contribute to
+    const localChunks = getLocalChunks(normA.domain, hereD, chunkShape);
+
+    for chunkIndices in localChunks {
+      // Get the part of the array that contributes to this chunk
+      const thisChunkDomain = getChunkDomain(chunkShape, chunkIndices);
+      const thisChunkHere = hereD[thisChunkDomain];
+      ref thisChunkSlice = hereA.localSlice(thisChunkHere);
+      const chunkPath = buildChunkPath(directoryPath, ".", chunkIndices);
+      locks[chunkIndices].writeEF(true);
+      writeChunk(dimCount, chunkPath, thisChunkDomain, thisChunkSlice);
+      locks[chunkIndices].readFE();
+    }
+  }
+}
+
+
+
+proc throughputTest(param dimCount: int, type dtype) {
+  var s: stopwatch;
+  var gbs = 1;
+  var name = "%igb%id%s%i".format(gbs, dimCount, (dtype:string).replace("(", "").replace(")",""));
+  {
+    writeln("Starting single-locale read of %iGB %iD store of %s".format(gbs, dimCount, dtype:string));
+    s.restart();
+    var A = readZarrArrayV2(name, dtype, dimCount);
+    var t = s.elapsed();
+    verifyCorrectness(A);
+    writeln("Throughput: ", gbs:real / t, " GB/s on ", 1, " locale");
+    
+  }
+  writeln();
+  {
+    writeln("Starting distributed read of %iGB %iD store of %s".format(gbs, dimCount, dtype:string));
+    s.restart();
+    var A = readZarrArrayV2Dist(name, dtype, dimCount);
+    var t = s.elapsed();
+    verifyCorrectness(A);
+    writeln("Throughput: ", gbs:real / t, " GB/s on ", numLocales, " locales");
+  }
+  writeln();
+}
+
+
+
diff --git a/test/library/draft/Zarr/Zarr.notest b/test/library/draft/Zarr/Zarr.notest
new file mode 100644
index 000000000000..e69de29bb2d1
diff --git a/test/library/draft/Zarr/ZarrPerf.chpl b/test/library/draft/Zarr/ZarrPerf.chpl
new file mode 100644
index 000000000000..8ced877450fe
--- /dev/null
+++ b/test/library/draft/Zarr/ZarrPerf.chpl
@@ -0,0 +1,64 @@
+use Zarr;
+use Time;
+use BlockDist;
+use FileSystem;
+use IO;
+proc throughputTest(type dtype, param dimCount: int, arrayShape: dimCount*int, chunkShape: dimCount*int, bloscThreads: int(32)) {
+  var ranges: dimCount*range(int);
+  for i in 0..<dimCount do
+    ranges[i] = 0..<arrayShape[i];
+  const undistD : domain(dimCount) = ranges;
+  const Dist = new blockDist(boundingBox=undistD);
+  const D = Dist.createDomain(undistD);
+  var A: [D] dtype;
+
+  coforall loc in Locales do on loc {
+    forall i in A.localSubdomain() do
+      A[i] = if isTuple(i) then i[0]:dtype else i:dtype;
+  }
+
+  if exists("PerfStore") then rmTree("PerfStore");
+  var s: stopwatch;
+  s.start();
+  writeZarrArray("PerfStore", A, chunkShape, bloscThreads=bloscThreads);
+  const writeTime = s.elapsed();
+
+  s.restart();
+  var B = readZarrArray("PerfStore", dtype, dimCount);
+  const readTime = s.elapsed();
+
+  assert(A.domain == B.domain);
+  coforall loc in Locales do on loc {
+    forall i in A.localSubdomain() do
+      assert(A[i] == B[i]);
+  }
+
+  const arrayBytes = A.size:real * numBits(dtype) / 8;
+  const arrayGBs = arrayBytes:real / (1000 ** 3);
+
+  const chunkBytes = (* reduce chunkShape) * numBits(dtype):real / 8;
+  const chunkMBs = chunkBytes:real / (1000 ** 2);
+  writeln("%?GB %?D array of %?, %?MB chunks. %n locales, %n bloscThreads".format(arrayGBs, dimCount, dtype:string, chunkMBs, Locales.size, bloscThreads));
+  writeln("%n GB/s writing".format(arrayGBs / writeTime));
+  writeln("%n GB/s reading".format(arrayGBs / readTime));
+}
+
+
+throughputTest(real(32), 1, (1000**3,), (1000**2,),1);
+throughputTest(real(32), 1, (1000**3,), (1000**2,),2);
+throughputTest(real(32), 1, (1000**3,), (1000**2,),4);
+throughputTest(real(32), 2, (30000,30000), (1000,1000),1);
+throughputTest(real(32), 2, (30000,30000), (1000,1000),2);
+throughputTest(real(32), 2, (30000,30000), (1000,1000),4);
+throughputTest(real(32), 3, (1000,1000,1000), (100,100,100),1);
+throughputTest(real(32), 3, (1000,1000,1000), (100,100,100),2);
+throughputTest(real(32), 3, (1000,1000,1000), (100,100,100),4);
+throughputTest(int(32), 1, (1000**3,), (1000**2,),1);
+throughputTest(int(32), 1, (1000**3,), (1000**2,),2);
+throughputTest(int(32), 1, (1000**3,), (1000**2,),4);
+throughputTest(int(32), 2, (30000,30000), (1000,1000),1);
+throughputTest(int(32), 2, (30000,30000), (1000,1000),2);
+throughputTest(int(32), 2, (30000,30000), (1000,1000),4);
+throughputTest(int(32), 3, (1000,1000,1000), (100,100,100),1);
+throughputTest(int(32), 3, (1000,1000,1000), (100,100,100),2);
+throughputTest(int(32), 3, (1000,1000,1000), (100,100,100),4);
diff --git a/test/library/draft/Zarr/ZarrPerf.notest b/test/library/draft/Zarr/ZarrPerf.notest
new file mode 100644
index 000000000000..e69de29bb2d1
diff --git a/test/library/draft/Zarr/ZarrTest.chpl b/test/library/draft/Zarr/ZarrTest.chpl
new file mode 100644
index 000000000000..c8a6d57eace5
--- /dev/null
+++ b/test/library/draft/Zarr/ZarrTest.chpl
@@ -0,0 +1,82 @@
+use Zarr;
+use IO;
+use FileSystem;
+use BlockDist;
+use Random;
+
+proc smallTest(type dtype) {
+  {
+    const N1 = 100;
+    const D1 = {0..<N1};
+    var A1: [D1] dtype;
+    for i in D1 do A1[i] = (i + 3):dtype;
+    if (exists("Test1D")) then rmTree("Test1D");
+    writeZarrArray("Test1D", A1, (7,));
+
+    var B1 = readZarrArray("Test1D", dtype, 1);
+    
+    assert(A1.domain == B1.domain, "Domain mismatch: %? %?".format(A1.domain, B1.domain));
+    forall i in A1.domain do 
+      assert(A1[i] == B1[i], "Mismatch for 1D real data on indices: %?.\nWritten: %?\nRead: %?".format(i, A1[i], B1[i]));
+    rmTree("Test1D");
+  }
+
+  {
+    const N2 = 100;
+    const D2 = {0..<N2,0..<N2};
+    var A2: [D2] dtype;
+    for (i,j) in D2 do A2[i,j] = (i:real(32) / (j+1)) : dtype;
+    if (exists("Test2D")) then rmTree("Test2D");
+    writeZarrArray("Test2D", A2, (7,18));
+
+    var B2 = readZarrArray("Test2D", dtype, 2);
+    
+    assert(A2.domain == B2.domain, "Domain mismatch: %? %?".format(A2.domain, B2.domain));
+    forall i in A2.domain do 
+      assert(A2[i] == B2[i], "Mismatch for 2D real data on indices: %?.\nWritten: %?\nRead: %?".format(i, A2[i], B2[i]));
+    rmTree("Test2D");
+  }
+
+  {
+    const N3 = 30;
+    const D3 = {0..<N3,0..<N3,0..<N3};
+    var A3: [D3] dtype;
+    for (i,j,k) in D3 do A3[i,j,k] = (i:real(32) / ((j+1) + (k+1))) : dtype;
+    if (exists("Test3D")) then rmTree("Test3D");
+    writeZarrArray("Test3D", A3, (7,18,22));
+
+    var B3 = readZarrArray("Test3D", dtype, 3);
+    
+    assert(A3.domain == B3.domain, "Domain mismatch: %? %?".format(A3.domain, B3.domain));
+    forall i in A3.domain do 
+      assert(A3[i] == B3[i], "Mismatch for 3D real data on indices: %?.\nWritten: %?\nRead: %?".format(i, A3[i], B3[i]));
+    rmTree("Test3D");
+  }
+}
+
+proc reindexTest(type dtype) {
+  const D = {1..20,14..<72,-10..10};
+  var A: [D] dtype;
+  forall (i,j,k) in D do A[i,j,k] = (i + j*i + k*j):dtype;
+  if exists("ReindexStore") then rmTree("ReindexStore");
+  writeZarrArray("ReindexStore", A, (5,8,12));
+
+  var B = readZarrArray("ReindexStore", dtype, 3);
+  ref viewB = B.reindex(D);
+  assert(A.domain == viewB.domain, "Domain mismatch for array with offset dimensions");
+  forall i in D do
+    assert(A[i] == viewB[i], "Mismatch on indices %?. Written: %?\n. Read: %?\n.".format(i, A[i], viewB[i]));
+  
+}
+
+
+proc main() {
+  const typeTuple = (0:int(32), 0:int(64), 0:real(32), 0:real(64));
+  for param i in 0..<typeTuple.size {
+    type dtype = typeTuple[i].type;
+    writeln("Testing ", dtype:string);
+    smallTest(dtype);
+    reindexTest(dtype);
+  }
+  writeln("Pass");
+}
diff --git a/test/library/draft/Zarr/ZarrTest.cleanfiles b/test/library/draft/Zarr/ZarrTest.cleanfiles
new file mode 100644
index 000000000000..59f15bc8a2fb
--- /dev/null
+++ b/test/library/draft/Zarr/ZarrTest.cleanfiles
@@ -0,0 +1,6 @@
+*.so
+c-blosc
+ReindexStore
+Test1D
+Test2D
+Test3D
diff --git a/test/library/draft/Zarr/ZarrTest.good b/test/library/draft/Zarr/ZarrTest.good
new file mode 100644
index 000000000000..6035a12a34e2
--- /dev/null
+++ b/test/library/draft/Zarr/ZarrTest.good
@@ -0,0 +1,5 @@
+Testing int(32)
+Testing int(64)
+Testing real(32)
+Testing real(64)
+Pass
diff --git a/test/library/draft/Zarr/ZarrTest.numlocales b/test/library/draft/Zarr/ZarrTest.numlocales
new file mode 100644
index 000000000000..d8263ee98605
--- /dev/null
+++ b/test/library/draft/Zarr/ZarrTest.numlocales
@@ -0,0 +1 @@
+2
\ No newline at end of file

From 59a1b1f1854f56bdd06d027078ed34103a915773 Mon Sep 17 00:00:00 2001
From: Brandon Neth <brandon.neth@hpe.com>
Date: Wed, 6 Mar 2024 15:00:50 -0700
Subject: [PATCH 2/3] updates from reviewer feedback. --- Signed-off-by:
 Brandon Neth <brandon.neth@hpe.com>

---
 test/library/draft/Zarr/.gitignore            |   3 +
 .../Zarr/{ZarrTest.cleanfiles => CLEANFILES}  |   0
 test/library/draft/Zarr/Zarr.chpl             | 695 +++++++++---------
 3 files changed, 333 insertions(+), 365 deletions(-)
 create mode 100644 test/library/draft/Zarr/.gitignore
 rename test/library/draft/Zarr/{ZarrTest.cleanfiles => CLEANFILES} (100%)

diff --git a/test/library/draft/Zarr/.gitignore b/test/library/draft/Zarr/.gitignore
new file mode 100644
index 000000000000..60fd84f41d78
--- /dev/null
+++ b/test/library/draft/Zarr/.gitignore
@@ -0,0 +1,3 @@
+c-blosc
+*.dylib
+*.so*
diff --git a/test/library/draft/Zarr/ZarrTest.cleanfiles b/test/library/draft/Zarr/CLEANFILES
similarity index 100%
rename from test/library/draft/Zarr/ZarrTest.cleanfiles
rename to test/library/draft/Zarr/CLEANFILES
diff --git a/test/library/draft/Zarr/Zarr.chpl b/test/library/draft/Zarr/Zarr.chpl
index 1b77e8b9290b..fb1ac9cedc60 100644
--- a/test/library/draft/Zarr/Zarr.chpl
+++ b/test/library/draft/Zarr/Zarr.chpl
@@ -1,403 +1,368 @@
-use IO;
-use FileSystem;
-use JSON;
-use Map;
-use List;
-use Path;
-use CTypes;
-use BlockDist;
-use Time;
-
-require "blosc.h";
-require "-lblosc";
-extern proc blosc_init();
-extern proc blosc_compress(clevel: c_int, doshuffle: c_int, typesize: c_size_t,
-                           nbytes: c_size_t, src: c_ptrConst(void), 
-                           dest: c_ptr(void), destsize: c_size_t): int;
-extern proc blosc_decompress(src: c_ptrConst(void), dest: c_ptr(void), destsize: c_size_t): int;
-extern proc blosc_destroy();
-extern proc blosc_set_nthreads(nthreads_new: c_int) : c_int;
-extern proc blosc_get_nthreads() : c_int;
-
-// checks values based on how the data is written
-proc verifyCorrectness(ref A: [?D]) {
-  param dimCount = D.rank;
-  if dimCount == 1 then
-    forall i in A.domain do assert(A[i] == i);
-
-  if dimCount == 2 {
-    forall (i,j) in A.domain {
-      if (A[i,j] != i*j) {
-        writeln("Failure for indices %i %i".format(i,j));
-        writeln("Expected: %i\nReceived: %s".format(i*j, A[i,j]:string));
-        
-        assert(A[i,j] == i*j);
-      }
-      
-    }
+/*
+  Support for distributed reading and writing of Zarr stores. Support is 
+  limited to v2 Zarr arrays stored on local filesystems. The module uses 
+  c-blosc to compress and decompress chunks. Zarr specification: 
+  https://zarr-specs.readthedocs.io/en/latest/v2/v2.0.html
+*/ 
+module Zarr {
+  use IO;
+  use FileSystem;
+  use JSON;
+  use Map;
+  use List;
+  use Path;
+  use CTypes;
+  use BlockDist;
+  use Time;
+
+  require "blosc.h";
+  require "-lblosc";
+
+  module Blosc {
+    extern proc blosc_init();
+    extern proc blosc_compress(clevel: c_int, doshuffle: c_int, typesize: c_size_t,
+                              nbytes: c_size_t, src: c_ptrConst(void), 
+                              dest: c_ptr(void), destsize: c_size_t): int;
+    extern proc blosc_decompress(src: c_ptrConst(void), dest: c_ptr(void), destsize: c_size_t): int;
+    extern proc blosc_destroy();
+    extern proc blosc_set_nthreads(nthreads_new: c_int) : c_int;
+    extern proc blosc_get_nthreads() : c_int;
   }
-    
-  if dimCount == 3 then
-    forall (i,j,k) in A.domain do
-      assert(A[i,j,k] == k + i*j);
-}
-
-record zarrMetadataV2 {
-  var zarr_format: int;
-  var chunks: list(int);
-  var dtype: string;
-  var shape: list(int);
-};
-
-record zarrMetadataV3 {
-  var zarr_format: int;
-  var node_type: string;
-  var shape: list(int);
-  var data_type: string;
-  var dimension_names: list(string);
-};
-
-proc dtypeString(type dtype) throws {
-  select dtype {
-    when real(32) do return "f4";
-    when real(64) do return "f8";
-    when int(32) do return "i4";
-    when int(64) do return "i8";
+  private use Blosc;
+
+  record zarrMetadataV2 {
+    var zarr_format: int;
+    var chunks: list(int);
+    var dtype: string;
+    var shape: list(int);
+  };
+
+  // Unused until support is added for v3.0 stores
+  record zarrMetadataV3 {
+    var zarr_format: int;
+    var node_type: string;
+    var shape: list(int);
+    var data_type: string;
+    var dimension_names: list(string);
+  };
+
+  proc dtypeString(type dtype) throws {
+    select dtype {
+      when real(32) do return "f4";
+      when real(64) do return "f8";
+      when int(32) do return "i4";
+      when int(64) do return "i8";
+    }
+    throw Error("Unexpected data type, only real and int types are supported.");
   }
-  throw Error("Unexpected data type, only real and int types are supported.");
-}
-
-proc getMetadata(directoryPath: string) {
-  var metadataPath = joinPath(directoryPath, ".zarray");
-  var r = openReader(metadataPath, deserializer = new jsonDeserializer(), locking=true);
-  var md: zarrMetadataV2;
-  r.readf("%?", md);
-  return md;
-}
-
-proc validateMetadata(metadata: zarrMetadataV2, type dtype, param dimCount) throws {
-  //dimensionality matches
-  if dimCount != metadata.shape.size then
-    throw new Error("Expected metadata shape field to have %i dimensions: %?".format(dimCount, metadata.shape));
-  if dimCount != metadata.chunks.size then
-    throw new Error("Expected metadata chunks field to have %i dimensions: %?".format(dimCount, metadata.chunks));
-  //positive, integer sizes
-  for i in 0..<dimCount {
-    if metadata.shape[i] <= 0 then
-      throw new Error("Metadata shape field must have positive side lengths: %?".format(metadata.shape));
-    if metadata.chunks[i] <= 0 then
-      throw new Error("Metadata chunks field must have positive side lengths: %?".format(metadata.chunks));
+
+  proc getMetadata(directoryPath: string) {
+    var metadataPath = joinPath(directoryPath, ".zarray");
+    var r = openReader(metadataPath, deserializer = new jsonDeserializer(), locking=false);
+    var md: zarrMetadataV2;
+    r.readf("%?", md);
+    return md;
   }
 
-  var chplType: string;
-  select metadata.dtype {
-    when "i4", "<i4" do chplType = "int(32)";
-    when "i8", "<i8" do chplType = "int(64)";
-    when "f4", "<f4" do chplType = "real(32)";
-    when "f8", "<f8" do chplType = "real(64)";
-    otherwise {
-      throw new Error("Only integer and floating point data types currently supported: %s".format(metadata.dtype));
+  proc validateMetadata(metadata: zarrMetadataV2, type dtype, param dimCount) throws {
+    //dimensionality matches
+    if dimCount != metadata.shape.size then
+      throw new Error("Expected metadata shape field to have %i dimensions: %?".format(dimCount, metadata.shape));
+    if dimCount != metadata.chunks.size then
+      throw new Error("Expected metadata chunks field to have %i dimensions: %?".format(dimCount, metadata.chunks));
+    //positive, integer sizes
+    for i in 0..<dimCount {
+      if metadata.shape[i] <= 0 then
+        throw new Error("Metadata shape field must have positive side lengths: %?".format(metadata.shape));
+      if metadata.chunks[i] <= 0 then
+        throw new Error("Metadata chunks field must have positive side lengths: %?".format(metadata.chunks));
     }
-  }
 
-  if chplType != dtype:string then
-    throw new Error("Expected entries of type %s. Found %s".format(dtype:string, chplType));
-}
+    var chplType: string;
+    select metadata.dtype {
+      when "i4", "<i4" do chplType = "int(32)";
+      when "i8", "<i8" do chplType = "int(64)";
+      when "f4", "<f4" do chplType = "real(32)";
+      when "f8", "<f8" do chplType = "real(64)";
+      otherwise {
+        throw new Error("Only integer and floating point data types currently supported: %s".format(metadata.dtype));
+      }
+    }
 
+    if chplType != dtype:string then
+      throw new Error("Expected entries of type %s. Found %s".format(dtype:string, chplType));
+  }
 
 
-proc buildChunkPath(directoryPath: string, delimiter: string, const chunkIndices: ?dimCount * int) {
-  var indexStrings: dimCount*string;
-  for i in 0..<dimCount do indexStrings[i] = chunkIndices[i] : string;
-  return joinPath(directoryPath, delimiter.join(indexStrings));
-}
-proc buildChunkPath(directoryPath: string, delimiter: string, chunkIndex: int) {
-  return joinPath(directoryPath, chunkIndex:string);
-}
 
-proc getLocalChunks(D: domain(?), localD: domain(?), chunkShape: ?dimCount*int): domain(dimCount) {
-  
-  const totalShape = D.shape;
-  var chunkCounts: dimCount*int;
-  for i in 0..<dimCount {
-    chunkCounts[i] = ceil(totalShape[i]:real / chunkShape[i]: real) : int;
+  proc buildChunkPath(directoryPath: string, delimiter: string, const chunkIndices: ?dimCount * int) {
+    var indexStrings: dimCount*string;
+    for i in 0..<dimCount do indexStrings[i] = chunkIndices[i] : string;
+    return joinPath(directoryPath, delimiter.join(indexStrings));
   }
-
-  var localChunks: dimCount*range(int);
-  for i in 0..<dimCount {
-    var l = if dimCount != 1 then localD.low[i] else localD.low;
-    var h = if dimCount != 1 then localD.high[i] else localD.high;
-    var low = floor(l:real / chunkShape[i]:real):int;
-    var high = ceil(h / chunkShape[i]:real):int;
-    localChunks[i] = max(low,0)..<min(high,chunkCounts[i]); 
+  proc buildChunkPath(directoryPath: string, delimiter: string, chunkIndex: int) {
+    return joinPath(directoryPath, chunkIndex:string);
   }
-  const localChunkDomain: domain(dimCount) = localChunks;
-  return localChunkDomain;
-}
 
+  proc getLocalChunks(D: domain(?), localD: domain(?), chunkShape: ?dimCount*int): domain(dimCount) {
+    
+    const totalShape = D.shape;
+    var chunkCounts: dimCount*int;
+    for i in 0..<dimCount {
+      chunkCounts[i] = ceil(totalShape[i]:real / chunkShape[i]: real) : int;
+    }
 
-/* Returns the domain of the `chunkIndices`-th chunk for chunks of size `chunkShape` */
-proc getChunkDomain(chunkShape: ?dimCount*int, chunkIndices: dimCount*int) {
-  var thisChunkRange: dimCount*range(int);
-  for i in 0..<dimCount {
-    const start = chunkIndices[i] * chunkShape[i];
-    thisChunkRange[i] = start..<start+chunkShape[i];
-  }
-  const thisChunkDomain: domain(dimCount) = thisChunkRange;
-  return thisChunkDomain;
-}
-proc getChunkDomain(chunkShape: ?dimCount*int, chunkIndices: int) {
-  return getChunkDomain(chunkShape, (chunkIndices,));
-}
-
-
-/* 
-   Reads a chunk from storage and fills `arraySlice` with its corresponding
-   values.
-   
-   :arg dimCount: Dimensionality of the array being read.
-
-   :arg chunkPath: Relative or absolute path to the chunk being read.
-
-   :arg chunkDomain: Array subdomain the chunk contains.
-
-   :arg arraySlice: Reference to the portion of the array the calling locale stores.
-
-   :throws Error: If the decompression fails
-*/
-proc readChunk(param dimCount: int, chunkPath: string, chunkDomain: domain(dimCount), ref arraySlice: [] ?t) throws {
-  const f: file;
-  // if the file does not exist, the chunk is empty
-  try {
-    f = open(chunkPath, ioMode.r);
-  } catch {
-    arraySlice[arraySlice.domain] = 0;
-    return;
+    var localChunks: dimCount*range(int);
+    for i in 0..<dimCount {
+      var l = if dimCount != 1 then localD.low[i] else localD.low;
+      var h = if dimCount != 1 then localD.high[i] else localD.high;
+      var low = floor(l:real / chunkShape[i]:real):int;
+      var high = ceil(h / chunkShape[i]:real):int;
+      localChunks[i] = max(low,0)..<min(high,chunkCounts[i]); 
+    }
+    const localChunkDomain: domain(dimCount) = localChunks;
+    return localChunkDomain;
   }
 
-  const r = f.reader(deserializer = new binaryDeserializer(), locking=true);
-  var compressedChunk = r.readBytes(f.size);
-  var readBytes = compressedChunk.size;
 
-  var copyIn: [chunkDomain] t;
-  var numRead = blosc_decompress(compressedChunk.c_str(), c_ptrTo(copyIn), copyIn.size*c_sizeof(t));
-  if numRead <= 0 {
-    throw new Error("Failed to decompress data from %?. Blosc error code: %?".format(chunkPath, numRead));
+  /* Returns the domain of the `chunkIndices`-th chunk for chunks of size `chunkShape` */
+  proc getChunkDomain(chunkShape: ?dimCount*int, chunkIndices: dimCount*int) {
+    var thisChunkRange: dimCount*range(int);
+    for i in 0..<dimCount {
+      const start = chunkIndices[i] * chunkShape[i];
+      thisChunkRange[i] = start..<start+chunkShape[i];
+    }
+    const thisChunkDomain: domain(dimCount) = thisChunkRange;
+    return thisChunkDomain;
+  }
+  proc getChunkDomain(chunkShape: ?dimCount*int, chunkIndices: int) {
+    return getChunkDomain(chunkShape, (chunkIndices,));
   }
 
-  arraySlice[arraySlice.domain] = copyIn[arraySlice.domain];
-}
 
-/* 
-   Updates a chunk in storage with a locale's contribution to that chunk.
-   The calling function is expected to manage synchronization among locales.
-   If the locale contributes the entire chunk, it will immediately compress
-   and write the chunk's data. If the contribution is partial, it decompresses
-   the chunk, updates the necessary values, then compresses and writes the
-   chunk to storage.
+  /* 
+    Reads a chunk from storage and fills `arraySlice` with its corresponding
+    values.
+    
+    :arg dimCount: Dimensionality of the array being read.
 
-   :arg dimCount: Dimensionality of the array being written.
+    :arg chunkPath: Relative or absolute path to the chunk being read.
 
-   :arg chunkPath: Relative or absolute path to the chunk being written.
+    :arg chunkDomain: Array subdomain the chunk contains.
 
-   :arg chunkDomain: Array subdomain that the chunk contains.
+    :arg arraySlice: Reference to the portion of the array the calling locale stores.
 
-   :arg arraySlice: The portion of the array that the calling locale 
-    contributes to this chunk.
+    :throws Error: If the decompression fails
+  */
+  proc readChunk(param dimCount: int, chunkPath: string, chunkDomain: domain(dimCount), ref arraySlice: [] ?t) throws {
+    const f: file;
+    // if the file does not exist, the chunk is empty
+    try {
+      f = open(chunkPath, ioMode.r);
+    } catch {
+      arraySlice[arraySlice.domain] = 0;
+      return;
+    }
 
-   :arg bloscLevel: Compression level to use. 0 indicates no compression,
-    9 (default) indicates maximum compression. Values outside of this range
-    will be clipped to a value between 0 and 9.
+    const r = f.reader(deserializer = new binaryDeserializer(), locking=true);
+    var compressedChunk = r.readAll(bytes); // TODO: stream straight through to blosc
+    var readBytes = compressedChunk.size;
 
-   :throws Error: If the compression fails
-*/
-proc writeChunk(param dimCount, chunkPath: string, chunkDomain: domain(dimCount), ref arraySlice: [] ?t, bloscLevel: int(32) = 9) throws {
-  //bloscLevel must be between 0 and 9
-  var _bloscLevel = min(9,max(0,bloscLevel));
+    var copyIn: [chunkDomain] t;
+    var numRead = blosc_decompress(compressedChunk.c_str(), c_ptrTo(copyIn), copyIn.size*c_sizeof(t));
+    if numRead <= 0 {
+      throw new Error("Failed to decompress data from %?. Blosc error code: %?".format(chunkPath, numRead));
+    }
 
-  // If this chunk is entirely contained in the array slice, we can write 
-  // it out immediately. Otherwise, we need to read in the chunk and update
-  // it with the partial data before writing
-  var copyOut: [chunkDomain] t;
-  if (chunkDomain != arraySlice.domain) {
-    readChunk(dimCount, chunkPath, chunkDomain, copyOut);
+    arraySlice[arraySlice.domain] = copyIn[arraySlice.domain];
   }
-  copyOut[arraySlice.domain] = arraySlice[arraySlice.domain];
-
-  // Create buffer for compressed bytes 
-  var compressedBuffer = allocate(t, copyOut.size + 16);
-
-  // Compress the chunk's data
-  var bytesCompressed = blosc_compress(_bloscLevel, 0, c_sizeof(t), copyOut.size*c_sizeof(t), c_ptrTo(copyOut), compressedBuffer, (copyOut.size + 16) * c_sizeof(t));
-  if bytesCompressed == 0 then
-    throw new Error("Failed to compress bytes");
-
-  // Write it to storage
-  const f = open(chunkPath, ioMode.cw);
-  const w = f.writer(serializer = new binarySerializer(),locking=true);
-  w.writeBinary(compressedBuffer: c_ptr(void),bytesCompressed);
-}
-
-/* 
-   Reads a zarr store from storage, returning a block distributed array. Each
-   locale reads and decompresses the chunks with elements in their subdomain.
-   
-   :arg directoryPath: Relative or absolute path to the root of the zarr store. 
-    The store is expected to contain a '.zarray' metadata file
-
-   :arg dtype: Chapel type of the store's data
-
-   :arg dimCount: Dimensionality of the zarr array
-
-   :arg bloscThreads: The number of threads to use during decompression (default=1)
-*/
-proc readZarrArray(directoryPath: string, type dtype, param dimCount: int, bloscThreads: int(32) = 1) {
-  var md = getMetadata(directoryPath);
-  validateMetadata(md, dtype, dimCount);
-  // Size and shape tuples
-  var totalShape, chunkShape : dimCount*int;
-  var chunkCounts: dimCount*int;
-  var totalRanges,chunkRanges: dimCount*range(int);
-  for i in 0..<dimCount {
-    totalShape[i] = md.shape[i];
-    chunkShape[i] = md.chunks[i];
-    chunkCounts[i] = ceil(totalShape[i]:real / chunkShape[i]:real) : int;
-    totalRanges[i] = 0..<totalShape[i];
-    chunkRanges[i] = 0..<chunkCounts[i];
-  }
-  const fullChunkDomain: domain(dimCount) = chunkRanges;
-
-  // Initialize the distributed domain and array
-  const undistD : domain(dimCount) = totalRanges;
-  const Dist = new blockDist(boundingBox=undistD);
-  const D = Dist.createDomain(undistD);
-  var A: [D] dtype;
-
-
-  coforall loc in Locales do on loc {
-    blosc_init();
-    blosc_set_nthreads(bloscThreads);
-    const hereD = A.localSubdomain();
-    ref hereA = A[hereD];
-
-    const localChunks = getLocalChunks(D, hereD, chunkShape);
-    for chunkIndices in localChunks do {
-      
-      const chunkPath = buildChunkPath(directoryPath, ".", chunkIndices);
-      
-      const thisChunkDomain = getChunkDomain(chunkShape, chunkIndices);
-      const thisChunkHere = hereD[thisChunkDomain];
-      
-      ref thisChunkSlice = hereA.localSlice(thisChunkHere);
-      readChunk(dimCount, chunkPath, thisChunkDomain, thisChunkSlice);
+
+  /* 
+    Updates a chunk in storage with a locale's contribution to that chunk.
+    The calling function is expected to manage synchronization among locales.
+    If the locale contributes the entire chunk, it will immediately compress
+    and write the chunk's data. If the contribution is partial, it decompresses
+    the chunk, updates the necessary values, then compresses and writes the
+    chunk to storage.
+
+    :arg dimCount: Dimensionality of the array being written.
+
+    :arg chunkPath: Relative or absolute path to the chunk being written.
+
+    :arg chunkDomain: Array subdomain that the chunk contains.
+
+    :arg arraySlice: The portion of the array that the calling locale 
+      contributes to this chunk.
+
+    :arg bloscLevel: Compression level to use. 0 indicates no compression,
+      9 (default) indicates maximum compression. Values outside of this range
+      will be clipped to a value between 0 and 9.
+
+    :throws Error: If the compression fails
+  */
+  proc writeChunk(param dimCount, chunkPath: string, chunkDomain: domain(dimCount), ref arraySlice: [] ?t, bloscLevel: int(32) = 9) throws {
+    //bloscLevel must be between 0 and 9
+    var _bloscLevel = min(9,max(0,bloscLevel));
+
+    // If this chunk is entirely contained in the array slice, we can write 
+    // it out immediately. Otherwise, we need to read in the chunk and update
+    // it with the partial data before writing
+    var copyOut: [chunkDomain] t;
+    if (chunkDomain != arraySlice.domain) {
+      readChunk(dimCount, chunkPath, chunkDomain, copyOut);
     }
-    blosc_destroy();
+    copyOut[arraySlice.domain] = arraySlice[arraySlice.domain];
+
+    // Create buffer for compressed bytes 
+    var compressedBuffer = allocate(t, copyOut.size + 16);
+
+    // Compress the chunk's data
+    var bytesCompressed = blosc_compress(_bloscLevel, 0, c_sizeof(t), copyOut.size*c_sizeof(t), c_ptrTo(copyOut), compressedBuffer, (copyOut.size + 16) * c_sizeof(t));
+    if bytesCompressed == 0 then
+      throw new Error("Failed to compress bytes");
+
+    // Write it to storage
+    const f = open(chunkPath, ioMode.cw);
+    const w = f.writer(serializer = new binarySerializer(),locking=true);
+    w.writeBinary(compressedBuffer: c_ptr(void),bytesCompressed);
   }
-  return A;
-}
 
-/*
-   Writes an array to storage as a v2.0 zarr store. The array metadata and
-   chunks will be stored within the `directoryPath` directory, which is created
-   if it does not yet exist. The chunks will have the dimensions given in the
-   `chunkShape` argument. This function writes chunks in parallel, and supports
-   distributed execution.
-
-   :arg directoryPath: Relative or absolute path to the root of the zarr store. 
-    The directory and all necessary parent directories will be created if it 
-    does not exist.
-
-   :arg A: The array to write to storage.
-
-   :arg chunkShape: The dimension extents to use when breaking A into chunks.
-
-   :arg bloscThreads: The number of threads to use during compression (default=1)
-
-   :arg bloscLevel: Compression level to use. 0 indicates no compression,
-    9 (default) indicates maximum compression. 
-*/
-proc writeZarrArray(directoryPath: string, ref A: [?domainType] ?dtype, chunkShape: ?dimCount*int, bloscThreads: int(32) = 1, bloscLevel: int(32) = 9) {
-
-  // Create the metadata record that is written before the chunks
-  var shape, chunks: list(int);
-  for size in A.shape do shape.pushBack(size);
-  for size in chunkShape do chunks.pushBack(size);
-  const md: zarrMetadataV2 = new zarrMetadataV2(2, chunks, dtypeString(dtype), shape);
-
-  // Clear the directory before writing
-  if exists(directoryPath) then rmTree(directoryPath);
-  mkdir(directoryPath, parents=true);
-
-  // Write the metadata
-  const metadataPath = joinPath(directoryPath, ".zarray");
-  const w = openWriter(metadataPath, serializer = new jsonSerializer(),locking=true);
-  w.writef("%?\n", md);
-
-  // Normalize the array's domain to be zero-indexed
-  var normalizedRanges: dimCount*range(int);
-  for i in 0..<dimCount do
-    normalizedRanges[i] = 0..<shape[i];
-  const undistD: domain(dimCount) = normalizedRanges;
-  const dist = new blockDist(boundingBox=undistD);
-  const D = dist.createDomain(undistD);
-  ref normA = A.reindex(D);
-
-  // Locks to synchronize locales writing to the same chunks
-  const allChunks = getLocalChunks(normA.domain, normA.domain, chunkShape);
-  var locks: [allChunks] sync bool;
-
-  // Write the chunks
-  coforall loc in Locales do on loc {
-    // Initialize blosc on each locale
-    blosc_init();
-    blosc_set_nthreads(bloscThreads);
-
-    // Get the part of the array that belongs to this locale
-    const hereD = normA.localSubdomain();
-    ref hereA = normA[hereD];
-
-    // Identify the range of chunks this locale will contribute to
-    const localChunks = getLocalChunks(normA.domain, hereD, chunkShape);
-
-    for chunkIndices in localChunks {
-      // Get the part of the array that contributes to this chunk
-      const thisChunkDomain = getChunkDomain(chunkShape, chunkIndices);
-      const thisChunkHere = hereD[thisChunkDomain];
-      ref thisChunkSlice = hereA.localSlice(thisChunkHere);
-      const chunkPath = buildChunkPath(directoryPath, ".", chunkIndices);
-      locks[chunkIndices].writeEF(true);
-      writeChunk(dimCount, chunkPath, thisChunkDomain, thisChunkSlice);
-      locks[chunkIndices].readFE();
+  /* 
+    Reads a v2.0 zarr store from storage, returning a block distributed array.
+    Each locale reads and decompresses the chunks with elements in its
+    subdomain. This method assumes a shared filesystem where all nodes can 
+    access the store directory.
+    
+    :arg directoryPath: Relative or absolute path to the root of the zarr 
+      store. The store is expected to contain a '.zarray' metadata file
+
+    :arg dtype: Chapel type of the store's data
+
+    :arg dimCount: Dimensionality of the zarr array
+
+    :arg bloscThreads: The number of threads to use during decompression 
+      (default=1)
+  */
+  proc readZarrArray(directoryPath: string, type dtype, param dimCount: int, bloscThreads: int(32) = 1) {
+    var md = getMetadata(directoryPath);
+    validateMetadata(md, dtype, dimCount);
+    // Size and shape tuples
+    var totalShape, chunkShape : dimCount*int;
+    var chunkCounts: dimCount*int;
+    var totalRanges,chunkRanges: dimCount*range(int);
+    for i in 0..<dimCount {
+      totalShape[i] = md.shape[i];
+      chunkShape[i] = md.chunks[i];
+      chunkCounts[i] = ceil(totalShape[i]:real / chunkShape[i]:real) : int;
+      totalRanges[i] = 0..<totalShape[i];
+      chunkRanges[i] = 0..<chunkCounts[i];
     }
-  }
-}
+    const fullChunkDomain: domain(dimCount) = chunkRanges;
 
+    // Initialize the distributed domain and array
+    const undistD : domain(dimCount) = totalRanges;
+    const Dist = new blockDist(boundingBox=undistD);
+    const D = Dist.createDomain(undistD);
+    var A: [D] dtype;
 
 
-proc throughputTest(param dimCount: int, type dtype) {
-  var s: stopwatch;
-  var gbs = 1;
-  var name = "%igb%id%s%i".format(gbs, dimCount, (dtype:string).replace("(", "").replace(")",""));
-  {
-    writeln("Starting single-locale read of %iGB %iD store of %s".format(gbs, dimCount, dtype:string));
-    s.restart();
-    var A = readZarrArrayV2(name, dtype, dimCount);
-    var t = s.elapsed();
-    verifyCorrectness(A);
-    writeln("Throughput: ", gbs:real / t, " GB/s on ", 1, " locale");
-    
-  }
-  writeln();
-  {
-    writeln("Starting distributed read of %iGB %iD store of %s".format(gbs, dimCount, dtype:string));
-    s.restart();
-    var A = readZarrArrayV2Dist(name, dtype, dimCount);
-    var t = s.elapsed();
-    verifyCorrectness(A);
-    writeln("Throughput: ", gbs:real / t, " GB/s on ", numLocales, " locales");
-  }
-  writeln();
-}
+    coforall loc in Locales do on loc {
+      blosc_init();
+      blosc_set_nthreads(bloscThreads);
+      const hereD = A.localSubdomain();
+      ref hereA = A[hereD];
 
+      const localChunks = getLocalChunks(D, hereD, chunkShape);
+      forall chunkIndices in localChunks do {
+        
+        const chunkPath = buildChunkPath(directoryPath, ".", chunkIndices);
+        
+        const thisChunkDomain = getChunkDomain(chunkShape, chunkIndices);
+        const thisChunkHere = hereD[thisChunkDomain];
+        
+        ref thisChunkSlice = hereA.localSlice(thisChunkHere);
+        readChunk(dimCount, chunkPath, thisChunkDomain, thisChunkSlice);
+      }
+      blosc_destroy();
+    }
+    return A;
+  }
 
+  /*
+    Writes an array to storage as a v2.0 zarr store. The array metadata and
+    chunks will be stored within the `directoryPath` directory, which is created
+    if it does not yet exist. The chunks will have the dimensions given in the
+    `chunkShape` argument. This function writes chunks in parallel, and supports
+    distributed execution. It assumes a shared filesystem where all nodes can 
+    access the store directory.
+
+    :arg directoryPath: Relative or absolute path to the root of the zarr store. 
+      The directory and all necessary parent directories will be created if it 
+      does not exist.
+
+    :arg A: The array to write to storage.
+
+    :arg chunkShape: The dimension extents to use when breaking A into chunks.
+
+    :arg bloscThreads: The number of threads to use during compression (default=1)
+
+    :arg bloscLevel: Compression level to use. 0 indicates no compression,
+      9 (default) indicates maximum compression. 
+  */
+  proc writeZarrArray(directoryPath: string, ref A: [?domainType] ?dtype, chunkShape: ?dimCount*int, bloscThreads: int(32) = 1, bloscLevel: int(32) = 9) {
+
+    // Create the metadata record that is written before the chunks
+    var shape, chunks: list(int);
+    for size in A.shape do shape.pushBack(size);
+    for size in chunkShape do chunks.pushBack(size);
+    const md: zarrMetadataV2 = new zarrMetadataV2(2, chunks, dtypeString(dtype), shape);
+
+    // Clear the directory before writing
+    if exists(directoryPath) then rmTree(directoryPath);
+    mkdir(directoryPath, parents=true);
+
+    // Write the metadata
+    const metadataPath = joinPath(directoryPath, ".zarray");
+    const w = openWriter(metadataPath, serializer = new jsonSerializer(),locking=true);
+    w.writef("%?\n", md);
+
+    // Normalize the array's domain to be zero-indexed
+    var normalizedRanges: dimCount*range(int);
+    for i in 0..<dimCount do
+      normalizedRanges[i] = 0..<shape[i];
+    const undistD: domain(dimCount) = normalizedRanges;
+    const dist = new blockDist(boundingBox=undistD);
+    const D = dist.createDomain(undistD);
+    ref normA = A.reindex(D);
+
+    // Locks to synchronize locales writing to the same chunks
+    const allChunks = getLocalChunks(normA.domain, normA.domain, chunkShape);
+    var locks: [allChunks] sync bool;
+
+    // Write the chunks
+    coforall loc in Locales do on loc {
+      // Initialize blosc on each locale
+      blosc_init();
+      blosc_set_nthreads(bloscThreads);
+
+      // Get the part of the array that belongs to this locale
+      const hereD = normA.localSubdomain();
+      ref hereA = normA[hereD];
+
+      // Identify the range of chunks this locale will contribute to
+      const localChunks = getLocalChunks(normA.domain, hereD, chunkShape);
+
+      forall chunkIndices in localChunks {
+        // Get the part of the array that contributes to this chunk
+        const thisChunkDomain = getChunkDomain(chunkShape, chunkIndices);
+        const thisChunkHere = hereD[thisChunkDomain];
+        ref thisChunkSlice = hereA.localSlice(thisChunkHere);
+        const chunkPath = buildChunkPath(directoryPath, ".", chunkIndices);
+        locks[chunkIndices].writeEF(true);
+        writeChunk(dimCount, chunkPath, thisChunkDomain, thisChunkSlice);
+        locks[chunkIndices].readFE();
+      }
+    }
+  }
 
+}
\ No newline at end of file

From 9fc7a1c01132311274a78a187b4f098691484231 Mon Sep 17 00:00:00 2001
From: Brandon Neth <brandon.neth@hpe.com>
Date: Wed, 6 Mar 2024 15:22:35 -0700
Subject: [PATCH 3/3] fixes for compilation errors --- Signed-off-by: Brandon
 Neth <brandon.neth@hpe.com>

---
 test/library/draft/Zarr/.gitignore | 4 ++++
 test/library/draft/Zarr/Zarr.chpl  | 7 ++++---
 2 files changed, 8 insertions(+), 3 deletions(-)

diff --git a/test/library/draft/Zarr/.gitignore b/test/library/draft/Zarr/.gitignore
index 60fd84f41d78..ecfb4da4741c 100644
--- a/test/library/draft/Zarr/.gitignore
+++ b/test/library/draft/Zarr/.gitignore
@@ -1,3 +1,7 @@
+ReindexStore
+Test1D
+Test2D
+Test3D
 c-blosc
 *.dylib
 *.so*
diff --git a/test/library/draft/Zarr/Zarr.chpl b/test/library/draft/Zarr/Zarr.chpl
index fb1ac9cedc60..33f1885781a6 100644
--- a/test/library/draft/Zarr/Zarr.chpl
+++ b/test/library/draft/Zarr/Zarr.chpl
@@ -19,6 +19,7 @@ module Zarr {
   require "-lblosc";
 
   module Blosc {
+    use CTypes;
     extern proc blosc_init();
     extern proc blosc_compress(clevel: c_int, doshuffle: c_int, typesize: c_size_t,
                               nbytes: c_size_t, src: c_ptrConst(void), 
@@ -56,7 +57,7 @@ module Zarr {
     throw Error("Unexpected data type, only real and int types are supported.");
   }
 
-  proc getMetadata(directoryPath: string) {
+  proc getMetadata(directoryPath: string) throws {
     var metadataPath = joinPath(directoryPath, ".zarray");
     var r = openReader(metadataPath, deserializer = new jsonDeserializer(), locking=false);
     var md: zarrMetadataV2;
@@ -243,7 +244,7 @@ module Zarr {
     :arg bloscThreads: The number of threads to use during decompression 
       (default=1)
   */
-  proc readZarrArray(directoryPath: string, type dtype, param dimCount: int, bloscThreads: int(32) = 1) {
+  proc readZarrArray(directoryPath: string, type dtype, param dimCount: int, bloscThreads: int(32) = 1) throws {
     var md = getMetadata(directoryPath);
     validateMetadata(md, dtype, dimCount);
     // Size and shape tuples
@@ -309,7 +310,7 @@ module Zarr {
     :arg bloscLevel: Compression level to use. 0 indicates no compression,
       9 (default) indicates maximum compression. 
   */
-  proc writeZarrArray(directoryPath: string, ref A: [?domainType] ?dtype, chunkShape: ?dimCount*int, bloscThreads: int(32) = 1, bloscLevel: int(32) = 9) {
+  proc writeZarrArray(directoryPath: string, ref A: [?domainType] ?dtype, chunkShape: ?dimCount*int, bloscThreads: int(32) = 1, bloscLevel: int(32) = 9) throws {
 
     // Create the metadata record that is written before the chunks
     var shape, chunks: list(int);