Refactor time series management (#78)

* Refactor time series management

Author: daniel-thom

.gitignore

@@ -130,3 +130,4 @@ dmypy.json
 # pixi environments
 .pixi
 
+.vscode

docs/how_tos/list_time_series.md

@@ -3,61 +3,52 @@
 Suppose that you have added multiple time series arrays to your components using differing
 names and attributes. How can you see what is present?
 
-This example assumes that a system with two generators and time series data has been serialized
-to a file.
+This example uses a test module in the `infrasys` repository.
+
+The call to `system.add_time_series` returns a key. You can store those keys yourself or look them
+up later with `system.list_time_series_keys`. Here's how to do it.
 
 ```python
-from infrasys import Component, System
-
-system = System.from_json("system.json")
-for component in system.get_components(Component):
-    for metadata in system.list_time_series_metadata(component):
-        print(f"{component.label}: {metadata.label} {metadata.user_attributes}")
-
-Generator.gen1: SingleTimeSeries.active_power {'scenario': 'high', 'model_year': '2030'}
-Generator.gen1: SingleTimeSeries.active_power {'scenario': 'high', 'model_year': '2035'}
-Generator.gen1: SingleTimeSeries.active_power {'scenario': 'low', 'model_year': '2030'}
-Generator.gen1: SingleTimeSeries.active_power {'scenario': 'low', 'model_year': '2035'}
-Generator.gen1: SingleTimeSeries.reactive_power {'scenario': 'high', 'model_year': '2030'}
-Generator.gen1: SingleTimeSeries.reactive_power {'scenario': 'high', 'model_year': '2035'}
-Generator.gen1: SingleTimeSeries.reactive_power {'scenario': 'low', 'model_year': '2030'}
-Generator.gen1: SingleTimeSeries.reactive_power {'scenario': 'low', 'model_year': '2035'}
-Generator.gen2: SingleTimeSeries.active_power {'scenario': 'high', 'model_year': '2030'}
-Generator.gen2: SingleTimeSeries.active_power {'scenario': 'high', 'model_year': '2035'}
-Generator.gen2: SingleTimeSeries.active_power {'scenario': 'low', 'model_year': '2030'}
-Generator.gen2: SingleTimeSeries.active_power {'scenario': 'low', 'model_year': '2035'}
-Generator.gen2: SingleTimeSeries.reactive_power {'scenario': 'high', 'model_year': '2030'}
-Generator.gen2: SingleTimeSeries.reactive_power {'scenario': 'high', 'model_year': '2035'}
-Generator.gen2: SingleTimeSeries.reactive_power {'scenario': 'low', 'model_year': '2030'}
-Generator.gen2: SingleTimeSeries.reactive_power {'scenario': 'low', 'model_year': '2035'}
+from datetime import datetime, timedelta
+
+import numpy as np
+
+from infrasys import SingleTimeSeries
+from tests.models.simple_system import SimpleSystem, SimpleGenerator, SimpleBus
+
+system = SimpleSystem()
+bus = SimpleBus(name="test-bus", voltage=1.1)
+gen = SimpleGenerator(name="gen", active_power=1.0, rating=1.0, bus=bus, available=True)
+system.add_components(bus, gen)
+
+length = 10
+initial_time = datetime(year=2020, month=1, day=1)
+timestamps = [initial_time + timedelta(hours=i) for i in range(length)]
+variable_name = "active_power"
+ts1 = SingleTimeSeries.from_time_array(np.random.rand(length), variable_name, timestamps)
+ts2 = SingleTimeSeries.from_time_array(np.random.rand(length), variable_name, timestamps)
+key1 = system.add_time_series(ts1, gen, scenario="low")
+key2 = system.add_time_series(ts2, gen, scenario="high")
+
+# Use the keys directly.
+ts1_b = system.get_time_series_by_key(gen, key1)
+ts2_b = system.get_time_series_by_key(gen, key2)
+
+# Identify the keys later.
+for key in system.list_time_series_keys(gen):
+    print(f"{gen.label}: {key}")
+```
+```
+SimpleGenerator.gen: variable_name='active_power' initial_time=datetime.datetime(2020, 1, 1, 0, 0) resolution=datetime.timedelta(seconds=3600) time_series_type=<class 'infrasys.time_series_models.SingleTimeSeries'> user_attributes={'scenario': 'high'} length=10
+SimpleGenerator.gen: variable_name='active_power' initial_time=datetime.datetime(2020, 1, 1, 0, 0) resolution=datetime.timedelta(seconds=3600) time_series_type=<class 'infrasys.time_series_models.SingleTimeSeries'> user_attributes={'scenario': 'low'} length=10
 ```
 
-Now you can retrieve the exact instance you want.
+You can also retrieve time series by specifying the parameters as shown here:
 
 ```python
-system.time_series.get(gen1, variable_name="active_power", scenario="high", model_year="2035").data
-<pyarrow.lib.Int64Array object at 0x107a38d60>
-[
-  0,
-  1,
-  2,
-  3,
-  4,
-  5,
-  6,
-  7,
-  8,
-  9,
-  ...
-  8774,
-  8775,
-  8776,
-  8777,
-  8778,
-  8779,
-  8780,
-  8781,
-  8782,
-  8783
-]
+system.time_series.get(gen, variable_name="active_power", scenario="high")
+```
+```
+SingleTimeSeries(variable_name='active_power', normalization=None, data=array([0.29276233, 0.97400382, 0.76499075, 0.95080431, 0.61749027,
+       0.73899945, 0.57877704, 0.3411286 , 0.80701393, 0.53051773]), resolution=datetime.timedelta(seconds=3600), initial_time=datetime.datetime(2020, 1, 1, 0, 0), length=10)
 ```

pyproject.toml

@@ -27,18 +27,16 @@ classifiers = [
     "Programming Language :: Python :: Implementation :: PyPy",
 ]
 dependencies = [
-    # This is what we want, but are currently be limited by pyarrow.
-    # Leave unbounded until pyarrow upgrades numpy.
-    #"numpy >= 2, < 3",
-    "numpy",
-    "pyarrow~=15.0.2",
-    "pint~=0.23",
     "loguru~=0.7.2",
-    "rich~=13.7.1",
+    "numpy >= 2, < 3",
+    "pint~=0.23",
+    "pyarrow~=19.0",
     "pydantic >= 2.7, < 3",
+    "rich~=13.7.1",
 ]
 [project.optional-dependencies]
 dev = [
+    "autodoc_pydantic~=2.0",
     "furo",
     "mypy >=1.13, < 2",
     "myst_parser",
@@ -49,7 +47,6 @@ dev = [
     "ruff",
     "sphinx",
     "sphinx-click",
-    "autodoc_pydantic~=2.0",
     "sphinx-copybutton",
     "sphinx-tabs~=3.4",
 ]

src/infrasys/__init__.py

@@ -11,7 +11,12 @@
 from .normalization import NormalizationModel
 from .supplemental_attribute import SupplementalAttribute
 from .system import System
-from .time_series_models import SingleTimeSeries
+from .time_series_models import (
+    SingleTimeSeries,
+    TimeSeriesStorageType,
+    TimeSeriesKey,
+    SingleTimeSeriesKey,
+)
 
 
 __all__ = (
@@ -21,6 +26,9 @@
     "Location",
     "NormalizationModel",
     "SingleTimeSeries",
+    "SingleTimeSeriesKey",
     "SupplementalAttribute",
     "System",
+    "TimeSeriesKey",
+    "TimeSeriesStorageType",
 )

src/infrasys/arrow_storage.py

@@ -19,6 +19,7 @@
     SingleTimeSeriesMetadata,
     TimeSeriesData,
     TimeSeriesMetadata,
+    TimeSeriesStorageType,
 )
 from infrasys.time_series_storage_base import TimeSeriesStorageBase
 
@@ -54,16 +55,15 @@ def add_time_series(
         self,
         metadata: TimeSeriesMetadata,
         time_series: TimeSeriesData,
+        connection: Any = None,
     ) -> None:
         if isinstance(time_series, SingleTimeSeries):
-            self.add_raw_single_time_series(metadata.time_series_uuid, time_series.data_array)
+            self._add_single_time_series(metadata.time_series_uuid, time_series.data_array)
         else:
             msg = f"Bug: need to implement add_time_series for {type(time_series)}"
             raise NotImplementedError(msg)
 
-    def add_raw_single_time_series(
-        self, time_series_uuid: UUID, time_series_data: NDArray
-    ) -> None:
+    def _add_single_time_series(self, time_series_uuid: UUID, time_series_data: NDArray) -> None:
         fpath = self._ts_directory.joinpath(f"{time_series_uuid}{EXTENSION}")
         if not fpath.exists():
             arrow_batch = self._convert_to_record_batch(time_series_data, str(time_series_uuid))
@@ -80,6 +80,7 @@ def get_time_series(
         metadata: TimeSeriesMetadata,
         start_time: datetime | None = None,
         length: int | None = None,
+        connection: Any = None,
     ) -> Any:
         if isinstance(metadata, SingleTimeSeriesMetadata):
             return self._get_single_time_series(
@@ -89,22 +90,29 @@ def get_time_series(
         msg = f"Bug: need to implement get_time_series for {type(metadata)}"
         raise NotImplementedError(msg)
 
-    def remove_time_series(self, uuid: UUID) -> None:
-        fpath = self._ts_directory.joinpath(f"{uuid}{EXTENSION}")
+    def remove_time_series(self, metadata: TimeSeriesMetadata, connection: Any = None) -> None:
+        fpath = self._ts_directory.joinpath(f"{metadata.time_series_uuid}{EXTENSION}")
         if not fpath.exists():
-            msg = f"No time series with {uuid} is stored"
+            msg = f"No time series with {metadata.time_series_uuid} is stored"
             raise ISNotStored(msg)
         fpath.unlink()
 
-    def serialize(self, dst: Path | str, src: Optional[Path | str] = None) -> None:
+    def serialize(
+        self, data: dict[str, Any], dst: Path | str, src: Path | str | None = None
+    ) -> None:
         # From the shutil documentation: the copying operation will continue if
         # it encounters existing directories, and files within the dst tree
         # will be overwritten by corresponding files from the src tree.
         if src is None:
             src = self._ts_directory
         shutil.copytree(src, dst, dirs_exist_ok=True)
+        self.add_serialized_data(data)
         logger.info("Copied time series data to {}", dst)
 
+    @staticmethod
+    def add_serialized_data(data: dict[str, Any]) -> None:
+        data["time_series_storage_type"] = TimeSeriesStorageType.ARROW.value
+
     def _get_single_time_series(
         self,
         metadata: SingleTimeSeriesMetadata,

src/infrasys/in_memory_time_series_storage.py

@@ -2,12 +2,13 @@
 
 from datetime import datetime
 from pathlib import Path
+from typing import Any
+
 from numpy.typing import NDArray
-from typing import Optional, TypeAlias
+from typing import TypeAlias
 from uuid import UUID
 
 from loguru import logger
-from infrasys.arrow_storage import ArrowTimeSeriesStorage
 
 from infrasys.exceptions import ISNotStored
 from infrasys.time_series_models import (
@@ -30,7 +31,12 @@ def __init__(self) -> None:
     def get_time_series_directory(self) -> None:
         return None
 
-    def add_time_series(self, metadata: TimeSeriesMetadata, time_series: TimeSeriesData) -> None:
+    def add_time_series(
+        self,
+        metadata: TimeSeriesMetadata,
+        time_series: TimeSeriesData,
+        connection: Any = None,
+    ) -> None:
         if isinstance(time_series, SingleTimeSeries):
             if metadata.time_series_uuid not in self._arrays:
                 self._arrays[metadata.time_series_uuid] = time_series.data_array
@@ -47,22 +53,23 @@ def get_time_series(
         metadata: TimeSeriesMetadata,
         start_time: datetime | None = None,
         length: int | None = None,
+        connection: Any = None,
     ) -> TimeSeriesData:
         if isinstance(metadata, SingleTimeSeriesMetadata):
             return self._get_single_time_series(metadata, start_time, length)
         raise NotImplementedError(str(metadata.get_time_series_data_type()))
 
-    def remove_time_series(self, uuid: UUID) -> None:
-        time_series = self._arrays.pop(uuid, None)
+    def remove_time_series(self, metadata: TimeSeriesMetadata, connection: Any = None) -> None:
+        time_series = self._arrays.pop(metadata.time_series_uuid, None)
         if time_series is None:
-            msg = f"No time series with {uuid} is stored"
+            msg = f"No time series with {metadata.time_series_uuid} is stored"
             raise ISNotStored(msg)
 
-    def serialize(self, dst: Path | str, _: Optional[Path | str] = None) -> None:
-        base_directory = dst if isinstance(dst, Path) else Path(dst)
-        storage = ArrowTimeSeriesStorage.create_with_permanent_directory(base_directory)
-        for ts_uuid, ts in self._arrays.items():
-            storage.add_raw_single_time_series(ts_uuid, ts)
+    def serialize(
+        self, data: dict[str, Any], dst: Path | str, src: Path | str | None = None
+    ) -> None:
+        msg = "Bug: InMemoryTimeSeriesStorage.serialize should never be called."
+        raise Exception(msg)
 
     def _get_single_time_series(
         self,