Merge pull request #77 from NREL/dt/convert-storage-generic

Make TimeSeriesStorage.convert_storage more generic

Author: daniel-thom

src/infrasys/arrow_storage.py

@@ -139,18 +139,6 @@ def _get_single_time_series(
             normalization=metadata.normalization,
         )
 
-    def get_raw_single_time_series(self, time_series_uuid: UUID) -> NDArray:
-        fpath = self._ts_directory.joinpath(f"{time_series_uuid}{EXTENSION}")
-        with pa.OSFile(str(fpath), "r") as source:  # type: ignore
-            base_ts = pa.ipc.open_file(source).get_record_batch(0)
-            logger.trace("Reading time series from {}", fpath)
-        columns = base_ts.column_names
-        if len(columns) != 1:
-            msg = f"Bug: expected a single column: {columns=}"
-            raise Exception(msg)
-        column = columns[0]
-        return base_ts[column].to_numpy()
-
     def _convert_to_record_batch(self, time_series_array: NDArray, column: str) -> pa.RecordBatch:
         """Create record batch to save array to disk."""
         pa_array = pa.array(time_series_array)

src/infrasys/in_memory_time_series_storage.py

@@ -2,15 +2,14 @@
 
 from datetime import datetime
 from pathlib import Path
-import numpy as np
 from numpy.typing import NDArray
 from typing import Optional, TypeAlias
 from uuid import UUID
 
 from loguru import logger
 from infrasys.arrow_storage import ArrowTimeSeriesStorage
 
-from infrasys.exceptions import ISNotStored, ISOperationNotAllowed
+from infrasys.exceptions import ISNotStored
 from infrasys.time_series_models import (
     SingleTimeSeries,
     SingleTimeSeriesMetadata,
@@ -43,15 +42,6 @@ def add_time_series(self, metadata: TimeSeriesMetadata, time_series: TimeSeriesD
             msg = f"add_time_series not implemented for {type(time_series)}"
             raise NotImplementedError(msg)
 
-    def add_raw_single_time_series(
-        self, time_series_uuid: UUID, time_series_data: DataStoreType
-    ) -> None:
-        if time_series_uuid not in self._arrays:
-            self._arrays[time_series_uuid] = time_series_data
-            logger.debug("Added {} to store", time_series_uuid)
-        else:
-            logger.debug("{} was already stored", time_series_uuid)
-
     def get_time_series(
         self,
         metadata: TimeSeriesMetadata,
@@ -62,13 +52,6 @@ def get_time_series(
             return self._get_single_time_series(metadata, start_time, length)
         raise NotImplementedError(str(metadata.get_time_series_data_type()))
 
-    def get_raw_single_time_series(self, time_series_uuid: UUID) -> NDArray:
-        data_array = self._arrays[time_series_uuid]
-        if not isinstance(data_array, np.ndarray):
-            msg = f"Can't retrieve type: {type(data_array)} as single_time_series"
-            raise ISOperationNotAllowed(msg)
-        return data_array
-
     def remove_time_series(self, uuid: UUID) -> None:
         time_series = self._arrays.pop(uuid, None)
         if time_series is None:

src/infrasys/time_series_manager.py

@@ -308,9 +308,15 @@ def convert_storage(self, **kwargs) -> None:
         for time_series_uuid in self.metadata_store.unique_uuids_by_type(
             SingleTimeSeries.__name__
         ):
-            new_storage.add_raw_single_time_series(
-                time_series_uuid, self._storage.get_raw_single_time_series(time_series_uuid)
+            metadata = self.metadata_store.list_metadata_with_time_series_uuid(
+                time_series_uuid, limit=1
             )
+            if len(metadata) != 1:
+                msg = f"Expected 1 metadata for {time_series_uuid}, got {len(metadata)}"
+                raise Exception(msg)
+
+            time_series = self._storage.get_time_series(metadata[0])
+            new_storage.add_time_series(metadata[0], time_series)
 
         self._storage = new_storage
         return None

src/infrasys/time_series_metadata_store.py

@@ -288,6 +288,25 @@ def list_metadata(
         rows = execute(cur, query, params=params).fetchall()
         return [_deserialize_time_series_metadata(x[0]) for x in rows]
 
+    def list_metadata_with_time_series_uuid(
+        self, time_series_uuid: UUID, limit: int | None = None
+    ) -> list[TimeSeriesMetadata]:
+        """Return metadata attached to the given time_series_uuid.
+
+        Parameters
+        ----------
+        time_series_uuid
+            The UUID of the time series.
+        limit
+            The maximum number of metadata to return. If None, all metadata are returned.
+        """
+        params = (str(time_series_uuid),)
+        limit_str = "" if limit is None else f"LIMIT {limit}"
+        query = f"SELECT metadata FROM {self.TABLE_NAME} WHERE time_series_uuid = ? {limit_str}"
+        cur = self._con.cursor()
+        rows = execute(cur, query, params=params).fetchall()
+        return [_deserialize_time_series_metadata(x[0]) for x in rows]
+
     def _list_metadata_no_sql_json(
         self,
         *components: Component | SupplementalAttribute,

src/infrasys/time_series_storage_base.py

@@ -3,7 +3,7 @@
 import abc
 from datetime import datetime
 from pathlib import Path
-from typing import Optional, Any
+from typing import Optional
 from uuid import UUID
 
 from infrasys.time_series_models import TimeSeriesData, TimeSeriesMetadata
@@ -16,14 +16,6 @@ class TimeSeriesStorageBase(abc.ABC):
     def add_time_series(self, metadata: TimeSeriesMetadata, time_series: TimeSeriesData) -> None:
         """Store a time series array."""
 
-    @abc.abstractmethod
-    def add_raw_single_time_series(self, time_series_uuid: UUID, time_series_data: Any) -> None:
-        """Store a time series array from raw data."""
-
-    @abc.abstractmethod
-    def get_raw_single_time_series(self, time_series_uuid: UUID) -> Any:
-        """Get the raw time series data for a given uuid."""
-
     @abc.abstractmethod
     def get_time_series_directory(self) -> Path | None:
         """Return the directory containing time series files. Will be none for in-memory time

tests/test_in_memory_storage.py

@@ -8,14 +8,6 @@
 import pytest
 
 
-def get_data_and_uuids(system):
-    uuids = system._time_series_mgr.metadata_store.unique_uuids_by_type(SingleTimeSeries.__name__)
-    data = {
-        uuid: system._time_series_mgr._storage.get_raw_single_time_series(uuid) for uuid in uuids
-    }
-    return uuids, data
-
-
 @pytest.mark.parametrize(
     "original_kwargs,new_kwargs,original_stype,new_stype",
     [
@@ -46,14 +38,11 @@ def test_memory_convert_storage_time_series(
     with pytest.raises(ISAlreadyAttached):
         system.add_time_series(test_time_series_data, test_generator)
 
-    original_uuids, original_data = get_data_and_uuids(system)
-
     system.convert_storage(**new_kwargs)
 
-    assert isinstance(system._time_series_mgr._storage, new_stype)
-    new_uuids, new_data = get_data_and_uuids(system)
+    assert isinstance(system.time_series.storage, new_stype)
 
-    assert set(original_uuids) == set(new_uuids)
-
-    for uuid in new_uuids:
-        assert np.array_equal(original_data[uuid], new_data[uuid])
+    ts2 = system.get_time_series(
+        test_generator, time_series_type=SingleTimeSeries, variable_name="load"
+    )
+    assert np.array_equal(ts2.data_array, test_time_series_data.data_array)