fix safe chunks validation

[josephnowak]

This MR fixes multiple cases on the logic of the safe_chunks option, for more info please see the test added, which includes multiple scenarios.

I know that some of the issues are already closed, but if you see the comments none of them was successfully closed, but with this fix a proper error is going to be raised.

• [Appending data to a dataset stored in Zarr format produce PermissonError or NaN values in the final result #5511, Is it unsafe to write on independent regions using the to_zarr method when one of chunks is not complete? #9072, Appending to existing zarr store writes mostly NaN from dask arrays, but not numpy arrays #7812] Closes #xxxx
• [test_zarr_safe_chunk] Tests added
• User visible changes (including notable bug fixes) are documented in whats-new.rst
• New functions/methods are listed in api.rst

[josephnowak]

Hi @shoyer, a long time ago I mentioned to you that it would be good to raise an error on certain scenarios where the dask and zarr chunks are not aligned, I was finally able to dedicate some time to this fix, and I think it would be useful to avoid many issues in the future related to incorrect writes on Zarr.

[max-sixty]

If the region kwarg is required now, should we make it mandatory? Or at least add a TODO in the function if we want to push that off to another PR?

[josephnowak]

I think that ideally, the validation should be outside the extraction of the encoding, and that would avoid the use of the region inside the extract_zarr_variable_encoding function, I added a comment with that option in the zarr.py file but, I tried to reduce the amount of modifications to the code so I decided to not apply it.

[josephnowak]

If needed I can add a TODO indicating that it would be good to apply that modification.

[max-sixty]

Are we sure we don't need to distinguish between append mode vs. others? If we have an array of (5,5,5) and we write to location (4,10) in chunks of [(4,5),(6,10)], then that's only OK if we're appending — otherwise we could be writing (1,4) in another process?

(To the extent the proposed code dominates the existing code — i.e. the existing code still has this problem — I would still vote to merge even if we can't solve that on this iteration, albeit with a TODO)

[josephnowak]

Could you specify that case using Xarray? I'm not able to understand correctly the chunk sizes that you are indicating, but I think that it is not necessary to distinguish between append and others because Xarray internally represents all as a region write, and for the update region it is always going to apply the logic of check if the first chunk is aligned with the border_size + N * zarr_chunk_size, for any integer N >= 0 and border_size being the amount of data between interval.start and the corresponding chunk end.

[max-sixty]

Take the example above:

This seems safe or unsafe depending on whether we're appending:

• If we're appending (mode="a"), then this is safe, because we can only append in a single process, so can only be writing the blue box
• But if we're doing a region write without append (mode="r+"), then we can be writing to the yellow and blue outlined box in two separate processes, which would cause corruption. So when mode="r+", writing half a chunk isn't safe.

(tbc, you seem to understand this overall better than I do, so I suspect I'm wrong, but also that this is important enough that it's worth clarifying)

[josephnowak]

I think that case is already covered by the validation, if the first chunk of the array (the blue box) is not of the size of the region that needs to be updated it is also going to raise an error.

For example, if you see the test on the image (test_zarr_region_chunk_partial_offset) the first chunk of the array is going to be written from position 5 until position 15, which means that it is going to touch only half of the first and second Zarr chunk, this would be completely valid if there is no other chunk in the array, but as there is another one it means that the second Zarr chunk and, only the second one is going to suffer the issue of more than one Dask chunk being written to the same Zarr chunk in parallel.

border_size = zchunk - interval.start % zchunk
if dchunks[0] % zchunk != border_size:

the translation of that condition is that the first chunk of Dask must be always greater or equal to the portion of the data that needs to be updated on the first Zarr chunk (the border size), and if it is greater then it must be exactly aligned with the other Zarr chunks unless there is no other Dask chunk

[josephnowak]

I hope these examples help to better understand my explanation.

import xarray as xr

path = "C:/Users/josep/Desktop/test_zarr3"
arr = xr.DataArray(list(range(10)), dims=["a"], coords={"a": list(range(10))}, name="foo").chunk(a=3)
arr.to_zarr(path, mode="w")

try:
    arr.isel(a=slice(0, 3)).chunk(a=(2, 1)).to_zarr(path, region="auto")
except ValueError:
    print("Chunk error, the first chunk is not completely covering the border size")

# This is valid the border size is covered by the first chunk of Dask
arr.isel(a=slice(1, 4)).chunk(a=(2, 1)).to_zarr(path, region="auto")


try:
    arr.isel(a=slice(1, 5)).chunk(a=(3, 1)).to_zarr(path, region="auto")
except ValueError:
    print("Chunk error, the first chunk is covering the border size but it is not completely covering the second chunk")

# This is valid because there is a single dask chunk, so it is not possible to write multiple chunks in parallel 
arr.isel(a=slice(1, 5)).chunk(a=(4, )).to_zarr(path, region="auto")

try:
    arr.isel(a=slice(0, 5)).chunk(a=(3, 1, 1)).to_zarr(path, region="auto")
except ValueError:
    print("Chunk error, the first chunk is correct but the other two chunks are overlapping the same Zarr chunk")

# This is the simplest case, the first chunk is completely covering the first Zarr chunk
arr.isel(a=slice(0, 5)).chunk(a=(3, 2)).to_zarr(path, region="auto")


# The first chunk is covering the border size (2 elements) and also the second chunk (3 elements), so it is valid
arr.isel(a=slice(1, 8)).chunk(a=(5, 2)).to_zarr(path, region="auto")

[josephnowak]

Oh ok, you are right it is hard to detect that alignment problem, and if that would be the meaning of the mode="a" it has sense, just that right now if I try to write on a region using that mode it raises the following error:

arr.isel(a=slice(1, 8)).chunk(a=(5, 2)).to_zarr(path, region="auto", mode="a")

ValueError: ``mode`` must be 'r+' when using ``region='auto'``, got 'a'

Do you think that we should allow to write on region using the "a" mode? if yes then I could proceed to change the validation code so it satisfy the condition that you specified before (partial writes for "a" and non-partial writes for "r+")

[max-sixty]

I think that would be a nice addition, yes! Though would need to refresh my memory to be confident, and others may have views.

I would encourage landing this PR with the improvements, and then another PR with that change, unless that's difficult!

Thank you!

[josephnowak]

Sounds fine, so would this PR be ready then? and just to clarify would you like that in the other PR I add the following features?

1. If the mode is "r+" then only allows full chunks.
2. If the mode is "a" allows writing using region.
3. If the mode is "a" allows partial writes.
4. Probably it would be good to also update the docs of the to_zarr method, honestly they do not indicate any of those behaviors.

[max-sixty]

Sounds fine, so would this PR be ready then?

I added one comment above, otherwise I think so!

[max-sixty]

1. If the mode is "r+" then only allows full chunks.

2. If the mode is "a" allows writing using region.

r+ can write using region — that's the default — as long as they are full chunks. (I think the 1&3 cover things, bullet 2 isn't needed)

3. If the mode is "a" allows partial writes.

4. Probably it would be good to also update the docs of the to_zarr method, honestly they do not indicate any of those behaviors.

That would be really good, we've fallen behind a bit on docs. Possibly an LLM can help if that makes improving the docs a bit more appealing :)

[max-sixty]

Very nice! Thank you v much @josephnowak . Added a couple of questions and comments

[josephnowak]

By the way, probably it would be good to create a utility function on Xarray (or a parameter on the to_zarr method) that automatically aligns the Dask chunks of a dataset with the Zarr chunks, this could be inefficient in some scenarios but it can be quite useful to avoid dealing with the alignment of chunks logic or using a synchronizer.

[max-sixty]

Re "so would this PR be ready then" — am I right in thinking the proposed code changes the behavior such that a single partial chunk can be written on mode="r+"? I don't think we want that yet (though am not confident, it seems unlikely to be a big problem, but still possibly confusing)

[josephnowak]

Yes, the code is going to allow writing on a single partial chunk on mode="r+".
From my perspective and based on the docs of the to_zarr method, the new behavior would not be confusing, "“r+” means modify existing array values only (raise an error if any metadata or shapes would change)", and writing in partial chunk means modify the existing array values.

[max-sixty]

But it breaks the rule we discussed above, re it being safe writing to different locations from distributed processes.

It's unlikely, because it's rare to write to a single partial chunk from distributed processes. But is it worth the complication, relative to forcing people to pass mode="a"?

[josephnowak]

@max-sixty I sent the following PR #9527 which contains all the changes requested, if you want I can add everything here or close this one, whatever you prefer.