From f231393a4d84654c1556a1b3d848e756a5cd60da Mon Sep 17 00:00:00 2001 From: jendrik Date: Thu, 24 Jan 2019 22:15:20 +0100 Subject: [PATCH 01/33] Initial version of appending to zarr store --- xarray/backends/api.py | 4 +- xarray/backends/zarr.py | 119 +++++++++++++++++++++++++--------- xarray/core/dataset.py | 8 +-- xarray/tests/test_backends.py | 6 +- 4 files changed, 96 insertions(+), 41 deletions(-) diff --git a/xarray/backends/api.py b/xarray/backends/api.py index 5f88783bb2e..7e9d3c6b15f 100644 --- a/xarray/backends/api.py +++ b/xarray/backends/api.py @@ -890,7 +890,7 @@ def save_mfdataset(datasets, paths, mode='w', format=None, groups=None, def to_zarr(dataset, store=None, mode='w-', synchronizer=None, group=None, - encoding=None, compute=True, consolidated=False): + encoding=None, compute=True, consolidated=False, append_dim=None): """This function creates an appropriate datastore for writing a dataset to a zarr ztore @@ -909,7 +909,7 @@ def to_zarr(dataset, store=None, mode='w-', synchronizer=None, group=None, synchronizer=synchronizer, group=group, consolidate_on_close=consolidated) - + zstore.append_dim = append_dim writer = ArrayWriter() # TODO: figure out how to properly handle unlimited_dims dump_to_store(dataset, zstore, writer, encoding=encoding) diff --git a/xarray/backends/zarr.py b/xarray/backends/zarr.py index feefaf1735f..c629311a08c 100644 --- a/xarray/backends/zarr.py +++ b/xarray/backends/zarr.py @@ -8,7 +8,7 @@ from ..core import indexing from ..core.pycompat import OrderedDict, integer_types, iteritems from ..core.utils import FrozenOrderedDict, HiddenKeyDict -from .common import AbstractWritableDataStore, BackendArray +from .common import AbstractWritableDataStore, BackendArray, _encode_variable_name # need some special secret attributes to tell us the dimensions _DIMENSION_KEY = '_ARRAY_DIMENSIONS' @@ -313,40 +313,97 @@ def encode_variable(self, variable): def encode_attribute(self, a): return _encode_zarr_attr_value(a) - def prepare_variable(self, name, variable, check_encoding=False, - unlimited_dims=None): - - attrs = variable.attrs.copy() - dims = variable.dims - dtype = variable.dtype - shape = variable.shape - - fill_value = attrs.pop('_FillValue', None) - if variable.encoding == {'_FillValue': None} and fill_value is None: - variable.encoding = {} - - encoding = _extract_zarr_variable_encoding( - variable, raise_on_invalid=check_encoding) - - encoded_attrs = OrderedDict() - # the magic for storing the hidden dimension data - encoded_attrs[_DIMENSION_KEY] = dims - for k, v in iteritems(attrs): - encoded_attrs[k] = self.encode_attribute(v) - - zarr_array = self.ds.create(name, shape=shape, dtype=dtype, - fill_value=fill_value, **encoding) - zarr_array.attrs.put(encoded_attrs) - - return zarr_array, variable.data - - def store(self, variables, attributes, *args, **kwargs): - AbstractWritableDataStore.store(self, variables, attributes, - *args, **kwargs) + def store(self, variables, attributes, check_encoding_set=frozenset(), + writer=None, unlimited_dims=None): + """ + Top level method for putting data on this store, this method: + - encodes variables/attributes + - sets dimensions + - sets variables + + Parameters + ---------- + variables : dict-like + Dictionary of key/value (variable name / xr.Variable) pairs + attributes : dict-like + Dictionary of key/value (attribute name / attribute) pairs + check_encoding_set : list-like + List of variables that should be checked for invalid encoding + values + writer : ArrayWriter + unlimited_dims : list-like + List of dimension names that should be treated as unlimited + dimensions. + dimension on which the zarray will be appended + only needed in append mode + """ + + variables, attributes = self.encode(variables, attributes) + + self.set_attributes(attributes) + self.set_dimensions(variables, unlimited_dims=unlimited_dims) + self.set_variables(variables, check_encoding_set, writer, + unlimited_dims=unlimited_dims) def sync(self): pass + def set_variables(self, variables, check_encoding_set, writer, + unlimited_dims=None, append_dim=None): + """ + This provides a centralized method to set the variables on the data + store. + + Parameters + ---------- + variables : dict-like + Dictionary of key/value (variable name / xr.Variable) pairs + check_encoding_set : list-like + List of variables that should be checked for invalid encoding + values + writer : + unlimited_dims : list-like + List of dimension names that should be treated as unlimited + dimensions. + append_dim: str + dimension on which the zarray will be appended + only needed in append mode + """ + for vn, v in iteritems(variables): + name = _encode_variable_name(vn) + check = vn in check_encoding_set + attrs = v.attrs.copy() + dims = v.dims + dtype = v.dtype + shape = v.shape + + fill_value = attrs.pop('_FillValue', None) + if v.encoding == {'_FillValue': None} and fill_value is None: + v.encoding = {} + append = False + try: + zarr_array = self.ds[name] + append = True + except KeyError: + encoding = _extract_zarr_variable_encoding( + v, raise_on_invalid=check) + encoded_attrs = OrderedDict() + # the magic for storing the hidden dimension data + encoded_attrs[_DIMENSION_KEY] = dims + for k2, v2 in iteritems(attrs): + encoded_attrs[k2] = self.encode_attribute(v2) + + zarr_array = self.ds.create(name, shape=shape, dtype=dtype, + fill_value=fill_value, **encoding) + zarr_array.attrs.put(encoded_attrs) + zarr_array[...] = v.data + if append: + if self.append_dim is None: + raise ValueError('The dimension which is appended has to be named.') + if self.append_dim not in dims: + continue + zarr_array.append(v.data) + def close(self): if self._consolidate_on_close: import zarr diff --git a/xarray/core/dataset.py b/xarray/core/dataset.py index 3f1c038fc11..622aecd7489 100644 --- a/xarray/core/dataset.py +++ b/xarray/core/dataset.py @@ -1246,7 +1246,7 @@ def to_netcdf(self, path=None, mode='w', format=None, group=None, compute=compute) def to_zarr(self, store=None, mode='w-', synchronizer=None, group=None, - encoding=None, compute=True, consolidated=False): + encoding=None, compute=True, consolidated=False, append_dim=None): """Write dataset contents to a zarr group. .. note:: Experimental @@ -1281,14 +1281,14 @@ def to_zarr(self, store=None, mode='w-', synchronizer=None, group=None, """ if encoding is None: encoding = {} - if mode not in ['w', 'w-']: - # TODO: figure out how to handle 'r+' and 'a' + if mode not in ['w', 'w-', 'a']: + # TODO: figure out how to handle 'r+' raise ValueError("The only supported options for mode are 'w' " "and 'w-'.") from ..backends.api import to_zarr return to_zarr(self, store=store, mode=mode, synchronizer=synchronizer, group=group, encoding=encoding, compute=compute, - consolidated=consolidated) + consolidated=consolidated, append_dim=append_dim) def __unicode__(self): return formatting.dataset_repr(self) diff --git a/xarray/tests/test_backends.py b/xarray/tests/test_backends.py index d3c8599b21b..f1b0c3366b0 100644 --- a/xarray/tests/test_backends.py +++ b/xarray/tests/test_backends.py @@ -1485,10 +1485,8 @@ def test_write_persistence_modes(self): self.save(original, store, mode='w-') # check that we can't use other persistence modes - # TODO: reconsider whether other persistence modes should be supported - with pytest.raises(ValueError): - with self.roundtrip(original, save_kwargs={'mode': 'a'}) as actual: - pass + with self.roundtrip(original, save_kwargs={'mode': 'a'}) as actual: + assert_identical(original, actual) def test_compressor_encoding(self): original = create_test_data() From f14f3b7e583cb9d925b29c4f23ec21306d3544f1 Mon Sep 17 00:00:00 2001 From: jendrik Date: Thu, 24 Jan 2019 22:40:33 +0100 Subject: [PATCH 02/33] Added docs --- doc/whats-new.rst | 1 + 1 file changed, 1 insertion(+) diff --git a/doc/whats-new.rst b/doc/whats-new.rst index c408306ffdb..ed48a5927a7 100644 --- a/doc/whats-new.rst +++ b/doc/whats-new.rst @@ -43,6 +43,7 @@ Enhancements report showing what exactly differs between the two objects (dimensions / coordinates / variables / attributes) (:issue:`1507`). By `Benoit Bovy `_. +- Added append capability to the zarr store. Bug fixes ~~~~~~~~~ From 928440d79b6025a29af2fec339f62aa94c7fdb6d Mon Sep 17 00:00:00 2001 From: jendrik Date: Thu, 24 Jan 2019 22:45:12 +0100 Subject: [PATCH 03/33] Resolve PEP8 incompliances --- xarray/backends/zarr.py | 6 ++++-- xarray/core/dataset.py | 10 +++++++--- 2 files changed, 11 insertions(+), 5 deletions(-) diff --git a/xarray/backends/zarr.py b/xarray/backends/zarr.py index c629311a08c..2586a8cec60 100644 --- a/xarray/backends/zarr.py +++ b/xarray/backends/zarr.py @@ -8,7 +8,8 @@ from ..core import indexing from ..core.pycompat import OrderedDict, integer_types, iteritems from ..core.utils import FrozenOrderedDict, HiddenKeyDict -from .common import AbstractWritableDataStore, BackendArray, _encode_variable_name +from .common import AbstractWritableDataStore, BackendArray, \ + _encode_variable_name # need some special secret attributes to tell us the dimensions _DIMENSION_KEY = '_ARRAY_DIMENSIONS' @@ -399,7 +400,8 @@ def set_variables(self, variables, check_encoding_set, writer, zarr_array[...] = v.data if append: if self.append_dim is None: - raise ValueError('The dimension which is appended has to be named.') + raise ValueError('The dimension on which the data is \ + appended has to be named.') if self.append_dim not in dims: continue zarr_array.append(v.data) diff --git a/xarray/core/dataset.py b/xarray/core/dataset.py index 622aecd7489..a9ba2962324 100644 --- a/xarray/core/dataset.py +++ b/xarray/core/dataset.py @@ -1246,7 +1246,8 @@ def to_netcdf(self, path=None, mode='w', format=None, group=None, compute=compute) def to_zarr(self, store=None, mode='w-', synchronizer=None, group=None, - encoding=None, compute=True, consolidated=False, append_dim=None): + encoding=None, compute=True, consolidated=False, + append_dim=None): """Write dataset contents to a zarr group. .. note:: Experimental @@ -1257,9 +1258,10 @@ def to_zarr(self, store=None, mode='w-', synchronizer=None, group=None, ---------- store : MutableMapping or str, optional Store or path to directory in file system. - mode : {'w', 'w-'} + mode : {'w', 'w-', 'a'} Persistence mode: 'w' means create (overwrite if exists); - 'w-' means create (fail if exists). + 'w-' means create (fail if exists); + 'a' means append (create if does not exist). synchronizer : object, optional Array synchronizer group : str, obtional @@ -1274,6 +1276,8 @@ def to_zarr(self, store=None, mode='w-', synchronizer=None, group=None, consolidated: bool, optional If True, apply zarr's `consolidate_metadata` function to the store after writing. + append_dim: str + If mode='a' hand the dimension on which the data will be appended References ---------- From 442e938c2c5dcc0f192f0db2348cd679d07c16cb Mon Sep 17 00:00:00 2001 From: jendrik Date: Sat, 26 Jan 2019 13:25:36 +0100 Subject: [PATCH 04/33] Added write and append test for mode 'a' --- xarray/backends/zarr.py | 3 ++- xarray/tests/test_backends.py | 12 ++++++++++-- xarray/tests/test_dataset.py | 37 +++++++++++++++++++++++++++++++++++ 3 files changed, 49 insertions(+), 3 deletions(-) diff --git a/xarray/backends/zarr.py b/xarray/backends/zarr.py index 2586a8cec60..353732101b8 100644 --- a/xarray/backends/zarr.py +++ b/xarray/backends/zarr.py @@ -404,7 +404,8 @@ def set_variables(self, variables, check_encoding_set, writer, appended has to be named.') if self.append_dim not in dims: continue - zarr_array.append(v.data) + axis = dims.index(self.append_dim) + zarr_array.append(v.data, axis=axis) def close(self): if self._consolidate_on_close: diff --git a/xarray/tests/test_backends.py b/xarray/tests/test_backends.py index f1b0c3366b0..81750e43047 100644 --- a/xarray/tests/test_backends.py +++ b/xarray/tests/test_backends.py @@ -37,7 +37,7 @@ requires_pathlib, requires_pseudonetcdf, requires_pydap, requires_pynio, requires_rasterio, requires_scipy, requires_scipy_or_netCDF4, requires_zarr) -from .test_dataset import create_test_data +from .test_dataset import create_test_data, create_append_test_data try: import netCDF4 as nc4 @@ -1484,10 +1484,18 @@ def test_write_persistence_modes(self): with pytest.raises(ValueError): self.save(original, store, mode='w-') - # check that we can't use other persistence modes + # check append mode for normal write with self.roundtrip(original, save_kwargs={'mode': 'a'}) as actual: assert_identical(original, actual) + # check append mode for append write + obj, obj2 = create_append_test_data() + with self.create_zarr_target() as store_target: + obj.to_zarr(store_target, mode='w') + obj2.to_zarr(store_target, mode='a', append_dim='dim1') + original = xr.concat([obj, obj2], dim='dim1') + assert_identical(original, xr.open_zarr(store_target)) + def test_compressor_encoding(self): original = create_test_data() # specify a custom compressor diff --git a/xarray/tests/test_dataset.py b/xarray/tests/test_dataset.py index 376c22104c5..1f10a700a90 100644 --- a/xarray/tests/test_dataset.py +++ b/xarray/tests/test_dataset.py @@ -54,6 +54,43 @@ def create_test_data(seed=None): return obj +def create_append_test_data(seed=None): + rs = np.random.RandomState(seed) + _vars = {'var1': ['dim1', 'dim2'], + 'var2': ['dim1', 'dim2'], + 'var3': ['dim3', 'dim1']} + _dims = {'dim1': 8, 'dim2': 9, 'dim3': 10} + + obj = Dataset() + obj['time'] = ('time', pd.date_range('2000-01-01', periods=20)) + obj['dim2'] = ('dim2', 0.5 * np.arange(_dims['dim2'])) + obj['dim3'] = ('dim3', list('abcdefghij')) + for v, dims in sorted(_vars.items()): + data = rs.normal(size=tuple(_dims[d] for d in dims)) + obj[v] = (dims, data, {'foo': 'variable'}) + obj.coords['numbers'] = ('dim3', np.array([0, 1, 2, 0, 0, 1, 1, 2, 2, 3], + dtype='int64')) + obj.encoding = {'foo': 'bar'} + assert all(objp.data.flags.writeable for objp in obj.variables.values()) + _vars = {'var1': ['dim1', 'dim2'], + 'var2': ['dim1', 'dim2'], + 'var3': ['dim3', 'dim1']} + _dims = {'dim1': 8, 'dim2': 9, 'dim3': 10} + + obj2 = Dataset() + obj2['time'] = ('time', pd.date_range('2000-01-01', periods=20)) + obj2['dim2'] = ('dim2', 0.5 * np.arange(_dims['dim2'])) + obj2['dim3'] = ('dim3', list('abcdefghij')) + for v, dims in sorted(_vars.items()): + data = rs.normal(size=tuple(_dims[d] for d in dims)) + obj2[v] = (dims, data, {'foo': 'variable'}) + obj2.coords['numbers'] = ('dim3', np.array([0, 1, 2, 0, 0, 1, 1, 2, 2, 3], + dtype='int64')) + obj2.encoding = {'foo': 'bar'} + assert all(objp.data.flags.writeable for objp in obj2.variables.values()) + return obj, obj2 + + def create_test_multiindex(): mindex = pd.MultiIndex.from_product([['a', 'b'], [1, 2]], names=('level_1', 'level_2')) From 390a792d18feb3ec1d58245ec241470be665d3c3 Mon Sep 17 00:00:00 2001 From: jendrik Date: Tue, 29 Jan 2019 21:32:17 +0100 Subject: [PATCH 05/33] Merged repaired master --- xarray/backends/zarr.py | 1 + 1 file changed, 1 insertion(+) diff --git a/xarray/backends/zarr.py b/xarray/backends/zarr.py index a528aa3a05b..4d286a24603 100644 --- a/xarray/backends/zarr.py +++ b/xarray/backends/zarr.py @@ -403,6 +403,7 @@ def set_variables(self, variables, check_encoding_set, writer, appended has to be named.') if self.append_dim not in dims: continue + print(name, v.data.shape) axis = dims.index(self.append_dim) zarr_array.append(v.data, axis=axis) From da9a962d7fdff259c68b40430e3e45a3851db944 Mon Sep 17 00:00:00 2001 From: jendrik Date: Wed, 30 Jan 2019 00:17:00 +0100 Subject: [PATCH 06/33] Resolved pep8 issue --- xarray/backends/zarr.py | 1 - xarray/tests/test_dataset.py | 2 +- 2 files changed, 1 insertion(+), 2 deletions(-) diff --git a/xarray/backends/zarr.py b/xarray/backends/zarr.py index 4d286a24603..a528aa3a05b 100644 --- a/xarray/backends/zarr.py +++ b/xarray/backends/zarr.py @@ -403,7 +403,6 @@ def set_variables(self, variables, check_encoding_set, writer, appended has to be named.') if self.append_dim not in dims: continue - print(name, v.data.shape) axis = dims.index(self.append_dim) zarr_array.append(v.data, axis=axis) diff --git a/xarray/tests/test_dataset.py b/xarray/tests/test_dataset.py index aba4e56002e..49627d117e2 100644 --- a/xarray/tests/test_dataset.py +++ b/xarray/tests/test_dataset.py @@ -83,7 +83,7 @@ def create_append_test_data(seed=None): data = rs.normal(size=tuple(_dims[d] for d in dims)) obj2[v] = (dims, data, {'foo': 'variable'}) obj2.coords['numbers'] = ('dim3', np.array([0, 1, 2, 0, 0, 1, 1, 2, 2, 3], - dtype='int64')) + dtype='int64')) obj2.encoding = {'foo': 'bar'} assert all(objp.data.flags.writeable for objp in obj2.variables.values()) return obj, obj2 From 6756b8f9b1396fcecf4cfe5f47ba071dc8df977d Mon Sep 17 00:00:00 2001 From: David Brochart Date: Wed, 3 Apr 2019 14:59:37 +0200 Subject: [PATCH 07/33] Put target store encoding in appended variable --- xarray/backends/zarr.py | 47 ++++++++++++++++++++++++++--------------- 1 file changed, 30 insertions(+), 17 deletions(-) diff --git a/xarray/backends/zarr.py b/xarray/backends/zarr.py index a528aa3a05b..ba49c4ab44e 100644 --- a/xarray/backends/zarr.py +++ b/xarray/backends/zarr.py @@ -257,6 +257,7 @@ def __init__(self, zarr_group, consolidate_on_close=False): self._synchronizer = self.ds.synchronizer self._group = self.ds.path self._consolidate_on_close = consolidate_on_close + self.append_dim = None def open_store_variable(self, name, zarr_array): data = indexing.LazilyOuterIndexedArray(ZarrArrayWrapper(name, self)) @@ -338,7 +339,24 @@ def store(self, variables, attributes, check_encoding_set=frozenset(), only needed in append mode """ - variables, attributes = self.encode(variables, attributes) + variables_to_encode = OrderedDict() + variables_encoded = OrderedDict() + ds = None + for vn, v in variables.items(): + name = _encode_variable_name(vn) + if name in self.ds: + # existing variable, get its encoding + # and apply it to appended variable + variables_encoded[vn] = v + if ds is None: + ds = open_zarr(self.ds.store, auto_chunk=False) + v.encoding = ds[vn].encoding + else: + # new variable, encode it + variables_to_encode[vn] = v + variables_encoded, _ = self.encode(variables_encoded, {}) + variables, attributes = self.encode(variables_to_encode, attributes) + variables.update(variables_encoded) self.set_attributes(attributes) self.set_dimensions(variables, unlimited_dims=unlimited_dims) @@ -349,7 +367,7 @@ def sync(self): pass def set_variables(self, variables, check_encoding_set, writer, - unlimited_dims=None, append_dim=None): + unlimited_dims=None): """ This provides a centralized method to set the variables on the data store. @@ -365,9 +383,6 @@ def set_variables(self, variables, check_encoding_set, writer, unlimited_dims : list-like List of dimension names that should be treated as unlimited dimensions. - append_dim: str - dimension on which the zarray will be appended - only needed in append mode """ for vn, v in variables.items(): name = _encode_variable_name(vn) @@ -380,11 +395,17 @@ def set_variables(self, variables, check_encoding_set, writer, fill_value = attrs.pop('_FillValue', None) if v.encoding == {'_FillValue': None} and fill_value is None: v.encoding = {} - append = False - try: + if name in self.ds: + # append to existing variable zarr_array = self.ds[name] - append = True - except KeyError: + if self.append_dim is None: + raise ValueError('dimension being appended is unknown; ' + 'did you forget to call to_zarr with append_dim argument?') + if self.append_dim in dims: + axis = dims.index(self.append_dim) + zarr_array.append(v.data, axis=axis) + else: + # new variable encoding = _extract_zarr_variable_encoding( v, raise_on_invalid=check) encoded_attrs = OrderedDict() @@ -397,14 +418,6 @@ def set_variables(self, variables, check_encoding_set, writer, fill_value=fill_value, **encoding) zarr_array.attrs.put(encoded_attrs) zarr_array[...] = v.data - if append: - if self.append_dim is None: - raise ValueError('The dimension on which the data is \ - appended has to be named.') - if self.append_dim not in dims: - continue - axis = dims.index(self.append_dim) - zarr_array.append(v.data, axis=axis) def close(self): if self._consolidate_on_close: From a750a9239e1b03991b153ec48be539d5dc6852f0 Mon Sep 17 00:00:00 2001 From: David Brochart Date: Thu, 4 Apr 2019 16:56:03 +0200 Subject: [PATCH 08/33] Rewrite test with appending along time dimension --- xarray/backends/zarr.py | 47 ++++++++++++++++-------------- xarray/tests/test_backends.py | 8 +++--- xarray/tests/test_dataset.py | 54 ++++++++++++----------------------- 3 files changed, 48 insertions(+), 61 deletions(-) diff --git a/xarray/backends/zarr.py b/xarray/backends/zarr.py index ba49c4ab44e..7ba1bb1a526 100644 --- a/xarray/backends/zarr.py +++ b/xarray/backends/zarr.py @@ -339,28 +339,29 @@ def store(self, variables, attributes, check_encoding_set=frozenset(), only needed in append mode """ - variables_to_encode = OrderedDict() - variables_encoded = OrderedDict() - ds = None - for vn, v in variables.items(): - name = _encode_variable_name(vn) - if name in self.ds: - # existing variable, get its encoding - # and apply it to appended variable - variables_encoded[vn] = v - if ds is None: - ds = open_zarr(self.ds.store, auto_chunk=False) - v.encoding = ds[vn].encoding - else: - # new variable, encode it - variables_to_encode[vn] = v - variables_encoded, _ = self.encode(variables_encoded, {}) - variables, attributes = self.encode(variables_to_encode, attributes) - variables.update(variables_encoded) + existing_variables = set([vn for vn in variables + if _encode_variable_name(vn) in self.ds]) + new_variables = set(variables).difference(existing_variables) + variables_without_encoding = OrderedDict([(vn, variables[vn]) + for vn in new_variables]) + variables_encoded, attributes = self.encode( + variables_without_encoding, attributes) + + if len(existing_variables) > 0: + # there are variables to append + # their encoding must be the same as in the store + ds = open_zarr(self.ds.store, auto_chunk=False) + variables_with_encoding = OrderedDict() + for vn in existing_variables: + variables_with_encoding[vn] = variables[vn] + variables_with_encoding[vn].encoding = ds[vn].encoding + variables_with_encoding, _ = self.encode(variables_with_encoding, + {}) + variables_encoded.update(variables_with_encoding) self.set_attributes(attributes) - self.set_dimensions(variables, unlimited_dims=unlimited_dims) - self.set_variables(variables, check_encoding_set, writer, + self.set_dimensions(variables_encoded, unlimited_dims=unlimited_dims) + self.set_variables(variables_encoded, check_encoding_set, writer, unlimited_dims=unlimited_dims) def sync(self): @@ -399,8 +400,10 @@ def set_variables(self, variables, check_encoding_set, writer, # append to existing variable zarr_array = self.ds[name] if self.append_dim is None: - raise ValueError('dimension being appended is unknown; ' - 'did you forget to call to_zarr with append_dim argument?') + raise ValueError( + 'dimension being appended is unknown; ' + 'did you forget to call to_zarr with append_dim ' + 'argument?') if self.append_dim in dims: axis = dims.index(self.append_dim) zarr_array.append(v.data, axis=axis) diff --git a/xarray/tests/test_backends.py b/xarray/tests/test_backends.py index 64f3cf29f93..ecff4946be5 100644 --- a/xarray/tests/test_backends.py +++ b/xarray/tests/test_backends.py @@ -1524,11 +1524,11 @@ def test_write_persistence_modes(self): assert_identical(original, actual) # check append mode for append write - obj, obj2 = create_append_test_data() + obj1, obj2 = create_append_test_data() with self.create_zarr_target() as store_target: - obj.to_zarr(store_target, mode='w') - obj2.to_zarr(store_target, mode='a', append_dim='dim1') - original = xr.concat([obj, obj2], dim='dim1') + obj1.to_zarr(store_target, mode='w') + obj2.to_zarr(store_target, mode='a', append_dim='time') + original = xr.concat([obj1, obj2], dim='time') assert_identical(original, xr.open_zarr(store_target)) def test_compressor_encoding(self): diff --git a/xarray/tests/test_dataset.py b/xarray/tests/test_dataset.py index 7fead6a800d..13ac5cd9c81 100644 --- a/xarray/tests/test_dataset.py +++ b/xarray/tests/test_dataset.py @@ -52,41 +52,25 @@ def create_test_data(seed=None): return obj -def create_append_test_data(seed=None): - rs = np.random.RandomState(seed) - _vars = {'var1': ['dim1', 'dim2'], - 'var2': ['dim1', 'dim2'], - 'var3': ['dim3', 'dim1']} - _dims = {'dim1': 8, 'dim2': 9, 'dim3': 10} - - obj = Dataset() - obj['time'] = ('time', pd.date_range('2000-01-01', periods=20)) - obj['dim2'] = ('dim2', 0.5 * np.arange(_dims['dim2'])) - obj['dim3'] = ('dim3', list('abcdefghij')) - for v, dims in sorted(_vars.items()): - data = rs.normal(size=tuple(_dims[d] for d in dims)) - obj[v] = (dims, data, {'foo': 'variable'}) - obj.coords['numbers'] = ('dim3', np.array([0, 1, 2, 0, 0, 1, 1, 2, 2, 3], - dtype='int64')) - obj.encoding = {'foo': 'bar'} - assert all(objp.data.flags.writeable for objp in obj.variables.values()) - _vars = {'var1': ['dim1', 'dim2'], - 'var2': ['dim1', 'dim2'], - 'var3': ['dim3', 'dim1']} - _dims = {'dim1': 8, 'dim2': 9, 'dim3': 10} - - obj2 = Dataset() - obj2['time'] = ('time', pd.date_range('2000-01-01', periods=20)) - obj2['dim2'] = ('dim2', 0.5 * np.arange(_dims['dim2'])) - obj2['dim3'] = ('dim3', list('abcdefghij')) - for v, dims in sorted(_vars.items()): - data = rs.normal(size=tuple(_dims[d] for d in dims)) - obj2[v] = (dims, data, {'foo': 'variable'}) - obj2.coords['numbers'] = ('dim3', np.array([0, 1, 2, 0, 0, 1, 1, 2, 2, 3], - dtype='int64')) - obj2.encoding = {'foo': 'bar'} - assert all(objp.data.flags.writeable for objp in obj2.variables.values()) - return obj, obj2 +def create_append_test_data(): + lat = [2, 1, 0] + lon = [0, 1, 2] + + nt1 = 3 + time1 = pd.date_range('2000-01-01', periods=nt1) + da1 = xr.DataArray(np.arange(3 * 3 * nt1).reshape(3, 3, nt1), coords=[lat, + lon, time1], dims=['lat', 'lon', 'time']) + ds1 = da1.to_dataset(name='da') + + nt2 = 2 + time2 = pd.date_range('2000-02-01', periods=nt2) + da2 = xr.DataArray((np.arange(3 * 3 * nt2) + 100).reshape(3, 3, nt2), + coords=[lat, lon, time2], dims=['lat', 'lon', 'time']) + ds2 = da2.to_dataset(name='da') + + assert all(objp.data.flags.writeable for objp in ds1.variables.values()) + assert all(objp.data.flags.writeable for objp in ds2.variables.values()) + return ds1, ds2 def create_test_multiindex(): From 295084b23ac0a32eabb030991bea4ea5cc299bba Mon Sep 17 00:00:00 2001 From: David Brochart Date: Mon, 22 Apr 2019 08:35:00 +0200 Subject: [PATCH 09/33] Add chunk_size parameter for rechunking appended coordinate --- doc/io.rst | 7 +++++++ xarray/backends/api.py | 4 +++- xarray/backends/zarr.py | 24 +++++++++++++++++++++--- xarray/core/dataset.py | 11 +++++++---- 4 files changed, 38 insertions(+), 8 deletions(-) diff --git a/doc/io.rst b/doc/io.rst index 51c747189da..be6a5e2d3e7 100644 --- a/doc/io.rst +++ b/doc/io.rst @@ -597,6 +597,13 @@ store is already present at that path, an error will be raised, preventing it from being overwritten. To override this behavior and overwrite an existing store, add ``mode='w'`` when invoking ``to_zarr``. +It is also possible to append to an existing store. For that, add ``mode='a'`` +and set ``append_dim`` to the name of the dimension along which to append. +Optionally, you can also pass ``chunk_dim`` to set the chunk size of the +appended dimension's coordinate. If you don't, this coordinate will keep the +same chunk size as before appending, which could be too small if this is e.g. a +1-D array (typically, the time coordinate). + To read back a zarr dataset that has been created this way, we use the :py:func:`~xarray.open_zarr` method: diff --git a/xarray/backends/api.py b/xarray/backends/api.py index 59326c0389a..0be7958398f 100644 --- a/xarray/backends/api.py +++ b/xarray/backends/api.py @@ -951,7 +951,8 @@ def save_mfdataset(datasets, paths, mode='w', format=None, groups=None, def to_zarr(dataset, store=None, mode='w-', synchronizer=None, group=None, - encoding=None, compute=True, consolidated=False, append_dim=None): + encoding=None, compute=True, consolidated=False, append_dim=None, + chunk_dim=None): """This function creates an appropriate datastore for writing a dataset to a zarr ztore @@ -971,6 +972,7 @@ def to_zarr(dataset, store=None, mode='w-', synchronizer=None, group=None, group=group, consolidate_on_close=consolidated) zstore.append_dim = append_dim + zstore.chunk_dim = chunk_dim writer = ArrayWriter() # TODO: figure out how to properly handle unlimited_dims dump_to_store(dataset, zstore, writer, encoding=encoding) diff --git a/xarray/backends/zarr.py b/xarray/backends/zarr.py index 7ba1bb1a526..ae228e3099d 100644 --- a/xarray/backends/zarr.py +++ b/xarray/backends/zarr.py @@ -258,6 +258,7 @@ def __init__(self, zarr_group, consolidate_on_close=False): self._group = self.ds.path self._consolidate_on_close = consolidate_on_close self.append_dim = None + self.chunk_dim = None def open_store_variable(self, name, zarr_array): data = indexing.LazilyOuterIndexedArray(ZarrArrayWrapper(name, self)) @@ -341,7 +342,7 @@ def store(self, variables, attributes, check_encoding_set=frozenset(), existing_variables = set([vn for vn in variables if _encode_variable_name(vn) in self.ds]) - new_variables = set(variables).difference(existing_variables) + new_variables = set(variables) - existing_variables variables_without_encoding = OrderedDict([(vn, variables[vn]) for vn in new_variables]) variables_encoded, attributes = self.encode( @@ -405,8 +406,25 @@ def set_variables(self, variables, check_encoding_set, writer, 'did you forget to call to_zarr with append_dim ' 'argument?') if self.append_dim in dims: - axis = dims.index(self.append_dim) - zarr_array.append(v.data, axis=axis) + if len(dims) == 1 and self.chunk_dim is not None: + # this is the coordinate along which to append + prev_coord = zarr_array[:] + attrs = zarr_array.attrs.asdict() + dtype = zarr_array.dtype + fill_value = zarr_array.fill_value + new_coord = np.hstack((prev_coord, v.data)) + self.ds.create_dataset(name, shape=new_coord.shape, + chunks=self.chunk_dim, + fill_value=fill_value, + dtype=dtype, overwrite=True) + zarr_array = self.ds[name] + zarr_array.attrs.put(attrs) + zarr_array[:] = new_coord + else: + # this is the DataArray that has append_dim as a + # dimension + axis = dims.index(self.append_dim) + zarr_array.append(v.data, axis=axis) else: # new variable encoding = _extract_zarr_variable_encoding( diff --git a/xarray/core/dataset.py b/xarray/core/dataset.py index c1d47fd2913..de1ab705f28 100644 --- a/xarray/core/dataset.py +++ b/xarray/core/dataset.py @@ -1325,7 +1325,7 @@ def to_netcdf(self, path=None, mode='w', format=None, group=None, def to_zarr(self, store=None, mode='w-', synchronizer=None, group=None, encoding=None, compute=True, consolidated=False, - append_dim=None): + append_dim=None, chunk_dim=None): """Write dataset contents to a zarr group. .. note:: Experimental @@ -1354,8 +1354,10 @@ def to_zarr(self, store=None, mode='w-', synchronizer=None, group=None, consolidated: bool, optional If True, apply zarr's `consolidate_metadata` function to the store after writing. - append_dim: str - If mode='a' hand the dimension on which the data will be appended + append_dim: str, optional + If mode='a', the dimension on which the data will be appended. + chunk_dim: int, optional + If mode='a', the chunk size of the `append_dim` coordinate. References ---------- @@ -1370,7 +1372,8 @@ def to_zarr(self, store=None, mode='w-', synchronizer=None, group=None, from ..backends.api import to_zarr return to_zarr(self, store=store, mode=mode, synchronizer=synchronizer, group=group, encoding=encoding, compute=compute, - consolidated=consolidated, append_dim=append_dim) + consolidated=consolidated, append_dim=append_dim, + chunk_dim=chunk_dim) def __repr__(self): return formatting.dataset_repr(self) From 519b39860c6c74af111f2fa830cfb48965e6daec Mon Sep 17 00:00:00 2001 From: David Brochart Date: Tue, 21 May 2019 12:40:42 +0200 Subject: [PATCH 10/33] Add chunk_dim test --- xarray/tests/test_backends.py | 10 ++++++++++ 1 file changed, 10 insertions(+) diff --git a/xarray/tests/test_backends.py b/xarray/tests/test_backends.py index f2ae4379d8f..9231920416b 100644 --- a/xarray/tests/test_backends.py +++ b/xarray/tests/test_backends.py @@ -1612,6 +1612,16 @@ def test_write_persistence_modes(self): obj2.to_zarr(store_target, mode='a', append_dim='time') original = xr.concat([obj1, obj2], dim='time') assert_identical(original, xr.open_zarr(store_target)) + # check chunk size of the append_dim coordinate when chunk_dim is set + with self.create_zarr_target() as store_target: + original = xr.concat([obj1, obj2], dim='time') + obj1.to_zarr(store_target, mode='w') + obj2.to_zarr(store_target, mode='a', append_dim='time', + chunk_dim=original['time'].shape) + import zarr + store = zarr.open(store_target) + assert_identical(original, xr.open_zarr(store_target)) + assert original['time'].shape == store['time'].chunks def test_compressor_encoding(self): original = create_test_data() From 3adfd49d07d36131f180593d01a4f245b52395ad Mon Sep 17 00:00:00 2001 From: jendrik Date: Mon, 17 Jun 2019 20:54:01 +0200 Subject: [PATCH 11/33] Add type check and tests for it. In append mode storing any datatype apart from number subtypes and fixed size strings raises and error. --- xarray/backends/api.py | 30 ++++++++++++++++++++++++++++++ xarray/tests/test_backends.py | 10 +++++++++- xarray/tests/test_dataset.py | 3 +++ 3 files changed, 42 insertions(+), 1 deletion(-) diff --git a/xarray/backends/api.py b/xarray/backends/api.py index c2106b59a7a..99f61890963 100644 --- a/xarray/backends/api.py +++ b/xarray/backends/api.py @@ -4,16 +4,21 @@ from io import BytesIO from numbers import Number from pathlib import Path +import re import numpy as np +import pandas as pd from .. import Dataset, backends, conventions from ..core import indexing from ..core.combine import ( _CONCAT_DIM_DEFAULT, _auto_combine, _infer_concat_order_from_positions) from ..core.utils import close_on_error, is_grib_path, is_remote_uri +from ..core.variable import Variable from .common import ArrayWriter from .locks import _get_scheduler +from ..coding.variables import safe_setitem, unpack_for_encoding +from ..coding.strings import encode_string_array DATAARRAY_NAME = '__xarray_dataarray_name__' DATAARRAY_VARIABLE = '__xarray_dataarray_variable__' @@ -1003,6 +1008,28 @@ def save_mfdataset(datasets, paths, mode='w', format=None, groups=None, for w, s in zip(writes, stores)]) +def encode_utf8(var, string_max_length): + dims, data, attrs, encoding = unpack_for_encoding(var) + missing = pd.isnull(data) + data[missing] = "" + data = encode_string_array(data, 'utf-8') + data = data.astype(np.dtype(f"S{string_max_length*2}")) + return Variable(dims, data, attrs, encoding) + + +def _validate_datatypes_for_zarr_append(dataset): + """DataArray.name and Dataset keys must be a string or None""" + def check_dtype(var): + if (not np.issubdtype(var.dtype, np.number) + and not np.issubdtype(var.dtype, np.string_)): + # and not re.match('^bytes[1-9]+$', var.dtype.name)): + raise ValueError('Invalid dtype for DataVariable: {} ' + 'dtype must be a subtype of number or ' + 'a fixed sized string'.format(var)) + for k in dataset.data_vars.values(): + check_dtype(k) + + def to_zarr(dataset, store=None, mode='w-', synchronizer=None, group=None, encoding=None, compute=True, consolidated=False, append_dim=None, chunk_dim=None): @@ -1020,6 +1047,9 @@ def to_zarr(dataset, store=None, mode='w-', synchronizer=None, group=None, _validate_dataset_names(dataset) _validate_attrs(dataset) + if mode == "a": + _validate_datatypes_for_zarr_append(dataset) + zstore = backends.ZarrStore.open_group(store=store, mode=mode, synchronizer=synchronizer, group=group, diff --git a/xarray/tests/test_backends.py b/xarray/tests/test_backends.py index 9231920416b..8e95741396e 100644 --- a/xarray/tests/test_backends.py +++ b/xarray/tests/test_backends.py @@ -20,6 +20,7 @@ from xarray import ( DataArray, Dataset, backends, open_dataarray, open_dataset, open_mfdataset, save_mfdataset, load_dataset, load_dataarray) +from xarray.backends.api import encode_utf8 from xarray.backends.common import robust_getitem from xarray.backends.netCDF4_ import _extract_nc4_variable_encoding from xarray.backends.pydap_ import PydapDataStore @@ -1607,6 +1608,13 @@ def test_write_persistence_modes(self): # check append mode for append write obj1, obj2 = create_append_test_data() + + with pytest.raises(ValueError): + with self.create_zarr_target() as store_target: + obj1.to_zarr(store_target, mode='a') + + obj1['string_var'] = encode_utf8(obj1.string_var, 2) + obj2['string_var'] = encode_utf8(obj2.string_var, 2) with self.create_zarr_target() as store_target: obj1.to_zarr(store_target, mode='w') obj2.to_zarr(store_target, mode='a', append_dim='time') @@ -1614,10 +1622,10 @@ def test_write_persistence_modes(self): assert_identical(original, xr.open_zarr(store_target)) # check chunk size of the append_dim coordinate when chunk_dim is set with self.create_zarr_target() as store_target: - original = xr.concat([obj1, obj2], dim='time') obj1.to_zarr(store_target, mode='w') obj2.to_zarr(store_target, mode='a', append_dim='time', chunk_dim=original['time'].shape) + original = xr.concat([obj1, obj2], dim='time') import zarr store = zarr.open(store_target) assert_identical(original, xr.open_zarr(store_target)) diff --git a/xarray/tests/test_dataset.py b/xarray/tests/test_dataset.py index 37a860a53aa..a86208e0058 100644 --- a/xarray/tests/test_dataset.py +++ b/xarray/tests/test_dataset.py @@ -55,18 +55,21 @@ def create_test_data(seed=None): def create_append_test_data(): lat = [2, 1, 0] lon = [0, 1, 2] + string_var = ["ae", "bc", "df"] nt1 = 3 time1 = pd.date_range('2000-01-01', periods=nt1) da1 = xr.DataArray(np.arange(3 * 3 * nt1).reshape(3, 3, nt1), coords=[lat, lon, time1], dims=['lat', 'lon', 'time']) ds1 = da1.to_dataset(name='da') + ds1['string_var'] = xr.DataArray(string_var, coords=[time1], dims=['time']) nt2 = 2 time2 = pd.date_range('2000-02-01', periods=nt2) da2 = xr.DataArray((np.arange(3 * 3 * nt2) + 100).reshape(3, 3, nt2), coords=[lat, lon, time2], dims=['lat', 'lon', 'time']) ds2 = da2.to_dataset(name='da') + ds2['string_var'] = xr.DataArray(string_var[:nt2], coords=[time2], dims=['time']) assert all(objp.data.flags.writeable for objp in ds1.variables.values()) assert all(objp.data.flags.writeable for objp in ds2.variables.values()) From 608813be5b8faa5b59ab048e75f08b8fe465506a Mon Sep 17 00:00:00 2001 From: jendrik Date: Mon, 17 Jun 2019 21:09:06 +0200 Subject: [PATCH 12/33] Add documentation --- doc/io.rst | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/doc/io.rst b/doc/io.rst index 36b4375114d..e6346801b49 100644 --- a/doc/io.rst +++ b/doc/io.rst @@ -606,6 +606,11 @@ store, add ``mode='w'`` when invoking ``to_zarr``. It is also possible to append to an existing store. For that, add ``mode='a'`` and set ``append_dim`` to the name of the dimension along which to append. +It is necessary that the data only contains types which are either a subclass +of ``np.number`` or ``np.string_``. It as well needs to be taken into account, +that the size of the ``np.string_`` in the first chunk sets the maximum +string size for all following ones. To encode the data consider using +:py:func:`~xarray.core.api.encode_utf8`. Optionally, you can also pass ``chunk_dim`` to set the chunk size of the appended dimension's coordinate. If you don't, this coordinate will keep the same chunk size as before appending, which could be too small if this is e.g. a From 2078838f4c88fc3fc9516b8719a3e28044c9e4ec Mon Sep 17 00:00:00 2001 From: jendrik Date: Mon, 17 Jun 2019 21:36:05 +0200 Subject: [PATCH 13/33] Add test for compute=False and commented it out --- xarray/tests/test_backends.py | 9 +++++++++ 1 file changed, 9 insertions(+) diff --git a/xarray/tests/test_backends.py b/xarray/tests/test_backends.py index 8e95741396e..8adccdfee99 100644 --- a/xarray/tests/test_backends.py +++ b/xarray/tests/test_backends.py @@ -1610,9 +1610,11 @@ def test_write_persistence_modes(self): obj1, obj2 = create_append_test_data() with pytest.raises(ValueError): + # check failure when having non-fixed-sized strings with self.create_zarr_target() as store_target: obj1.to_zarr(store_target, mode='a') + # encode strings to fixed sized arrays obj1['string_var'] = encode_utf8(obj1.string_var, 2) obj2['string_var'] = encode_utf8(obj2.string_var, 2) with self.create_zarr_target() as store_target: @@ -1630,6 +1632,13 @@ def test_write_persistence_modes(self): store = zarr.open(store_target) assert_identical(original, xr.open_zarr(store_target)) assert original['time'].shape == store['time'].chunks + # TODO: implement that compute = False keeps zarr uninitialized + # with self.create_zarr_target() as store_target: + # obj1.to_zarr(store_target, mode='w', compute=False) + # with pytest.raises(ValueError): + # assert_identical(obj1, xr.open_zarr(store_target)) + # print(xr.open_zarr(store_target).lon.values) + def test_compressor_encoding(self): original = create_test_data() From b8af5bd77f0de3aa46d6d81dc5f8b69a9f31d5e3 Mon Sep 17 00:00:00 2001 From: jendrik Date: Mon, 17 Jun 2019 23:47:06 +0200 Subject: [PATCH 14/33] Remove python 3.7 string formatting --- xarray/backends/api.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/xarray/backends/api.py b/xarray/backends/api.py index 99f61890963..c5e1f255d0f 100644 --- a/xarray/backends/api.py +++ b/xarray/backends/api.py @@ -1013,7 +1013,7 @@ def encode_utf8(var, string_max_length): missing = pd.isnull(data) data[missing] = "" data = encode_string_array(data, 'utf-8') - data = data.astype(np.dtype(f"S{string_max_length*2}")) + data = data.astype(np.dtype("S{}".format(string_max_length*2))) return Variable(dims, data, attrs, encoding) From 5bee0dc58edf7d09839cab4d9570ed237ac12d24 Mon Sep 17 00:00:00 2001 From: jendrik Date: Mon, 17 Jun 2019 23:49:57 +0200 Subject: [PATCH 15/33] Fix PEP8 incompliance --- xarray/backends/api.py | 2 +- xarray/tests/test_backends.py | 1 - xarray/tests/test_dataset.py | 3 ++- 3 files changed, 3 insertions(+), 3 deletions(-) diff --git a/xarray/backends/api.py b/xarray/backends/api.py index 99f61890963..c5e1f255d0f 100644 --- a/xarray/backends/api.py +++ b/xarray/backends/api.py @@ -1013,7 +1013,7 @@ def encode_utf8(var, string_max_length): missing = pd.isnull(data) data[missing] = "" data = encode_string_array(data, 'utf-8') - data = data.astype(np.dtype(f"S{string_max_length*2}")) + data = data.astype(np.dtype("S{}".format(string_max_length*2))) return Variable(dims, data, attrs, encoding) diff --git a/xarray/tests/test_backends.py b/xarray/tests/test_backends.py index 5e89583ce43..11cefa664a4 100644 --- a/xarray/tests/test_backends.py +++ b/xarray/tests/test_backends.py @@ -1651,7 +1651,6 @@ def test_write_persistence_modes(self): # assert_identical(obj1, xr.open_zarr(store_target)) # print(xr.open_zarr(store_target).lon.values) - def test_compressor_encoding(self): original = create_test_data() # specify a custom compressor diff --git a/xarray/tests/test_dataset.py b/xarray/tests/test_dataset.py index acd4372e93d..cee570222d9 100644 --- a/xarray/tests/test_dataset.py +++ b/xarray/tests/test_dataset.py @@ -69,7 +69,8 @@ def create_append_test_data(): da2 = xr.DataArray((np.arange(3 * 3 * nt2) + 100).reshape(3, 3, nt2), coords=[lat, lon, time2], dims=['lat', 'lon', 'time']) ds2 = da2.to_dataset(name='da') - ds2['string_var'] = xr.DataArray(string_var[:nt2], coords=[time2], dims=['time']) + ds2['string_var'] = xr.DataArray(string_var[:nt2], + coords=[time2], dims=['time']) assert all(objp.data.flags.writeable for objp in ds1.variables.values()) assert all(objp.data.flags.writeable for objp in ds2.variables.values()) From b4ff1c7cf7cf6683fa8e264b19ad4ef36514227d Mon Sep 17 00:00:00 2001 From: jendrik Date: Tue, 18 Jun 2019 09:07:11 +0200 Subject: [PATCH 16/33] Add missing whitespaces --- xarray/backends/api.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/xarray/backends/api.py b/xarray/backends/api.py index c5e1f255d0f..a8099f5f6b5 100644 --- a/xarray/backends/api.py +++ b/xarray/backends/api.py @@ -1013,7 +1013,7 @@ def encode_utf8(var, string_max_length): missing = pd.isnull(data) data[missing] = "" data = encode_string_array(data, 'utf-8') - data = data.astype(np.dtype("S{}".format(string_max_length*2))) + data = data.astype(np.dtype("S{}".format(string_max_length * 2))) return Variable(dims, data, attrs, encoding) From 5b3f8ea6c6244ce6bf55935bc040d4eb8c68b888 Mon Sep 17 00:00:00 2001 From: Shikhar Goenka Date: Thu, 20 Jun 2019 15:19:07 +0100 Subject: [PATCH 17/33] allowed for compute=False when appending to a zarr store --- xarray/backends/common.py | 17 ++++++++++--- xarray/backends/zarr.py | 13 ++++++---- xarray/tests/test_backends.py | 45 ++++++++++++++++++++++++++++++----- 3 files changed, 62 insertions(+), 13 deletions(-) diff --git a/xarray/backends/common.py b/xarray/backends/common.py index 5819292bb8a..38ada5405e9 100644 --- a/xarray/backends/common.py +++ b/xarray/backends/common.py @@ -158,14 +158,20 @@ class ArrayWriter: def __init__(self, lock=None): self.sources = [] self.targets = [] + self.regions = [] self.lock = lock - def add(self, source, target): + def add(self, source, target, region=None): if isinstance(source, dask_array_type): self.sources.append(source) self.targets.append(target) + if region: + self.regions.append(region) else: - target[...] = source + if region: + target[region] = source + else: + target[...] = source def sync(self, compute=True): if self.sources: @@ -173,9 +179,14 @@ def sync(self, compute=True): # TODO: consider wrapping targets with dask.delayed, if this makes # for any discernable difference in perforance, e.g., # targets = [dask.delayed(t) for t in self.targets] + + if not self.regions: + regions = None + else: + regions = self.regions delayed_store = da.store(self.sources, self.targets, lock=self.lock, compute=compute, - flush=True) + flush=True, regions=regions) self.sources = [] self.targets = [] return delayed_store diff --git a/xarray/backends/zarr.py b/xarray/backends/zarr.py index da7fcb1adad..260fcfba1f9 100644 --- a/xarray/backends/zarr.py +++ b/xarray/backends/zarr.py @@ -420,12 +420,17 @@ def set_variables(self, variables, check_encoding_set, writer, dtype=dtype, overwrite=True) zarr_array = self.ds[name] zarr_array.attrs.put(attrs) - zarr_array[:] = new_coord + writer.add(new_coord, zarr_array) else: # this is the DataArray that has append_dim as a # dimension - axis = dims.index(self.append_dim) - zarr_array.append(v.data, axis=axis) + append_axis = dims.index(self.append_dim) + new_shape = list(zarr_array.shape) + new_shape[append_axis] += v.shape[append_axis] + new_region = [slice(None),] * len(new_shape) + new_region[append_axis] = slice(zarr_array.shape[append_axis],None) + zarr_array.resize(new_shape) + writer.add(v.data, zarr_array, region=tuple(new_region)) else: # new variable encoding = _extract_zarr_variable_encoding( @@ -439,7 +444,7 @@ def set_variables(self, variables, check_encoding_set, writer, zarr_array = self.ds.create(name, shape=shape, dtype=dtype, fill_value=fill_value, **encoding) zarr_array.attrs.put(encoded_attrs) - zarr_array[...] = v.data + writer.add(v.data, zarr_array) def close(self): if self._consolidate_on_close: diff --git a/xarray/tests/test_backends.py b/xarray/tests/test_backends.py index 11cefa664a4..4a965f0e189 100644 --- a/xarray/tests/test_backends.py +++ b/xarray/tests/test_backends.py @@ -1644,12 +1644,6 @@ def test_write_persistence_modes(self): store = zarr.open(store_target) assert_identical(original, xr.open_zarr(store_target)) assert original['time'].shape == store['time'].chunks - # TODO: implement that compute = False keeps zarr uninitialized - # with self.create_zarr_target() as store_target: - # obj1.to_zarr(store_target, mode='w', compute=False) - # with pytest.raises(ValueError): - # assert_identical(obj1, xr.open_zarr(store_target)) - # print(xr.open_zarr(store_target).lon.values) def test_compressor_encoding(self): original = create_test_data() @@ -1699,11 +1693,50 @@ def test_to_zarr_compute_false_roundtrip(self): with self.create_zarr_target() as store: delayed_obj = self.save(original, store, compute=False) assert isinstance(delayed_obj, Delayed) + + # make sure target store has not been written to yet + with pytest.raises(AssertionError): + with self.open(store) as actual: + assert_identical(original, actual) + delayed_obj.compute() with self.open(store) as actual: assert_identical(original, actual) + def test_to_zarr_append_compute_false_roundtrip(self): + from dask.delayed import Delayed + + obj1, obj2 = create_append_test_data() + obj1['string_var'] = encode_utf8(obj1.string_var, 2) + obj2['string_var'] = encode_utf8(obj2.string_var, 2) + obj1, obj2 = obj1.chunk(), obj2.chunk() + + with self.create_zarr_target() as store: + delayed_obj = self.save(obj1, store, compute=False, mode='w') + assert isinstance(delayed_obj, Delayed) + + with pytest.raises(AssertionError): + with self.open(store) as actual: + assert_identical(obj1, actual) + + delayed_obj.compute() + + with self.open(store) as actual: + assert_identical(obj1, actual) + + delayed_obj = self.save(obj2, store, compute=False, mode='a', append_dim='time') + assert isinstance(delayed_obj, Delayed) + + with pytest.raises(AssertionError): + with self.open(store) as actual: + assert_identical(xr.concat([obj1, obj2], dim='time'), actual) + + delayed_obj.compute() + + with self.open(store) as actual: + assert_identical(xr.concat([obj1, obj2], dim='time'), actual) + def test_encoding_chunksizes(self): # regression test for GH2278 # see also test_encoding_chunksizes_unlimited From 756432955e249c4f6ac8a99d406d5db3fbff16b8 Mon Sep 17 00:00:00 2001 From: Shikhar Goenka Date: Thu, 20 Jun 2019 17:40:14 +0100 Subject: [PATCH 18/33] Fixed empty array data error When using create_append_test_data we used np.arange(...), which was incidently also the default value of the zarr array when fill_value is set to None. So appending to the data with compute=False and then expecting an error when asserting the source and target to be the same failed the tests. Using random data passes the tests --- xarray/tests/test_dataset.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/xarray/tests/test_dataset.py b/xarray/tests/test_dataset.py index cee570222d9..754eddaba29 100644 --- a/xarray/tests/test_dataset.py +++ b/xarray/tests/test_dataset.py @@ -52,14 +52,15 @@ def create_test_data(seed=None): return obj -def create_append_test_data(): +def create_append_test_data(seed=None): + rs = np.random.RandomState(seed) lat = [2, 1, 0] lon = [0, 1, 2] string_var = ["ae", "bc", "df"] nt1 = 3 time1 = pd.date_range('2000-01-01', periods=nt1) - da1 = xr.DataArray(np.arange(3 * 3 * nt1).reshape(3, 3, nt1), coords=[lat, + da1 = xr.DataArray(rs.rand(3, 3, nt1), coords=[lat, lon, time1], dims=['lat', 'lon', 'time']) ds1 = da1.to_dataset(name='da') ds1['string_var'] = xr.DataArray(string_var, coords=[time1], dims=['time']) From 93be7906b21958614ba0c79ac5e603fd41454cd3 Mon Sep 17 00:00:00 2001 From: Shikhar Goenka Date: Thu, 20 Jun 2019 18:02:19 +0100 Subject: [PATCH 19/33] flake8 fixes --- xarray/backends/zarr.py | 10 +++++++--- xarray/tests/test_backends.py | 6 ++++-- 2 files changed, 11 insertions(+), 5 deletions(-) diff --git a/xarray/backends/zarr.py b/xarray/backends/zarr.py index 260fcfba1f9..fe44c7640cf 100644 --- a/xarray/backends/zarr.py +++ b/xarray/backends/zarr.py @@ -427,10 +427,14 @@ def set_variables(self, variables, check_encoding_set, writer, append_axis = dims.index(self.append_dim) new_shape = list(zarr_array.shape) new_shape[append_axis] += v.shape[append_axis] - new_region = [slice(None),] * len(new_shape) - new_region[append_axis] = slice(zarr_array.shape[append_axis],None) + new_region = [slice(None)] * len(new_shape) + new_region[append_axis] = slice( + zarr_array.shape[append_axis], + None + ) zarr_array.resize(new_shape) - writer.add(v.data, zarr_array, region=tuple(new_region)) + writer.add(v.data, zarr_array, + region=tuple(new_region)) else: # new variable encoding = _extract_zarr_variable_encoding( diff --git a/xarray/tests/test_backends.py b/xarray/tests/test_backends.py index 4a965f0e189..53872c24274 100644 --- a/xarray/tests/test_backends.py +++ b/xarray/tests/test_backends.py @@ -1725,12 +1725,14 @@ def test_to_zarr_append_compute_false_roundtrip(self): with self.open(store) as actual: assert_identical(obj1, actual) - delayed_obj = self.save(obj2, store, compute=False, mode='a', append_dim='time') + delayed_obj = self.save(obj2, store, compute=False, mode='a', + append_dim='time') assert isinstance(delayed_obj, Delayed) with pytest.raises(AssertionError): with self.open(store) as actual: - assert_identical(xr.concat([obj1, obj2], dim='time'), actual) + assert_identical(xr.concat([obj1, obj2], dim='time'), + actual) delayed_obj.compute() From 58c4b783916603d95f6fefe6dda81c0e3cf1c7c3 Mon Sep 17 00:00:00 2001 From: Shikhar Goenka Date: Fri, 21 Jun 2019 16:13:47 +0100 Subject: [PATCH 20/33] removed chunk_dim argument to to_zarr function --- doc/io.rst | 4 ---- xarray/backends/api.py | 4 +--- xarray/backends/zarr.py | 42 +++++++++++------------------------ xarray/core/dataset.py | 7 ++---- xarray/tests/test_backends.py | 10 --------- 5 files changed, 16 insertions(+), 51 deletions(-) diff --git a/doc/io.rst b/doc/io.rst index e6346801b49..2a2d3b3cb33 100644 --- a/doc/io.rst +++ b/doc/io.rst @@ -611,10 +611,6 @@ of ``np.number`` or ``np.string_``. It as well needs to be taken into account, that the size of the ``np.string_`` in the first chunk sets the maximum string size for all following ones. To encode the data consider using :py:func:`~xarray.core.api.encode_utf8`. -Optionally, you can also pass ``chunk_dim`` to set the chunk size of the -appended dimension's coordinate. If you don't, this coordinate will keep the -same chunk size as before appending, which could be too small if this is e.g. a -1-D array (typically, the time coordinate). To read back a zarr dataset that has been created this way, we use the :py:func:`~xarray.open_zarr` method: diff --git a/xarray/backends/api.py b/xarray/backends/api.py index a8099f5f6b5..e74280b5d9b 100644 --- a/xarray/backends/api.py +++ b/xarray/backends/api.py @@ -1031,8 +1031,7 @@ def check_dtype(var): def to_zarr(dataset, store=None, mode='w-', synchronizer=None, group=None, - encoding=None, compute=True, consolidated=False, append_dim=None, - chunk_dim=None): + encoding=None, compute=True, consolidated=False, append_dim=None): """This function creates an appropriate datastore for writing a dataset to a zarr ztore @@ -1055,7 +1054,6 @@ def to_zarr(dataset, store=None, mode='w-', synchronizer=None, group=None, group=group, consolidate_on_close=consolidated) zstore.append_dim = append_dim - zstore.chunk_dim = chunk_dim writer = ArrayWriter() # TODO: figure out how to properly handle unlimited_dims dump_to_store(dataset, zstore, writer, encoding=encoding) diff --git a/xarray/backends/zarr.py b/xarray/backends/zarr.py index fe44c7640cf..e1330406aca 100644 --- a/xarray/backends/zarr.py +++ b/xarray/backends/zarr.py @@ -259,7 +259,6 @@ def __init__(self, zarr_group, consolidate_on_close=False): self._group = self.ds.path self._consolidate_on_close = consolidate_on_close self.append_dim = None - self.chunk_dim = None def open_store_variable(self, name, zarr_array): data = indexing.LazilyOuterIndexedArray(ZarrArrayWrapper(name, self)) @@ -407,34 +406,19 @@ def set_variables(self, variables, check_encoding_set, writer, 'did you forget to call to_zarr with append_dim ' 'argument?') if self.append_dim in dims: - if len(dims) == 1 and self.chunk_dim is not None: - # this is the coordinate along which to append - prev_coord = zarr_array[:] - attrs = zarr_array.attrs.asdict() - dtype = zarr_array.dtype - fill_value = zarr_array.fill_value - new_coord = np.hstack((prev_coord, v.data)) - self.ds.create_dataset(name, shape=new_coord.shape, - chunks=self.chunk_dim, - fill_value=fill_value, - dtype=dtype, overwrite=True) - zarr_array = self.ds[name] - zarr_array.attrs.put(attrs) - writer.add(new_coord, zarr_array) - else: - # this is the DataArray that has append_dim as a - # dimension - append_axis = dims.index(self.append_dim) - new_shape = list(zarr_array.shape) - new_shape[append_axis] += v.shape[append_axis] - new_region = [slice(None)] * len(new_shape) - new_region[append_axis] = slice( - zarr_array.shape[append_axis], - None - ) - zarr_array.resize(new_shape) - writer.add(v.data, zarr_array, - region=tuple(new_region)) + # this is the DataArray that has append_dim as a + # dimension + append_axis = dims.index(self.append_dim) + new_shape = list(zarr_array.shape) + new_shape[append_axis] += v.shape[append_axis] + new_region = [slice(None)] * len(new_shape) + new_region[append_axis] = slice( + zarr_array.shape[append_axis], + None + ) + zarr_array.resize(new_shape) + writer.add(v.data, zarr_array, + region=tuple(new_region)) else: # new variable encoding = _extract_zarr_variable_encoding( diff --git a/xarray/core/dataset.py b/xarray/core/dataset.py index e2bd696ef1d..3568bc2d947 100644 --- a/xarray/core/dataset.py +++ b/xarray/core/dataset.py @@ -1323,7 +1323,7 @@ def to_netcdf(self, path=None, mode='w', format=None, group=None, def to_zarr(self, store=None, mode='w-', synchronizer=None, group=None, encoding=None, compute=True, consolidated=False, - append_dim=None, chunk_dim=None): + append_dim=None): """Write dataset contents to a zarr group. .. note:: Experimental @@ -1354,8 +1354,6 @@ def to_zarr(self, store=None, mode='w-', synchronizer=None, group=None, after writing. append_dim: str, optional If mode='a', the dimension on which the data will be appended. - chunk_dim: int, optional - If mode='a', the chunk size of the `append_dim` coordinate. References ---------- @@ -1370,8 +1368,7 @@ def to_zarr(self, store=None, mode='w-', synchronizer=None, group=None, from ..backends.api import to_zarr return to_zarr(self, store=store, mode=mode, synchronizer=synchronizer, group=group, encoding=encoding, compute=compute, - consolidated=consolidated, append_dim=append_dim, - chunk_dim=chunk_dim) + consolidated=consolidated, append_dim=append_dim) def __repr__(self): return formatting.dataset_repr(self) diff --git a/xarray/tests/test_backends.py b/xarray/tests/test_backends.py index 53872c24274..b69af8dc9bf 100644 --- a/xarray/tests/test_backends.py +++ b/xarray/tests/test_backends.py @@ -1634,16 +1634,6 @@ def test_write_persistence_modes(self): obj2.to_zarr(store_target, mode='a', append_dim='time') original = xr.concat([obj1, obj2], dim='time') assert_identical(original, xr.open_zarr(store_target)) - # check chunk size of the append_dim coordinate when chunk_dim is set - with self.create_zarr_target() as store_target: - obj1.to_zarr(store_target, mode='w') - obj2.to_zarr(store_target, mode='a', append_dim='time', - chunk_dim=original['time'].shape) - original = xr.concat([obj1, obj2], dim='time') - import zarr - store = zarr.open(store_target) - assert_identical(original, xr.open_zarr(store_target)) - assert original['time'].shape == store['time'].chunks def test_compressor_encoding(self): original = create_test_data() From 5316593e2b103092c0e2d5094fa73894fa13664b Mon Sep 17 00:00:00 2001 From: Shikhar Goenka Date: Mon, 24 Jun 2019 22:50:07 +0100 Subject: [PATCH 21/33] implemented requested changes --- xarray/backends/common.py | 4 ++-- xarray/backends/zarr.py | 6 +++--- xarray/core/dataset.py | 4 ++-- xarray/tests/test_backends.py | 6 ++++++ 4 files changed, 13 insertions(+), 7 deletions(-) diff --git a/xarray/backends/common.py b/xarray/backends/common.py index 38ada5405e9..9c46bea7787 100644 --- a/xarray/backends/common.py +++ b/xarray/backends/common.py @@ -165,8 +165,7 @@ def add(self, source, target, region=None): if isinstance(source, dask_array_type): self.sources.append(source) self.targets.append(target) - if region: - self.regions.append(region) + self.regions.append(region) else: if region: target[region] = source @@ -189,6 +188,7 @@ def sync(self, compute=True): flush=True, regions=regions) self.sources = [] self.targets = [] + self.regions = [] return delayed_store diff --git a/xarray/backends/zarr.py b/xarray/backends/zarr.py index e1330406aca..64862e3cc95 100644 --- a/xarray/backends/zarr.py +++ b/xarray/backends/zarr.py @@ -402,9 +402,9 @@ def set_variables(self, variables, check_encoding_set, writer, zarr_array = self.ds[name] if self.append_dim is None: raise ValueError( - 'dimension being appended is unknown; ' - 'did you forget to call to_zarr with append_dim ' - 'argument?') + "variable '{}' already exists, but append_dim " + "was not set".format(name) + ) if self.append_dim in dims: # this is the DataArray that has append_dim as a # dimension diff --git a/xarray/core/dataset.py b/xarray/core/dataset.py index 3568bc2d947..a3065778a83 100644 --- a/xarray/core/dataset.py +++ b/xarray/core/dataset.py @@ -1363,8 +1363,8 @@ def to_zarr(self, store=None, mode='w-', synchronizer=None, group=None, encoding = {} if mode not in ['w', 'w-', 'a']: # TODO: figure out how to handle 'r+' - raise ValueError("The only supported options for mode are 'w' " - "and 'w-'.") + raise ValueError("The only supported options for mode are 'w'," + "'w-' and 'a'.") from ..backends.api import to_zarr return to_zarr(self, store=store, mode=mode, synchronizer=synchronizer, group=group, encoding=encoding, compute=compute, diff --git a/xarray/tests/test_backends.py b/xarray/tests/test_backends.py index b69af8dc9bf..ee85c60b43b 100644 --- a/xarray/tests/test_backends.py +++ b/xarray/tests/test_backends.py @@ -1635,6 +1635,12 @@ def test_write_persistence_modes(self): original = xr.concat([obj1, obj2], dim='time') assert_identical(original, xr.open_zarr(store_target)) + # check failure when append_dim not set + with pytest.raises(ValueError): + with self.create_zarr_target() as store_target: + obj1.to_zarr(store_target, mode='w') + obj2.to_zarr(store_target, mode='a') + def test_compressor_encoding(self): original = create_test_data() # specify a custom compressor From ad08c7311fa7f2b6656b2de4fd9a8ddcb4e9004d Mon Sep 17 00:00:00 2001 From: shikharsg Date: Tue, 25 Jun 2019 10:48:04 +0100 Subject: [PATCH 22/33] Update xarray/backends/api.py Co-Authored-By: Stephan Hoyer --- xarray/backends/api.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/xarray/backends/api.py b/xarray/backends/api.py index e74280b5d9b..b6d7a264894 100644 --- a/xarray/backends/api.py +++ b/xarray/backends/api.py @@ -1023,7 +1023,7 @@ def check_dtype(var): if (not np.issubdtype(var.dtype, np.number) and not np.issubdtype(var.dtype, np.string_)): # and not re.match('^bytes[1-9]+$', var.dtype.name)): - raise ValueError('Invalid dtype for DataVariable: {} ' + raise ValueError('Invalid dtype for data variable: {} ' 'dtype must be a subtype of number or ' 'a fixed sized string'.format(var)) for k in dataset.data_vars.values(): From 4d2122b5946174c2cbe4a54dd7452004cd4f8057 Mon Sep 17 00:00:00 2001 From: Shikhar Goenka Date: Tue, 25 Jun 2019 14:38:40 +0100 Subject: [PATCH 23/33] added contributors and example of using append to zarr --- doc/io.rst | 19 +++++++++++++++++++ doc/whats-new.rst | 4 ++++ 2 files changed, 23 insertions(+) diff --git a/doc/io.rst b/doc/io.rst index 2a2d3b3cb33..02000c843fb 100644 --- a/doc/io.rst +++ b/doc/io.rst @@ -606,6 +606,25 @@ store, add ``mode='w'`` when invoking ``to_zarr``. It is also possible to append to an existing store. For that, add ``mode='a'`` and set ``append_dim`` to the name of the dimension along which to append. + +.. ipython:: python + :suppress: + + ! rm -rf path/to/directory.zarr + +.. ipython:: python + + ds1 = xr.Dataset({'foo': (('x', 'y'), np.random.rand(4, 5, 2))}, + coords={'x': [10, 20, 30, 40], + 'y': [1,2,3,4,5], + 't': pd.date_range('2001-01-01', periods=2)}) + ds1.to_zarr('path/to/directory.zarr') + ds2 = xr.Dataset({'foo': (('x', 'y'), np.random.rand(4, 5, 2))}, + coords={'x': [10, 20, 30, 40], + 'y': [1,2,3,4,5], + 't': pd.date_range('2001-01-03', periods=2)}) + ds2.to_zarr('path/to/directory.zarr', mode='a', append_dim='time') + It is necessary that the data only contains types which are either a subclass of ``np.number`` or ``np.string_``. It as well needs to be taken into account, that the size of the ``np.string_`` in the first chunk sets the maximum diff --git a/doc/whats-new.rst b/doc/whats-new.rst index eb1379d856a..500269d1720 100644 --- a/doc/whats-new.rst +++ b/doc/whats-new.rst @@ -189,6 +189,10 @@ Other enhancements coordinates / variables / attributes) (:issue:`1507`). By `Benoit Bovy `_. - Added append capability to the zarr store. + By `Jendrik Jördening `_, + `David Brochart `_, + `Ryan Abernathey `_ and + `Shikhar Goenka`_. - Add ``tolerance`` option to ``resample()`` methods ``bfill``, ``pad``, ``nearest``. (:issue:`2695`) By `Hauke Schulz `_. From af4a5a52c6ed5cdc5990dc21e3dab40083c68d72 Mon Sep 17 00:00:00 2001 From: Shikhar Goenka Date: Tue, 25 Jun 2019 15:04:24 +0100 Subject: [PATCH 24/33] fixed docs fail --- doc/io.rst | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/doc/io.rst b/doc/io.rst index 02000c843fb..ffb6fed6ea2 100644 --- a/doc/io.rst +++ b/doc/io.rst @@ -614,16 +614,16 @@ and set ``append_dim`` to the name of the dimension along which to append. .. ipython:: python - ds1 = xr.Dataset({'foo': (('x', 'y'), np.random.rand(4, 5, 2))}, + ds1 = xr.Dataset({'foo': (('x', 'y', 't'), np.random.rand(4, 5, 2))}, coords={'x': [10, 20, 30, 40], 'y': [1,2,3,4,5], 't': pd.date_range('2001-01-01', periods=2)}) ds1.to_zarr('path/to/directory.zarr') - ds2 = xr.Dataset({'foo': (('x', 'y'), np.random.rand(4, 5, 2))}, + ds2 = xr.Dataset({'foo': (('x', 'y', 't'), np.random.rand(4, 5, 2))}, coords={'x': [10, 20, 30, 40], 'y': [1,2,3,4,5], 't': pd.date_range('2001-01-03', periods=2)}) - ds2.to_zarr('path/to/directory.zarr', mode='a', append_dim='time') + ds2.to_zarr('path/to/directory.zarr', mode='a', append_dim='t') It is necessary that the data only contains types which are either a subclass of ``np.number`` or ``np.string_``. It as well needs to be taken into account, From 62d4f521d6354627850f311a353c16acf4d57d01 Mon Sep 17 00:00:00 2001 From: Shikhar Goenka Date: Tue, 25 Jun 2019 15:57:22 +0100 Subject: [PATCH 25/33] fixed docs --- xarray/backends/zarr.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/xarray/backends/zarr.py b/xarray/backends/zarr.py index 64862e3cc95..7548a064649 100644 --- a/xarray/backends/zarr.py +++ b/xarray/backends/zarr.py @@ -351,7 +351,7 @@ def store(self, variables, attributes, check_encoding_set=frozenset(), if len(existing_variables) > 0: # there are variables to append # their encoding must be the same as in the store - ds = open_zarr(self.ds.store, auto_chunk=False) + ds = open_zarr(self.ds.store, chunks=None) variables_with_encoding = OrderedDict() for vn in existing_variables: variables_with_encoding[vn] = variables[vn] From 9d70e02b467c9b21368a7db178bdfd8b68619f31 Mon Sep 17 00:00:00 2001 From: Shikhar Goenka Date: Tue, 25 Jun 2019 17:14:08 +0100 Subject: [PATCH 26/33] removed unnecessary condition --- xarray/backends/common.py | 6 +----- 1 file changed, 1 insertion(+), 5 deletions(-) diff --git a/xarray/backends/common.py b/xarray/backends/common.py index 9c46bea7787..a433c80505a 100644 --- a/xarray/backends/common.py +++ b/xarray/backends/common.py @@ -179,13 +179,9 @@ def sync(self, compute=True): # for any discernable difference in perforance, e.g., # targets = [dask.delayed(t) for t in self.targets] - if not self.regions: - regions = None - else: - regions = self.regions delayed_store = da.store(self.sources, self.targets, lock=self.lock, compute=compute, - flush=True, regions=regions) + flush=True, regions=self.regions) self.sources = [] self.targets = [] self.regions = [] From a6ff49436a620eb4071c0529d785370b2eb739b0 Mon Sep 17 00:00:00 2001 From: Shikhar Goenka Date: Wed, 26 Jun 2019 17:33:56 +0100 Subject: [PATCH 27/33] attempt at clean string encoding and variable length strings --- doc/io.rst | 6 +----- xarray/backends/api.py | 19 ++++++------------- xarray/backends/zarr.py | 4 +++- xarray/tests/test_backends.py | 11 ----------- xarray/tests/test_dataset.py | 10 ++++++++-- 5 files changed, 18 insertions(+), 32 deletions(-) diff --git a/doc/io.rst b/doc/io.rst index 51a002e8a8d..857066899c8 100644 --- a/doc/io.rst +++ b/doc/io.rst @@ -625,11 +625,7 @@ and set ``append_dim`` to the name of the dimension along which to append. 't': pd.date_range('2001-01-03', periods=2)}) ds2.to_zarr('path/to/directory.zarr', mode='a', append_dim='t') -It is necessary that the data only contains types which are either a subclass -of ``np.number`` or ``np.string_``. It as well needs to be taken into account, -that the size of the ``np.string_`` in the first chunk sets the maximum -string size for all following ones. To encode the data consider using -:py:func:`~xarray.core.api.encode_utf8`. +To store variable length strings use ``dtype=object``. To read back a zarr dataset that has been created this way, we use the :py:func:`~xarray.open_zarr` method: diff --git a/xarray/backends/api.py b/xarray/backends/api.py index dd96e91cd2b..ddaed443885 100644 --- a/xarray/backends/api.py +++ b/xarray/backends/api.py @@ -19,7 +19,6 @@ from .common import ArrayWriter from .locks import _get_scheduler from ..coding.variables import safe_setitem, unpack_for_encoding -from ..coding.strings import encode_string_array DATAARRAY_NAME = '__xarray_dataarray_name__' DATAARRAY_VARIABLE = '__xarray_dataarray_variable__' @@ -1028,24 +1027,18 @@ def save_mfdataset(datasets, paths, mode='w', format=None, groups=None, for w, s in zip(writes, stores)]) -def encode_utf8(var, string_max_length): - dims, data, attrs, encoding = unpack_for_encoding(var) - missing = pd.isnull(data) - data[missing] = "" - data = encode_string_array(data, 'utf-8') - data = data.astype(np.dtype("S{}".format(string_max_length * 2))) - return Variable(dims, data, attrs, encoding) - - def _validate_datatypes_for_zarr_append(dataset): """DataArray.name and Dataset keys must be a string or None""" def check_dtype(var): if (not np.issubdtype(var.dtype, np.number) - and not np.issubdtype(var.dtype, np.string_)): + and not np.issubdtype(var.dtype, np.string_) + and not np.issubdtype(var.dtype, np.unicode_) + and not var.dtype == object): # and not re.match('^bytes[1-9]+$', var.dtype.name)): raise ValueError('Invalid dtype for data variable: {} ' - 'dtype must be a subtype of number or ' - 'a fixed sized string'.format(var)) + 'dtype must be a subtype of number, ' + 'a fixed sized string, a fixed size ' + 'unicode string or an object'.format(var)) for k in dataset.data_vars.values(): check_dtype(k) diff --git a/xarray/backends/zarr.py b/xarray/backends/zarr.py index 7548a064649..42b60a64a44 100644 --- a/xarray/backends/zarr.py +++ b/xarray/backends/zarr.py @@ -213,7 +213,7 @@ def encode_zarr_variable(var, needs_copy=True, name=None): # zarr allows unicode, but not variable-length strings, so it's both # simpler and more compact to always encode as UTF-8 explicitly. # TODO: allow toggling this explicitly via dtype in encoding. - coder = coding.strings.EncodedStringCoder(allows_unicode=False) + coder = coding.strings.EncodedStringCoder(allows_unicode=True) var = coder.encode(var, name=name) var = coding.strings.ensure_fixed_length_bytes(var) @@ -429,6 +429,8 @@ def set_variables(self, variables, check_encoding_set, writer, for k2, v2 in attrs.items(): encoded_attrs[k2] = self.encode_attribute(v2) + if dtype == object: + dtype = str zarr_array = self.ds.create(name, shape=shape, dtype=dtype, fill_value=fill_value, **encoding) zarr_array.attrs.put(encoded_attrs) diff --git a/xarray/tests/test_backends.py b/xarray/tests/test_backends.py index 6bcfccea1de..47987de0e6a 100644 --- a/xarray/tests/test_backends.py +++ b/xarray/tests/test_backends.py @@ -20,7 +20,6 @@ from xarray import ( DataArray, Dataset, backends, load_dataarray, load_dataset, open_dataarray, open_dataset, open_mfdataset, save_mfdataset) -from xarray.backends.api import encode_utf8 from xarray.backends.common import robust_getitem from xarray.backends.netCDF4_ import _extract_nc4_variable_encoding from xarray.backends.pydap_ import PydapDataStore @@ -1621,14 +1620,6 @@ def test_write_persistence_modes(self): # check append mode for append write obj1, obj2 = create_append_test_data() - with pytest.raises(ValueError): - # check failure when having non-fixed-sized strings - with self.create_zarr_target() as store_target: - obj1.to_zarr(store_target, mode='a') - - # encode strings to fixed sized arrays - obj1['string_var'] = encode_utf8(obj1.string_var, 2) - obj2['string_var'] = encode_utf8(obj2.string_var, 2) with self.create_zarr_target() as store_target: obj1.to_zarr(store_target, mode='w') obj2.to_zarr(store_target, mode='a', append_dim='time') @@ -1704,8 +1695,6 @@ def test_to_zarr_append_compute_false_roundtrip(self): from dask.delayed import Delayed obj1, obj2 = create_append_test_data() - obj1['string_var'] = encode_utf8(obj1.string_var, 2) - obj2['string_var'] = encode_utf8(obj2.string_var, 2) obj1, obj2 = obj1.chunk(), obj2.chunk() with self.create_zarr_target() as store: diff --git a/xarray/tests/test_dataset.py b/xarray/tests/test_dataset.py index 88e3579259a..c7d6b2c9482 100644 --- a/xarray/tests/test_dataset.py +++ b/xarray/tests/test_dataset.py @@ -56,7 +56,9 @@ def create_append_test_data(seed=None): rs = np.random.RandomState(seed) lat = [2, 1, 0] lon = [0, 1, 2] - string_var = ["ae", "bc", "df"] + string_var = np.array(["ae", "bc", "df"], dtype=object) + string_var_to_append = np.array(['asdf', 'asdfg'], dtype=object) + unicode_var = ["áó", "áó", "áó"] nt1 = 3 time1 = pd.date_range('2000-01-01', periods=nt1) @@ -64,14 +66,18 @@ def create_append_test_data(seed=None): lon, time1], dims=['lat', 'lon', 'time']) ds1 = da1.to_dataset(name='da') ds1['string_var'] = xr.DataArray(string_var, coords=[time1], dims=['time']) + ds1['unicode_var'] = xr.DataArray(unicode_var, coords=[time1], + dims=['time']).astype(np.unicode_) nt2 = 2 time2 = pd.date_range('2000-02-01', periods=nt2) da2 = xr.DataArray((np.arange(3 * 3 * nt2) + 100).reshape(3, 3, nt2), coords=[lat, lon, time2], dims=['lat', 'lon', 'time']) ds2 = da2.to_dataset(name='da') - ds2['string_var'] = xr.DataArray(string_var[:nt2], + ds2['string_var'] = xr.DataArray(string_var_to_append, coords=[time2], dims=['time']) + ds2['unicode_var'] = xr.DataArray(unicode_var[:nt2], coords=[time2], + dims=['time']).astype(np.unicode_) assert all(objp.data.flags.writeable for objp in ds1.variables.values()) assert all(objp.data.flags.writeable for objp in ds2.variables.values()) From 34b700f07afa0c7cd2c0f0b810cdc379efab93fa Mon Sep 17 00:00:00 2001 From: Shikhar Goenka Date: Wed, 26 Jun 2019 19:23:04 +0100 Subject: [PATCH 28/33] implemented suggestions --- xarray/backends/api.py | 5 ++--- xarray/backends/zarr.py | 2 +- 2 files changed, 3 insertions(+), 4 deletions(-) diff --git a/xarray/backends/api.py b/xarray/backends/api.py index ddaed443885..861379e1288 100644 --- a/xarray/backends/api.py +++ b/xarray/backends/api.py @@ -9,7 +9,7 @@ import numpy as np import pandas as pd -from .. import Dataset, DataArray, backends, conventions +from .. import Dataset, DataArray, backends, conventions, coding from ..core import indexing from .. import auto_combine from ..core.combine import (combine_auto, _manual_combine, @@ -1031,8 +1031,7 @@ def _validate_datatypes_for_zarr_append(dataset): """DataArray.name and Dataset keys must be a string or None""" def check_dtype(var): if (not np.issubdtype(var.dtype, np.number) - and not np.issubdtype(var.dtype, np.string_) - and not np.issubdtype(var.dtype, np.unicode_) + and not coding.strings.is_unicode_dtype(var.dtype) and not var.dtype == object): # and not re.match('^bytes[1-9]+$', var.dtype.name)): raise ValueError('Invalid dtype for data variable: {} ' diff --git a/xarray/backends/zarr.py b/xarray/backends/zarr.py index 42b60a64a44..d1eda484928 100644 --- a/xarray/backends/zarr.py +++ b/xarray/backends/zarr.py @@ -429,7 +429,7 @@ def set_variables(self, variables, check_encoding_set, writer, for k2, v2 in attrs.items(): encoded_attrs[k2] = self.encode_attribute(v2) - if dtype == object: + if coding.strings.check_vlen_dtype(dtype) == str: dtype = str zarr_array = self.ds.create(name, shape=shape, dtype=dtype, fill_value=fill_value, **encoding) From 3e54cb90703ca6c0e6ed9534abe106ed09788e9f Mon Sep 17 00:00:00 2001 From: Shikhar Goenka Date: Thu, 27 Jun 2019 14:14:45 +0100 Subject: [PATCH 29/33] * append_dim does not need to be specified if creating a new array with Dataset.to_zarr(store, mode='a') * cleand up to_zarr append mode tests --- xarray/backends/zarr.py | 32 ++++++++++++++---- xarray/tests/test_backends.py | 39 +++++++++++++--------- xarray/tests/test_dataset.py | 62 +++++++++++++++++++++++------------ 3 files changed, 90 insertions(+), 43 deletions(-) diff --git a/xarray/backends/zarr.py b/xarray/backends/zarr.py index d1eda484928..f4fc780e5fa 100644 --- a/xarray/backends/zarr.py +++ b/xarray/backends/zarr.py @@ -368,6 +368,19 @@ def store(self, variables, attributes, check_encoding_set=frozenset(), def sync(self): pass + def _new_variable_dimensions(self, variables): + """ + Returns a list of all existing dimensions that + belong to new variables + """ + new_var_dims = list() + for vn, v in variables.items(): + if _encode_variable_name(vn) not in self.ds: + for dim in v.dims: + if _encode_variable_name(dim) in self.ds: + new_var_dims.append(dim) + return new_var_dims + def set_variables(self, variables, check_encoding_set, writer, unlimited_dims=None): """ @@ -386,6 +399,9 @@ def set_variables(self, variables, check_encoding_set, writer, List of dimension names that should be treated as unlimited dimensions. """ + + new_variable_dimensions = self._new_variable_dimensions(variables) + for vn, v in variables.items(): name = _encode_variable_name(vn) check = vn in check_encoding_set @@ -398,13 +414,17 @@ def set_variables(self, variables, check_encoding_set, writer, if v.encoding == {'_FillValue': None} and fill_value is None: v.encoding = {} if name in self.ds: - # append to existing variable zarr_array = self.ds[name] - if self.append_dim is None: - raise ValueError( - "variable '{}' already exists, but append_dim " - "was not set".format(name) - ) + """ + If variable is a dimension of an existing array + append_dim does not need to be specified + """ + if name not in new_variable_dimensions: + if self.append_dim is None: + raise ValueError( + "variable '{}' already exists, but append_dim " + "was not set".format(name) + ) if self.append_dim in dims: # this is the DataArray that has append_dim as a # dimension diff --git a/xarray/tests/test_backends.py b/xarray/tests/test_backends.py index 47987de0e6a..5843b7dd5fc 100644 --- a/xarray/tests/test_backends.py +++ b/xarray/tests/test_backends.py @@ -1617,20 +1617,26 @@ def test_write_persistence_modes(self): with self.roundtrip(original, save_kwargs={'mode': 'a'}) as actual: assert_identical(original, actual) - # check append mode for append write - obj1, obj2 = create_append_test_data() + ds, ds_to_append, ds_with_new_var = create_append_test_data() + # check append mode for append write with self.create_zarr_target() as store_target: - obj1.to_zarr(store_target, mode='w') - obj2.to_zarr(store_target, mode='a', append_dim='time') - original = xr.concat([obj1, obj2], dim='time') + ds.to_zarr(store_target, mode='w') + ds_to_append.to_zarr(store_target, mode='a', append_dim='time') + original = xr.concat([ds, ds_to_append], dim='time') assert_identical(original, xr.open_zarr(store_target)) # check failure when append_dim not set with pytest.raises(ValueError): with self.create_zarr_target() as store_target: - obj1.to_zarr(store_target, mode='w') - obj2.to_zarr(store_target, mode='a') + ds.to_zarr(store_target, mode='w') + ds_to_append.to_zarr(store_target, mode='a') + + # check append mode for new variable + with self.create_zarr_target() as store_target: + xr.concat([ds, ds_to_append], dim='time').to_zarr(store_target, + mode='w') + ds_with_new_var.to_zarr(store_target, mode='a') def test_compressor_encoding(self): original = create_test_data() @@ -1694,35 +1700,36 @@ def test_to_zarr_compute_false_roundtrip(self): def test_to_zarr_append_compute_false_roundtrip(self): from dask.delayed import Delayed - obj1, obj2 = create_append_test_data() - obj1, obj2 = obj1.chunk(), obj2.chunk() + ds, ds_to_append, ds_with_new_var = create_append_test_data() + ds, ds_to_append = ds.chunk(), ds_to_append.chunk() with self.create_zarr_target() as store: - delayed_obj = self.save(obj1, store, compute=False, mode='w') + delayed_obj = self.save(ds, store, compute=False, mode='w') assert isinstance(delayed_obj, Delayed) with pytest.raises(AssertionError): with self.open(store) as actual: - assert_identical(obj1, actual) + assert_identical(ds, actual) delayed_obj.compute() with self.open(store) as actual: - assert_identical(obj1, actual) + assert_identical(ds, actual) - delayed_obj = self.save(obj2, store, compute=False, mode='a', - append_dim='time') + delayed_obj = self.save(ds_to_append, store, compute=False, + mode='a', append_dim='time') assert isinstance(delayed_obj, Delayed) with pytest.raises(AssertionError): with self.open(store) as actual: - assert_identical(xr.concat([obj1, obj2], dim='time'), + assert_identical(xr.concat([ds, ds_to_append], dim='time'), actual) delayed_obj.compute() with self.open(store) as actual: - assert_identical(xr.concat([obj1, obj2], dim='time'), actual) + assert_identical(xr.concat([ds, ds_to_append], dim='time'), + actual) def test_encoding_chunksizes(self): # regression test for GH2278 diff --git a/xarray/tests/test_dataset.py b/xarray/tests/test_dataset.py index c7d6b2c9482..6b10900d9c5 100644 --- a/xarray/tests/test_dataset.py +++ b/xarray/tests/test_dataset.py @@ -54,34 +54,54 @@ def create_test_data(seed=None): def create_append_test_data(seed=None): rs = np.random.RandomState(seed) + lat = [2, 1, 0] lon = [0, 1, 2] + nt1 = 3 + nt2 = 2 + time1 = pd.date_range('2000-01-01', periods=nt1) + time2 = pd.date_range('2000-02-01', periods=nt2) string_var = np.array(["ae", "bc", "df"], dtype=object) string_var_to_append = np.array(['asdf', 'asdfg'], dtype=object) unicode_var = ["áó", "áó", "áó"] - nt1 = 3 - time1 = pd.date_range('2000-01-01', periods=nt1) - da1 = xr.DataArray(rs.rand(3, 3, nt1), coords=[lat, - lon, time1], dims=['lat', 'lon', 'time']) - ds1 = da1.to_dataset(name='da') - ds1['string_var'] = xr.DataArray(string_var, coords=[time1], dims=['time']) - ds1['unicode_var'] = xr.DataArray(unicode_var, coords=[time1], - dims=['time']).astype(np.unicode_) + ds = xr.Dataset( + data_vars={ + 'da': xr.DataArray(rs.rand(3, 3, nt1), coords=[lat, lon, time1], + dims=['lat', 'lon', 'time']), + 'string_var': xr.DataArray(string_var, coords=[time1], + dims=['time']), + 'unicode_var': xr.DataArray(unicode_var, coords=[time1], + dims=['time']).astype(np.unicode_) + } + ) - nt2 = 2 - time2 = pd.date_range('2000-02-01', periods=nt2) - da2 = xr.DataArray((np.arange(3 * 3 * nt2) + 100).reshape(3, 3, nt2), - coords=[lat, lon, time2], dims=['lat', 'lon', 'time']) - ds2 = da2.to_dataset(name='da') - ds2['string_var'] = xr.DataArray(string_var_to_append, - coords=[time2], dims=['time']) - ds2['unicode_var'] = xr.DataArray(unicode_var[:nt2], coords=[time2], - dims=['time']).astype(np.unicode_) - - assert all(objp.data.flags.writeable for objp in ds1.variables.values()) - assert all(objp.data.flags.writeable for objp in ds2.variables.values()) - return ds1, ds2 + ds_to_append = xr.Dataset( + data_vars={ + 'da': xr.DataArray(rs.rand(3, 3, nt2), coords=[lat, lon, time2], + dims=['lat', 'lon', 'time']), + 'string_var': xr.DataArray(string_var_to_append, coords=[time2], + dims=['time']), + 'unicode_var': xr.DataArray(unicode_var[:nt2], coords=[time2], + dims=['time']).astype(np.unicode_) + } + ) + + ds_with_new_var = xr.Dataset( + data_vars={ + 'new_var': xr.DataArray( + rs.rand(3, 3, nt1 + nt2), + coords=[lat, lon, time1.append(time2)], + dims=['lat', 'lon', 'time'] + ), + } + ) + + assert all(objp.data.flags.writeable for objp in ds.variables.values()) + assert all( + objp.data.flags.writeable for objp in ds_to_append.variables.values() + ) + return ds, ds_to_append, ds_with_new_var def create_test_multiindex(): From beb12e5d199dbb1c6249cf764bd5a5067fe7760f Mon Sep 17 00:00:00 2001 From: Shikhar Goenka Date: Thu, 27 Jun 2019 17:22:15 +0100 Subject: [PATCH 30/33] raise ValueError when append_dim is not a valid dimension --- xarray/backends/api.py | 21 +++++++++++++++++++++ xarray/backends/zarr.py | 21 --------------------- xarray/tests/test_backends.py | 9 +++++++++ 3 files changed, 30 insertions(+), 21 deletions(-) diff --git a/xarray/backends/api.py b/xarray/backends/api.py index ead9fef4d7f..f26ef8d167d 100644 --- a/xarray/backends/api.py +++ b/xarray/backends/api.py @@ -1043,6 +1043,25 @@ def check_dtype(var): check_dtype(k) +def _validate_append_dim(ds_to_append, store, append_dim, **open_kwargs): + try: + ds = backends.zarr.open_zarr(store, **open_kwargs) + except ValueError: # store empty + return + if append_dim: + if append_dim not in ds.dims: + raise ValueError( + "{} not a valid dimension in the Dataset".format(append_dim) + ) + for data_var in ds_to_append: + if data_var in ds: + if append_dim is None: + raise ValueError( + "variable '{}' already exists, but append_dim " + "was not set".format(data_var) + ) + + def to_zarr(dataset, store=None, mode='w-', synchronizer=None, group=None, encoding=None, compute=True, consolidated=False, append_dim=None): """This function creates an appropriate datastore for writing a dataset to @@ -1061,6 +1080,8 @@ def to_zarr(dataset, store=None, mode='w-', synchronizer=None, group=None, if mode == "a": _validate_datatypes_for_zarr_append(dataset) + _validate_append_dim(dataset, store, append_dim, + consolidated=consolidated, group=group) zstore = backends.ZarrStore.open_group(store=store, mode=mode, synchronizer=synchronizer, diff --git a/xarray/backends/zarr.py b/xarray/backends/zarr.py index f4fc780e5fa..d734c68fd88 100644 --- a/xarray/backends/zarr.py +++ b/xarray/backends/zarr.py @@ -368,19 +368,6 @@ def store(self, variables, attributes, check_encoding_set=frozenset(), def sync(self): pass - def _new_variable_dimensions(self, variables): - """ - Returns a list of all existing dimensions that - belong to new variables - """ - new_var_dims = list() - for vn, v in variables.items(): - if _encode_variable_name(vn) not in self.ds: - for dim in v.dims: - if _encode_variable_name(dim) in self.ds: - new_var_dims.append(dim) - return new_var_dims - def set_variables(self, variables, check_encoding_set, writer, unlimited_dims=None): """ @@ -400,8 +387,6 @@ def set_variables(self, variables, check_encoding_set, writer, dimensions. """ - new_variable_dimensions = self._new_variable_dimensions(variables) - for vn, v in variables.items(): name = _encode_variable_name(vn) check = vn in check_encoding_set @@ -419,12 +404,6 @@ def set_variables(self, variables, check_encoding_set, writer, If variable is a dimension of an existing array append_dim does not need to be specified """ - if name not in new_variable_dimensions: - if self.append_dim is None: - raise ValueError( - "variable '{}' already exists, but append_dim " - "was not set".format(name) - ) if self.append_dim in dims: # this is the DataArray that has append_dim as a # dimension diff --git a/xarray/tests/test_backends.py b/xarray/tests/test_backends.py index f3b9ea961f2..fc28d93eeae 100644 --- a/xarray/tests/test_backends.py +++ b/xarray/tests/test_backends.py @@ -1632,11 +1632,20 @@ def test_write_persistence_modes(self): ds.to_zarr(store_target, mode='w') ds_to_append.to_zarr(store_target, mode='a') + # check failure when append_dim not valid + with pytest.raises(ValueError): + with self.create_zarr_target() as store_target: + ds.to_zarr(store_target, mode='w') + ds_to_append.to_zarr(store_target, mode='a', append_dim='notvalid') + # check append mode for new variable with self.create_zarr_target() as store_target: xr.concat([ds, ds_to_append], dim='time').to_zarr(store_target, mode='w') ds_with_new_var.to_zarr(store_target, mode='a') + combined = xr.concat([ds, ds_to_append], dim='time') + combined['new_var'] = ds_with_new_var['new_var'] + assert_identical(combined, xr.open_zarr(store_target)) def test_compressor_encoding(self): original = create_test_data() From 41a6ca384ddf5eb9a481b743b50cbbe72480614c Mon Sep 17 00:00:00 2001 From: Shikhar Goenka Date: Thu, 27 Jun 2019 17:31:21 +0100 Subject: [PATCH 31/33] flake8 fix --- xarray/tests/test_backends.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/xarray/tests/test_backends.py b/xarray/tests/test_backends.py index fc28d93eeae..76ed7ad57f8 100644 --- a/xarray/tests/test_backends.py +++ b/xarray/tests/test_backends.py @@ -1636,7 +1636,8 @@ def test_write_persistence_modes(self): with pytest.raises(ValueError): with self.create_zarr_target() as store_target: ds.to_zarr(store_target, mode='w') - ds_to_append.to_zarr(store_target, mode='a', append_dim='notvalid') + ds_to_append.to_zarr(store_target, mode='a', + append_dim='notvalid') # check append mode for new variable with self.create_zarr_target() as store_target: From 321aec1822181bc06d90888f9bbf557bdc063714 Mon Sep 17 00:00:00 2001 From: Shikhar Goenka Date: Thu, 27 Jun 2019 17:36:55 +0100 Subject: [PATCH 32/33] removed unused comment --- xarray/backends/zarr.py | 4 ---- 1 file changed, 4 deletions(-) diff --git a/xarray/backends/zarr.py b/xarray/backends/zarr.py index d734c68fd88..9d223289271 100644 --- a/xarray/backends/zarr.py +++ b/xarray/backends/zarr.py @@ -400,10 +400,6 @@ def set_variables(self, variables, check_encoding_set, writer, v.encoding = {} if name in self.ds: zarr_array = self.ds[name] - """ - If variable is a dimension of an existing array - append_dim does not need to be specified - """ if self.append_dim in dims: # this is the DataArray that has append_dim as a # dimension From 2b130ff3628b16cea6717a1bb97fafa14964742e Mon Sep 17 00:00:00 2001 From: Shikhar Goenka Date: Sat, 29 Jun 2019 14:33:14 +0100 Subject: [PATCH 33/33] * raise error when appending with encoding provided for existing variable * add test for encoding consistency when appending * implemented: https://github.com/pydata/xarray/pull/2706#discussion_r298267776 * refactored tests --- xarray/backends/api.py | 14 +++-- xarray/backends/zarr.py | 2 +- xarray/tests/test_backends.py | 99 +++++++++++++++++++++++++---------- 3 files changed, 83 insertions(+), 32 deletions(-) diff --git a/xarray/backends/api.py b/xarray/backends/api.py index f26ef8d167d..7a8c022b87d 100644 --- a/xarray/backends/api.py +++ b/xarray/backends/api.py @@ -1043,7 +1043,8 @@ def check_dtype(var): check_dtype(k) -def _validate_append_dim(ds_to_append, store, append_dim, **open_kwargs): +def _validate_append_dim_and_encoding(ds_to_append, store, append_dim, + encoding, **open_kwargs): try: ds = backends.zarr.open_zarr(store, **open_kwargs) except ValueError: # store empty @@ -1060,6 +1061,11 @@ def _validate_append_dim(ds_to_append, store, append_dim, **open_kwargs): "variable '{}' already exists, but append_dim " "was not set".format(data_var) ) + if data_var in encoding.keys(): + raise ValueError( + "variable '{}' already exists, but encoding was" + "provided".format(data_var) + ) def to_zarr(dataset, store=None, mode='w-', synchronizer=None, group=None, @@ -1080,8 +1086,10 @@ def to_zarr(dataset, store=None, mode='w-', synchronizer=None, group=None, if mode == "a": _validate_datatypes_for_zarr_append(dataset) - _validate_append_dim(dataset, store, append_dim, - consolidated=consolidated, group=group) + _validate_append_dim_and_encoding(dataset, store, append_dim, + group=group, + consolidated=consolidated, + encoding=encoding) zstore = backends.ZarrStore.open_group(store=store, mode=mode, synchronizer=synchronizer, diff --git a/xarray/backends/zarr.py b/xarray/backends/zarr.py index 9d223289271..c0634fff009 100644 --- a/xarray/backends/zarr.py +++ b/xarray/backends/zarr.py @@ -354,7 +354,7 @@ def store(self, variables, attributes, check_encoding_set=frozenset(), ds = open_zarr(self.ds.store, chunks=None) variables_with_encoding = OrderedDict() for vn in existing_variables: - variables_with_encoding[vn] = variables[vn] + variables_with_encoding[vn] = variables[vn].copy(deep=False) variables_with_encoding[vn].encoding = ds[vn].encoding variables_with_encoding, _ = self.encode(variables_with_encoding, {}) diff --git a/xarray/tests/test_backends.py b/xarray/tests/test_backends.py index 76ed7ad57f8..638ab3584a6 100644 --- a/xarray/tests/test_backends.py +++ b/xarray/tests/test_backends.py @@ -1617,7 +1617,7 @@ def test_write_persistence_modes(self): with self.roundtrip(original, save_kwargs={'mode': 'a'}) as actual: assert_identical(original, actual) - ds, ds_to_append, ds_with_new_var = create_append_test_data() + ds, ds_to_append, _ = create_append_test_data() # check append mode for append write with self.create_zarr_target() as store_target: @@ -1626,28 +1626,6 @@ def test_write_persistence_modes(self): original = xr.concat([ds, ds_to_append], dim='time') assert_identical(original, xr.open_zarr(store_target)) - # check failure when append_dim not set - with pytest.raises(ValueError): - with self.create_zarr_target() as store_target: - ds.to_zarr(store_target, mode='w') - ds_to_append.to_zarr(store_target, mode='a') - - # check failure when append_dim not valid - with pytest.raises(ValueError): - with self.create_zarr_target() as store_target: - ds.to_zarr(store_target, mode='w') - ds_to_append.to_zarr(store_target, mode='a', - append_dim='notvalid') - - # check append mode for new variable - with self.create_zarr_target() as store_target: - xr.concat([ds, ds_to_append], dim='time').to_zarr(store_target, - mode='w') - ds_with_new_var.to_zarr(store_target, mode='a') - combined = xr.concat([ds, ds_to_append], dim='time') - combined['new_var'] = ds_with_new_var['new_var'] - assert_identical(combined, xr.open_zarr(store_target)) - def test_compressor_encoding(self): original = create_test_data() # specify a custom compressor @@ -1676,17 +1654,82 @@ def test_encoding_kwarg_fixed_width_string(self): def test_dataset_caching(self): super(CFEncodedBase, self).test_dataset_caching() - @pytest.mark.xfail(reason="Zarr stores can not be appended to") def test_append_write(self): - super(CFEncodedBase, self).test_append_write() + ds, ds_to_append, _ = create_append_test_data() + with self.create_zarr_target() as store_target: + ds.to_zarr(store_target, mode='w') + ds_to_append.to_zarr(store_target, mode='a', append_dim='time') + original = xr.concat([ds, ds_to_append], dim='time') + assert_identical(original, xr.open_zarr(store_target)) @pytest.mark.xfail(reason="Zarr stores can not be appended to") def test_append_overwrite_values(self): super(CFEncodedBase, self).test_append_overwrite_values() - @pytest.mark.xfail(reason="Zarr stores can not be appended to") def test_append_with_invalid_dim_raises(self): - super(CFEncodedBase, self).test_append_with_invalid_dim_raises() + + ds, ds_to_append, _ = create_append_test_data() + + # check failure when append_dim not valid + with pytest.raises(ValueError): + with self.create_zarr_target() as store_target: + ds.to_zarr(store_target, mode='w') + ds_to_append.to_zarr(store_target, mode='a', + append_dim='notvalid') + + def test_append_with_append_dim_not_set_raises(self): + + ds, ds_to_append, _ = create_append_test_data() + + # check failure when append_dim not set + with pytest.raises(ValueError): + with self.create_zarr_target() as store_target: + ds.to_zarr(store_target, mode='w') + ds_to_append.to_zarr(store_target, mode='a') + + def test_append_with_existing_encoding_raises(self): + + ds, ds_to_append, _ = create_append_test_data() + + # check failure when providing encoding to existing variable + with pytest.raises(ValueError): + with self.create_zarr_target() as store_target: + ds.to_zarr(store_target, mode='w') + ds_to_append.to_zarr(store_target, mode='a', + append_dim='time', + encoding={'da': {'compressor': None}}) + + def test_check_encoding_is_consistent_after_append(self): + + ds, ds_to_append, _ = create_append_test_data() + + # check encoding consistency + with self.create_zarr_target() as store_target: + import zarr + compressor = zarr.Blosc() + encoding = {'da': {'compressor': compressor}} + ds.to_zarr(store_target, mode='w', encoding=encoding) + ds_to_append.to_zarr(store_target, mode='a', append_dim='time') + actual_ds = xr.open_zarr(store_target) + actual_encoding = actual_ds['da'].encoding['compressor'] + assert actual_encoding.get_config() == compressor.get_config() + assert_identical( + xr.open_zarr(store_target).compute(), + xr.concat([ds, ds_to_append], dim='time') + ) + + def test_append_with_new_variable(self): + + ds, ds_to_append, ds_with_new_var = create_append_test_data() + + # check append mode for new variable + with self.create_zarr_target() as store_target: + xr.concat([ds, ds_to_append], dim='time').to_zarr(store_target, + mode='w') + ds_with_new_var.to_zarr(store_target, mode='a') + combined = xr.concat([ds, ds_to_append], dim='time') + combined['new_var'] = ds_with_new_var['new_var'] + assert_identical(combined, xr.open_zarr(store_target)) def test_to_zarr_compute_false_roundtrip(self): from dask.delayed import Delayed @@ -1710,7 +1753,7 @@ def test_to_zarr_compute_false_roundtrip(self): def test_to_zarr_append_compute_false_roundtrip(self): from dask.delayed import Delayed - ds, ds_to_append, ds_with_new_var = create_append_test_data() + ds, ds_to_append, _ = create_append_test_data() ds, ds_to_append = ds.chunk(), ds_to_append.chunk() with self.create_zarr_target() as store: