Writing icechunk group piecewise leads to exponentially slower iterations

[mtvector]

Hi again,
In the process of attempting to switch my code over from pure zarr to icechunk, I've hit a weird snag in the following function and I figure it's because I don't quite understand the sequential writing for icechunk. Basically after a few iterations, each write takes exponentially longer to the point that it would take years to finish.

The purpose of the function is to sequentially read files and write their data to the store along the appropriate dimension (the files are too large to load all into memory and then write the group in one big piece). It works fine in the old version of the function which just used pure zarr3 with xarray.

#old version
def _add_x_to_ds(adata, bw_files, consensus_peaks, target, target_region_width,
                      obs_index, var_index, chunk_size=2048,n_bins=None):
    #I've removed some identical stuff here to focus on the zarr part
    
    store_path = adata.encoding['source']
    #Ensure the store is up to date
    adata.to_zarr(store_path,mode='a')

    # Open the Zarr store directly for chunk-wise writes
    store = zarr.open(store_path, mode='a')##############

    # Explicitly create the "X" dataset if not already present
    if 'X' not in store:
        store.create_dataset(##############
            'X', shape=(n_obs, n_var, seq_len),
            chunks=(n_obs, chunk_size, seq_len),
            dtype='float32',
            fill_value=np.nan
        )

    # Write data chunk-wise directly to Zarr store
    for i, bw_file in tqdm(enumerate(bw_files), total=n_obs):
        temp_bed = _create_temp_bed_file(consensus_peaks, target_region_width)
        result = _extract_values_from_bigwig(bw_file, temp_bed, target=target, n_bins=n_bins)
        os.remove(temp_bed)

        if result.ndim == 1:
            result = result.reshape(n_var, 1)
        # Directly assign the chunk into the Zarr array
        store['X'][i, :, :] = result.astype('float32')##############

Here I'm committing after running this function, but I have tried another version where you commit on each loop iteration. Both have the same problem (and so does a version that calls to_icechunk on each iteration).

#new icechunk version
def _add_x_to_ds(adata, bw_files, consensus_peaks, target, target_region_width,
                 obs_index, var_index, chunk_size=2048, n_bins=None):
    #I've removed some identical stuff here to focus on the zarr part
    # Use the icechunk-backed store from the session.
    adata.session = adata.repo.writable_session("main")##############
    store = adata.session.store##############
    # Open (or create) a Zarr group on the store in append mode.
    group = zarr.open_group(store, mode='a')##############

    # Create the "X" array in the group if it doesn't already exist.
    if 'X' not in group:
        array = group.create(##############
            # store=store,
            name='X',
            shape=(n_obs, n_var, seq_len),
            chunks=(n_obs, chunk_size, seq_len),
            dtype='float32',
            fill_value=np.nan
        )
        array.attrs["_ARRAY_DIMENSIONS"] = ["obs", "var", "seq_bins"]
    else:
        array = group['X']

    # Write data chunk-wise directly to the Zarr store.
    for i, bw_file in tqdm(enumerate(bw_files), total=n_obs):
        temp_bed = _create_temp_bed_file(consensus_peaks, target_region_width)
        result = _extract_values_from_bigwig(bw_file, temp_bed, target=target, n_bins=n_bins)
        os.remove(temp_bed)

        if result.ndim == 1:
            result = result.reshape(n_var, 1)
        # Write the computed result to the corresponding slice of the Zarr array.
        array[i, :, :] = result.astype('float32')##############

Any chance you can see where I'm going wrong?
Thanks!

notes:
I marked the zarr-relevant lines with ###
adata is basically an xarray DataSet subclass that holds onto the current repo and session (is this a bad idea?)

[dcherian]

Thanks for opening an issue. Are you able to make this a minimal reproducible example please? Can you also tell us the output of icechunk.print_debug_info()

At the very least it would be good to know actual values for shape=(n_obs, n_var, seq_len), chunks=(n_obs, chunk_size, seq_len)`.

[mtvector]

Here are the minimal examples, which are writing to a local store. They all grow in write times, less dramatically than with the large files so perhaps it is a general feature of how the files will be structured causing the increase in write times as the files grow. Also it appears the icechunk write times are much slower than the pure zarr, is this expected?

First, without committing in each:

import numpy as np
import icechunk
import zarr
import psutil
import os
from tqdm import tqdm

# === Parameters for demonstration ===
OBS = 10
VAR = int(2e4)      # 2,000,000 variables per observation
SEQ_LEN = 2114      # Sequence length per variable
CHUNK_SIZE = 512    # Chunk size for the 'var' dimension

# Create local zarr
storage = icechunk.local_filesystem_storage('/home/matthew.schmitz/Matthew/test.zarr')
repo = icechunk.Repository.create(storage)
session = repo.writable_session("main")
store = session.store

group = zarr.group(store=store, overwrite=True)
array = group.create(
    name='X',
    shape=(OBS, VAR, SEQ_LEN),
    chunks=(OBS, CHUNK_SIZE, SEQ_LEN),
    dtype='float32',
    fill_value=np.nan
)
array.attrs["_ARRAY_DIMENSIONS"] = ["obs", "var", "seq_bins"]

# Display the initial memory usage.
process = psutil.Process(os.getpid())
print("Initial memory usage (bytes):", process.memory_info().rss)

# === Loop to write random data into each observation ===
for i in tqdm(range(OBS), desc="Writing observations"):
    # Generate random data for a single observation.
    data = np.random.random((VAR, SEQ_LEN)).astype('float32')
    
    array[i, :, :] = data

    # Report memory usage after this observation.
    current_mem = process.memory_info().rss
    print(f"After writing observation {i+1}, memory usage: {current_mem} bytes")
    
# Commit the changes for this observation.
print(session.commit('dummy commit'))
print("Finished writing all observations.")
icechunk.print_debug_info()

Initial memory usage (bytes): 3963637760
Writing observations:  10%|█         | 1[/10](https://aidc-ood-prd.corp.alleninstitute.org/10) [00:08<01:19,  8.85s[/it](https://aidc-ood-prd.corp.alleninstitute.org/it)]
After writing observation 1, memory usage: 4134625280 bytes
Writing observations:  20%|██        | 2[/10](https://aidc-ood-prd.corp.alleninstitute.org/10) [00:30<02:10, 16.26s[/it](https://aidc-ood-prd.corp.alleninstitute.org/it)]
After writing observation 2, memory usage: 4152950784 bytes
Writing observations:  30%|███       | 3[/10](https://aidc-ood-prd.corp.alleninstitute.org/10) [01:03<02:46, 23.77s[/it](https://aidc-ood-prd.corp.alleninstitute.org/it)]
After writing observation 3, memory usage: 4159426560 bytes
Writing observations:  40%|████      | 4[/10](https://aidc-ood-prd.corp.alleninstitute.org/10) [01:57<03:34, 35.77s[/it](https://aidc-ood-prd.corp.alleninstitute.org/it)]
After writing observation 4, memory usage: 4211236864 bytes
Writing observations:  50%|█████     | 5[/10](https://aidc-ood-prd.corp.alleninstitute.org/10) [02:37<03:06, 37.39s[/it](https://aidc-ood-prd.corp.alleninstitute.org/it)]
After writing observation 5, memory usage: 4157231104 bytes
Writing observations:  60%|██████    | 6[/10](https://aidc-ood-prd.corp.alleninstitute.org/10) [03:31<02:52, 43.22s[/it](https://aidc-ood-prd.corp.alleninstitute.org/it)]
After writing observation 6, memory usage: 4331810816 bytes
Writing observations:  70%|███████   | 7[/10](https://aidc-ood-prd.corp.alleninstitute.org/10) [04:41<02:35, 51.83s[/it](https://aidc-ood-prd.corp.alleninstitute.org/it)]
After writing observation 7, memory usage: 4463869952 bytes
Writing observations:  80%|████████  | 8[/10](https://aidc-ood-prd.corp.alleninstitute.org/10) [06:15<02:10, 65.23s[/it](https://aidc-ood-prd.corp.alleninstitute.org/it)]
After writing observation 8, memory usage: 4745273344 bytes
Writing observations:  90%|█████████ | 9[/10](https://aidc-ood-prd.corp.alleninstitute.org/10) [07:24<01:06, 66.28s[/it](https://aidc-ood-prd.corp.alleninstitute.org/it)]
After writing observation 9, memory usage: 4913094656 bytes
Writing observations: 100%|██████████| 10[/10](https://aidc-ood-prd.corp.alleninstitute.org/10) [09:16<00:00, 55.61s[/it](https://aidc-ood-prd.corp.alleninstitute.org/it)]
After writing observation 10, memory usage: 5011529728 bytes
FBC3FSDPD1TV47Z7R8J0
Finished writing all observations.
platform:  Linux-5.14.0-427.42.1.el9_4.x86_64-x86_64-with-glibc2.34
python:  3.12.7
icechunk:  0.2.12
zarr:  3.0.0
numcodecs:  0.15.1

xarray:  2025.1.2

Secondly, committing in each iteration:

import numpy as np
import icechunk
import zarr
import psutil
import os
from tqdm import tqdm

# === Parameters for demonstration ===
OBS = 10
VAR = int(2e4)      # 20,000 variables per observation
SEQ_LEN = 2114      # Sequence length per variable
CHUNK_SIZE = 512    # Chunk size for the 'var' dimension

# === Create an Icechunk repository on the local filesystem ===
# Note: replace the path with a location you can write to.
storage = icechunk.local_filesystem_storage('/home/matthew.schmitz/Matthew/test.zarr')
repo = icechunk.Repository.create(storage)
session = repo.writable_session("main")
store = session.store

# Create the Zarr group and array on the Icechunk store.
group = zarr.group(store=store, overwrite=True)
array = group.create(
    name='X',
    shape=(OBS, VAR, SEQ_LEN),
    chunks=(OBS, CHUNK_SIZE, SEQ_LEN),
    dtype='float32',
    fill_value=np.nan
)
array.attrs["_ARRAY_DIMENSIONS"] = ["obs", "var", "seq_bins"]

# Display the initial memory usage.
process = psutil.Process(os.getpid())
print("Initial memory usage (bytes):", process.memory_info().rss)

# === Loop to write random data into each observation ===
for i in tqdm(range(OBS), desc="Writing observations"):
    # Generate random data for a single observation.
    data = np.random.random((VAR, SEQ_LEN)).astype('float32')
    
    # Write the random data to the corresponding slice in the Zarr array.
    array[i, :, :] = data
    
    # Commit the changes for this observation.
    commit_id = session.commit(f'dummy commit_{i}')
    print(f"Committed observation {i+1} with id: {commit_id}")
    
    # Re-open a new writable session for further writes.
    session = repo.writable_session("main")
    store = session.store
    # Re-open the updated Zarr array from the new store.
    group = zarr.open_group(store, mode='a')
    array = group['X']
    current_mem = process.memory_info().rss
    print(f"After writing observation {i+1}, memory usage: {current_mem} bytes")

final_commit_id = session.commit('final dummy commit')
print(f"Final commit id: {final_commit_id}")
print("Finished writing all observations.")

icechunk.print_debug_info()

Initial memory usage (bytes): 3762229248
Writing observations:  10%|█         | 1[/10](https://aidc-ood-prd.corp.alleninstitute.org/10) [00:08<01:18,  8.75s[/it](https://aidc-ood-prd.corp.alleninstitute.org/it)]
Committed observation 1 with id: V4638RB2KFB5EJYSPQ0G
After writing observation 1, memory usage: 3934769152 bytes
Writing observations:  20%|██        | 2[/10](https://aidc-ood-prd.corp.alleninstitute.org/10) [01:05<04:54, 36.86s[/it](https://aidc-ood-prd.corp.alleninstitute.org/it)]
Committed observation 2 with id: NNQ6WTAF3FAE46E1J2R0
After writing observation 2, memory usage: 3412848640 bytes
Writing observations:  30%|███       | 3[/10](https://aidc-ood-prd.corp.alleninstitute.org/10) [01:29<03:37, 31.10s[/it](https://aidc-ood-prd.corp.alleninstitute.org/it)]
Committed observation 3 with id: M5ATGZ0G1J1Z9KKENR7G
After writing observation 3, memory usage: 3388350464 bytes
Writing observations:  40%|████      | 4[/10](https://aidc-ood-prd.corp.alleninstitute.org/10) [02:14<03:39, 36.60s[/it](https://aidc-ood-prd.corp.alleninstitute.org/it)]
Committed observation 4 with id: N7NTR5FQPGXN7MMKQ1H0
After writing observation 4, memory usage: 3486969856 bytes
Writing observations:  50%|█████     | 5[/10](https://aidc-ood-prd.corp.alleninstitute.org/10) [03:16<03:47, 45.60s[/it](https://aidc-ood-prd.corp.alleninstitute.org/it)]
Committed observation 5 with id: DGM4BD24MCPJ8J6JGT50
After writing observation 5, memory usage: 3025842176 bytes
Writing observations:  60%|██████    | 6[/10](https://aidc-ood-prd.corp.alleninstitute.org/10) [04:18<03:25, 51.41s[/it](https://aidc-ood-prd.corp.alleninstitute.org/it)]
Committed observation 6 with id: 3DCAPHPQ4P222297PJ4G
After writing observation 6, memory usage: 3045105664 bytes
Writing observations:  70%|███████   | 7[/10](https://aidc-ood-prd.corp.alleninstitute.org/10) [05:40<03:03, 61.30s[/it](https://aidc-ood-prd.corp.alleninstitute.org/it)]
Committed observation 7 with id: GKQY747MM33GE1Q61SGG
After writing observation 7, memory usage: 3195998208 bytes
Writing observations:  80%|████████  | 8[/10](https://aidc-ood-prd.corp.alleninstitute.org/10) [06:48<02:06, 63.32s[/it](https://aidc-ood-prd.corp.alleninstitute.org/it)]
Committed observation 8 with id: DG9Y13M1GY6S9T4GS7C0
After writing observation 8, memory usage: 3324370944 bytes
Writing observations:  90%|█████████ | 9[/10](https://aidc-ood-prd.corp.alleninstitute.org/10) [07:56<01:04, 64.87s[/it](https://aidc-ood-prd.corp.alleninstitute.org/it)]
Committed observation 9 with id: TM6CMCZQ9JYKE61ZCHQG
After writing observation 9, memory usage: 3345354752 bytes
Writing observations: 100%|██████████| 10[/10](https://aidc-ood-prd.corp.alleninstitute.org/10) [09:11<00:00, 55.12s[/it](https://aidc-ood-prd.corp.alleninstitute.org/it)]
Committed observation 10 with id: WNZJ0SP4ZFAMDNGWJ4QG
After writing observation 10, memory usage: 3365064704 bytes

And finally using pure zarr, which is significantly faster but also increases with each iteration:

import numpy as np
import zarr
import psutil
import os
from tqdm import tqdm

# === Parameters for demonstration ===
OBS = 10
VAR = int(2e4)      # 20,000 variables per observation (adjust as needed)
SEQ_LEN = 2114      # Sequence length per variable
CHUNK_SIZE = 512    # Chunk size along the 'var' dimension

# === Create a local Zarr DirectoryStore ===
store_dir = '/home/matthew.schmitz/Matthew/test.zarr'
print("Using Zarr store directory:", store_dir)
store = zarr.storage.LocalStore(store_dir)

# Create a Zarr group and array with the specified dimensions and chunking.
group = zarr.group(store=store, overwrite=True)
array = group.create(
    name='X',
    shape=(OBS, VAR, SEQ_LEN),
    chunks=(OBS, CHUNK_SIZE, SEQ_LEN),
    dtype='float32',
    fill_value=np.nan
)
# (Optional) Set dimension metadata for downstream tools.
array.attrs["_ARRAY_DIMENSIONS"] = ["obs", "var", "seq_bins"]

# Display the initial memory usage.
process = psutil.Process(os.getpid())
print("Initial memory usage (bytes):", process.memory_info().rss)

# === Loop to write random data into each observation ===
for i in tqdm(range(OBS), desc="Writing observations"):
    # Generate random data for the observation.
    data = np.random.random((VAR, SEQ_LEN)).astype('float32')
    
    # Write the random data to the corresponding slice of the Zarr array.
    array[i, :, :] = data

    # Report memory usage after this observation.
    current_mem = process.memory_info().rss
    print(f"After writing observation {i+1}, memory usage: {current_mem} bytes")

print("Finished writing all observations.")

Using Zarr store directory: /home/matthew.schmitz/Matthew/test.zarr
Initial memory usage (bytes): 3242545152
Writing observations:  10%|█         | 1[/10](https://aidc-ood-prd.corp.alleninstitute.org/10) [00:02<00:18,  2.03s[/it](https://aidc-ood-prd.corp.alleninstitute.org/it)]
After writing observation 1, memory usage: 3413397504 bytes
Writing observations:  20%|██        | 2[/10](https://aidc-ood-prd.corp.alleninstitute.org/10) [00:05<00:21,  2.70s[/it](https://aidc-ood-prd.corp.alleninstitute.org/it)]
After writing observation 2, memory usage: 3420995584 bytes
Writing observations:  30%|███       | 3[/10](https://aidc-ood-prd.corp.alleninstitute.org/10) [00:08<00:21,  3.13s[/it](https://aidc-ood-prd.corp.alleninstitute.org/it)]
After writing observation 3, memory usage: 3430469632 bytes
Writing observations:  40%|████      | 4[/10](https://aidc-ood-prd.corp.alleninstitute.org/10) [00:12<00:21,  3.53s[/it](https://aidc-ood-prd.corp.alleninstitute.org/it)]
After writing observation 4, memory usage: 3447918592 bytes
Writing observations:  50%|█████     | 5[/10](https://aidc-ood-prd.corp.alleninstitute.org/10) [00:18<00:20,  4.09s[/it](https://aidc-ood-prd.corp.alleninstitute.org/it)]
After writing observation 5, memory usage: 3281653760 bytes
Writing observations:  60%|██████    | 6[/10](https://aidc-ood-prd.corp.alleninstitute.org/10) [00:23<00:18,  4.71s[/it](https://aidc-ood-prd.corp.alleninstitute.org/it)]
After writing observation 6, memory usage: 3469582336 bytes
Writing observations:  70%|███████   | 7[/10](https://aidc-ood-prd.corp.alleninstitute.org/10) [00:30<00:15,  5.29s[/it](https://aidc-ood-prd.corp.alleninstitute.org/it)]
After writing observation 7, memory usage: 3719208960 bytes
Writing observations:  80%|████████  | 8[/10](https://aidc-ood-prd.corp.alleninstitute.org/10) [00:37<00:11,  5.91s[/it](https://aidc-ood-prd.corp.alleninstitute.org/it)]
After writing observation 8, memory usage: 3823661056 bytes
Writing observations:  90%|█████████ | 9[/10](https://aidc-ood-prd.corp.alleninstitute.org/10) [00:46<00:06,  6.73s[/it](https://aidc-ood-prd.corp.alleninstitute.org/it)]
After writing observation 9, memory usage: 3718164480 bytes
Writing observations: 100%|██████████| 10[/10](https://aidc-ood-prd.corp.alleninstitute.org/10) [00:54<00:00,  5.45s[/it](https://aidc-ood-prd.corp.alleninstitute.org/it)]
After writing observation 10, memory usage: 3718926336 bytes
Finished writing all observations.

[dcherian]

Thank you for the very good reproducible example!

Yesterday we released icechunk 0.2.13 with a bugfix that substantially reduces the overhead for icechunk.
Couple of comments:

1. This workflow has chunksize > 1 along the dimension of appending. This means Zarr has to read back the previously written chunk, merge in the new update, and the write the whole chunk back. So the amount of bytes being written in each iteration increases linearly.
2. At the moment, the icechunk workflow is still slower than Zarr. We are looking into it more.

PS: Note that the memory measurements aren't very helpful because of python's garbage collector. If you are interested in memory, you're better off using memray, scalene, or similar to track allocations.

[mtvector]

Yeah... I haven't tested the performance with chunk size of 1 along the dimension I was appending. For my application I need all of the values along that dimension every time, so I figured I'd speed up reading by keeping it all together. From what I had done originally, it worked fine with normal zarr so I figured it would work just as well with icechunk, look forward to testing the new version! I also am not opposed to writing a normal zarr store if this ends up being more straightforward and then converting it to an icechunk store (which I raised a feature request asking about that utility).

Thanks for the point about the memory, I'll use an actual profiler going forward. The reason it's in there as you may have surmised is that I was wondering if once elements from memory get written to the store whether they stay in memory indefinitely or if they get managed by the cache even without reopening the store.

[dcherian]

For my application I need all of the values along that dimension every time, so I figured I'd speed up reading by keeping it all together.

In that case I'd append along some other dimension that is chunked.

---

In my tests Icechunk worked well with a chunk size of 1. With the current setup, there are interleaved sets and gets, each of which holds the GIL while sending data across the python-rust bridge. This GIL contention is the slowdown you see.