Improve safe chunk validation

[josephnowak]

@max-sixty Here is the PR with all the features requested #9513 (comment), also I added more tests and improved the logic of the validation algorithm, I think now it is much more simple and covers certain cases that were not validated before, for example, the "r+" mode only allows full chunk writes even on the last chunk.

• [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]
• [test_zarr_safe_chunk_region ] Tests added
• User visible changes (including notable bug fixes) are documented in whats-new.rst
• New functions/methods are listed in api.rst

[max-sixty]

This looks really good @josephnowak ! The tests look great.

I left a couple of small comments and will have another read through of the code later, but overall 👍 . Thank you.

[max-sixty]

Couple of typing errors, then we can merge!

[josephnowak]

Hi @max-sixty, sorry for being insistent with the review but I would like to avoid any misunderstood with the validation logic and affect the usability of Xarray. The following case is going to raise an error using the mode "r+":

    arr = xr.DataArray(
        list(range(10)), dims=["a"], coords={"a": list(range(10))}, name="foo"
    ).chunk(a=3)
    arr.to_zarr(store, mode="w")
    with pytest.raises(ValueError):
        # It is not possible to write on a chunk that is not full, even when it is the last one which can be of a smaller size
        arr.isel(a=slice(9, 10)).to_zarr(store, region="auto", mode="r+")

Is it the expected behavior? if it is, it probably would be good to make "a" the default mode, because it allows partial writes on the first and last chunk inside the region which I think is more convenient in more scenarios

[max-sixty]

Hi @max-sixty, sorry for being insistent with the review but I would like to avoid any misunderstood with the validation logic and affect the usability of Xarray.

(tbc, thank you very much for being insistent, this is a complicated and important problem which we haven't spent enough time on, so need exactly this sort of help!)

The following case is going to raise an error using the mode "r+":

OK, that's a very good point — that writing to the full final chunk will raise an error with the proposed code. The rule should be that partial chunk writes are not allowed, but that is a full chunk write, so should be allowed. Otherwise people can't write to the full array!

Could we adjust the logic so writing to the final chunk is allowed?

if it is, it probably would be good to make "a" the default mode, because it allows partial writes on the first and last chunk inside the region which I think is more convenient in more scenarios

I weigh the cost of unsafe writes very highly, so my initial impulse is that writing to partial chunks should fail by default.

Possibly making this strict could break some existing code; possibly we should explain ourselves better. For users who want to write to a single partial chunk from a single process, it would be a fairly small change for people to pass mode="a", or allow unsafe writes.

But open to feedback from others; CC @pydata/xarray. One compromise would be to have a new mode, like rd+, for distributed writes, which would be strict like this. (though we should only do that if people feel strongly, otherwise we're at risk of config proliferation)

[max-sixty]

nit: we could do an early return outside this loop

[max-sixty]

OK, very nice! (I haven't re-reviewed all the code, just the most recent commit).

There are still a couple of type errors; lmk if you need a hand with those.

[josephnowak]

Hi @max-sixty, sorry for being insistent with the review but I would like to avoid any misunderstood with the validation logic and affect the usability of Xarray.

(tbc, thank you very much for being insistent, this is a complicated and important problem which we haven't spent enough time on, so need exactly this sort of help!)

The following case is going to raise an error using the mode "r+":

OK, that's a very good point — that writing to the full final chunk will raise an error with the proposed code. The rule should be that partial chunk writes are not allowed, but that is a full chunk write, so should be allowed. Otherwise people can't write to the full array!

Could we adjust the logic so writing to the final chunk is allowed?

if it is, it probably would be good to make "a" the default mode, because it allows partial writes on the first and last chunk inside the region which I think is more convenient in more scenarios

I weigh the cost of unsafe writes very highly, so my initial impulse is that writing to partial chunks should fail by default.

Possibly making this strict could break some existing code; possibly we should explain ourselves better. For users who want to write to a single partial chunk from a single process, it would be a fairly small change for people to pass mode="a", or allow unsafe writes.

But open to feedback from others; CC @pydata/xarray. One compromise would be to have a new mode, like rd+, for distributed writes, which would be strict like this. (though we should only do that if people feel strongly, otherwise we're at risk of config proliferation)

Yes, I can change the logic to avoid raising an error, I will have to send the shape of the Zarr array to validate if the region is covering the last chunk.

Hi @max-sixty, sorry for being insistent with the review but I would like to avoid any misunderstood with the validation logic and affect the usability of Xarray.

(tbc, thank you very much for being insistent, this is a complicated and important problem which we haven't spent enough time on, so need exactly this sort of help!)

The following case is going to raise an error using the mode "r+":

OK, that's a very good point — that writing to the full final chunk will raise an error with the proposed code. The rule should be that partial chunk writes are not allowed, but that is a full chunk write, so should be allowed. Otherwise people can't write to the full array!

Could we adjust the logic so writing to the final chunk is allowed?

if it is, it probably would be good to make "a" the default mode, because it allows partial writes on the first and last chunk inside the region which I think is more convenient in more scenarios

I weigh the cost of unsafe writes very highly, so my initial impulse is that writing to partial chunks should fail by default.

Possibly making this strict could break some existing code; possibly we should explain ourselves better. For users who want to write to a single partial chunk from a single process, it would be a fairly small change for people to pass mode="a", or allow unsafe writes.

But open to feedback from others; CC @pydata/xarray. One compromise would be to have a new mode, like rd+, for distributed writes, which would be strict like this. (though we should only do that if people feel strongly, otherwise we're at risk of config proliferation)

I have already push

OK, very nice! (I haven't re-reviewed all the code, just the most recent commit).

There are still a couple of type errors; lmk if you need a hand with those.

I will check the type error, I thought that it was related to the code of someone else

[max-sixty]

Great! I'll merge, and if anyone has feedback we can make any adjustments later.

Thank you very much @josephnowak !

[shoyer]

Thanks Max and Joseph! These checks will save Xarray users a lot of grief

…

On Sat, Sep 21, 2024 at 6:30 PM Maximilian Roos ***@***.***> wrote: Merged #9527 <#9527> into main. — Reply to this email directly, view it on GitHub <#9527 (comment)>, or unsubscribe <https://github.com/notifications/unsubscribe-auth/AAJJFVV5RUSTENY4ZNH3GCTZXYMU5AVCNFSM6AAAAABOSZLXESVHI2DSMVQWIX3LMV45UABCJFZXG5LFIV3GK3TUJZXXI2LGNFRWC5DJN5XDWMJUGM2TAOJRGQZDOMA> . You are receiving this because you are on a team that was mentioned.Message ID: ***@***.***>

[shoyer]

Unfortunately this error seems to be triggered incorrectly in some cases: #9557

Please take a look if you have the time.

[shoyer]

We reverted this in #9558.

Would love to get this back in again after fixing #9557!