apache-hamilton 1.90.0.dev0__py3-none-any.whl

hamilton/caching/stores/file.py

@@ -0,0 +1,140 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+import inspect
+import shutil
+from pathlib import Path
+from typing import Any
+
+try:
+    from typing import override
+except ImportError:
+    override = lambda x: x  # noqa E731
+
+from hamilton.io.data_adapters import DataLoader, DataSaver
+
+from .base import ResultStore, StoredResult
+
+
+class FileResultStore(ResultStore):
+    def __init__(self, path: str, create_dir: bool = True) -> None:
+        self.path = Path(path)
+        self.create_dir = create_dir
+
+        if self.create_dir:
+            self.path.mkdir(exist_ok=True, parents=True)
+
+    def __getstate__(self) -> dict:
+        """Serialize the `__init__` kwargs to pass in Parallelizable branches
+        when using multiprocessing.
+        """
+        return {"path": str(self.path)}
+
+    @staticmethod
+    def _write_result(file_path: Path, stored_result: StoredResult) -> None:
+        file_path.write_bytes(stored_result.save())
+
+    @staticmethod
+    def _load_result_from_path(path: Path) -> StoredResult | None:
+        try:
+            data = path.read_bytes()
+            return StoredResult.load(data)
+        except FileNotFoundError:
+            return None
+
+    def _path_from_data_version(self, data_version: str) -> Path:
+        return self.path.joinpath(data_version)
+
+    def _materialized_path(self, data_version: str, saver_cls: DataSaver) -> Path:
+        # TODO allow a more flexible mechanism to specify file path extension
+        return self._path_from_data_version(data_version).with_suffix(f".{saver_cls.name()}")
+
+    @override
+    def exists(self, data_version: str) -> bool:
+        result_path = self._path_from_data_version(data_version)
+        return result_path.exists()
+
+    @override
+    def set(
+        self,
+        data_version: str,
+        result: Any,
+        saver_cls: DataSaver | None = None,
+        loader_cls: DataLoader | None = None,
+    ) -> None:
+        # != operator on boolean is XOR
+        if bool(saver_cls is not None) != bool(loader_cls is not None):
+            raise ValueError(
+                "Must pass both `saver` and `loader` or neither. Currently received: "
+                f"`saver`: `{saver_cls}`; `loader`: `{loader_cls}`"
+            )
+
+        if saver_cls is not None:
+            # materialized_path
+            materialized_path = self._materialized_path(data_version, saver_cls)
+            saver_argspec = inspect.getfullargspec(saver_cls.__init__)
+            loader_argspec = inspect.getfullargspec(loader_cls.__init__)
+            if "file" in saver_argspec.args:
+                saver = saver_cls(file=str(materialized_path.absolute()))
+            elif "path" in saver_argspec.args:
+                saver = saver_cls(path=str(materialized_path.absolute()))
+            else:
+                raise ValueError(
+                    f"Saver [{saver_cls.name()}] must have either `file` or `path` as an argument."
+                )
+            if "file" in loader_argspec.args:
+                loader = loader_cls(file=str(materialized_path.absolute()))
+            elif "path" in loader_argspec.args:
+                loader = loader_cls(path=str(materialized_path.absolute()))
+            else:
+                raise ValueError(
+                    f"Loader [{loader_cls.name()}] must have either `file` or `path` as an argument."
+                )
+        else:
+            saver = None
+            loader = None
+
+        self.path.mkdir(exist_ok=True)
+        result_path = self._path_from_data_version(data_version)
+        stored_result = StoredResult.new(value=result, saver=saver, loader=loader)
+        self._write_result(result_path, stored_result)
+
+    @override
+    def get(self, data_version: str) -> Any | None:
+        result_path = self._path_from_data_version(data_version)
+        stored_result = self._load_result_from_path(result_path)
+
+        if stored_result is None:
+            return None
+
+        return stored_result.value
+
+    @override
+    def delete(self, data_version: str) -> None:
+        result_path = self._path_from_data_version(data_version)
+        result_path.unlink(missing_ok=True)
+
+    @override
+    def delete_all(self) -> None:
+        shutil.rmtree(self.path)
+        self.path.mkdir(exist_ok=True)
+
+    def delete_expired(self) -> None:
+        for file_path in self.path.iterdir():
+            stored_result = self._load_result_from_path(file_path)
+            if stored_result and stored_result.expired:
+                file_path.unlink(missing_ok=True)

hamilton/caching/stores/memory.py

@@ -0,0 +1,297 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+from collections.abc import Sequence
+from typing import Any
+
+try:
+    from typing import override
+except ImportError:
+    override = lambda x: x  # noqa E731
+
+from hamilton.caching.cache_key import decode_key
+
+from .base import MetadataStore, ResultStore, StoredResult
+from .file import FileResultStore
+from .sqlite import SQLiteMetadataStore
+
+
+class InMemoryMetadataStore(MetadataStore):
+    def __init__(self) -> None:
+        self._data_versions: dict[str, str] = {}  # {cache_key: data_version}
+        self._cache_keys_by_run: dict[str, list[str]] = {}  # {run_id: [cache_key]}
+        self._run_ids: list[str] = []
+
+    @override
+    def __len__(self) -> int:
+        """Number of unique ``cache_key`` values."""
+        return len(self._data_versions.keys())
+
+    @override
+    def exists(self, cache_key: str) -> bool:
+        """Indicate if ``cache_key`` exists and it can retrieve a ``data_version``."""
+        return cache_key in self._data_versions.keys()
+
+    @override
+    def initialize(self, run_id: str) -> None:
+        """Set up and log the beginning of the run."""
+        self._cache_keys_by_run[run_id] = []
+        self._run_ids.append(run_id)
+
+    @override
+    def set(self, cache_key: str, data_version: str, run_id: str, **kwargs) -> Any | None:
+        """Set the ``data_version`` for ``cache_key`` and associate it with the ``run_id``."""
+        self._data_versions[cache_key] = data_version
+        self._cache_keys_by_run[run_id].append(cache_key)
+
+    @override
+    def get(self, cache_key: str) -> str | None:
+        """Retrieve the ``data_version`` for ``cache_key``."""
+        return self._data_versions.get(cache_key, None)
+
+    @override
+    def delete(self, cache_key: str) -> None:
+        """Delete the ``data_version`` for ``cache_key``."""
+        del self._data_versions[cache_key]
+
+    @override
+    def delete_all(self) -> None:
+        """Delete all stored metadata."""
+        self._data_versions.clear()
+
+    def persist_to(self, metadata_store: MetadataStore | None = None) -> None:
+        """Persist in-memory metadata using another MetadataStore implementation.
+
+        :param metadata_store: MetadataStore implementation to use for persistence.
+            If None, a SQLiteMetadataStore is created with the default path "./.hamilton_cache".
+
+        .. code-block:: python
+
+            from hamilton import driver
+            from hamilton.caching.stores.sqlite import SQLiteMetadataStore
+            from hamilton.caching.stores.memory import InMemoryMetadataStore
+            import my_dataflow
+
+            dr = (
+              driver.Builder()
+              .with_modules(my_dataflow)
+              .with_cache(metadata_store=InMemoryMetadataStore())
+              .build()
+            )
+
+            # execute the Driver several time. This will populate the in-memory metadata store
+            dr.execute(...)
+
+            # persist to disk in-memory metadata
+            dr.cache.metadata_store.persist_to(SQLiteMetadataStore(path="./.hamilton_cache"))
+
+        """
+        if metadata_store is None:
+            metadata_store = SQLiteMetadataStore(path="./.hamilton_cache")
+
+        for run_id in self._run_ids:
+            metadata_store.initialize(run_id)
+
+        for run_id, cache_keys in self._cache_keys_by_run.items():
+            for cache_key in cache_keys:
+                data_version = self._data_versions[cache_key]
+                metadata_store.set(
+                    cache_key=cache_key,
+                    data_version=data_version,
+                    run_id=run_id,
+                )
+
+    @classmethod
+    def load_from(cls, metadata_store: MetadataStore) -> "InMemoryMetadataStore":
+        """Load in-memory metadata from another MetadataStore instance.
+
+        :param metadata_store: MetadataStore instance to load from.
+        :return: InMemoryMetadataStore copy of the ``metadata_store``.
+
+        .. code-block:: python
+
+            from hamilton import driver
+            from hamilton.caching.stores.sqlite import SQLiteMetadataStore
+            from hamilton.caching.stores.memory import InMemoryMetadataStore
+            import my_dataflow
+
+            sqlite_metadata_store = SQLiteMetadataStore(path="./.hamilton_cache")
+            in_memory_metadata_store = InMemoryMetadataStore.load_from(sqlite_metadata_store)
+
+            # create the Driver with the in-memory metadata store
+            dr = (
+              driver.Builder()
+              .with_modules(my_dataflow)
+              .with_cache(metadata_store=in_memory_metadata_store)
+              .build()
+            )
+
+        """
+        in_memory_metadata_store = InMemoryMetadataStore()
+
+        for run_id in metadata_store.get_run_ids():
+            in_memory_metadata_store.initialize(run_id)
+
+            for node_metadata in metadata_store.get_run(run_id):
+                in_memory_metadata_store.set(
+                    cache_key=node_metadata["cache_key"],
+                    data_version=node_metadata["data_version"],
+                    run_id=run_id,
+                )
+
+        return in_memory_metadata_store
+
+    @override
+    def get_run_ids(self) -> list[str]:
+        """Return a list of all ``run_id`` values stored."""
+        return self._run_ids
+
+    @override
+    def get_run(self, run_id: str) -> list[dict[str, str]]:
+        """Return a list of node metadata associated with a run."""
+        if self._cache_keys_by_run.get(run_id, None) is None:
+            raise IndexError(f"Run ID not found: {run_id}")
+
+        nodes_metadata = []
+        for cache_key in self._cache_keys_by_run[run_id]:
+            decoded_key = decode_key(cache_key)
+            nodes_metadata.append(
+                dict(
+                    cache_key=cache_key,
+                    data_version=self._data_versions[cache_key],
+                    node_name=decoded_key["node_name"],
+                    code_version=decoded_key["code_version"],
+                    dependencies_data_versions=decoded_key["dependencies_data_versions"],
+                )
+            )
+
+        return nodes_metadata
+
+
+class InMemoryResultStore(ResultStore):
+    def __init__(self, persist_on_exit: bool = False) -> None:
+        self._results: dict[str, StoredResult] = {}  # {data_version: result}
+
+    @override
+    def exists(self, data_version: str) -> bool:
+        return data_version in self._results.keys()
+
+    # TODO handle materialization
+    @override
+    def set(self, data_version: str, result: Any, **kwargs) -> None:
+        self._results[data_version] = StoredResult.new(value=result)
+
+    @override
+    def get(self, data_version: str) -> Any | None:
+        stored_result = self._results.get(data_version, None)
+        if stored_result is None:
+            return None
+
+        return stored_result.value
+
+    @override
+    def delete(self, data_version: str) -> None:
+        del self._results[data_version]
+
+    @override
+    def delete_all(self) -> None:
+        self._results.clear()
+
+    def delete_expired(self) -> None:
+        to_delete = [
+            data_version
+            for data_version, stored_result in self._results.items()
+            if stored_result.expired
+        ]
+
+        # first collect keys then delete because you can delete from dictionary
+        # as you iterate through it
+        for data_version in to_delete:
+            self.delete(data_version)
+
+    def persist_to(self, result_store: ResultStore | None = None) -> None:
+        """Persist in-memory results using another ``ResultStore`` implementation.
+
+        :param result_store: ResultStore implementation to use for persistence.
+            If None, a FileResultStore is created with the default path "./.hamilton_cache".
+        """
+        if result_store is None:
+            result_store = FileResultStore(path="./.hamilton_cache")
+
+        for data_version, stored_result in self._results.items():
+            result_store.set(data_version, stored_result.value)
+
+    @classmethod
+    def load_from(
+        cls,
+        result_store: ResultStore,
+        metadata_store: MetadataStore | None = None,
+        data_versions: Sequence[str] | None = None,
+    ) -> "InMemoryResultStore":
+        """Load in-memory results from another ResultStore instance.
+
+        Since result stores do not store an index of their keys, you must provide a
+        ``MetadataStore`` instance or a list of ``data_version`` for which results
+        should be loaded in memory.
+
+        :param result_store: ``ResultStore`` instance to load results from.
+        :param metadata_store: ``MetadataStore`` instance from which all ``data_version`` are retrieved.
+        :return: InMemoryResultStore copy of the ``result_store``.
+
+        .. code-block:: python
+
+            from hamilton import driver
+            from hamilton.caching.stores.sqlite import SQLiteMetadataStore
+            from hamilton.caching.stores.memory import InMemoryMetadataStore
+            import my_dataflow
+
+            sqlite_metadata_store = SQLiteMetadataStore(path="./.hamilton_cache")
+            in_memory_metadata_store = InMemoryMetadataStore.load_from(sqlite_metadata_store)
+
+            # create the Driver with the in-memory metadata store
+            dr = (
+              driver.Builder()
+              .with_modules(my_dataflow)
+              .with_cache(metadata_store=in_memory_metadata_store)
+              .build()
+            )
+
+
+        """
+        if metadata_store is None and data_versions is None:
+            raise ValueError(
+                "A `metadata_store` or `data_versions` must be provided to load results."
+            )
+
+        in_memory_result_store = InMemoryResultStore()
+
+        data_versions_to_retrieve = set()
+        if data_versions is not None:
+            data_versions_to_retrieve.update(data_versions)
+
+        if metadata_store is not None:
+            for run_id in metadata_store.get_run_ids():
+                for node_metadata in metadata_store.get_run(run_id):
+                    data_versions_to_retrieve.add(node_metadata["data_version"])
+
+        for data_version in data_versions_to_retrieve:
+            # TODO disambiguate "result is None" from the sentinel value when `data_version`
+            # is not found in `result_store`.
+            result = result_store.get(data_version)
+            in_memory_result_store.set(data_version, result)
+
+        return in_memory_result_store