Skip to content

Commit

Permalink
Remove unpack_chunks code
Browse files Browse the repository at this point in the history
dask.array.core.normalize_chunks is a better way
to do the same thing.
  • Loading branch information
Kirill888 committed May 29, 2024
1 parent cac88f9 commit 50518e5
Show file tree
Hide file tree
Showing 2 changed files with 1 addition and 28 deletions.
3 changes: 1 addition & 2 deletions odc/loader/_builder.py
Original file line number Diff line number Diff line change
Expand Up @@ -32,7 +32,6 @@
from odc.geo.geobox import GeoBox, GeoBoxBase, GeoboxTiles
from odc.geo.xr import xr_coords

from ._dask import unpack_chunks
from ._reader import nodata_mask, resolve_dst_fill_value, resolve_src_nodata
from ._utils import SizedIterable, pmap
from .types import (
Expand Down Expand Up @@ -195,7 +194,7 @@ def __call__(
*postfix_dims,
)
assert len(chunk_shape) == len(shape)
chunks = unpack_chunks(chunk_shape, shape)
chunks: tuple[tuple[int, ...], ...] = normalize_chunks(chunk_shape, shape)
tchunk_range = [
range(last - n, last) for last, n in zip(np.cumsum(chunks[0]), chunks[0])
]
Expand Down
26 changes: 0 additions & 26 deletions odc/loader/_dask.py
Original file line number Diff line number Diff line change
Expand Up @@ -34,29 +34,3 @@ def tokenize_stream(
for k, x in kx:
dsk[k] = x
yield k, x


def unpack_chunksize(chunk: int, N: int) -> Tuple[int, ...]:
"""
Compute chunk sizes
Example: 4, 11 -> (4, 4, 3)
"""
if chunk >= N or chunk < 0:
return (N,)

nb = N // chunk
last_chunk = N - chunk * nb
if last_chunk == 0:
return tuple(chunk for _ in range(nb))

return tuple(chunk for _ in range(nb)) + (last_chunk,)


def unpack_chunks(
chunks: Tuple[int, ...], shape: Tuple[int, ...]
) -> Tuple[Tuple[int, ...], ...]:
"""
Expand chunks
"""
assert len(chunks) == len(shape)
return tuple(unpack_chunksize(ch, n) for ch, n in zip(chunks, shape))

0 comments on commit 50518e5

Please sign in to comment.