Merge pull request #158 from csiro-coasts/geometry-hash

Dataset Cache Keys

Author: mx-moth

docs/api/operations/cache.rst

@@ -0,0 +1,10 @@
+=========================
+emsarray.operations.cache
+=========================
+
+.. automodule:: emsarray.operations.cache
+
+Functions
+=========
+
+.. autofunction:: make_cache_key

docs/api/operations/index.rst

@@ -13,6 +13,10 @@ and behave the same across all supported conventions.
 
    ./*
 
+:doc:`cache`
+    These operations create hash keys of datasets for use
+    in caching data.
+
 :doc:`depth`
     These operations manipulate datasets with a depth axis,
     such as the output of ocean models.

docs/releases/development.rst

@@ -11,3 +11,6 @@ Next release (in development)
   but the trade off is worth the added security
   after the invalid polygons found in :pr:`154`
   (:pr:`156`).
+* Added new :mod:`emsarray.operations.cache` module
+  for generating cache keys based on dataset geometry.
+  (:issue:`153`, :pr:`158`).

src/emsarray/conventions/_base.py

@@ -1,6 +1,7 @@
 import abc
 import dataclasses
 import enum
+import hashlib
 import logging
 import warnings
 from collections.abc import Callable, Hashable, Iterable, Sequence
@@ -18,6 +19,7 @@
 from emsarray.compat.shapely import SpatialIndex
 from emsarray.exceptions import InvalidPolygonWarning, NoSuchCoordinateError
 from emsarray.operations import depth, point_extraction
+from emsarray.operations.cache import hash_attributes, hash_int, hash_string
 from emsarray.plot import (
     _requires_plot, animate_on_figure, make_plot_title, plot_on_figure,
     polygons_to_collection
@@ -1950,6 +1952,39 @@ def normalize_depth_variables(
             self.dataset, self.depth_coordinates,
             positive_down=positive_down, deep_to_shallow=deep_to_shallow)
 
+    def hash_geometry(self, hash: "hashlib._Hash") -> None:
+        """
+        Updates the provided hash with all of the relevant geometry data for this dataset.
+
+        Parameters
+        ----------
+        hash : hashlib-style hash instance
+            The hash instance to update with geometry data.
+            This must follow the interface defined in :mod:`hashlib`.
+        """
+        geometry_names = self.get_all_geometry_names()
+
+        for geometry_name in geometry_names:
+            data_array = self.dataset[geometry_name]
+
+            # Include the variable name in the digest.
+            hash_string(hash, str(geometry_name))
+
+            # Include the dtype of the data array.
+            # A float array and an int array mean very different things,
+            # but could have identical byte patterns.
+            hash_string(hash, data_array.encoding['dtype'].name)
+
+            # Include the size and shape of the data.
+            # 1D coordinate arrays are very different to 2D coordinate arrays,
+            # but could have identical byte patterns.
+            hash_int(hash, data_array.size)
+            hash.update(numpy.array(data_array.shape, dtype='int32').tobytes('C'))
+            hash.update(data_array.to_numpy().tobytes('C'))
+
+            # Hash dataset attributes
+            hash_attributes(hash, data_array.attrs)
+
 
 class DimensionConvention(Convention[GridKind, Index]):
     """

src/emsarray/operations/cache.py

@@ -0,0 +1,158 @@
+"""
+Operations for making cache keys based on dataset geometry.
+
+Some operations such as :func:`~.operations.triangulate.triangulate_dataset`
+only depend on the dataset geometry and are expensive to compute.
+For applications that need to derive data from the dataset geometry
+it would be useful if the derived data could be reused between different runs of the same application
+or between multiple time slices of the same geometry distributed across multiple files.
+This module provides :func:`.make_cache_key` to assist in this process
+by deriving a cache key from the important parts of a dataset geometry.
+Applications can use this cache key
+as part of a filename when save derived geometry data to disk
+or as a key to an in-memory cache of derived geometry.
+
+The derived cache keys will be identical between different instances of an application,
+and between different files in multi-file datasets split over an unlimited dimension.
+
+This module does not provide an actual cache implementation.
+"""
+import hashlib
+import marshal
+
+import numpy
+import xarray
+
+import emsarray
+
+
+def hash_attributes(hash: "hashlib._Hash", attributes: dict) -> None:
+    """
+    Adds the contents of an :attr:`attributes dictionary <xarray.DataArray.attrs>`
+    to a hash.
+
+    Parameters
+    ----------
+    hash : hashlib-style hash instance
+        The hash instance to add the attribute dictionary to.
+        This must follow the interface defined in :mod:`hashlib`.
+    attributes : dict
+        A dictionary of attributes from a :class:`~xarray.Dataset` or :class:`~xarray.DataArray`.
+
+    Notes
+    -----
+    The attribute dictionary is serialized to bytes using :func:`marshal.dumps`.
+    This is an implementation detail that may change in future releases.
+    """
+    # Prepend the marshal encoding version
+    marshal_version = 4
+    hash_int(hash, marshal_version)
+    # Specify marshal encoding version when serialising
+    attribute_dict_marshal_bytes = marshal.dumps(attributes, marshal_version)
+    # Prepend the number of attributes
+    hash_int(hash, len(attributes))
+    # Prepend the size of the pickled attributes
+    hash_int(hash, len(attribute_dict_marshal_bytes))
+    hash.update(attribute_dict_marshal_bytes)
+
+
+def hash_string(hash: "hashlib._Hash", value: str) -> None:
+    """
+    Adds a :class:`string <str>` to a hash.
+
+    Parameters
+    ----------
+    hash : hashlib-style hash instance
+        The hash instance to add the string to.
+        This must follow the interface defined in :mod:`hashlib`.
+    value : str
+        Any unicode string.
+
+    Notes
+    -----
+    The string is UTF-8 encoded as part of being added to the hash.
+    This is an implementation detail that may change in future releases.
+    """
+    # Prepend the length of the string to the hash
+    # to prevent malicious datasets generating overlapping string hashes.
+    hash_int(hash, len(value))
+    hash.update(value.encode('utf-8'))
+
+
+def hash_int(hash: "hashlib._Hash", value: int) -> None:
+    """
+    Adds an :class:`int` to a hash.
+
+    Parameters
+    ----------
+    hash : hashlib-style hash instance
+        The hash instance to add the integer to.
+        This must follow the interface defined in :mod:`hashlib`.
+    value : int
+        Any int representable as an :data:`numpy.int32`
+
+    Notes
+    -----
+    The int is cast to a :data:`numpy.int32` as part of being added to the hash.
+    This is an implementation detail that may change in the future
+    if larger integers are required.
+    """
+    with numpy.errstate(over='raise'):
+        # Manual overflow check as older numpy versions dont throw the exception
+        if numpy.iinfo("int32").min <= value <= numpy.iinfo("int32").max:
+            hash.update(numpy.int32(value).tobytes())
+        else:
+            raise OverflowError
+
+
+def make_cache_key(dataset: xarray.Dataset, hash: "hashlib._Hash | None" = None) -> str:
+    """
+    Derive a cache key from the geometry of a dataset.
+
+    Parameters
+    ----------
+    dataset : xarray.Dataset
+        The dataset to generate a cache key from.
+    hash : :mod:`hashlib`-compatible hash instance, optional
+        An instance of a hashlib hash class.
+        Defaults to :func:`hashlib.blake2b` with a digest size of 32,
+        which is secure enough and fast enough for most purposes.
+
+    Returns
+    -------
+    cache_key : str
+        A string suitable for use as a cache key.
+        The string will be safe for use as part of a filename if data is to be cached to disk.
+
+    Examples
+    --------
+
+    .. code-block:: python
+
+        import emsarray
+        from emsarray.operations.cache import make_cache_key
+
+        # Make a cache key from the dataset
+        dataset = emsarray.tuorial.open_dataset("austen")
+        cache_key = make_cache_key(dataset)
+        >>> cache_key
+        '580853c44e732878937598e86d0b26cb81e18d986072c0790a122244e9d3f480'
+
+    Notes
+    -----
+    The cache key will depend on the Convention class,
+    the emsarray version, and a hash of the geometry of the dataset.
+    The specific structure of the cache key may change between emsarray and python versions,
+    and should not be relied upon.
+    """
+    if hash is None:
+        hash = hashlib.blake2b(digest_size=32)
+
+    dataset.ems.hash_geometry(hash)
+
+    # Hash convention name, convention module path and emsarray version
+    hash_string(hash, dataset.ems.__class__.__module__)
+    hash_string(hash, dataset.ems.__class__.__name__)
+    hash_string(hash, emsarray.__version__)
+
+    return hash.hexdigest()