Index variables loaded from dask can be computed twice

[shoyer]

as reported by @crusaderky in #1522

[stale]

In order to maintain a list of currently relevant issues, we mark issues as stale after a period of inactivity

If this issue remains relevant, please comment here or remove the stale label; otherwise it will be marked as closed automatically

[crusaderky]

Still valid in xarray 0.13

import xarray
import numpy
import dask.array

def kernel():
    print("Kernel invoked!")
    return numpy.array([100, 200])

data = dask.array.Array(name='foo', dask={('foo', 0): (kernel, )}, chunks=((2,),), dtype=float)
ds = xarray.Dataset(coords={'z': ('z', data)})

output:

Kernel invoked!
Kernel invoked!

[dcherian]

Nice example!

I tracked this down to

1. as_variable being called when extracting the index. This seems sensible.
   

   xarray/xarray/core/merge.py

   Lines 413 to 420 in 21705e6

   	def merge_data_and_coords(data, coords, compat="broadcast_equals", join="outer"):
   	"""Used in Dataset.__init__."""
   	objs = [data, coords]
   	explicit_coords = coords.keys()
   	indexes = dict(extract_indexes(coords))
   	return merge_core(
   	objs, compat, join, explicit_coords=explicit_coords, indexes=indexes
   	)

2. as_variable being called again in the expand_variable_dicts call in merge_core (should be avoided):
   

   xarray/xarray/core/merge.py

   Lines 493 to 497 in 21705e6

   	coerced = coerce_pandas_values(objs)
   	aligned = deep_align(
   	coerced, join=join, copy=False, indexes=indexes, fill_value=fill_value
   	)
   	expanded = expand_variable_dicts(aligned)

We need to replace the appropriate objects in aligned with IndexVariables from indexes if indexes is not None.

[shoyer]

We should probably merge my merging refactor before doing much more on this...

…

On Wed, Oct 2, 2019 at 7:28 AM Deepak Cherian ***@***.***> wrote: Nice example! I tracked this down to 1. as_variable being called when extracting the index. This seems sensible. https://github.com/pydata/xarray/blob/21705e61503fb49f000186c0d556e5623bd5ac82/xarray/core/merge.py#L413-L420 2. as_variable being called again in the expand_variable_dicts call in merge_core (should be avoided): https://github.com/pydata/xarray/blob/21705e61503fb49f000186c0d556e5623bd5ac82/xarray/core/merge.py#L493-L497 We need to replace the appropriate objects in aligned with IndexVariables from indexes if indexes is not None. — You are receiving this because you authored the thread. Reply to this email directly, view it on GitHub <#1533?email_source=notifications&email_token=AAJJFVQSTUPMQV6FAQGHX5LQMSVXFA5CNFSM4DYTHFMKYY3PNVWWK3TUL52HS4DFVREXG43VMVBW63LNMVXHJKTDN5WW2ZLOORPWSZGOEAE6GKI#issuecomment-537518889>, or mute the thread <https://github.com/notifications/unsubscribe-auth/AAJJFVW6FY37NVRUPXSMQU3QMSVXFANCNFSM4DYTHFMA> .

[dcherian]

Haha I thought it was already in!

[dcherian]

Adding

if not isinstance(coords, Frozen):
    coords.update(indexes)

before the merge_core call fixes this. Not sure if its right.

xarray/xarray/core/merge.py

Lines 455 to 462 in 283b4fe

	def merge_data_and_coords(data, coords, compat="broadcast_equals", join="outer"):
	"""Used in Dataset.__init__."""
	objects = [data, coords]
	explicit_coords = coords.keys()
	indexes = dict(_extract_indexes_from_coords(coords))
	return merge_core(
	objects, compat, join, explicit_coords=explicit_coords, indexes=indexes
	)