Normalize IndexSet.data DB storage

ixmp4/core/optimization/indexset.py

@@ -23,17 +23,17 @@ def name(self) -> str:
         return self._model.name
 
     @property
-    def elements(self) -> list[float | int | str]:
-        return self._model.elements
+    def data(self) -> list[float | int | str]:
+        return self._model.data
 
-    def add(self, elements: float | int | list[float | int | str] | str) -> None:
-        """Adds elements to an existing IndexSet."""
-        self.backend.optimization.indexsets.add_elements(
-            indexset_id=self._model.id, elements=elements
+    def add(self, data: float | int | list[float | int | str] | str) -> None:
+        """Adds data to an existing IndexSet."""
+        self.backend.optimization.indexsets.add_data(
+            indexset_id=self._model.id, data=data
         )
-        self._model.elements = self.backend.optimization.indexsets.get(
+        self._model.data = self.backend.optimization.indexsets.get(
             run_id=self._model.run__id, name=self._model.name
-        ).elements
+        ).data
 
     @property
     def run_id(self) -> int:

ixmp4/data/abstract/optimization/equation.py

@@ -177,7 +177,7 @@ def add_data(self, equation_id: int, data: dict[str, Any] | pd.DataFrame) -> Non
         The data will be validated with the linked constrained
         :class:`ixmp4.data.abstract.optimization.IndexSet`s. For that, `data.keys()`
         must correspond to the names of the Equation's columns. Each column can only
-        contain values that are in the linked `IndexSet.elements`. Each row of entries
+        contain values that are in the linked `IndexSet.data`. Each row of entries
         must be unique. No values can be missing, `None`, or `NaN`. If `data.keys()`
         contains names already present in `Equation.data`, existing values will be
         overwritten.

ixmp4/data/abstract/optimization/indexset.py

@@ -17,8 +17,8 @@ class IndexSet(base.BaseModel, Protocol):
     """The id of the :class:`ixmp4.data.abstract.Run` for which this IndexSet is
     defined. """
 
-    elements: types.JsonList
-    """Unique list of str or int."""
+    data: types.OptimizationDataList
+    """Unique list of str, int, or float."""
 
     created_at: types.DateTime
     "Creation date/time. TODO"
@@ -120,24 +120,24 @@ def tabulate(self, *, name: str | None = None, **kwargs) -> pd.DataFrame:
             A data frame with the columns:
                 - id
                 - name
-                - elements
+                - data
                 - run__id
                 - created_at
                 - created_by
         """
         ...
 
-    def add_elements(
-        self, indexset_id: int, elements: float | int | List[float | int | str] | str
+    def add_data(
+        self, indexset_id: int, data: float | int | List[float | int | str] | str
     ) -> None:
-        """Adds elements to an existing IndexSet.
+        """Adds data to an existing IndexSet.
 
         Parameters
         ----------
         indexset_id : int
             The id of the target IndexSet.
-        elements : float | int | List[float | int | str] | str
-            The elements to be added to the IndexSet.
+        data : float | int | List[float | int | str] | str
+            The data to be added to the IndexSet.
 
         Returns
         -------

ixmp4/data/abstract/optimization/parameter.py

@@ -176,7 +176,7 @@ def add_data(self, parameter_id: int, data: dict[str, Any] | pd.DataFrame) -> No
         The data will be validated with the linked constrained
         :class:`ixmp4.data.abstract.optimization.IndexSet`s. For that, `data.keys()`
         must correspond to the names of the Parameter's columns. Each column can only
-        contain values that are in the linked `IndexSet.elements`. Each row of entries
+        contain values that are in the linked `IndexSet.data`. Each row of entries
         must be unique. No values can be missing, `None`, or `NaN`. If `data.keys()`
         contains names already present in `Parameter.data`, existing values will be
         overwritten.

ixmp4/data/abstract/optimization/table.py

@@ -176,7 +176,7 @@ def add_data(self, table_id: int, data: dict[str, Any] | pd.DataFrame) -> None:
         The data will be validated with the linked constrained
         :class:`ixmp4.data.abstract.optimization.IndexSet`s. For that, `data.keys()`
         must correspond to the names of the Table's columns. Each column can only
-        contain values that are in the linked `IndexSet.elements`. Each row of entries
+        contain values that are in the linked `IndexSet.data`. Each row of entries
         must be unique. No values can be missing, `None`, or `NaN`. If `data.keys()`
         contains names already present in `Table.data`, existing values will be
         overwritten.

ixmp4/data/abstract/optimization/variable.py

@@ -179,7 +179,7 @@ def add_data(self, variable_id: int, data: dict[str, Any] | pd.DataFrame) -> Non
         The data will be validated with the linked constrained
         :class:`ixmp4.data.abstract.optimization.IndexSet`s. For that, `data.keys()`
         must correspond to the names of the Variable's columns. Each column can only
-        contain values that are in the linked `IndexSet.elements`. Each row of entries
+        contain values that are in the linked `IndexSet.data`. Each row of entries
         must be unique. No values can be missing, `None`, or `NaN`. If `data.keys()`
         contains names already present in `Variable.data`, existing values will be
         overwritten.

ixmp4/data/api/optimization/indexset.py

@@ -17,7 +17,7 @@ class IndexSet(base.BaseModel):
 
     id: int
     name: str
-    elements: (
+    data: (
         StrictFloat
         | StrictInt
         | StrictStr
@@ -67,13 +67,13 @@ def list(self, **kwargs) -> list[IndexSet]:
     def tabulate(self, **kwargs) -> pd.DataFrame:
         return super()._tabulate(json=kwargs)
 
-    def add_elements(
+    def add_data(
         self,
         indexset_id: int,
-        elements: StrictFloat
+        data: StrictFloat
         | StrictInt
         | List[StrictFloat | StrictInt | StrictStr]
         | StrictStr,
     ) -> None:
-        kwargs = {"indexset_id": indexset_id, "elements": elements}
+        kwargs = {"indexset_id": indexset_id, "data": data}
         self._request("PATCH", self.prefix + str(indexset_id) + "/", json=kwargs)

ixmp4/data/db/optimization/equation/repository.py

@@ -65,7 +65,7 @@ def _add_column(
         self.columns.create(
             name=column_name,
             constrained_to_indexset=indexset.id,
-            dtype=pd.Series(indexset.elements).dtype.name,
+            dtype=pd.Series(indexset.data).dtype.name,
             equation_id=equation_id,
             unique=True,
             **kwargs,

ixmp4/data/db/optimization/indexset/model.py

@@ -1,6 +1,4 @@
-from typing import ClassVar
-
-from sqlalchemy.orm import validates
+from typing import ClassVar, Literal
 
 from ixmp4 import db
 from ixmp4.core.exceptions import OptimizationDataValidationError
@@ -10,26 +8,50 @@
 from .. import base
 
 
+# TODO Feels like there ought to be this kind of functionality already
+def cast_data_as_type(
+    data: "IndexSetData", type: Literal["float", "int", "str"] | None
+) -> float | int | str:
+    if type == "str":
+        return data.value
+    elif type == "int":
+        return int(data.value)
+    elif type == "float":
+        return float(data.value)
+    else:  # type is None
+        return 0
+
+
 class IndexSet(base.BaseModel):
     NotFound: ClassVar = abstract.IndexSet.NotFound
     NotUnique: ClassVar = abstract.IndexSet.NotUnique
     DataInvalid: ClassVar = OptimizationDataValidationError
     DeletionPrevented: ClassVar = abstract.IndexSet.DeletionPrevented
 
-    elements: types.JsonList = db.Column(db.JsonType, nullable=False, default=[])
+    data_type: types.OptimizationDataType
+
+    _data: types.Mapped[list["IndexSetData"]] = db.relationship(
+        back_populates="indexset"
+    )
 
-    @validates("elements")
-    def validate_elements(self, key, value: list[float | int | str]):
-        unique = set()
-        for element in value:
-            if element in unique:
-                raise self.DataInvalid(
-                    f"{element} already defined for IndexSet {self.name}!"
-                )
-            else:
-                unique.add(element)
-        return value
+    @db.hybrid_property
+    def data(self) -> list[float | int | str]:
+        return [cast_data_as_type(data, self.data_type) for data in self._data]
+
+    # NOTE For the core layer (setting and retrieving) to work, the property needs a
+    # setter method
+    @data.inplace.setter
+    def _data_setter(self, value: list[float | int | str]) -> None:
+        return None
 
     run__id: types.RunId
 
     __table_args__ = (db.UniqueConstraint("name", "run__id"),)
+
+
+class IndexSetData(base.RootBaseModel):
+    indexset: types.Mapped["IndexSet"] = db.relationship(back_populates="_data")
+    indexset__id: types.IndexSetId
+    value: types.String = db.Column(db.String, nullable=False)
+
+    __table_args__ = (db.UniqueConstraint("indexset__id", "value"),)

ixmp4/data/db/optimization/indexset/repository.py

@@ -8,7 +8,7 @@
 
 from .. import base
 from .docs import IndexSetDocsRepository
-from .model import IndexSet
+from .model import IndexSet, IndexSetData
 
 
 class IndexSetRepository(
@@ -61,21 +61,36 @@ def list(self, *args, **kwargs) -> list[IndexSet]:
 
     @guard("view")
     def tabulate(self, *args, **kwargs) -> pd.DataFrame:
-        return super().tabulate(*args, **kwargs)
+        result = super().tabulate(*args, **kwargs).drop(labels="data_type", axis=1)
+        result.insert(
+            loc=0,
+            column="data",
+            value=[self.get_by_id(id=indexset_id).data for indexset_id in result.id],
+        )
+        return result
 
     @guard("edit")
-    def add_elements(
+    def add_data(
         self,
         indexset_id: int,
-        elements: float | int | List[float | int | str] | str,
+        data: float | int | List[float | int | str] | str,
     ) -> None:
         indexset = self.get_by_id(id=indexset_id)
-        if not isinstance(elements, list):
-            elements = [elements]
-        if indexset.elements is None:
-            indexset.elements = elements
-        else:
-            indexset.elements = indexset.elements + elements
+        if not isinstance(data, list):
+            data = [data]
+        # TODO If adding rows one by one is too expensive, look into executemany pattern
+        for value in data:
+            self.session.add(
+                IndexSetData(indexset=indexset, indexset__id=indexset_id, value=value)
+            )
+
+        try:
+            self.session.flush()
+        except db.IntegrityError as e:
+            self.session.rollback()
+            raise indexset.DataInvalid from e
+
+        indexset.data_type = type(data[0]).__name__
 
         self.session.add(indexset)
         self.session.commit()

ixmp4/data/db/optimization/parameter/repository.py

@@ -66,7 +66,7 @@ def _add_column(
         self.columns.create(
             name=column_name,
             constrained_to_indexset=indexset.id,
-            dtype=pd.Series(indexset.elements).dtype.name,
+            dtype=pd.Series(indexset.data).dtype.name,
             parameter_id=parameter_id,
             unique=True,
             **kwargs,

ixmp4/data/db/optimization/table/repository.py

@@ -65,7 +65,7 @@ def _add_column(
         self.columns.create(
             name=column_name,
             constrained_to_indexset=indexset.id,
-            dtype=pd.Series(indexset.elements).dtype.name,
+            dtype=pd.Series(indexset.data).dtype.name,
             table_id=table_id,
             unique=True,
             **kwargs,

ixmp4/data/db/optimization/utils.py

@@ -12,10 +12,10 @@ def collect_indexsets_to_check(
     columns: list["Column"],
 ) -> dict[str, Any]:
     """Creates a {key:value} dict from linked Column.names and their
-    IndexSet.elements."""
+    IndexSet.data."""
     collection: dict[str, Any] = {}
     for column in columns:
-        collection[column.name] = column.indexset.elements
+        collection[column.name] = column.indexset.data
     return collection
 
 

ixmp4/data/db/optimization/variable/repository.py

@@ -65,7 +65,7 @@ def _add_column(
         self.columns.create(
             name=column_name,
             constrained_to_indexset=indexset.id,
-            dtype=pd.Series(indexset.elements).dtype.name,
+            dtype=pd.Series(indexset.data).dtype.name,
             variable_id=variable_id,
             unique=True,
             **kwargs,

ixmp4/data/types.py

@@ -1,5 +1,5 @@
 from datetime import datetime
-from typing import Any
+from typing import Any, Literal
 
 from sqlalchemy.orm import Mapped as Mapped
 
@@ -8,9 +8,11 @@
 Boolean = Mapped[bool]
 DateTime = Mapped[datetime]
 Float = Mapped[float]
+IndexSetId = Mapped[db.IndexSetIdType]
 Integer = Mapped[int]
-JsonList = Mapped[list[float | int | str]]
+OptimizationDataList = Mapped[list[float | int | str]]
 JsonDict = Mapped[dict[str, Any]]
+OptimizationDataType = Mapped[Literal["float", "int", "str"] | None]
 String = Mapped[str]
 Name = Mapped[db.NameType]
 UniqueName = Mapped[db.UniqueNameType]

ixmp4/db/__init__.py

@@ -50,7 +50,8 @@
     update,
 )
 from sqlalchemy.dialects.postgresql import JSONB
-from sqlalchemy.exc import MultipleResultsFound
+from sqlalchemy.exc import IntegrityError, MultipleResultsFound
+from sqlalchemy.ext.hybrid import hybrid_property
 from sqlalchemy.orm import (
     Relationship,
     Session,
@@ -65,17 +66,16 @@
 from . import utils
 
 Column = mapped_column
+IndexSetIdType = Annotated[
+    int,
+    Column(Integer, ForeignKey("optimization_indexset.id"), nullable=False, index=True),
+]
 JsonType = JSON()
 JsonType = JsonType.with_variant(JSONB(), "postgresql")
 NameType = Annotated[str, Column(String(255), nullable=False, unique=False)]
 RunIdType = Annotated[
     int,
-    Column(
-        Integer,
-        ForeignKey("run.id"),
-        nullable=False,
-        index=True,
-    ),
+    Column(Integer, ForeignKey("run.id"), nullable=False, index=True),
 ]
 UniqueNameType = Annotated[str, Column(String(255), nullable=False, unique=True)]
 UsernameType = Annotated[str, Column(String(255), nullable=True)]

ixmp4/server/rest/optimization/indexset.py

@@ -20,8 +20,8 @@ class IndexSetInput(BaseModel):
     name: str
 
 
-class ElementsInput(BaseModel):
-    elements: (
+class DataInput(BaseModel):
+    data: (
         StrictFloat | StrictInt | StrictStr | list[StrictFloat | StrictInt | StrictStr]
     )
 
@@ -57,11 +57,11 @@ def create(
 
 @autodoc
 @router.patch("/{indexset_id}/")
-def add_elements(
+def add_data(
     indexset_id: int,
-    elements: ElementsInput,
+    data: DataInput,
     backend: Backend = Depends(deps.get_backend),
 ):
-    backend.optimization.indexsets.add_elements(
-        indexset_id=indexset_id, **elements.model_dump()
+    backend.optimization.indexsets.add_data(
+        indexset_id=indexset_id, **data.model_dump()
     )