Read a chunk from all sub-groups of single zarr store

[sanjaysrikakulam]

Hi,

I am attempting to read n rows or 1 chunk at a time from millions of 1D array sub-group stored under a single Zarr store. Doing this in a for loop or multiprocessing would be really slow. Is there a way to do this in an efficient way using Dask? I am a complete novice to Dask, so any pointers or help would be really great.

Here is what I am trying to do:

import os
import zarr
from numcodecs import Blosc, PackBits
import numpy as np
from multiprocessing import Process

# Zarr store related #
store_file = "data/bloomfilters.zarr"
store = zarr.DirectoryStore(store_file)
sync_dir = os.path.splitext(store_file)[0] + ".sync"
synchronizer = zarr.ProcessSynchronizer(sync_dir)
Blosc.use_threads = False
compressor = Blosc(cname="zstd", clevel=9, shuffle=Blosc.BITSHUFFLE)
root_group = zarr.group(store=store, synchronizer=synchronizer)
bf_group = root_group.create_group(name="BloomFilterIndex")
max_memory = 64420000000

# Parallelize bloom filter writes to zarr #
# samples_list size is 215M sometimes even more
bytes_per_row = 215000000
row_chunk = 250                # (row_chunk * btyes_per_row) <= ((max_memory) * 90/100)
bloomfilter_size = 120000
n_cores = 64
processes = []
for samp_list in np.array_split(samples_list, n_cores):
    proc = Process(target=write_bf_to_zarr, args=(bf_group, compressor, bloomfilter_size, samples_list, row_chunk))
    proc.start()
    processes.append(proc)

for proc in processes:
    proc.join()

# Write function #
def write_bf_to_zarr(bf_group_handle, compressor, bloomfilter_size, samples_list, row_chunk):
    for sample_info in samples_list:
        col_idx = sample_info[0]
        file_paths = sample_info[1]
        bloomfilter = bf_group_handle.zeros(
            f"sample_{col_idx}",
            shape=(bloomfilter_size),
            chunks=(row_chunk, ),
            dtype="bool",
            filters=[PackBits()],
            compressor=compressor,
        )
        # Note: bloomfilter_cython will create bloom filter as NumPy boolean array and will return it
        bloomfilter[:] = bloomfilter_cython(file_paths, bloomfilter_size)

# To slice 'n' rows (example 10 rows or even a chunk ) from each of the sub-group
m = len(bf_group.keys())
n = row_chunk
n_rows = np.zeros((n, m), dtype=bool)

start = 0
stop = n
for col_idx, key in enumerate(bf_group.keys()):
    n_rows[:, col_idx] = bf_group.get(key)[start:stop]

What am I trying to do with the concatenated chunk:
I would like to extract the non-zero index position from this concatenated 2D array row-wise and store it in a DB for fast querying.

Thank you!