Source code for

This module contains ``CachedDataset``, a dataset wrapper which caches in memory the data saved,
so that the user avoids io operations with slow storage media
from __future__ import annotations

import logging
from typing import Any

from import VERSIONED_FLAG_KEY, AbstractDataset, Version
from import MemoryDataset

[docs] class CachedDataset(AbstractDataset): """``CachedDataset`` is a dataset wrapper which caches in memory the data saved, so that the user avoids io operations with slow storage media. You can also specify a ``CachedDataset`` in catalog.yml: :: >>> test_ds: >>> type: CachedDataset >>> versioned: true >>> dataset: >>> type: pandas.CSVDataset >>> filepath: example.csv Please note that if your dataset is versioned, this should be indicated in the wrapper class as shown above. """ # this dataset cannot be used with ``ParallelRunner``, # therefore it has the attribute ``_SINGLE_PROCESS = True`` # for parallelism please consider ``ThreadRunner`` instead _SINGLE_PROCESS = True
[docs] def __init__( self, dataset: AbstractDataset | dict, version: Version | None = None, copy_mode: str | None = None, metadata: dict[str, Any] | None = None, ): """Creates a new instance of ``CachedDataset`` pointing to the provided Python object. Args: dataset: A Kedro Dataset object or a dictionary to cache. version: If specified, should be an instance of ````. If its ``load`` attribute is None, the latest version will be loaded. If its ``save`` attribute is None, save version will be autogenerated. copy_mode: The copy mode used to copy the data. Possible values are: "deepcopy", "copy" and "assign". If not provided, it is inferred based on the data type. metadata: Any arbitrary metadata. This is ignored by Kedro, but may be consumed by users or external plugins. Raises: ValueError: If the provided dataset is not a valid dict/YAML representation of a dataset or an actual dataset. """ self._EPHEMERAL = True if isinstance(dataset, dict): self._dataset = self._from_config(dataset, version) elif isinstance(dataset, AbstractDataset): self._dataset = dataset else: raise ValueError( "The argument type of 'dataset' should be either a dict/YAML " "representation of the dataset, or the actual dataset object." ) self._cache = MemoryDataset(copy_mode=copy_mode) self.metadata = metadata
def _release(self) -> None: self._cache.release() self._dataset.release() @staticmethod def _from_config(config: dict, version: Version | None) -> AbstractDataset: if VERSIONED_FLAG_KEY in config: raise ValueError( "Cached datasets should specify that they are versioned in the " "'CachedDataset', not in the wrapped dataset." ) if version: config[VERSIONED_FLAG_KEY] = True return AbstractDataset.from_config( "_cached", config, version.load, ) return AbstractDataset.from_config("_cached", config) def _describe(self) -> dict[str, Any]: return { "dataset": self._dataset._describe(), "cache": self._cache._describe(), } def _load(self) -> Any: data = self._cache.load() if self._cache.exists() else self._dataset.load() if not self._cache.exists(): return data def _save(self, data: Any) -> None: def _exists(self) -> bool: return self._cache.exists() or self._dataset.exists() def __getstate__(self) -> dict[str, Any]: # clearing the cache can be prevented by modifying # how parallel runner handles datasets (not trivial!) logging.getLogger(__name__).warning("%s: clearing cache to pickle.", str(self)) self._cache.release() return self.__dict__