Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Cache the result of DaskManager.normalize_chunks #9897

Draft
wants to merge 2 commits into
base: main
Choose a base branch
from
Draft

Cache the result of DaskManager.normalize_chunks #9897

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
0e34528
Cache the result of `DaskManager.normalize_chunks`
dcherian Dec 3, 2024
de010ea
Try using uuids
dcherian Dec 3, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
5 changes: 3 additions & 2 deletions xarray/core/dataset.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -6,6 +6,7 @@
import itertools
import math
import sys
import uuid
import warnings
from collections import defaultdict
from collections.abc import (
Expand Down Expand Up @@ -257,7 +258,6 @@ def _get_chunk(var: Variable, chunks, chunkmanager: ChunkManagerEntrypoint):
chunk_shape = chunkmanager.normalize_chunks(
chunk_shape, shape=shape, dtype=var.dtype, previous_chunks=preferred_chunk_shape
)

# Warn where requested chunks break preferred chunks, provided that the variable
# contains data.
if var.size:
Expand Down Expand Up @@ -344,7 +344,8 @@ def _maybe_chunk(
# by providing chunks as an input to tokenize.
# subtle bugs result otherwise. see GH3350
# we use str() for speed, and use the name for the final array name on the next line
token2 = tokenize(token if token else var._data, str(chunks))
mixin = uuid.uuid4()
token2 = tokenize(token if token else var._data, mixin)
name2 = f"{name_prefix}{name}-{token2}"

from_array_kwargs = utils.consolidate_dask_from_array_kwargs(
Expand Down
11 changes: 9 additions & 2 deletions xarray/namedarray/daskmanager.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,6 +1,7 @@
from __future__ import annotations

from collections.abc import Callable, Iterable, Sequence
from functools import lru_cache
from typing import TYPE_CHECKING, Any

import numpy as np
Expand All @@ -19,12 +20,20 @@

try:
from dask.array import Array as DaskArray

except ImportError:
DaskArray = np.ndarray[Any, Any]


dask_available = module_available("dask")

if dask_available:
from dask.array.core import normalize_chunks

normalize_chunks = lru_cache(normalize_chunks)
else:
normalize_chunks = None


class DaskManager(ChunkManagerEntrypoint["DaskArray"]):
array_cls: type[DaskArray]
Expand Down Expand Up @@ -52,8 +61,6 @@ def normalize_chunks(
previous_chunks: _NormalizedChunks | None = None,
) -> Any:
"""Called by open_dataset"""
from dask.array.core import normalize_chunks

return normalize_chunks(
chunks,
shape=shape,
Expand Down
Loading