Skip to content

rechunk() sometimes lead to arrays with non-regular chunks (dask==2025.1.0) #11713

New issue
New issue
Closed
Closed
rechunk() sometimes lead to arrays with non-regular chunks (dask==2025.1.0)#11713
Labels
needs triageNeeds a response from a contributor

Description

@LucaMarconato
LucaMarconato
opened on Feb 4, 2025

Describe the issue:

Rechunking with 'auto' sometimes lead to arrays whose chunks are not regular, making dask.array.to_zarr() fail (the validation _check_regular_chunks() from dask.array.core fails. This is shown in the example below. Passing explicit chunk dimensions instead of 'auto' fixes the problem. Bug reproducible with dask==2025.1.0.

Related to: #11544. The bug was originally reported from @christinedien in the context of the spatialdata package. I kept the dimensionalities of the original bug report in the code example below.

Minimal Complete Verifiable Example:

import dask.array as da
from dask.array.core import _check_regular_chunks
import pytest


# Function to create a lazy zero-filled dask array
def create_lazy_zeros(shape, chunks, dtype='uint16'):
    return da.zeros(shape, chunks=chunks, dtype=dtype)


# Create the mock array with lazy zeros
def create_mock_array():
    # Define the shape and chunks for scale0
    shape = (1, 25778, 35416)
    chunks = (1, 4096, 4096)

    # Create a lazy zero-filled dask array
    data = create_lazy_zeros(shape, chunks)

    return data


# Create the mock array
mock_array = create_mock_array()

# Slice the array
sliced_array = mock_array[:, :13000, 15000:35000]
print('Chunks before rechunking')
print(sliced_array.chunks)

# Rechunk the array with 'auto'
sliced_array_auto = sliced_array.rechunk('auto')
print('Chunks after rechunking with auto')
print(sliced_array_auto.chunks)

# Check for regular chunks
with pytest.raises(AssertionError):
    assert _check_regular_chunks(sliced_array_auto.chunks)

# Rechunk the array with a fixed chunk size
FIXED_CHUNK = 4096
sliced_array_fixed = sliced_array.rechunk((1, FIXED_CHUNK, FIXED_CHUNK))
print(f'Chunks after rechunking with a fixed chunk size (FIXED_CHUNK)')
print(sliced_array_fixed.chunks)

# Assert that chunks are regular
assert _check_regular_chunks(sliced_array_fixed.chunks)

Anything else we need to know?:

Environment:

  • Dask version: 2025.1.0
  • Python version: 3.11.0
  • Operating System: macOS 14.6.1
  • Install method (conda, pip, source): conda (conda-forge)

Activity

  1. github-actions
    added
    needs triageNeeds a response from a contributor
    on Feb 4, 2025

  2. AAYUSH-SPIDEY-SHARMA commented on Feb 5, 2025

    @AAYUSH-SPIDEY-SHARMA
    AAYUSH-SPIDEY-SHARMA
    on Feb 5, 2025

    hi @LucaMarconato i would like to work on this issue , please can you assign it to me.

  3. LucaMarconato commented on Feb 6, 2025

    @LucaMarconato
    LucaMarconato
    on Feb 6, 2025
    Author

    Hi @AAYUSH-SPIDEY-SHARMA, I am not a maintainer of dask so I do not have triaging rights. We need to wait for a maintainer. Thanks for the interest meanwhile! 😊

  4. AAYUSH-SPIDEY-SHARMA commented on Feb 7, 2025

    @AAYUSH-SPIDEY-SHARMA
    AAYUSH-SPIDEY-SHARMA
    on Feb 7, 2025

    @LucaMarconato thank you bro, for this information.....

  5. phofl commented on Feb 10, 2025

    @phofl
    phofl
    on Feb 10, 2025
    Collaborator

    Deviating from regular chunks was an intentional decision to improve efficiency of computations. So that part works as expected. I agree though that the default experience when writing to zarr should be better and just work through injecting a rechunk to get the required chunking

  6. LucaMarconato commented on Feb 10, 2025

    @LucaMarconato
    LucaMarconato
    on Feb 10, 2025
    Author

    Thank you for your comment. What I find confusing is that calling sliced_array.rechunk('auto') seems to have no effects. I think it could be made cleared that the new default behavior for 'auto' is to keep the chunks irregular.

  7. phofl commented on Feb 11, 2025

    @phofl
    phofl
    on Feb 11, 2025
    Collaborator

    Do you have a reproducer that is actually hitting this error?

            if not _check_regular_chunks(arr.chunks):
            # We almost certainly get here because auto chunking has been used
            # on irregular chunks. The max will then be smaller than auto, so using
            # max is a safe choice
            arr = arr.rechunk(tuple(map(max, arr.chunks)))

    this should protect against the problem you are seeing

  8. LucaMarconato commented on Feb 11, 2025

    @LucaMarconato
    LucaMarconato
    on Feb 11, 2025
    Author

    Yes, I implemented a workaround here https://github.com/scverse/spatialdata/pull/861/files, I could use the code you suggest instead. In this way I don't have to hardcode (or have the user decide): CHUNK_SIZE = 1024.

  9. LucaMarconato commented on Feb 11, 2025

    @LucaMarconato
    LucaMarconato
    on Feb 11, 2025
    Author

    I tracked this task here scverse/spatialdata#872 (not urgent). Meanwhile if you think your approach is enough, I think this issue can be closed. Thanks again!

  10. phofl commented on Feb 11, 2025

    @phofl
    phofl
    on Feb 11, 2025
    Collaborator

    Yeah I saw this, but do you have a unit test or something similar that explains how you got to this exception? I'd like to fix this in Dask itself so that you don't have to handle this in your library

  11. LucaMarconato commented on Feb 11, 2025

    @LucaMarconato
    LucaMarconato
    on Feb 11, 2025
    Author

    Ah got it. The real code triggering this was case basically the one that I posted as first post, with the exception of replacing _check_regular_chunks(sliced_array_auto.chunks) with sliced_array_auto.to_zarr().

    I'll make a snippet now, simplifying the code above.

  12. LucaMarconato commented on Feb 11, 2025

    @LucaMarconato
    LucaMarconato
    on Feb 11, 2025
    Author

    Here is it:

    import dask.array as da
from tempfile import TemporaryDirectory


shape = (1, 25778, 35416)
chunks = (1, 4096, 4096)
mock_array = da.zeros(shape, chunks=chunks, dtype="uint16")
sliced_array = mock_array[:, :13000, 15000:35000]
sliced_array_auto = sliced_array.rechunk("auto")

with TemporaryDirectory() as tmpdir:
    sliced_array_auto.to_zarr(tmpdir)

  13. phofl commented on Feb 11, 2025

    @phofl
    phofl
    on Feb 11, 2025
    Collaborator

    This completes for me without raising because the rechunk we added in Dask hits and smoothes out the chunks. Does this raise for you? If yes, I'd be curious about the traceback

  14. phofl commented on Feb 11, 2025

    @phofl
    phofl
    on Feb 11, 2025
    Collaborator

    Any chance you are on 2024.11.2? I think we had one release without this automatic rechunking functionality.

  15. LucaMarconato commented on Feb 11, 2025

    @LucaMarconato
    LucaMarconato
    on Feb 11, 2025
    Author

    Good point! My first code snipped (the one using

    with pytest.raises(AssertionError):
    assert _check_regular_chunks(sliced_array_auto.chunks)

    fails with 2025.1.0. Instead, the one using to_zarr() works on 2025.1.0. I could not try the real world example on 2025.1.0 because spatialdata does not support the latest dask (due to the .attrs missing).

    So it seems that no further changes on the dask size for this bug are needed. Wdyt?

  16. phofl commented on Feb 11, 2025

    @phofl
    phofl
    on Feb 11, 2025
    Collaborator

    Yes, correct. I think the bug has been fixed in 2024.12.0. I am closing this one

  17. phofl
    closed this as completedon Feb 11, 2025
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Metadata

Metadata

Assignees

No one assigned

    Labels

    needs triageNeeds a response from a contributor

    Type

    No type

    Projects

    No projects

      Milestone

      No milestone

      Relationships

      None yet

      Development

      No branches or pull requests

      Issue actions