junifer 0.0.7.dev332__py3-none-any.whl → 0.0.8__py3-none-any.whl

junifer/storage/pandas_base.py

@@ -6,14 +6,13 @@
 
 import json
 from collections.abc import Sequence
-from pathlib import Path
-from typing import ClassVar, Optional, Union
+from typing import ClassVar
 
 import numpy as np
 import pandas as pd
 
 from ..utils import raise_error
-from .base import BaseFeatureStorage
+from .base import BaseFeatureStorage, StorageType
 
 
 __all__ = ["PandasBaseFeatureStorage"]
@@ -38,19 +37,12 @@ class PandasBaseFeatureStorage(BaseFeatureStorage):
 
     """
 
-    _STORAGE_TYPES: ClassVar[Sequence[str]] = [
-        "vector",
-        "timeseries",
-        "matrix",
+    _STORAGE_TYPES: ClassVar[Sequence[StorageType]] = [
+        StorageType.Vector,
+        StorageType.Timeseries,
+        StorageType.Matrix,
     ]
 
-    def __init__(
-        self,
-        uri: Union[str, Path],
-        single_output: bool = True,
-    ) -> None:
-        super().__init__(uri=uri, single_output=single_output)
-
     def _meta_row(self, meta: dict, meta_md5: str) -> pd.DataFrame:
         """Convert the metadata to a pandas DataFrame.
 
@@ -75,8 +67,8 @@ class PandasBaseFeatureStorage(BaseFeatureStorage):
 
     @staticmethod
     def element_to_index(
-        element: dict, n_rows: int = 1, rows_col_name: Optional[str] = None
-    ) -> Union[pd.Index, pd.MultiIndex]:
+        element: dict, n_rows: int = 1, rows_col_name: str | None = None
+    ) -> pd.Index | pd.MultiIndex:
         """Convert the element metadata to index.
 
         Parameters
@@ -124,7 +116,11 @@ class PandasBaseFeatureStorage(BaseFeatureStorage):
         return index
 
     def store_df(
-        self, meta_md5: str, element: dict, df: Union[pd.DataFrame, pd.Series]
+        self,
+        meta_md5: str,
+        element: dict,
+        df: pd.DataFrame | pd.Series,
+        processed_element: dict | None = None,
     ) -> None:
         """Implement pandas DataFrame storing.
 
@@ -136,6 +132,10 @@ class PandasBaseFeatureStorage(BaseFeatureStorage):
             The element as a dictionary.
         df : pandas.DataFrame or pandas.Series
             The pandas DataFrame or Series to store.
+        processed_element : dict or None, optional
+            The element being processed when the data was computed, used to
+            choose the file of each element if ``single_output=False``. If
+            None, ``element`` is used (default None).
 
         Raises
         ------
@@ -153,9 +153,10 @@ class PandasBaseFeatureStorage(BaseFeatureStorage):
         self,
         meta_md5: str,
         element: dict,
-        data: Union[np.ndarray, list],
-        col_names: Optional[Sequence[str]] = None,
-        rows_col_name: Optional[str] = None,
+        data: np.ndarray | list,
+        col_names: Sequence[str] | None = None,
+        rows_col_name: str | None = None,
+        processed_element: dict | None = None,
     ) -> None:
         """Store 2D data.
 
@@ -173,6 +174,10 @@ class PandasBaseFeatureStorage(BaseFeatureStorage):
             The column name to use in case number of rows greater than 1.
             If None and number of rows greater than 1, then the name will be
             "idx" (default None).
+        processed_element : dict or None, optional
+            The element being processed when the data was computed, used to
+            choose the file of each element if ``single_output=False``. If
+            None, ``element`` is used (default None).
 
         """
         # Convert element metadata to index
@@ -186,14 +191,20 @@ class PandasBaseFeatureStorage(BaseFeatureStorage):
             index=idx,  # type: ignore
         )
         # Store dataframe
-        self.store_df(meta_md5=meta_md5, element=element, df=df)
+        self.store_df(
+            meta_md5=meta_md5,
+            element=element,
+            df=df,
+            processed_element=processed_element,
+        )
 
     def store_vector(
         self,
         meta_md5: str,
         element: dict,
-        data: Union[np.ndarray, list],
-        col_names: Optional[Sequence[str]] = None,
+        data: np.ndarray | list,
+        col_names: Sequence[str] | None = None,
+        processed_element: dict | None = None,
     ) -> None:
         """Store vector.
 
@@ -207,6 +218,10 @@ class PandasBaseFeatureStorage(BaseFeatureStorage):
             The vector data to store.
         col_names : list-like of str, optional
             The column labels (default None).
+        processed_element : dict or None, optional
+            The element being processed when the data was computed, used to
+            choose the file of each element if ``single_output=False``. If
+            None, ``element`` is used (default None).
 
         """
         if isinstance(data, list):
@@ -224,6 +239,7 @@ class PandasBaseFeatureStorage(BaseFeatureStorage):
             element=element,
             data=data,
             col_names=col_names,
+            processed_element=processed_element,
         )
 
     def store_timeseries(
@@ -231,7 +247,8 @@ class PandasBaseFeatureStorage(BaseFeatureStorage):
         meta_md5: str,
         element: dict,
         data: np.ndarray,
-        col_names: Optional[Sequence[str]] = None,
+        col_names: Sequence[str] | None = None,
+        processed_element: dict | None = None,
     ) -> None:
         """Store timeseries.
 
@@ -245,6 +262,10 @@ class PandasBaseFeatureStorage(BaseFeatureStorage):
             The timeseries data to store.
         col_names : list-like of str, optional
             The column labels (default None).
+        processed_element : dict or None, optional
+            The element being processed when the data was computed, used to
+            choose the file of each element if ``single_output=False``. If
+            None, ``element`` is used (default None).
 
         """
         self._store_2d(
@@ -253,4 +274,5 @@ class PandasBaseFeatureStorage(BaseFeatureStorage):
             data=data,
             col_names=col_names,
             rows_col_name="timepoint",
+            processed_element=processed_element,
         )

junifer/storage/sqlite.py

@@ -5,28 +5,40 @@
 # License: AGPL
 
 import json
-from collections.abc import Sequence
-from pathlib import Path
-from typing import TYPE_CHECKING, Any, Optional, Union
+from enum import StrEnum
+from typing import TYPE_CHECKING, Any, Optional
 
 import numpy as np
 import pandas as pd
 from pandas.core.base import NoNewAttributesMixin
 from pandas.io.sql import pandasSQL_builder
 from sqlalchemy import create_engine, inspect
+from sqlalchemy.pool import NullPool
 from tqdm import tqdm
 
 from ..api.decorators import register_storage
-from ..utils import logger, raise_error, warn_with_log
+from ..utils import raise_error, warn_with_log
+from .base import MatrixKind, logger
 from .pandas_base import PandasBaseFeatureStorage
-from .utils import element_to_prefix, matrix_to_vector, store_matrix_checks
+from .utils import (
+    element_to_prefix,
+    matrix_to_vector,
+    store_matrix_checks,
+)
 
 
 if TYPE_CHECKING:
     from sqlalchemy.engine import Engine
 
 
-__all__ = ["SQLiteFeatureStorage"]
+__all__ = ["SQLiteFeatureStorage", "Upsert"]
+
+
+class Upsert(StrEnum):
+    """Accepted upsert value."""
+
+    Update = "update"
+    Ignore = "ignore"
 
 
 @register_storage
@@ -44,9 +56,11 @@ class SQLiteFeatureStorage(PandasBaseFeatureStorage):
         ``uri`` and store all the elements in the same file. This behaviour
         is only suitable for non-parallel executions. SQLite does not
         support concurrency (default True).
-    upsert : {"ignore", "update"}, optional
-        Upsert mode. If "ignore" is used, the existing elements are ignored.
-        If "update", the existing elements are updated (default "update").
+    upsert : :enum:`.Upsert`, optional
+        Upsert mode. If ``Upsert.Ignore`` is used, the existing elements are
+        ignored, so the elements that are already stored are not stored
+        again. If ``Upsert.Update``, the existing elements are updated
+        (default ``Upsert.Ignore``).
 
     See Also
     --------
@@ -55,27 +69,9 @@ class SQLiteFeatureStorage(PandasBaseFeatureStorage):
 
     """
 
-    def __init__(
-        self,
-        uri: Union[str, Path],
-        single_output: bool = True,
-        upsert: str = "update",
-    ) -> None:
-        # Check and set upsert argument value
-        if upsert not in ["update", "ignore"]:
-            raise_error(
-                msg=(
-                    "Invalid choice for `upsert`. "
-                    "Must be either 'update' or 'ignore'."
-                )
-            )
-        self.upsert = upsert
-        super().__init__(
-            uri=uri,
-            single_output=single_output,
-        )
+    upsert: Upsert = Upsert.Ignore
 
-    def get_engine(self, element: Optional[dict] = None) -> "Engine":
+    def get_engine(self, element: dict | None = None) -> "Engine":
         """Get engine.
 
         Parameters
@@ -105,11 +101,11 @@ class SQLiteFeatureStorage(PandasBaseFeatureStorage):
                 prefix = element_to_prefix(element)
         # Format URI for engine creation
         uri = f"sqlite:///{self.uri.parent}/{prefix}{self.uri.name}"
-        return create_engine(uri, echo=False)
+        return create_engine(uri, echo=False, poolclass=NullPool)
 
     def _save_upsert(
         self,
-        df: Union[pd.DataFrame, pd.Series],
+        df: pd.DataFrame | pd.Series,
         name: str,
         engine: Optional["Engine"] = None,
         if_exists: str = "append",
@@ -225,11 +221,9 @@ class SQLiteFeatureStorage(PandasBaseFeatureStorage):
 
     def read(
         self,
-        feature_name: Optional[str] = None,
-        feature_md5: Optional[str] = None,
-    ) -> dict[
-        str, Union[str, list[Union[int, str, dict[str, str]]], np.ndarray]
-    ]:
+        feature_name: str | None = None,
+        feature_md5: str | None = None,
+    ) -> dict[str, str | list[int | str | dict[str, str]] | np.ndarray]:
         """Read stored feature.
 
         Parameters
@@ -261,8 +255,8 @@ class SQLiteFeatureStorage(PandasBaseFeatureStorage):
 
     def read_df(
         self,
-        feature_name: Optional[str] = None,
-        feature_md5: Optional[str] = None,
+        feature_name: str | None = None,
+        feature_md5: str | None = None,
     ) -> pd.DataFrame:
         """Implement feature reading into a pandas DataFrame.
 
@@ -347,7 +341,13 @@ class SQLiteFeatureStorage(PandasBaseFeatureStorage):
         df = df.set_index(index_names)
         return df
 
-    def store_metadata(self, meta_md5: str, element: dict, meta: dict) -> None:
+    def store_metadata(
+        self,
+        meta_md5: str,
+        element: dict,
+        meta: dict,
+        processed_element: dict | None = None,
+    ) -> None:
         """Implement metadata storing in the storage.
 
         Parameters
@@ -358,10 +358,15 @@ class SQLiteFeatureStorage(PandasBaseFeatureStorage):
             The element as a dictionary.
         meta : dict
             The metadata as a dictionary.
+        processed_element : dict or None, optional
+            The element being processed when the data was computed, used to
+            choose the file of each element if ``single_output=False``. If
+            None, ``element`` is used (default None).
 
         """
         # Get sqlalchemy engine
-        engine = self.get_engine(element=element)
+        # The file of the element being processed (if single_output=False)
+        engine = self.get_engine(element=processed_element or element)
         table_name = f"meta_{meta_md5}"
         if table_name not in inspect(engine).get_table_names():
             # Convert metadata to dataframe
@@ -370,7 +375,11 @@ class SQLiteFeatureStorage(PandasBaseFeatureStorage):
             self._save_upsert(meta_df, "meta", engine)
 
     def store_df(
-        self, meta_md5: str, element: dict, df: Union[pd.DataFrame, pd.Series]
+        self,
+        meta_md5: str,
+        element: dict,
+        df: pd.DataFrame | pd.Series,
+        processed_element: dict | None = None,
     ) -> None:
         """Implement pandas DataFrame storing.
 
@@ -382,6 +391,10 @@ class SQLiteFeatureStorage(PandasBaseFeatureStorage):
             The element as a dictionary.
         df : pandas.DataFrame or pandas.Series
             The pandas DataFrame or Series to store.
+        processed_element : dict or None, optional
+            The element being processed when the data was computed, used to
+            choose the file of each element if ``single_output=False``. If
+            None, ``element`` is used (default None).
 
         Raises
         ------
@@ -416,7 +429,8 @@ class SQLiteFeatureStorage(PandasBaseFeatureStorage):
 
         table_name = f"meta_{meta_md5}"
         # Get sqlalchemy engine
-        engine = self.get_engine(element)
+        # The file of the element being processed (if single_output=False)
+        engine = self.get_engine(element=processed_element or element)
         # Save data
         self._save_upsert(df, table_name, engine)
 
@@ -425,10 +439,11 @@ class SQLiteFeatureStorage(PandasBaseFeatureStorage):
         meta_md5: str,
         element: dict,
         data: np.ndarray,
-        col_names: Optional[Sequence[str]] = None,
-        row_names: Optional[Sequence[str]] = None,
-        matrix_kind: str = "full",
+        col_names: list[str] | None = None,
+        row_names: list[str] | None = None,
+        matrix_kind: MatrixKind = MatrixKind.Full,
         diagonal: bool = True,
+        processed_element: dict | None = None,
     ) -> None:
         """Store matrix.
 
@@ -440,23 +455,19 @@ class SQLiteFeatureStorage(PandasBaseFeatureStorage):
             The element as a dictionary.
         data : numpy.ndarray
             The matrix data to store.
-        meta : dict
-            The metadata as a dictionary.
         col_names : list-like of str, optional
             The column labels (default None).
         row_names : list-like of str, optional
             The row labels (optional None).
-        matrix_kind : str, optional
-            The kind of matrix:
-
-            * ``triu`` : store upper triangular only
-            * ``tril`` : store lower triangular
-            * ``full`` : full matrix
-
-            (default "full").
+        matrix_kind : :enum:`.MatrixKind`, optional
+            The matrix kind (default ``MatrixKind.Full``).
         diagonal : bool, optional
-            Whether to store the diagonal. If ``matrix_kind = full``, setting
-            this to False will raise an error (default True).
+            Whether to store the diagonal. If ``matrix_kind=MatrixKind.Full``,
+            setting this to False will raise an error (default True).
+        processed_element : dict or None, optional
+            The element being processed when the data was computed, used to
+            choose the file of each element if ``single_output=False``. If
+            None, ``element`` is used (default None).
 
         """
         # Row data validation
@@ -508,7 +519,12 @@ class SQLiteFeatureStorage(PandasBaseFeatureStorage):
             data_df.index.names = new_names
 
         # Store dataframe
-        self.store_df(meta_md5=meta_md5, element=element, df=data_df)
+        self.store_df(
+            meta_md5=meta_md5,
+            element=element,
+            df=data_df,
+            processed_element=processed_element,
+        )
 
     def collect(self) -> None:
         """Implement data collection.
@@ -526,9 +542,10 @@ class SQLiteFeatureStorage(PandasBaseFeatureStorage):
             )
         logger.info(f"Collecting data from {self.uri.parent}/*{self.uri.name}")
         # Create new instance
-        out_storage = SQLiteFeatureStorage(uri=self.uri, upsert="ignore")
+        out_storage = SQLiteFeatureStorage(uri=self.uri, upsert=Upsert.Ignore)
         # Glob files
-        files = self.uri.parent.glob(f"*{self.uri.name}")
+        # Only the files of the elements, not the collected file
+        files = self.uri.parent.glob(f"*_{self.uri.name}")
         for elem in tqdm(files, desc="file"):
             logger.debug(f"Reading from {elem.absolute()!s}")
             in_storage = SQLiteFeatureStorage(uri=elem)
@@ -539,6 +556,7 @@ class SQLiteFeatureStorage(PandasBaseFeatureStorage):
             )
             # Save metadata
             out_storage._save_upsert(t_meta_df, "meta")
+            in_features = in_storage.list_features()
             # Save dataframes
             for meta_md5 in tqdm(t_meta_df.index, desc="feature"):
                 logger.debug(f"Collecting feature {meta_md5}")
@@ -546,8 +564,27 @@ class SQLiteFeatureStorage(PandasBaseFeatureStorage):
                 # properly
                 table_name = f"meta_{meta_md5}"
                 t_df = in_storage.read_df(feature_md5=meta_md5)
+                # Only the files of single elements can be collected
+                element_keys = in_features[meta_md5]["_element_keys"]
+                n_elements = len(
+                    t_df.index.to_frame(index=False)[
+                        element_keys
+                    ].drop_duplicates()
+                )
+                if n_elements != 1:
+                    raise_error(
+                        msg=(
+                            f"The file {elem} has {n_elements} elements for "
+                            f"the feature {meta_md5}, but only the files with "
+                            "one element (as stored for each element with "
+                            "`single_output=False`) can be collected"
+                        ),
+                        klass=RuntimeError,
+                    )
                 # Save data
-                out_storage._save_upsert(t_df, table_name, if_exists="nocheck")
+                # Skip the elements already collected (e.g., the data of a
+                # subject in the files of each of its tasks)
+                out_storage._save_upsert(t_df, table_name)
 
 
 # TODO: refactor
@@ -582,7 +619,7 @@ def _generate_update_statements(table, index_col, rows_to_update):
     for i, (_, keys) in enumerate(pk_indb.iterrows()):
         stmt = (
             table.update()
-            .where(and_(col == keys[j] for j, col in enumerate(pk_cols)))  # type: ignore
+            .where(and_(col == keys.iloc[j] for j, col in enumerate(pk_cols)))  # type: ignore
             .values(new_records[i])
         )
         stmts.append(stmt)

junifer/storage/tests/test_hdf5.py

@@ -10,6 +10,7 @@ from pathlib import Path
 import h5py
 import numpy as np
 import pytest
+from h5io import read_hdf5
 from numpy.testing import assert_array_equal
 from pandas.testing import assert_frame_equal
 
@@ -21,18 +22,6 @@ from junifer.storage.utils import (
 )
 
 
-def test_get_valid_inputs() -> None:
-    """Test valid inputs."""
-    storage = HDF5FeatureStorage(uri="/tmp")
-    assert set(storage.get_valid_inputs()) == {
-        "matrix",
-        "vector",
-        "timeseries",
-        "scalar_table",
-        "timeseries_2d",
-    }
-
-
 def test_single_output(tmp_path: Path) -> None:
     """Test single output setup.
 
@@ -1178,3 +1167,149 @@ def test_collect_error_single_output() -> None:
     ):
         storage = HDF5FeatureStorage(uri="/tmp", single_output=True)
         storage.collect()
+
+
+def test_read_collected_with_h5io_fork() -> None:
+    """Test reading features collected with junifer's fork of h5io.
+
+    The file was collected from three elements (``sub-01`` to ``sub-03``),
+    each with a vector, a matrix and a timeseries computed from the element
+    index ``i``.
+
+    """
+    uri = Path(__file__).parent / "data" / "collected_h5io_fork.hdf5"
+    storage = HDF5FeatureStorage(uri=uri)
+    md5s = {meta["name"]: md5 for md5, meta in storage.list_features().items()}
+    elements = [{"subject": f"sub-0{i}"} for i in range(1, 4)]
+    # The vector data is stored as a "chunkedarray"
+    vector = storage.read(feature_md5=md5s["BOLD_vec"])
+    assert vector["element"] == elements
+    assert vector["data"].dtype == np.float32
+    for i in range(3):
+        assert_array_equal(
+            vector["data"][:, i], np.float32(np.array([0.1, 0.2]) + i)
+        )
+    # The matrix data is stored as a "chunkedarray"
+    matrix = storage.read(feature_md5=md5s["BOLD_mat"])
+    assert matrix["element"] == elements
+    for i in range(3):
+        assert_array_equal(
+            matrix["data"][..., i], np.array([[1.0, 0.5], [0.5, 1.0]]) + i
+        )
+    # The timeseries data is stored as a "chunkedlist"
+    timeseries = storage.read(feature_md5=md5s["BOLD_ts"])
+    assert timeseries["element"] == elements
+    assert len(timeseries["data"]) == 3
+    for i in range(3):
+        assert_array_equal(
+            timeseries["data"][i], np.arange(6.0).reshape(3, 2) + i
+        )
+    assert storage.read_df(feature_md5=md5s["BOLD_ts"]).shape == (9, 2)
+
+
+def _store_and_collect(uri: Path, force_float32: bool) -> HDF5FeatureStorage:
+    """Store a vector and a timeseries for three elements and collect them.
+
+    Parameters
+    ----------
+    uri : pathlib.Path
+        The path to the collected file.
+    force_float32 : bool
+        Whether to store the data as float32.
+
+    Returns
+    -------
+    HDF5FeatureStorage
+        The storage of the collected file.
+
+    """
+    storage = HDF5FeatureStorage(
+        uri=uri,
+        single_output=False,
+        force_float32=force_float32,
+        chunk_size=2,
+    )
+    for i in range(3):
+        meta = {
+            "element": {"subject": f"sub-0{i}"},
+            "dependencies": [],
+            "type": "BOLD",
+        }
+        storage.store(
+            kind="vector",
+            meta={**meta, "marker": {"name": "vec"}},
+            data=np.array([[0.1, 0.2]]) + i,
+            col_names=["a", "b"],
+        )
+        storage.store(
+            kind="timeseries",
+            meta={**meta, "marker": {"name": "ts"}},
+            data=np.arange(6.0).reshape(3, 2) + i,
+            col_names=["a", "b"],
+        )
+    storage.collect()
+    return HDF5FeatureStorage(uri=uri)
+
+
+@pytest.mark.parametrize(
+    "force_float32, dtype", [(True, np.float32), (False, np.float64)]
+)
+def test_collect_dtype(
+    tmp_path: Path, force_float32: bool, dtype: np.dtype
+) -> None:
+    """Test collected data keeps its dtype.
+
+    Parameters
+    ----------
+    tmp_path : pathlib.Path
+        The path to the test directory.
+    force_float32 : bool
+        The parametrized casting of the data to float32.
+    dtype : numpy.dtype
+        The parametrized dtype of the collected data.
+
+    """
+    storage = _store_and_collect(tmp_path / "out.hdf5", force_float32)
+    md5s = {meta["name"]: md5 for md5, meta in storage.list_features().items()}
+    vector = storage.read(feature_md5=md5s["BOLD_vec"])
+    timeseries = storage.read(feature_md5=md5s["BOLD_ts"])
+    assert vector["data"].dtype == dtype
+    assert all(x.dtype == dtype for x in timeseries["data"])
+    # The elements are collected in any order
+    for i_element, element in enumerate(vector["element"]):
+        i = int(element["subject"][-1])
+        expected = np.array([0.1, 0.2]) + i
+        assert_array_equal(
+            vector["data"][:, i_element], expected.astype(dtype)
+        )
+
+
+def test_collect_h5io_types(tmp_path: Path) -> None:
+    """Test collected data is stored with h5io types and can be re-collected.
+
+    Parameters
+    ----------
+    tmp_path : pathlib.Path
+        The path to the test directory.
+
+    """
+    uri = tmp_path / "out.hdf5"
+    storage = _store_and_collect(uri, force_float32=False)
+    collected = {
+        md5: read_hdf5(uri, title=md5) for md5 in storage.list_features()
+    }
+    # The data is only in the features
+    with h5py.File(uri, mode="r") as fid:
+        assert set(fid.keys()) == {"meta", *collected.keys()}
+    # Collect again
+    _store_and_collect(uri, force_float32=False)
+    for md5, feature in collected.items():
+        recollected = read_hdf5(uri, title=md5)
+        assert recollected["element"] == feature["element"]
+        if isinstance(feature["data"], list):
+            for new, old in zip(
+                recollected["data"], feature["data"], strict=True
+            ):
+                assert_array_equal(new, old)
+        else:
+            assert_array_equal(recollected["data"], feature["data"])