Fix Zarr region transpose

[max-sixty]

This wasn't working on an unregion-ed write; I think because new_var was being lost.

[max-sixty]

CC @slevang , feedback welcome!

[max-sixty]

FYI I was getting a tb like:

File /opt/homebrew/lib/python3.9/site-packages/xarray/backends/api.py:1306, in dump_to_store(dataset, store, writer, encoder, encoding, unlimited_dims)
   1303 if encoder:
   1304     variables, attrs = encoder(variables, attrs)
-> 1306 store.store(variables, attrs, check_encoding, writer, unlimited_dims=unlimited_dims)

File /opt/homebrew/lib/python3.9/site-packages/xarray/backends/zarr.py:641, in ZarrStore.store(self, variables, attributes, check_encoding_set, writer, unlimited_dims)
    638     self.set_attributes(attributes)
    639     self.set_dimensions(variables_encoded, unlimited_dims=unlimited_dims)
--> 641 self.set_variables(
    642     variables_encoded, check_encoding_set, writer, unlimited_dims=unlimited_dims
    643 )
    644 if self._consolidate_on_close:
    645     zarr.consolidate_metadata(self.zarr_group.store)

File /opt/homebrew/lib/python3.9/site-packages/xarray/backends/zarr.py:761, in ZarrStore.set_variables(self, variables, check_encoding_set, writer, unlimited_dims)
    758     zarr_array.resize(new_shape)
    760 region = tuple(write_region[dim] for dim in dims)
--> 761 writer.add(v.data, zarr_array, region)

File /opt/homebrew/lib/python3.9/site-packages/xarray/backends/common.py:241, in ArrayWriter.add(self, source, target, region)
    239 else:
    240     if region:
--> 241         target[region] = source
    242     else:
    243         target[...] = source

File /opt/homebrew/lib/python3.9/site-packages/zarr/core.py:1495, in Array.__setitem__(self, selection, value)
   1493     self.vindex[selection] = value
   1494 elif is_pure_orthogonal_indexing(pure_selection, self.ndim):
-> 1495     self.set_orthogonal_selection(pure_selection, value, fields=fields)
   1496 else:
   1497     self.set_basic_selection(pure_selection, value, fields=fields)

File /opt/homebrew/lib/python3.9/site-packages/zarr/core.py:1684, in Array.set_orthogonal_selection(self, selection, value, fields)
   1681 # setup indexer
   1682 indexer = OrthogonalIndexer(selection, self)
-> 1684 self._set_selection(indexer, value, fields=fields)

File /opt/homebrew/lib/python3.9/site-packages/zarr/core.py:2011, in Array._set_selection(self, indexer, value, fields)
   2009     if not hasattr(value, "shape"):
   2010         value = np.asanyarray(value, like=self._meta_array)
-> 2011     check_array_shape("value", value, sel_shape)
   2013 # iterate over chunks in range
   2014 if (
   2015     not hasattr(self.chunk_store, "setitems")
   2016     or self._synchronizer is not None
   2017     or any(map(lambda x: x == 0, self.shape))
   2018 ):
   2019     # iterative approach

File /opt/homebrew/lib/python3.9/site-packages/zarr/util.py:560, in check_array_shape(param, array, shape)
    556     raise TypeError(
    557         "parameter {!r}: expected an array-like object, got {!r}".format(param, type(array))
    558     )
    559 if array.shape != shape:
--> 560     raise ValueError(
    561         "parameter {!r}: expected array with shape {!r}, got {!r}".format(
    562             param, shape, array.shape
    563         )
    564     )

ValueError: parameter 'value': expected array with shape (5, 10), got (10, 5)

I'm not really sure how the region kwarg was causing it to work — or possibly it's something different. Regardless, the test previously failed and now passes, so I think this is an improvement...

[slevang]

Good catch, thanks. I am also confused about how this was working even with region specified?

Only thing is you may want to rename the test, since this (unintentionally but correctly) suggested that the feature only works when explicitly passing region.

[max-sixty]

Only thing is you may want to rename the test, since this (unintentionally but correctly) suggested

:)

Done!