Repository: tomwhite/cubed (python)

Issue:
Don't materialize `full` or offsets arrays
Arrays whose values are a function of the index do not need to be materialized - either in memory or on disk. These include arrays where every entry is the same (created with `full`), and arrays of offsets (used in `map_blocks` to provide block IDs to the map function).

Zarr arrays created with `full` are not expensive to create since they specify `write_empty_chunks=False` so no chunks are written - just the metadata. However, offsets are materialized to disk: https://github.com/tomwhite/cubed/blob/e21786591d9832a85f2e492641ad2061bdb8c14a/cubed/core/ops.py#L431-L432

Although the array is small since its size is the number of chunks, there is a scalability issue due to the fact that every chunk is stored in a separate file in the Zarr store.

These problems could be solved by writing an array implementation that implements indexing on the fly: either by returning a subarray of fill values for the case of `full`, or `np.ravel_multi_index` to turn an index into an offset for an array of offsets.

Pull request description:
Introduce `VirtualFullArray` 
…and use it in `full` so no Zarr arrays are materialized

Fixes #247

Patch:
diff --git a/cubed/array_api/creation_functions.py b/cubed/array_api/creation_functions.py
index 1a717211..8a4cf7be 100644
--- a/cubed/array_api/creation_functions.py
+++ b/cubed/array_api/creation_functions.py
@@ -6,8 +6,8 @@ from zarr.util import normalize_shape
 from cubed.core import Plan, gensym, map_blocks
 from cubed.core.ops import map_direct
 from cubed.core.plan import new_temp_path
-from cubed.storage.virtual import virtual_empty, virtual_offsets
-from cubed.storage.zarr import lazy_from_array, lazy_full
+from cubed.storage.virtual import virtual_empty, virtual_full, virtual_offsets
+from cubed.storage.zarr import lazy_from_array
 from cubed.utils import to_chunksize
 from cubed.vendor.dask.array.core import normalize_chunks
 
@@ -79,9 +79,9 @@ def asarray(
 
 
 def empty(shape, *, dtype=None, device=None, chunks="auto", spec=None) -> "Array":
-    if dtype is None:
-        dtype = np.float64
-    return full(shape, None, dtype=dtype, device=device, chunks=chunks, spec=spec)
+    return empty_virtual_array(
+        shape, dtype=dtype, device=device, chunks=chunks, spec=spec, hidden=False
+    )
 
 
 def empty_like(x, /, *, dtype=None, device=None, chunks=None, spec=None) -> "Array":
@@ -89,7 +89,7 @@ def empty_like(x, /, *, dtype=None, device=None, chunks=None, spec=None) -> "Arr
 
 
 def empty_virtual_array(
-    shape, *, dtype=None, device=None, chunks="auto", spec=None
+    shape, *, dtype=None, device=None, chunks="auto", spec=None, hidden=True
 ) -> "Array":
     if dtype is None:
         dtype = np.float64
@@ -100,7 +100,7 @@ def empty_virtual_array(
 
     from .array_object import Array
 
-    plan = Plan._new(name, "empty", target, hidden=True)
+    plan = Plan._new(name, "empty", target, hidden=hidden)
     return Array(name, target, spec, plan)
 
 
@@ -140,8 +140,6 @@ def _eye(x, *arrays, k=None, chunksize=None, block_id=None):
 def full(
     shape, fill_value, *, dtype=None, device=None, chunks="auto", spec=None
 ) -> "Array":
-    # write to zarr
-    # note that write_empty_chunks=False means no chunks are written to disk, so it is very efficient to create large arrays
     shape = normalize_shape(shape)
     if dtype is None:
         # check bool first since True/False are instances of int and float
@@ -155,15 +153,7 @@ def full(
             raise TypeError("Invalid input to full")
     chunksize = to_chunksize(normalize_chunks(chunks, shape=shape, dtype=dtype))
     name = gensym()
-    zarr_path = new_temp_path(name=name, spec=spec)
-    target = lazy_full(
-        shape,
-        fill_value,
-        dtype=dtype,
-        chunks=chunksize,
-        store=zarr_path,
-        write_empty_chunks=False,
-    )
+    target = virtual_full(shape, fill_value, dtype=dtype, chunks=chunksize)
 
     from .array_object import Array
 
diff --git a/cubed/storage/virtual.py b/cubed/storage/virtual.py
index ee93e6c7..e1408bd6 100644
--- a/cubed/storage/virtual.py
+++ b/cubed/storage/virtual.py
@@ -1,4 +1,5 @@
 from numbers import Integral
+from typing import Any
 
 import numpy as np
 import zarr
@@ -31,6 +32,45 @@ class VirtualEmptyArray:
         indexer = BasicIndexer(key, self.template)
         return np.empty(indexer.shape, dtype=self.dtype)
 
+    @property
+    def oindex(self):
+        return self.template.oindex
+
+
+class VirtualFullArray:
+    """An array that is never materialized (in memory or on disk) and contains a single fill value."""
+
+    def __init__(
+        self,
+        shape: T_Shape,
+        dtype: T_DType,
+        chunks: T_RegularChunks,
+        fill_value: Any = None,
+    ):
+        # use an empty in-memory Zarr array as a template since it normalizes its properties
+        template = zarr.full(
+            shape,
+            fill_value,
+            dtype=dtype,
+            chunks=chunks,
+            store=zarr.storage.MemoryStore(),
+        )
+        self.shape = template.shape
+        self.dtype = template.dtype
+        self.chunks = template.chunks
+        self.template = template
+        self.fill_value = fill_value
+
+    def __getitem__(self, key):
+        if not isinstance(key, tuple):
+            key = (key,)
+        indexer = BasicIndexer(key, self.template)
+        return np.full(indexer.shape, fill_value=self.fill_value, dtype=self.dtype)
+
+    @property
+    def oindex(self):
+        return self.template.oindex
+
 
 class VirtualOffsetsArray:
     """An array that is never materialized (in memory or on disk) and contains sequentially incrementing integers."""
@@ -74,5 +114,16 @@ def virtual_empty(
     return VirtualEmptyArray(shape, dtype, chunks, **kwargs)
 
 
+def virtual_full(
+    shape: T_Shape,
+    fill_value: Any,
+    *,
+    dtype: T_DType,
+    chunks: T_RegularChunks,
+    **kwargs,
+) -> VirtualFullArray:
+    return VirtualFullArray(shape, dtype, chunks, fill_value, **kwargs)
+
+
 def virtual_offsets(shape: T_Shape) -> VirtualOffsetsArray:
     return VirtualOffsetsArray(shape)


Test patch:
diff --git a/cubed/tests/test_optimization.py b/cubed/tests/test_optimization.py
index b772ce1d..b232077d 100644
--- a/cubed/tests/test_optimization.py
+++ b/cubed/tests/test_optimization.py
@@ -61,7 +61,7 @@ def test_no_fusion(spec):
     c = xp.positive(b)
     d = xp.equal(b, c)
 
-    num_created_arrays = 4  # a, b, c, d
+    num_created_arrays = 3  # b, c, d
     assert d.plan.num_tasks(optimize_graph=False) == num_created_arrays + 3
     assert d.plan.num_tasks(optimize_graph=True) == num_created_arrays + 3
 
@@ -81,7 +81,7 @@ def test_no_fusion_multiple_edges(spec):
     # this should not be fused under the current logic
     d = xp.equal(b, c)
 
-    num_created_arrays = 3  # a, c, d
+    num_created_arrays = 2  # c, d
     assert d.plan.num_tasks(optimize_graph=False) == num_created_arrays + 2
     assert d.plan.num_tasks(optimize_graph=True) == num_created_arrays + 2