From 1933e0b9befbb696ef4829f239f4de8164cf83d8 Mon Sep 17 00:00:00 2001 From: Deepak Cherian Date: Mon, 18 Dec 2023 20:13:33 -0700 Subject: [PATCH] Small cleanup --- xarray/core/parallel.py | 18 +++++++----------- 1 file changed, 7 insertions(+), 11 deletions(-) diff --git a/xarray/core/parallel.py b/xarray/core/parallel.py index f40e0e0f3a8..7635d77d8fd 100644 --- a/xarray/core/parallel.py +++ b/xarray/core/parallel.py @@ -491,8 +491,6 @@ def subset_dataset_to_block( for dim in variable.dims } if set(variable.dims) < chunk_dims_set: - # We are including a dimension coordinate, - # minimize duplication by not copying it in the graph for every chunk. this_var_chunk_tuple = tuple( chunk_index[dim] for dim in variable.dims ) @@ -502,6 +500,8 @@ def subset_dataset_to_block( chunk_variable_task = ( f"{name}-{gname}-{dask.base.tokenize(subsetter)}", ) + this_var_chunk_tuple + # We are including a dimension coordinate, + # minimize duplication by not copying it in the graph for every chunk. if variable.ndim == 0 or chunk_variable_task not in graph: subset = variable.isel(subsetter) graph[chunk_variable_task] = ( @@ -539,7 +539,6 @@ def subset_dataset_to_block( } expected["data_vars"] = set(template.data_vars.keys()) # type: ignore[assignment] expected["coords"] = set(template.coords.keys()) # type: ignore[assignment] - # Minimize duplication due to broadcasting by only including any new or modified indexes # Others can be inferred by inputs to wrapper (GH8412) expected["indexes"] = { @@ -560,14 +559,11 @@ def subset_dataset_to_block( gname_l = f"{name}-{gname}" var_key_map[name] = gname_l - key: tuple[Any, ...] = (gname_l,) - for dim in variable.dims: - if dim in chunk_index: - key += (chunk_index[dim],) - else: - # unchunked dimensions in the input have one chunk in the result - # output can have new dimensions with exactly one chunk - key += (0,) + # unchunked dimensions in the input have one chunk in the result + # output can have new dimensions with exactly one chunk + key: tuple[Any, ...] = (gname_l,) + tuple( + chunk_index[dim] if dim in chunk_index else 0 for dim in variable.dims + ) # We're adding multiple new layers to the graph: # The first new layer is the result of the computation on