Optimize setitem with chunk equal to fill_value, round 2

zarr/core.py

@@ -33,6 +33,7 @@
 from zarr.util import (
     InfoReporter,
     check_array_shape,
+    flatten,
     human_readable_size,
     is_total_slice,
     nolock,
@@ -74,6 +75,12 @@ class Array:
         If True and while the chunk_store is a FSStore and the compresion used
         is Blosc, when getting data from the array chunks will be partially
         read and decompressed when possible.
+    write_empty_chunks : bool, optional
+        Determines chunk writing behavior for chunks filled with `fill_value` ("empty" chunks).
+        If True (default), all chunks will be written regardless of their contents.
+        If False, empty chunks will not be written, and the `store` entry for
+        the chunk key of an empty chunk will be deleted. Note that setting this option to False
+        will incur additional overhead per chunk write.
 
         .. versionadded:: 2.7
 
@@ -138,6 +145,7 @@ def __init__(
         cache_metadata=True,
         cache_attrs=True,
         partial_decompress=False,
+        write_empty_chunks=True,
     ):
         # N.B., expect at this point store is fully initialized with all
         # configuration metadata fully specified and normalized
@@ -154,6 +162,7 @@ def __init__(
         self._cache_metadata = cache_metadata
         self._is_view = False
         self._partial_decompress = partial_decompress
+        self._write_empty_chunks = write_empty_chunks
 
         # initialize metadata
         self._load_metadata()
@@ -1586,6 +1595,17 @@ def _set_basic_selection_zd(self, selection, value, fields=None):
         else:
             chunk[selection] = value
 
+        # clear chunk if it only contains the fill value
+        if self._chunk_isempty(chunk):
+            try:
+                del self.chunk_store[ckey]
+                return
+            except KeyError:
+                return
+            except Exception:
+                # deleting failed, fallback to overwriting
+                pass
+
         # encode and store
         cdata = self._encode_chunk(chunk)
         self.chunk_store[ckey] = cdata
@@ -1859,9 +1879,38 @@ def _chunk_setitems(self, lchunk_coords, lchunk_selection, values, fields=None):
         ckeys = [self._chunk_key(co) for co in lchunk_coords]
         cdatas = [self._process_for_setitem(key, sel, val, fields=fields)
                   for key, sel, val in zip(ckeys, lchunk_selection, values)]
-        values = {k: v for k, v in zip(ckeys, cdatas)}
+        values = {}
+        if not self._write_empty_chunks:
+            for ckey, cdata in zip(ckeys, cdatas):
+                if self._chunk_isempty(cdata) and not self._chunk_delitem(ckey):
+                    values[ckey] = self._encode_chunk(cdata)
+        else:
+            values = dict(zip(ckeys, map(self._encode_chunk, cdatas)))
         self.chunk_store.setitems(values)
 
+    def _chunk_isempty(self, chunk):
+        if self.dtype == 'object':
+            # we have to flatten the result of np.equal to handle outputs like
+            # [np.array([True,True]), True, True]
+            is_empty = all(flatten(np.equal(chunk, self.fill_value, dtype='object')))
+        else:
+            is_empty = np.all(chunk == self._fill_value)
+        return is_empty
+
+    def _chunk_delitem(self, ckey):
+        """
+        Attempt to delete the value associated with ckey.
+        Returns True if deletion succeeds or KeyError is raised.
+        Returns False if any other exception is raised.
+        """
+        try:
+            del self.chunk_store[ckey]
+            return True
+        except KeyError:
+            return True
+        except Exception:
+            return False
+
     def _chunk_setitem(self, chunk_coords, chunk_selection, value, fields=None):
         """Replace part or whole of a chunk.
 
@@ -1889,10 +1938,17 @@ def _chunk_setitem(self, chunk_coords, chunk_selection, value, fields=None):
                                        fields=fields)
 
     def _chunk_setitem_nosync(self, chunk_coords, chunk_selection, value, fields=None):
+        do_store = True
         ckey = self._chunk_key(chunk_coords)
         cdata = self._process_for_setitem(ckey, chunk_selection, value, fields=fields)
+
+        # clear chunk if it only contains the fill value
+        if (not self._write_empty_chunks) and self._chunk_isempty(cdata):
+            do_store = not self._chunk_delitem(ckey)
+
         # store
-        self.chunk_store[ckey] = cdata
+        if do_store:
+            self.chunk_store[ckey] = self._encode_chunk(cdata)
 
     def _process_for_setitem(self, ckey, chunk_selection, value, fields=None):
         if is_total_slice(chunk_selection, self._chunks) and not fields:
@@ -1948,8 +2004,7 @@ def _process_for_setitem(self, ckey, chunk_selection, value, fields=None):
             else:
                 chunk[chunk_selection] = value
 
-        # encode chunk
-        return self._encode_chunk(chunk)
+        return chunk
 
     def _chunk_key(self, chunk_coords):
         return self._key_prefix + '.'.join(map(str, chunk_coords))
@@ -2169,7 +2224,8 @@ def hexdigest(self, hashname="sha1"):
 
     def __getstate__(self):
         return (self._store, self._path, self._read_only, self._chunk_store,
-                self._synchronizer, self._cache_metadata, self._attrs.cache)
+                self._synchronizer, self._cache_metadata, self._attrs.cache,
+                self._partial_decompress, self._write_empty_chunks)
 
     def __setstate__(self, state):
         self.__init__(*state)

zarr/creation.py

@@ -20,7 +20,7 @@ def create(shape, chunks=True, dtype=None, compressor='default',
            fill_value=0, order='C', store=None, synchronizer=None,
            overwrite=False, path=None, chunk_store=None, filters=None,
            cache_metadata=True, cache_attrs=True, read_only=False,
-           object_codec=None, dimension_separator=None, **kwargs):
+           object_codec=None, dimension_separator=None, write_empty_chunks=True, **kwargs):
     """Create an array.
 
     Parameters
@@ -70,6 +70,12 @@ def create(shape, chunks=True, dtype=None, compressor='default',
     dimension_separator : {'.', '/'}, optional
         Separator placed between the dimensions of a chunk.
         .. versionadded:: 2.8
+    write_empty_chunks : bool, optional
+        Determines chunk writing behavior for chunks filled with `fill_value` ("empty" chunks).
+        If True (default), all chunks will be written regardless of their contents.
+        If False, empty chunks will not be written, and the `store` entry for
+        the chunk key of an empty chunk will be deleted. Note that setting this option to False
+        will incur additional overhead per chunk write.
 
     Returns
     -------
@@ -140,7 +146,8 @@ def create(shape, chunks=True, dtype=None, compressor='default',
 
     # instantiate array
     z = Array(store, path=path, chunk_store=chunk_store, synchronizer=synchronizer,
-              cache_metadata=cache_metadata, cache_attrs=cache_attrs, read_only=read_only)
+              cache_metadata=cache_metadata, cache_attrs=cache_attrs, read_only=read_only,
+              write_empty_chunks=write_empty_chunks)
 
     return z
 
@@ -396,6 +403,7 @@ def open_array(
     chunk_store=None,
     storage_options=None,
     partial_decompress=False,
+    write_empty_chunks=True,
     **kwargs
 ):
     """Open an array using file-mode-like semantics.
@@ -450,8 +458,12 @@ def open_array(
         If True and while the chunk_store is a FSStore and the compresion used
         is Blosc, when getting data from the array chunks will be partially
         read and decompressed when possible.
-
-        .. versionadded:: 2.7
+    write_empty_chunks : bool, optional
+        Determines chunk writing behavior for chunks filled with `fill_value` ("empty" chunks).
+        If True (default), all chunks will be written regardless of their contents.
+        If False, empty chunks will not be written, and the `store` entry for
+        the chunk key of an empty chunk will be deleted. Note that setting this option to False
+        will incur additional overhead per chunk write.
 
     Returns
     -------
@@ -541,7 +553,7 @@ def open_array(
     # instantiate array
     z = Array(store, read_only=read_only, synchronizer=synchronizer,
               cache_metadata=cache_metadata, cache_attrs=cache_attrs, path=path,
-              chunk_store=chunk_store)
+              chunk_store=chunk_store, write_empty_chunks=write_empty_chunks)
 
     return z
 

zarr/tests/test_core.py

@@ -1597,6 +1597,37 @@ def test_nbytes_stored(self):
         assert -1 == z.nbytes_stored
 
 
+class TestArrayWithoutEmptyWrites(TestArray):
+
+    @staticmethod
+    def create_array(read_only=False, **kwargs):
+        path = mkdtemp()
+        atexit.register(shutil.rmtree, path)
+        store = DirectoryStore(path)
+        cache_metadata = kwargs.pop('cache_metadata', True)
+        cache_attrs = kwargs.pop('cache_attrs', True)
+        kwargs.setdefault('compressor', Zlib(1))
+        init_array(store, **kwargs)
+        return Array(store, read_only=read_only, cache_metadata=cache_metadata,
+                     cache_attrs=cache_attrs, write_empty_chunks=False)
+
+        def test_nchunks_initialized(self):
+            for fill_value in -1, 0, 1, 10:
+                z = self.create_array(shape=100, chunks=10, fill_value=fill_value)
+                assert 0 == z.nchunks_initialized
+                # manually put something into the store to confuse matters
+                z.store['foo'] = b'bar'
+                assert 0 == z.nchunks_initialized
+                z[:] = 42
+                assert 10 == z.nchunks_initialized
+                z[:] = fill_value
+                assert 0 == z.nchunks_initialized
+                z[0] = 42
+                assert 1 == z.nchunks_initialized
+                if hasattr(z.store, 'close'):
+                    z.store.close()
+
+
 class TestArrayWithDirectoryStore(TestArray):
 
     @staticmethod

zarr/util.py

@@ -9,13 +9,22 @@
 import numpy as np
 from asciitree import BoxStyle, LeftAligned
 from asciitree.traversal import Traversal
+from collections.abc import Iterable
 from numcodecs.compat import ensure_ndarray, ensure_text
 from numcodecs.registry import codec_registry
 from numcodecs.blosc import cbuffer_sizes, cbuffer_metainfo
 
 from typing import Any, Callable, Dict, Optional, Tuple, Union
 
 
+def flatten(arg: Iterable) -> Iterable:
+    for element in arg:
+        if isinstance(element, Iterable) and not isinstance(element, (str, bytes)):
+            yield from flatten(element)
+        else:
+            yield element
+
+
 # codecs to use for object dtype convenience API
 object_codecs = {
     str.__name__: 'vlen-utf8',