resolves #2 by adding support for axis argument to resize; also fixes bug in >1D resize

Author: alimanfoo

zarr/core.py

@@ -45,7 +45,8 @@ def empty(shape, chunks, dtype=None, cname=None, clevel=None, shuffle=None,
 
 def zeros(shape, chunks, dtype=None, cname=None, clevel=None, shuffle=None,
           synchronized=True):
-    """Create an array filled with zeros.
+    """Create an array, with zero being used as the default value for
+    uninitialised portions of the array.
 
     Parameters
     ----------
@@ -80,7 +81,8 @@ def zeros(shape, chunks, dtype=None, cname=None, clevel=None, shuffle=None,
 
 def ones(shape, chunks, dtype=None, cname=None, clevel=None, shuffle=None,
          synchronized=True):
-    """Create an array filled with ones.
+    """Create an array, with one being used as the default value for
+    uninitialised portions of the array.
 
     Parameters
     ----------
@@ -116,7 +118,8 @@ def ones(shape, chunks, dtype=None, cname=None, clevel=None, shuffle=None,
 
 def full(shape, chunks, fill_value, dtype=None, cname=None, clevel=None,
          shuffle=None, synchronized=True):
-    """Create an array filled with `fill_value`.
+    """Create an array, with `fill_value` being used as the default value for
+    uninitialised portions of the array.
 
     Parameters
     ----------
@@ -174,6 +177,8 @@ def array(data, chunks=None, dtype=None, cname=None, clevel=None,
     synchronized : bool, optional
         If True, each chunk will be protected with a lock to prevent data
         collision during write operations.
+    fill_value : object
+        Default value to use for uninitialised portions of the array.
 
     Returns
     -------

zarr/ext.pyx

@@ -583,7 +583,33 @@ cdef class Array:
         return r
 
     def resize(self, *args):
-        """TODO"""
+        """Resize the array by growing or shrinking one or more dimensions.
+
+        Parameters
+        ----------
+        args : int or sequence of ints
+            New shape to resize to.
+
+        Notes
+        -----
+        This function can only be used to change the size of existing
+        dimensions, it cannot add or drop dimensions.
+
+        N.B., this function does *not* behave in the same way as the numpy
+        resize method on the ndarray class. Existing data are *not*
+        reorganised in any way. Axes are simply grown or shrunk. When
+        growing an axis, uninitialised portions of the array will appear to
+        contain the value of the `fill_value` attribute on this array,
+        and when shrinking an array any data beyond the new shape will be
+        lost (although see note below).
+
+        N.B., because of the way the underlying chunks are organised,
+        and in particular the fact that chunks may overhang the edge of the
+        array, the value of uninitialised portions of this array is not
+        guaranteed to respect the setting of the `fill_value` attribute when
+        shrinking then regrowing an array.
+
+        """
 
         # normalise new shape argument
         if len(args) == 1:
@@ -601,15 +627,24 @@ cdef class Array:
         new_shape = tuple(s if n is None else n
                           for s, n in zip(self.shape, new_shape))
 
-        # set new shape
-        self.shape = new_shape
+        # work-around Cython problems with accessing .shape attribute as tuple
+        old_cdata = np.asarray(self.cdata)
+
+        # remember old cdata shape
+        old_cdata_shape = old_cdata.shape
 
         # determine the new number and arrangement of chunks
         new_cdata_shape = tuple(int(np.ceil(s / c))
                                 for s, c in zip(new_shape, self.chunks))
 
-        # resize chunks array
-        self.cdata.resize(new_cdata_shape, refcheck=False)
+        # setup new chunks array
+        new_cdata = np.empty(new_cdata_shape, dtype=object)
+
+        # copy across any chunks to be kept
+        cdata_overlap = tuple(
+            slice(min(o, n)) for o, n in zip(old_cdata_shape, new_cdata_shape)
+        )
+        new_cdata[cdata_overlap] = old_cdata[cdata_overlap]
 
         # determine function for instantiating chunks
         if self.synchronized:
@@ -618,33 +653,64 @@ cdef class Array:
             create_chunk = Chunk
 
         # instantiate any new chunks as needed
-        self.cdata.flat = [create_chunk(self.chunks, dtype=self.dtype,
-                                        cname=self.cname,
-                                        clevel=self.clevel,
-                                        shuffle=self.shuffle,
-                                        fill_value=self.fill_value)
-                           if c == 0 else c
-                           for c in self.cdata.flat]
+        new_cdata.flat = [create_chunk(self.chunks, dtype=self.dtype,
+                                       cname=self.cname,
+                                       clevel=self.clevel,
+                                       shuffle=self.shuffle,
+                                       fill_value=self.fill_value)
+                          if c is None else c
+                          for c in new_cdata.flat]
 
-    def append(self, data):
-        """TODO"""
+        # set new shape
+        self.shape = new_shape
+
+        # set new chunks
+        self.cdata = new_cdata
+
+    def append(self, data, axis=0):
+        """Append `data` to `axis`.
+
+        Parameters
+        ----------
+        data : array_like
+            Data to be appended.
+        axis : int
+            Axis along which to append.
+
+        Notes
+        -----
+        The size of all dimensions other than `axis` must match between this
+        array and `data`.
+
+        """
 
         # ensure data is array-like
         if not hasattr(data, 'shape') or not hasattr(data, 'dtype'):
             data = np.asanyarray(data)
 
-        # ensure shapes are compatible for trailing dimensions
-        if self.shape[1:] != data.shape[1:]:
-            raise ValueError('shape not compatible')
+        # ensure shapes are compatible for non-append dimensions
+        self_shape_preserved = tuple(s for i, s in enumerate(self.shape)
+                                     if i != axis)
+        data_shape_preserved = tuple(s for i, s in enumerate(data.shape)
+                                     if i != axis)
+        if self_shape_preserved != data_shape_preserved:
+            raise ValueError('shapes not compatible')
 
         # remember old shape
         old_shape = self.shape
 
         # determine new shape
-        new_shape = (self.shape[0] + data.shape[0],) + self.shape[1:]
+        new_shape = tuple(
+            self.shape[i] if i != axis else self.shape[i] + data.shape[i]
+            for i in range(len(self.shape))
+        )
 
         # resize
         self.resize(new_shape)
 
         # store data
-        self[old_shape[0]:] = data
+        append_selection = tuple(
+            slice(None) if i != axis else slice(old_shape[i], new_shape[i])
+            for i in range(len(self.shape))
+        )
+        self[append_selection] = data

zarr/tests/test_ext_array.py

@@ -147,57 +147,71 @@ def test_array_2d():
 
 def test_resize_1d():
 
-    z = Array(105, chunks=10, dtype='i4')
+    z = Array(105, chunks=10, dtype='i4', fill_value=0)
+    a = np.arange(105, dtype='i4')
+    z[:] = a
     eq((105,), z.shape)
     eq((105,), z[:].shape)
     eq(np.dtype('i4'), z.dtype)
     eq(np.dtype('i4'), z[:].dtype)
     eq((10,), z.chunks)
+    assert_array_equal(a, z[:])
 
     z.resize(205)
     eq((205,), z.shape)
     eq((205,), z[:].shape)
     eq(np.dtype('i4'), z.dtype)
     eq(np.dtype('i4'), z[:].dtype)
     eq((10,), z.chunks)
+    assert_array_equal(a, z[:105])
+    assert_array_equal(np.zeros(100, dtype='i4'), z[105:])
 
     z.resize(55)
     eq((55,), z.shape)
     eq((55,), z[:].shape)
     eq(np.dtype('i4'), z.dtype)
     eq(np.dtype('i4'), z[:].dtype)
     eq((10,), z.chunks)
+    assert_array_equal(a[:55], z[:55])
 
 
 def test_resize_2d():
 
-    z = Array((105, 105), chunks=(10, 10), dtype='i4')
+    z = Array((105, 105), chunks=(10, 10), dtype='i4', fill_value=0)
+    a = np.arange(105*105, dtype='i4').reshape((105, 105))
+    z[:] = a
     eq((105, 105), z.shape)
     eq((105, 105), z[:].shape)
     eq(np.dtype('i4'), z.dtype)
     eq(np.dtype('i4'), z[:].dtype)
     eq((10, 10), z.chunks)
+    assert_array_equal(a, z[:])
 
     z.resize((205, 205))
     eq((205, 205), z.shape)
     eq((205, 205), z[:].shape)
     eq(np.dtype('i4'), z.dtype)
     eq(np.dtype('i4'), z[:].dtype)
     eq((10, 10), z.chunks)
+    assert_array_equal(a, z[:105, :105])
+    assert_array_equal(np.zeros((100, 205), dtype='i4'), z[105:, :])
+    assert_array_equal(np.zeros((205, 100), dtype='i4'), z[:, 105:])
 
     z.resize((55, 55))
     eq((55, 55), z.shape)
     eq((55, 55), z[:].shape)
     eq(np.dtype('i4'), z.dtype)
     eq(np.dtype('i4'), z[:].dtype)
     eq((10, 10), z.chunks)
+    assert_array_equal(a[:55, :55], z[:])
 
     z.resize((55, 1))
     eq((55, 1), z.shape)
     eq((55, 1), z[:].shape)
     eq(np.dtype('i4'), z.dtype)
     eq(np.dtype('i4'), z[:].dtype)
     eq((10, 10), z.chunks)
+    assert_array_equal(a[:55, :1], z[:])
 
 
 def test_append_1d():
@@ -236,3 +250,22 @@ def test_append_2d():
     eq(e.dtype, z.dtype)
     eq((10, 10), z.chunks)
     assert_array_equal(e, z[:])
+
+
+def test_append_2d_axis():
+
+    a = np.arange(105*105, dtype='i4').reshape((105, 105))
+    z = Array(a.shape, chunks=(10, 10), dtype=a.dtype)
+    z[:] = a
+    eq(a.shape, z.shape)
+    eq(a.dtype, z.dtype)
+    eq((10, 10), z.chunks)
+    assert_array_equal(a, z[:])
+
+    b = np.arange(105*105, 2*105*105, dtype='i4').reshape((105, 105))
+    e = np.append(a, b, axis=1)
+    z.append(b, axis=1)
+    eq(e.shape, z.shape)
+    eq(e.dtype, z.dtype)
+    eq((10, 10), z.chunks)
+    assert_array_equal(e, z[:])