Repository: Aarhus-Psychiatry-Research/timeseriesflattener

Issue:
Refactor: Split out cache functionality to abstractclass

Discussion:

Patch:
diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml
index 7b3ce38..af1c9de 100644
--- a/.pre-commit-config.yaml
+++ b/.pre-commit-config.yaml
@@ -18,12 +18,6 @@ repos:
     hooks:
       - id: ssort
 
-  - repo: https://github.com/myint/docformatter
-    rev: v1.5.0
-    hooks:
-      - id: docformatter
-        args: [--in-place]
-
   - repo: https://github.com/asottile/add-trailing-comma
     rev: v2.2.3
     hooks:
diff --git a/src/timeseriesflattener/feature_cache/__init__.py b/src/timeseriesflattener/feature_cache/__init__.py
new file mode 100644
index 0000000..e69de29
diff --git a/src/timeseriesflattener/feature_cache/abstract_feature_cache.py b/src/timeseriesflattener/feature_cache/abstract_feature_cache.py
new file mode 100644
index 0000000..d5493e5
--- /dev/null
+++ b/src/timeseriesflattener/feature_cache/abstract_feature_cache.py
@@ -0,0 +1,46 @@
+"""Base method that defines a feature cache."""
+from abc import ABCMeta, abstractmethod
+from typing import Optional
+
+import pandas as pd
+
+from timeseriesflattener.feature_spec_objects import TemporalSpec
+
+
+class FeatureCache(metaclass=ABCMeta):
+    """Base class that defines a feature cache."""
+
+    @abstractmethod
+    def __init__(
+        self,
+        *args,
+        prediction_times_df: Optional[pd.DataFrame] = None,
+        pred_time_uuid_col_name: str = "pred_time_uuid",
+        id_col_name: str = "entity_id",
+        timestamp_col_name: str = "timestamp",
+    ):
+        """Initialize a FeatureCache.
+
+        Args:
+            prediction_times_df (Optional[pd.DataFrame], optional): DataFrame containing prediction times.
+                Must be set at some point, but doesn't have to be set at init.
+                Useful when e.g. used as a component in TimeseriesFlattener, which already knows the prediction_times_df and can set it as a pointer during initialization. Defaults to None. Defaults to None.
+            pred_time_uuid_col_name (str, optional): Name of column containing prediction time uuids.
+            id_col_name (str, optional): Name of column containing entity ids. Defaults to "entity_id".
+            timestamp_col_name (str, optional): Name of column containing timestamps. Defaults to "timestamp".
+            Defaults to "pred_time_uuid".
+        """
+        self.prediction_times_df = prediction_times_df
+        self.pred_time_uuid_col_name = pred_time_uuid_col_name
+
+    @abstractmethod
+    def feature_exists(self, feature_spec: TemporalSpec) -> bool:
+        """Check if feature exists in cache."""
+
+    @abstractmethod
+    def read_feature(self, feature_spec: TemporalSpec) -> pd.DataFrame:
+        """Read feature from cache."""
+
+    @abstractmethod
+    def write_feature(self, feature_spec: TemporalSpec, df: pd.DataFrame) -> None:
+        """Write feature to cache."""
diff --git a/src/timeseriesflattener/feature_cache/cache_to_disk.py b/src/timeseriesflattener/feature_cache/cache_to_disk.py
new file mode 100644
index 0000000..0a40ad7
--- /dev/null
+++ b/src/timeseriesflattener/feature_cache/cache_to_disk.py
@@ -0,0 +1,185 @@
+"""Cache module for writing features to disk."""
+import datetime as dt
+import os
+from pathlib import Path
+from typing import Optional
+
+import pandas as pd
+
+from timeseriesflattener.feature_cache.abstract_feature_cache import FeatureCache
+from timeseriesflattener.feature_spec_objects import TemporalSpec
+from timeseriesflattener.utils import load_dataset_from_file, write_df_to_file
+
+
+class DiskCache(FeatureCache):
+    """Cache module for writing features to disk."""
+
+    def __init__(
+        self,
+        feature_cache_dir: Path,
+        prediction_times_df: Optional[pd.DataFrame] = None,
+        pred_time_uuid_col_name: str = "pred_time_uuid",
+        id_col_name: str = "entity_id",
+        timestamp_col_name: str = "timestamp",
+        cache_file_suffix: str = ".parquet",
+    ):
+        """Initialize DiskCache.
+
+        Args:
+            feature_cache_dir (Path): Path to directory where features are cached
+            prediction_times_df (Optional[pd.DataFrame], optional): DataFrame containing prediction times.
+                Must be set at some point, but doesn't have to be set at init.
+                Useful when e.g. used as a component in TimeseriesFlattener, which already knows the prediction_times_df and can set it as a pointer. Defaults to None.
+            pred_time_uuid_col_name (str, optional): Name of column containing prediction time uuids. Defaults to "pred_time_uuid".
+            cache_file_suffix (str, optional): File suffix for cache files. Defaults to ".parquet".
+        """
+
+        super().__init__(
+            prediction_times_df=prediction_times_df,
+            pred_time_uuid_col_name=pred_time_uuid_col_name,
+        )
+
+        self.feature_cache_dir = feature_cache_dir
+        self.cache_file_suffix = cache_file_suffix
+        self.entity_id_col_name = id_col_name
+        self.timestamp_col_name = timestamp_col_name
+
+        if not self.feature_cache_dir.exists():
+            self.feature_cache_dir.mkdir()
+
+    def _load_most_recent_df_matching_pattern(
+        self,
+        file_pattern: str,
+    ) -> pd.DataFrame:
+        """Load most recent df matching pattern.
+
+        Args:
+            file_pattern (str): Pattern to match
+            file_suffix (str, optional): File suffix to match.
+
+        Returns:
+            DataFrame: DataFrame matching pattern
+
+        Raises:
+            FileNotFoundError: If no file matching pattern is found
+        """
+        files_with_suffix = list(self.feature_cache_dir.glob(file_pattern))
+
+        if len(files_with_suffix) == 0:
+            raise FileNotFoundError(f"No files matching pattern {file_pattern} found")
+
+        path_of_most_recent_file = max(files_with_suffix, key=os.path.getctime)
+
+        return load_dataset_from_file(
+            file_path=path_of_most_recent_file,
+        )
+
+    def _get_file_name(
+        self,
+        feature_spec: TemporalSpec,
+    ) -> str:
+        """Get file name for feature spec.
+
+        Args:
+            feature_spec (AnySpec): Feature spec
+
+        Returns:
+            str: File name
+        """
+        n_rows = feature_spec.values_df.shape[0]  # type: ignore
+
+        return f"{feature_spec.get_col_str()}_{n_rows}_uuids"
+
+    def _get_file_pattern(
+        self,
+        feature_spec: TemporalSpec,
+    ) -> str:
+        """Get file pattern for feature spec.
+
+        Args:
+            feature_spec (AnySpec): Feature spec
+
+        Returns:
+            str: File pattern
+        """
+        file_name = self._get_file_name(feature_spec=feature_spec)
+
+        return f"*{file_name}*.{self.cache_file_suffix}*"
+
+    def read_feature(self, feature_spec: TemporalSpec) -> pd.DataFrame:
+        """Load most recent df matching pattern, and expand fallback column.
+
+        Args:
+            feature_spec (AnySpec): Feature spec
+
+        Returns:
+            DataFrame: DataFrame with fallback column expanded
+        """
+        df = self._load_most_recent_df_matching_pattern(
+            file_pattern=self._get_file_pattern(feature_spec=feature_spec),
+        )
+
+        # Expand fallback column
+        df = pd.merge(
+            left=self.prediction_times_df[self.pred_time_uuid_col_name],
+            right=df,
+            how="left",
+            on=self.pred_time_uuid_col_name,
+            validate="m:1",
+        )
+
+        # Replace NaNs with fallback
+        df[feature_spec.get_col_str()] = df[feature_spec.get_col_str()].fillna(
+            feature_spec.fallback,
+        )
+
+        return df
+
+    def write_feature(
+        self,
+        feature_spec: TemporalSpec,
+        df: pd.DataFrame,
+    ):
+        """Write feature to cache."""
+        n_uuids = df[self.pred_time_uuid_col_name].nunique()
+        file_name = f"{feature_spec.get_col_str()}_{n_uuids}_uuids"
+
+        # Drop rows containing fallback, since it's non-informative
+        df = df[df[feature_spec.get_col_str()] != feature_spec.fallback].dropna()
+
+        # Drop entity and timestamp columns if they exists
+        for col in [self.entity_id_col_name, self.timestamp_col_name]:
+            if col in df.columns:
+                df = df.drop(columns=col)
+
+        # Write df to cache
+        timestamp = dt.datetime.now().strftime("%Y-%m-%d_%H-%M-%S")
+        write_df_to_file(
+            df=df,
+            file_path=self.feature_cache_dir
+            / f"{file_name}_{timestamp}.{self.cache_file_suffix}",
+        )
+
+    def feature_exists(
+        self,
+        feature_spec: TemporalSpec,
+    ) -> bool:
+        """Check if cache is hit.
+
+        Args:
+            feature_spec (AnySpec): Feature spec
+
+        Returns:
+            bool: True if cache is hit, False otherwise
+        """
+        file_pattern = self._get_file_pattern(feature_spec=feature_spec)
+
+        # Check that file exists
+        file_pattern_hits = list(
+            self.feature_cache_dir.glob(file_pattern),
+        )
+
+        if len(file_pattern_hits) == 0:
+            return False
+
+        return True
diff --git a/src/timeseriesflattener/flattened_dataset.py b/src/timeseriesflattener/flattened_dataset.py
index cc200ac..cf7d46a 100644
--- a/src/timeseriesflattener/flattened_dataset.py
+++ b/src/timeseriesflattener/flattened_dataset.py
@@ -1,14 +1,14 @@
-"""Takes a time-series and flattens it into a set of prediction times with
-describing values."""
+"""Takes a time-series and flattens it into a set of prediction times with.
+
+describing values.
+"""
 import datetime as dt
-import os
 import random
 import time
 from collections.abc import Callable
 from datetime import timedelta
 from multiprocessing import Pool
-from pathlib import Path
-from typing import Any, Optional
+from typing import Optional
 
 import numpy as np
 import pandas as pd
@@ -18,6 +18,7 @@ from dask.diagnostics import ProgressBar
 from pandas import DataFrame
 from wasabi import Printer, msg
 
+from timeseriesflattener.feature_cache.abstract_feature_cache import FeatureCache
 from timeseriesflattener.feature_spec_objects import (
     AnySpec,
     OutcomeSpec,
@@ -26,7 +27,6 @@ from timeseriesflattener.feature_spec_objects import (
 )
 from timeseriesflattener.flattened_ds_validator import ValidateInitFlattenedDataset
 from timeseriesflattener.resolve_multiple_functions import resolve_multiple_fns
-from timeseriesflattener.utils import load_dataset_from_file, write_df_to_file
 
 ProgressBar().register()
 
@@ -45,15 +45,42 @@ class TimeseriesFlattener:  # pylint: disable=too-many-instance-attributes
         feature_cache_dir (Path): Path to cache directory for feature dataframes.
     """
 
+    def _override_cache_attributes_with_self_attributes(
+        self,
+        prediction_times_df: DataFrame,
+    ):
+        """Make cache inherit attributes from flattened dataset.
+
+        Avoids duplicate specification.
+        """
+        if (
+            not hasattr(self.cache, "prediction_times_df")
+            or self.cache.prediction_times_df is None
+        ):
+            self.cache.prediction_times_df = prediction_times_df
+        elif not self.cache.prediction_times_df.equals(prediction_times_df):
+            msg.warn(
+                "Overriding prediction_times_df in cache with prediction_times_df passed to init",
+            )
+            self.cache.prediction_times_df = prediction_times_df
+
+        for attr in ["pred_time_uuid_col_name", "timestamp_col_name", "id_col_name"]:
+            if hasattr(self.cache, attr) and getattr(self.cache, attr) is not None:
+                if getattr(self.cache, attr) != getattr(self, attr):
+                    msg.warn(
+                        f"Overriding {attr} in cache with {attr} passed to init of flattened dataset",
+                    )
+                    setattr(self.cache, attr, getattr(self, attr))
+
     def __init__(  # pylint: disable=too-many-arguments
         self,
         prediction_times_df: DataFrame,
+        cache: Optional[FeatureCache] = None,
         id_col_name: str = "dw_ek_borger",
         timestamp_col_name: str = "timestamp",
         predictor_col_name_prefix: str = "pred",
         outcome_col_name_prefix: str = "outc",
         n_workers: int = 60,
-        feature_cache_dir: Optional[Path] = None,
     ):
         """Class containing a time-series, flattened.
 
@@ -80,12 +107,12 @@ class TimeseriesFlattener:  # pylint: disable=too-many-instance-attributes
 
         Args:
             prediction_times_df (DataFrame): Dataframe with prediction times, required cols: patient_id, .
-            timestamp_col_name (str, optional): Column name name for timestamps. Is used across outcomes and predictors. Defaults to "timestamp".
+            cache (Optional[FeatureCache], optional): Object for feature caching. Should be initialised when passed to init. Defaults to None.
             id_col_name (str, optional): Column namn name for patients ids. Is used across outcome and predictors. Defaults to "dw_ek_borger".
+            timestamp_col_name (str, optional): Column name name for timestamps. Is used across outcomes and predictors. Defaults to "timestamp".
             predictor_col_name_prefix (str, optional): Prefix for predictor col names. Defaults to "pred_".
             outcome_col_name_prefix (str, optional): Prefix for outcome col names. Defaults to "outc_".
             n_workers (int): Number of subprocesses to spawn for parallelization. Defaults to 60.
-            feature_cache_dir (Path): Path to cache directory for feature dataframes. Defaults to None.
 
         Raises:
             ValueError: If timestamp_col_name or id_col_name is not in prediction_times_df
@@ -98,18 +125,19 @@ class TimeseriesFlattener:  # pylint: disable=too-many-instance-attributes
         self.predictor_col_name_prefix = predictor_col_name_prefix
         self.outcome_col_name_prefix = outcome_col_name_prefix
         self.pred_time_uuid_col_name = "prediction_time_uuid"
+        self.cache = cache
 
-        if feature_cache_dir:
-            self.feature_cache_dir = feature_cache_dir
-            if not self.feature_cache_dir.exists():
-                self.feature_cache_dir.mkdir()
+        if self.cache:
+            self._override_cache_attributes_with_self_attributes(prediction_times_df)
 
-        self.n_uuids = len(prediction_times_df)
+        self.n_uuids = prediction_times_df.shape[0]
 
         self.msg = Printer(timestamp=True)
 
         if "value" in prediction_times_df.columns:
-            prediction_times_df.drop("value", axis=1, inplace=True)
+            raise ValueError(
+                "Column 'value' should not occur in prediction_times_df, only timestamps and ids.",
+            )
 
         self._df = prediction_times_df
 
@@ -124,83 +152,17 @@ class TimeseriesFlattener:  # pylint: disable=too-many-instance-attributes
             str,
         ) + self._df[self.timestamp_col_name].dt.strftime("-%Y-%m-%d-%H-%M-%S")
 
-    def _load_most_recent_df_matching_pattern(
-        self,
-        dir_path: Path,
-        file_pattern: str,
-        file_suffix: str,
-    ) -> DataFrame:
-        """Load most recent df matching pattern.
-
-        Args:
-            file_pattern (str): Pattern to match
-            file_suffix (str, optional): File suffix to match.
-
-        Returns:
-            DataFrame: DataFrame matching pattern
-
-        Raises:
-            FileNotFoundError: If no file matching pattern is found
-        """
-        files_with_suffix = list(dir_path.glob(f"*{file_pattern}*.{file_suffix}"))
-
-        if len(files_with_suffix) == 0:
-            raise FileNotFoundError(f"No files matching pattern {file_pattern} found")
-
-        path_of_most_recent_file = max(files_with_suffix, key=os.path.getctime)
-
-        return load_dataset_from_file(
-            file_path=path_of_most_recent_file,
-        )
-
-    def _load_cached_df_and_expand_fallback(
-        self,
-        file_pattern: str,
-        file_suffix: str,
-        fallback: Any,
-        full_col_str: str,
-    ) -> pd.DataFrame:
-        """Load most recent df matching pattern, and expand fallback column.
-
-        Args:
-            file_pattern (str): File pattern to search for
-            file_suffix (str): File suffix to search for
-            fallback (Any): Fallback value
-            full_col_str (str): Full column name for values
-
-        Returns:
-            DataFrame: DataFrame with fallback column expanded
-        """
-        df = self._load_most_recent_df_matching_pattern(
-            dir_path=self.feature_cache_dir,
-            file_pattern=file_pattern,
-            file_suffix=file_suffix,
-        )
-
-        # Expand fallback column
-        df = pd.merge(
-            left=self._df[self.pred_time_uuid_col_name],
-            right=df,
-            how="left",
-            on=self.pred_time_uuid_col_name,
-            validate="m:1",
-        )
-
-        df[full_col_str] = df[full_col_str].fillna(fallback)
-
-        return df
-
     @staticmethod
-    def _flatten_temporal_values_to_df(  # noqa pylint: disable=too-many-locals
+    def flatten_temporal_values_to_df(  # noqa pylint: disable=too-many-locals
         prediction_times_with_uuid_df: DataFrame,
-        output_spec: TemporalSpec,
+        output_spec: AnySpec,
         id_col_name: str,
         timestamp_col_name: str,
         pred_time_uuid_col_name: str,
         verbose: bool = False,
     ) -> DataFrame:
+        """Create a dataframe with flattened values (either predictor or.
 
-        """Create a dataframe with flattened values (either predictor or
         outcome depending on the value of "direction").
 
         Args:
@@ -288,159 +250,9 @@ class TimeseriesFlattener:  # pylint: disable=too-many-instance-attributes
 
         return df[[pred_time_uuid_col_name, output_spec.get_col_str()]]
 
-    def _generate_values_for_cache_checking(
-        self,
-        output_spec: TemporalSpec,
-        value_col_str: str,
-        n_to_generate: int = 100_000,
-    ):
-        generated_df = pd.DataFrame({value_col_str: []})
-
-        # Check that some values in generated_df differ from fallback
-        # Otherwise, comparison to cache is meaningless
-        n_trials = 0
-
-        while not any(
-            generated_df[value_col_str] != output_spec.fallback,
-        ):
-            if n_trials != 0:
-                self.msg.info(
-                    f"{value_col_str[20]}, {n_trials}: Generated_df was all fallback values, regenerating",
-                )
-
-            n_to_generate = int(min(n_to_generate, len(self._df)))
-
-            generated_df = self._flatten_temporal_values_to_df(
-                prediction_times_with_uuid_df=self._df.sample(
-                    n=n_to_generate,
-                    replace=False,
-                ),
-                id_col_name=self.id_col_name,
-                timestamp_col_name=self.timestamp_col_name,
-                pred_time_uuid_col_name=self.pred_time_uuid_col_name,
-                output_spec=output_spec,
-            ).dropna()
-
-            # Fallback values are not interesting for cache hit. If they exist in generated_df, they should be dropped
-            # in the cache. Saves on storage. Don't use them to check if cache is hit.
-            if not np.isnan(output_spec.fallback):  # type: ignore
-                generated_df = generated_df[
-                    generated_df[value_col_str] != output_spec.fallback
-                ]
-
-            n_to_generate = (
-                n_to_generate**1.5
-            )  # Increase n_to_generate by 1.5x each time to increase chance of non_fallback values
-
-            n_trials += 1
-
-        return generated_df
-
-    def _cache_is_hit(
-        self,
-        output_spec: TemporalSpec,
-        file_pattern: str,
-        file_suffix: str,
-    ) -> bool:
-        """Check if cache is hit.
-
-        Args:
-            kwargs_dict (dict): dictionary of kwargs
-            file_pattern (str): File pattern to match. Looks for *file_pattern* in cache dir.
-            e.g. "*feature_name*_uuids*.file_suffix"
-            full_col_str (str): Full column string. e.g. "feature_name_ahead_interval_days_resolve_multiple_fallback"
-            file_suffix (str): File suffix to match. e.g. "csv"
-
-        Returns:
-            bool: True if cache is hit, False otherwise
-        """
-        # Check that file exists
-        file_pattern_hits = list(
-            self.feature_cache_dir.glob(f"*{file_pattern}*.{file_suffix}"),
-        )
-
-        if len(file_pattern_hits) == 0:
-            self.msg.info(f"Cache miss, {file_pattern} didn't exist")
-            return False
-
-        value_col_str = output_spec.get_col_str()
-
-        # Check that file contents match expected
-        # NAs are not interesting when comparing if computed values are identical
-        cache_df = self._load_most_recent_df_matching_pattern(
-            dir_path=self.feature_cache_dir,
-            file_pattern=file_pattern,
-            file_suffix=file_suffix,
-        )
-
-        generated_df = self._generate_values_for_cache_checking(
-            output_spec=output_spec,
-            value_col_str=value_col_str,
-        )
-
-        cached_suffix = "_c"
-        generated_suffix = "_g"
-
-        # We frequently hit rounding errors with cache hits, so we round to 3 decimal places
-        generated_df[value_col_str] = generated_df[value_col_str].round(3)
-        cache_df[value_col_str] = cache_df[value_col_str].round(3)
-
-        # Merge cache_df onto generated_df
-        merged_df = pd.merge(
-            left=generated_df,
-            right=cache_df,
-            how="left",
-            on=self.pred_time_uuid_col_name,
-            suffixes=(generated_suffix, cached_suffix),
-            validate="1:1",
-            indicator=True,
-        )
-
-        # Check that all rows in generated_df are in cache_df
-        if not merged_df[value_col_str + generated_suffix].equals(
-            merged_df[value_col_str + cached_suffix],
-        ):
-            self.msg.info(f"Cache miss, computed values didn't match {file_pattern}")
-
-            # Keep this variable for easier inspection
-            unequal_rows = merged_df[  # pylint: disable=unused-variable
-                merged_df[value_col_str + generated_suffix]
-                != merged_df[value_col_str + cached_suffix]
-            ]
-
-            return False
-
-        # If all checks passed, return true
-        msg.good(f"Cache hit for {value_col_str}")
-        return True
-
-    def _write_feature_to_cache(
-        self,
-        values_df: pd.DataFrame,
-        predictor_spec: TemporalSpec,
-        file_name: str,
-    ):
-        """Write feature to cache."""
-        out_df = values_df
-
-        # Drop rows containing fallback, since it's non-informative
-        out_df = out_df[
-            out_df[predictor_spec.get_col_str()] != predictor_spec.fallback
-        ].dropna()
-
-        # Write df to cache
-        timestamp = dt.datetime.now().strftime("%Y-%m-%d_%H-%M-%S")
-
-        # Write df to cache
-        write_df_to_file(
-            df=out_df,
-            file_path=self.feature_cache_dir / f"{file_name}_{timestamp}.parquet",
-        )
-
     def _get_temporal_feature(
         self,
         feature_spec: TemporalSpec,
-        file_suffix: str = "parquet",
     ) -> pd.DataFrame:
         """Get feature. Either load from cache, or generate if necessary.
 
@@ -450,26 +262,14 @@ class TimeseriesFlattener:  # pylint: disable=too-many-instance-attributes
         Returns:
             pd.DataFrame: Feature
         """
-        file_name = f"{feature_spec.get_col_str()}_{self.n_uuids}_uuids"
-
-        if hasattr(self, "feature_cache_dir"):
-            if self._cache_is_hit(
-                file_pattern=file_name,
-                output_spec=feature_spec,
-                file_suffix="parquet",
-            ):
-                df = self._load_cached_df_and_expand_fallback(
-                    file_pattern=file_name,
-                    full_col_str=feature_spec.get_col_str(),
-                    fallback=feature_spec.fallback,
-                    file_suffix=file_suffix,
-                )
 
-                return df.set_index(keys=self.pred_time_uuid_col_name).sort_index()
-        else:
-            msg.info("No cache dir specified, not attempting load")
+        if self.cache and self.cache.feature_exists(feature_spec=feature_spec):
+            df = self.cache.read_feature(feature_spec=feature_spec)
+            return df.set_index(keys=self.pred_time_uuid_col_name).sort_index()
+        elif not self.cache:
+            msg.info("No cache specified, not attempting load")
 
-        df = self._flatten_temporal_values_to_df(
+        df = self.flatten_temporal_values_to_df(
             prediction_times_with_uuid_df=self._df[
                 [
                     self.pred_time_uuid_col_name,
@@ -484,11 +284,10 @@ class TimeseriesFlattener:  # pylint: disable=too-many-instance-attributes
         )
 
         # Write df to cache if exists
-        if hasattr(self, "feature_cache_dir"):
-            self._write_feature_to_cache(
-                predictor_spec=feature_spec,
-                file_name=file_name,
-                values_df=df,
+        if self.cache:
+            self.cache.write_feature(
+                feature_spec=feature_spec,
+                df=df,
             )
 
         return (
@@ -509,7 +308,10 @@ class TimeseriesFlattener:  # pylint: disable=too-many-instance-attributes
                     raise ValueError("Dataframes are not of equal length")
 
     def _check_dfs_have_identical_indexes(self, dfs: list[pd.DataFrame]):
-        """Randomly sample 50 positions in each df and check that their indeces
+        """Randomly sample 50 positions in each df and check that their.
+
+        indeces.
+
         are identical.
 
         This checks that all the dataframes are aligned before
@@ -585,7 +387,8 @@ class TimeseriesFlattener:  # pylint: disable=too-many-instance-attributes
         output_prefix: str = "pred",
         birth_year_as_predictor: bool = False,
     ):
-        """Add age at prediction time and patient's birth year to each
+        """Add age at prediction time and patient's birth year to each.
+
         prediction time.
 
         Args:
@@ -756,6 +559,7 @@ class TimeseriesFlattener:  # pylint: disable=too-many-instance-attributes
         output_spec: PredictorSpec,
     ):
         """Add a column with predictor values to the flattened dataset (e.g.
+
         "average value of bloodsample within n days").
 
         Args:
@@ -767,9 +571,10 @@ class TimeseriesFlattener:  # pylint: disable=too-many-instance-attributes
 
     def add_temporal_col_to_flattened_dataset(
         self,
-        output_spec: TemporalSpec,
+        output_spec: AnySpec,
     ):
-        """Add a column to the dataset (either predictor or outcome depending
+        """Add a column to the dataset (either predictor or outcome depending.
+
         on the value of "direction").
 
         Args:
@@ -783,7 +588,7 @@ class TimeseriesFlattener:  # pylint: disable=too-many-instance-attributes
                 f"{self.timestamp_col_name} is of type {timestamp_col_type}, not 'Timestamp' from Pandas. Will cause problems. Convert before initialising FlattenedDataset.",
             )
 
-        df = TimeseriesFlattener._flatten_temporal_values_to_df(
+        df = TimeseriesFlattener.flatten_temporal_values_to_df(
             prediction_times_with_uuid_df=self._df[
                 [
                     self.id_col_name,
@@ -809,7 +614,8 @@ class TimeseriesFlattener:  # pylint: disable=too-many-instance-attributes
         pred_times_with_uuid: DataFrame,
         pred_time_uuid_colname: str,
     ) -> DataFrame:
-        """Ensure all prediction times are represented in the returned
+        """Ensure all prediction times are represented in the returned.
+
         dataframe.
 
         Args:
@@ -834,7 +640,8 @@ class TimeseriesFlattener:  # pylint: disable=too-many-instance-attributes
         df: DataFrame,
         pred_time_uuid_colname: str,
     ) -> DataFrame:
-        """Apply the resolve_multiple function to prediction_times where there
+        """Apply the resolve_multiple function to prediction_times where there.
+
         are multiple values within the interval_days lookahead.
 
         Args:
@@ -875,7 +682,8 @@ class TimeseriesFlattener:  # pylint: disable=too-many-instance-attributes
         timestamp_pred_colname: str,
         timestamp_value_colname: str,
     ) -> DataFrame:
-        """Keep only rows where timestamp_value is within interval_days in
+        """Keep only rows where timestamp_value is within interval_days in.
+
         direction of timestamp_pred.
 
         Args:


Test patch:
diff --git a/src/timeseriesflattener/testing/utils_for_testing.py b/src/timeseriesflattener/testing/utils_for_testing.py
index 13d28bc..0f9e2ef 100644
--- a/src/timeseriesflattener/testing/utils_for_testing.py
+++ b/src/timeseriesflattener/testing/utils_for_testing.py
@@ -11,7 +11,7 @@ from pandas import DataFrame
 from pandas.testing import assert_series_equal
 
 from timeseriesflattener import TimeseriesFlattener
-from timeseriesflattener.feature_spec_objects import TemporalSpec
+from timeseriesflattener.feature_spec_objects import AnySpec
 from timeseriesflattener.testing.load_synth_data import (
     load_synth_outcome,
     load_synth_prediction_times,
@@ -77,7 +77,7 @@ def str_to_df(
 
 def assert_flattened_data_as_expected(
     prediction_times_df: Union[pd.DataFrame, str],
-    output_spec: TemporalSpec,
+    output_spec: AnySpec,
     expected_df: Optional[pd.DataFrame] = None,
     expected_values: Optional[Sequence[Any]] = None,
 ):
diff --git a/tests/test_feature_cache/test_cache_to_disk.py b/tests/test_feature_cache/test_cache_to_disk.py
new file mode 100644
index 0000000..d95a315
--- /dev/null
+++ b/tests/test_feature_cache/test_cache_to_disk.py
@@ -0,0 +1,119 @@
+"""Testing of the DiskCache class."""
+
+import numpy as np
+import pandas as pd
+from pandas.testing import assert_frame_equal
+
+from timeseriesflattener.feature_cache.cache_to_disk import DiskCache
+from timeseriesflattener.feature_spec_objects import PredictorSpec
+from timeseriesflattener.resolve_multiple_functions import latest
+
+
+def test_write_and_check_feature(tmp_path):
+    """Test that write_feature writes a feature to disk."""
+
+    cache = DiskCache(
+        feature_cache_dir=tmp_path,
+        pred_time_uuid_col_name="pred_time_uuid",
+        id_col_name="dw_ek_borger",
+        cache_file_suffix="csv",
+        prediction_times_df=pd.DataFrame(
+            {"uuid": [1, 2, 3], "pred_time_uuid": [1, 2, 3]},
+        ),
+    )
+
+    values_df = pd.DataFrame(
+        {
+            "dw_ek_borger": [1, 2, 3],
+            "pred_time_uuid": [1, 2, 3],
+            "timestamp": [1, 2, 3],
+        },
+    )
+
+    test_spec = PredictorSpec(
+        values_df=values_df,
+        interval_days=5,
+        resolve_multiple_fn=latest,
+        key_for_resolve_multiple="latest",
+        fallback=np.nan,
+        feature_name="test_feature",
+    )
+
+    generated_df = pd.DataFrame(
+        {
+            "dw_ek_borger": [1, 2, 3],
+            "pred_time_uuid": [1, 2, 3],
+            "timestamp": [1, 2, 3],
+            f"{test_spec.get_col_str()}": [1, 2, 3],
+        },
+    )
+
+    assert cache.feature_exists(feature_spec=test_spec) is False
+
+    cache.write_feature(feature_spec=test_spec, df=generated_df)
+
+    assert cache.feature_exists(feature_spec=test_spec) is True
+
+
+def test_read_feature(tmp_path):
+    """Test that read_feature reads a feature from disk.
+
+    Important that one row contains the fallback because we then test
+    removing fallback vals when saving and expanding them again when
+    reading.
+    """
+
+    # Note that initialisation is much simpler i flattened dataset, since
+    # many of the col names are specified in the instantiation of the
+    # flattened dataset, and passed along to the cache.
+    cache = DiskCache(
+        feature_cache_dir=tmp_path,
+        pred_time_uuid_col_name="pred_time_uuid",
+        id_col_name="dw_ek_borger",
+        timestamp_col_name="timestamp",
+        cache_file_suffix="csv",
+        prediction_times_df=pd.DataFrame(
+            {"pred_time_uuid": [1, 2, 3], "dw_ek_borger": [1, 2, 3]},
+        ),
+    )
+
+    values_df = pd.DataFrame(
+        {
+            "dw_ek_borger": [1, 2, 3],
+            "pred_time_uuid": [1, 2, 3],
+            "timestamp": [1, 2, 3],
+        },
+    )
+
+    test_spec = PredictorSpec(
+        values_df=values_df,
+        interval_days=5,
+        resolve_multiple_fn=latest,
+        key_for_resolve_multiple="latest",
+        fallback=np.nan,
+        feature_name="test_feature",
+    )
+
+    generated_df = pd.DataFrame(
+        {
+            "dw_ek_borger": [
+                1,
+                2,
+                3,
+            ],  # The merge recasts to the most general type, int -> float
+            "pred_time_uuid": [1, 2, 3],
+            "timestamp": [1, 2, 3],
+            f"{test_spec.get_col_str()}": [1, 2, np.nan],
+        },
+    )
+
+    cache.write_feature(feature_spec=test_spec, df=generated_df)
+
+    df = cache.read_feature(feature_spec=test_spec)
+
+    # For each column in df, check that the values are equal to generated_df
+    for col in df.columns:
+        assert_frame_equal(
+            df[col].to_frame(),
+            generated_df[col].to_frame(),
+        )
diff --git a/tests/test_timeseriesflattener/test_flattened_dataset_generation.py b/tests/test_timeseriesflattener/test_flattened_dataset_generation.py
index 7c1ad59..c4e801c 100644
--- a/tests/test_timeseriesflattener/test_flattened_dataset_generation.py
+++ b/tests/test_timeseriesflattener/test_flattened_dataset_generation.py
@@ -3,17 +3,19 @@
 # pylint: disable=unused-import, redefined-outer-name
 
 from pathlib import Path
-from typing import Iterable, List
+from typing import Iterable, List, Optional
 
 import numpy as np
 import pandas as pd
 import pytest
 
+from timeseriesflattener.feature_cache.abstract_feature_cache import FeatureCache
+from timeseriesflattener.feature_cache.cache_to_disk import DiskCache
 from timeseriesflattener.feature_spec_objects import (
+    AnySpec,
     OutcomeSpec,
     PredictorGroupSpec,
     PredictorSpec,
-    TemporalSpec,
 )
 from timeseriesflattener.flattened_dataset import TimeseriesFlattener
 from timeseriesflattener.testing.load_synth_data import (
@@ -89,15 +91,15 @@ def check_dfs_have_same_contents_by_column(df1, df2):
 
 
 def create_flattened_df(
-    cache_dir: Path,
     predictor_specs: list[PredictorSpec],
     prediction_times_df: pd.DataFrame,
+    cache: Optional[FeatureCache] = None,
 ):
     """Create a dataset df for testing."""
     flat_ds = TimeseriesFlattener(
         prediction_times_df=prediction_times_df,
         n_workers=1,
-        feature_cache_dir=cache_dir,
+        cache=cache,
     )
 
     flat_ds.add_temporal_predictors_from_pred_specs(
@@ -118,18 +120,23 @@ def test_cache_hitting(
 ):
     """Test that cache hits."""
 
+    cache = DiskCache(
+        feature_cache_dir=tmp_path,
+        id_col_name="dw_ek_borger",
+    )
+
     # Create the cache
     first_df = create_flattened_df(
-        cache_dir=tmp_path,
         predictor_specs=predictor_specs.copy(),
         prediction_times_df=synth_prediction_times,
+        cache=cache,
     )
 
     # Load the cache
     cache_df = create_flattened_df(
-        cache_dir=tmp_path,
         predictor_specs=predictor_specs.copy(),
         prediction_times_df=synth_prediction_times,
+        cache=cache,
     )
 
     # Assert that each column has the same contents