✨ Using the Dask arrays is now optional.

Author: fbriol

.pre-commit-config.yaml

@@ -15,7 +15,7 @@ repos:
       - id: trailing-whitespace
         exclude: conda/meta.yaml
   - repo: https://github.com/asottile/pyupgrade
-    rev: "v3.3.1"
+    rev: "v3.4.0"
     hooks:
     - id: pyupgrade
       args: [--py38-plus]
@@ -43,7 +43,7 @@ repos:
         additional_dependencies:
           - toml
   - repo: https://github.com/myint/docformatter
-    rev: "v1.6.3"
+    rev: "v1.6.5"
     hooks:
     - id: docformatter
   - repo: https://github.com/codespell-project/codespell

conftest.py

@@ -9,7 +9,7 @@
 """
 
 
-def pytest_addoption(parser):
+def pytest_addoption(parser) -> None:
     """Add command line options to pytest."""
     parser.addoption(
         '--s3',

docs/source/api.rst

@@ -31,6 +31,19 @@ Merging of existing datasets in a partition.
   zcollection.merging.time_series
   zcollection.merging.period
 
+Variable
+========
+
+Variables handled by the datasets. These objects manage access to the data
+stored in the collection.
+
+.. autosummary::
+  :toctree: _generated/
+
+  zcollection.variable.abc
+  zcollection.variable.array
+  zcollection.variable.delayed_array
+
 Collection
 ==========
 
@@ -45,7 +58,6 @@ Collection
   zcollection.meta
   zcollection.sync
   zcollection.type_hints
-  zcollection.variable
   zcollection.view
 
 Indexing

examples/ex_collection.py

@@ -7,6 +7,7 @@
 """
 from __future__ import annotations
 
+from typing import Iterator
 import datetime
 import pprint
 
@@ -23,19 +24,21 @@
 # ---------------------------------
 #
 # Before we create our first collection, we will create a dataset to record.
-def create_dataset():
+def create_dataset() -> zcollection.Dataset:
     """Create a dataset to record."""
-    generator = zcollection.tests.data.create_test_dataset_with_fillvalue()
+    generator: Iterator[zcollection.Dataset] = \
+        zcollection.tests.data.create_test_dataset_with_fillvalue()
     return next(generator)
 
 
-ds = create_dataset()
-ds.to_xarray()
+zds: zcollection.Dataset | None = create_dataset()
+assert zds is not None
+zds.to_xarray()
 
 # %%
 # We will create the file system that we will use. In this example, a file
 # system in memory.
-fs = fsspec.filesystem('memory')
+fs: fsspec.AbstractFileSystem = fsspec.filesystem('memory')
 
 # %%
 # Finally we create a local dask cluster using only threads in order to work
@@ -54,11 +57,8 @@ def create_dataset():
 
 # %%
 # Finally, we create our collection:
-collection = zcollection.create_collection('time',
-                                           ds,
-                                           partition_handler,
-                                           '/my_collection',
-                                           filesystem=fs)
+collection: zcollection.Collection = zcollection.create_collection(
+    'time', zds, partition_handler, '/my_collection', filesystem=fs)
 
 # %%
 # .. note::
@@ -75,7 +75,7 @@ def create_dataset():
 
 # %%
 # Now that the collection has been created, we can insert new records.
-collection.insert(ds)
+collection.insert(zds)
 
 # %%
 # .. note::
@@ -103,9 +103,15 @@ def create_dataset():
 # To load the dataset call the method
 # :py:meth:`load<zcollection.collection.Collection.load>` on the instance.  By
 # default, the method loads all partitions stored in the collection.
-collection.load()
+collection.load(delayed=True)
 
 # %%
+# .. note::
+#
+#   By default, the data is loaded as a :py:class:`dask.array<da.Array>`. It is
+#   possible to load the data as a :py:class:`numpy.ndarray` by specifying the
+#   parameter ``delayed=False``.
+#
 # You can also filter the partitions to be considered by filtering the
 # partitions using keywords used for partitioning in a valid Python expression.
 collection.load(filters='year == 2000 and month == 2')
@@ -145,13 +151,13 @@ def create_dataset():
 # %%
 # The :py:meth:`add_variable<zcollection.collection.Collection.add_variable>`
 # method allows you to add a new variable to the collection.
-collection.add_variable(ds.metadata().variables['var2'])
+collection.add_variable(zds.metadata().variables['var2'])
 
 # %%
 # The newly created variable is initialized with its default value.
-ds = collection.load()
-assert ds is not None
-ds.variables['var2'].values
+zds = collection.load()
+assert zds is not None
+zds.variables['var2'].values
 
 
 # %%
@@ -161,19 +167,28 @@ def create_dataset():
 #
 # In this example, we will alter the variable ``var2`` by setting it to 1
 # anywhere the variable ``var1`` is defined.
-def ones(ds):
+def ones(zds) -> dict[str, numpy.ndarray]:
     """Returns a variable with ones everywhere."""
-    return dict(var2=ds.variables['var1'].values * 0 + 1)
+    return dict(var2=zds.variables['var1'].values * 0 + 1)
 
 
 collection.update(ones)  # type: ignore[arg-type]
 
-ds = collection.load()
-assert ds is not None
-ds.variables['var2'].values
+zds = collection.load()
+assert zds is not None
+zds.variables['var2'].values
 
 
 # %%
+# ..note::
+#
+#   The method :py:meth:`update<zcollection.collection.Collection.update>`
+#   supports the ``delayed`` parameter. If ``delayed=True``, the function
+#   ``ones`` is applied to each partition using a Dask array as container
+#   for the variables data stored in the provided dataset. This is the default
+#   behavior. If ``delayed=False``, the function ``ones`` is applied to each
+#   partition using a Numpy array as container.
+#
 # Sometime is it important to know the values of the neighboring partitions.
 # This can be done using the
 # :py:meth:`update<zcollection.collection.Collection.update>` method with the
@@ -188,7 +203,7 @@ def ones(ds):
 #   start of the slice is 0, it means that the left partition is missing. If the
 #   stop of the slice is equal to the length of the given dataset, it means that
 #   the right partition is missing.
-def twos(ds, partition_info: tuple[str, slice]):
+def twos(ds, partition_info: tuple[str, slice]) -> dict[str, numpy.ndarray]:
     """Returns a variable with twos everywhere if the partition is surrounded
     by partitions on both sides, -1 if the left partition is missing and -2 if
     the right partition is missing."""
@@ -206,9 +221,9 @@ def twos(ds, partition_info: tuple[str, slice]):
 
 collection.update(twos, depth=1)  # type: ignore[arg-type]
 
-ds = collection.load()
-assert ds is not None
-ds.variables['var2'].values
+zds = collection.load()
+assert zds is not None
+zds.variables['var2'].values
 
 # %%
 # Map a function over the collection

examples/ex_indexing.py

@@ -14,7 +14,7 @@
 
 import zcollection
 import zcollection.indexing
-import zcollection.tests.data
+import zcollection.partitioning.tests.data
 
 # %%
 # Initialization of the environment
@@ -32,18 +32,18 @@
 #
 # For this latest example, we will index another data set. This one contains
 # measurements of a fictitious satellite on several half-orbits.
-ds = zcollection.Dataset.from_xarray(
-    zcollection.tests.data.create_test_sequence(5, 20, 10))
-ds
+zds: zcollection.Dataset = zcollection.Dataset.from_xarray(
+    zcollection.partitioning.tests.data.create_test_sequence(5, 20, 10))
+print(zds)
 
 # %%
-collection = zcollection.create_collection(
+collection: zcollection.Collection = zcollection.create_collection(
     'time',
-    ds,
+    zds,
     zcollection.partitioning.Date(('time', ), 'M'),
     partition_base_dir='/one_other_collection',
     filesystem=fs)
-collection.insert(ds, merge_callable=zcollection.merging.merge_time_series)
+collection.insert(zds, merge_callable=zcollection.merging.merge_time_series)
 
 # %%
 # Here we have created a collection partitioned by month.
@@ -87,7 +87,7 @@ def split_half_orbit(
 # Now we will compute these constant parts from a dataset contained in a
 # partition.
 def _half_orbit(
-    ds: zcollection.Dataset,
+    zds: zcollection.Dataset,
     *args,
     **kwargs,
 ) -> numpy.ndarray:
@@ -100,8 +100,8 @@ def _half_orbit(
     """
     pass_number_varname = kwargs.pop('pass_number', 'pass_number')
     cycle_number_varname = kwargs.pop('cycle_number', 'cycle_number')
-    pass_number = ds.variables[pass_number_varname].values
-    cycle_number = ds.variables[cycle_number_varname].values
+    pass_number = zds.variables[pass_number_varname].values
+    cycle_number = zds.variables[cycle_number_varname].values
 
     generator = ((
         i0,
@@ -141,7 +141,7 @@ def dtype(cls, /, **kwargs) -> List[Tuple[str, str]]:
     def create(
         cls,
         path: Union[pathlib.Path, str],
-        ds: zcollection.Collection,
+        zds: zcollection.Collection,
         filesystem: Optional[fsspec.AbstractFileSystem] = None,
         **kwargs,
     ) -> 'HalfOrbitIndexer':
@@ -155,13 +155,13 @@ def create(
             The created index.
         """
         return super()._create(path,
-                               ds,
+                               zds,
                                meta=dict(attribute=b'value'),
                                filesystem=filesystem)  # type: ignore
 
     def update(
         self,
-        ds: zcollection.Collection,
+        zds: zcollection.Collection,
         partition_size: Optional[int] = None,
         npartitions: Optional[int] = None,
         **kwargs,
@@ -177,15 +177,18 @@ def update(
             pass_number: The name of the pass number variable stored in the
                 collection. Defaults to "pass_number".
         """
-        super()._update(ds, _half_orbit, partition_size, npartitions, **kwargs)
+        super()._update(zds, _half_orbit, partition_size, npartitions,
+                        **kwargs)
 
 
 # %%
 # Using the index
 # ---------------
 #
 # Now we can create our index and fill it.
-indexer = HalfOrbitIndexer.create('/index.parquet', collection, filesystem=fs)
+indexer: HalfOrbitIndexer = HalfOrbitIndexer.create('/index.parquet',
+                                                    collection,
+                                                    filesystem=fs)
 indexer.update(collection)
 
 # The following command allows us to view the information stored in our index:
@@ -195,9 +198,12 @@ def update(
 
 # %%
 # This index can now be used to load a part of a collection.
-selection = collection.load(indexer=indexer.query(dict(pass_number=[1, 2])))
+selection: zcollection.Dataset | None = collection.load(
+    indexer=indexer.query(dict(pass_number=[1, 2])),
+    delayed=False,
+)
 assert selection is not None
-selection.to_xarray().compute()
+selection.to_xarray()
 
 # %%
 # Close the local cluster to avoid printing warning messages in the other