From 50518e5b7c51933774e4f0571bfba603f278381f Mon Sep 17 00:00:00 2001 From: Kirill Kouzoubov Date: Wed, 29 May 2024 14:08:21 +1000 Subject: [PATCH] Remove unpack_chunks code dask.array.core.normalize_chunks is a better way to do the same thing. --- odc/loader/_builder.py | 3 +-- odc/loader/_dask.py | 26 -------------------------- 2 files changed, 1 insertion(+), 28 deletions(-) diff --git a/odc/loader/_builder.py b/odc/loader/_builder.py index 08e8689..2768c45 100644 --- a/odc/loader/_builder.py +++ b/odc/loader/_builder.py @@ -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 ( @@ -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]) ] diff --git a/odc/loader/_dask.py b/odc/loader/_dask.py index 1751668..5cc79b4 100644 --- a/odc/loader/_dask.py +++ b/odc/loader/_dask.py @@ -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))