Accept mixed time column in IAMC data (#251)

meksor · web-flow · commit c63f39f863df · 2026-06-08T18:48:49.000+02:00
* allow supplying years and datetime values in one "time" column

* add mixed time format core tests

* fix comment length

* fix docstring indentation

* delete doubled `.drop()` call in `RunIamcData.remove()`

* fix pandas edgecase when only pd.Timestamps are supplied

* skip creating the year and datetime columns if no ANNUAL/DATETIME values exist
diff --git a/ixmp4/core/iamc/data.py b/ixmp4/core/iamc/data.py
@@ -1,5 +1,6 @@
 from typing import TYPE_CHECKING, TypeVar
 
+import numpy as np
 import pandas as pd
 
 # TODO Import this from typing when dropping Python 3.11
@@ -34,10 +35,23 @@ def _rename_arg_cols(df: pd.DataFrame) -> pd.DataFrame:
                 "category": "step_category",
                 "subannual": "step_category",
                 "datetime": "step_datetime",
-                "time": "step_datetime",
             }
         )
 
+    @staticmethod
+    def _split_time_col(df: pd.DataFrame) -> pd.DataFrame:
+        time = df["time"]
+        is_year = time.apply(lambda x: isinstance(x, (int, np.integer)))
+
+        if is_year.any():
+            df["year"] = pd.Series(np.nan, index=df.index, dtype="object")
+            df.loc[is_year, "year"] = time.loc[is_year]
+
+        if not is_year.all():
+            df["datetime"] = pd.to_datetime(time.where(~is_year), errors="coerce")
+
+        return df.drop(columns=["time"])
+
     @classmethod
     def _convert_to_std_format(
         cls, df: pd.DataFrame, join_runs: bool, join_run_id: bool
@@ -104,8 +118,8 @@ def _get_or_create_ts(self, df: pd.DataFrame) -> pd.DataFrame:
         ts_df = ts_df.rename(columns={"id": "time_series__id"})
 
         # merge on the identity columns
-        return pd.merge(
-            df, ts_df, how="left", on=id_cols, suffixes=(None, "_y")
+        return pd.merge(df, ts_df, how="left", on=id_cols, suffixes=(None, "_y")).drop(
+            columns=id_cols
         )  # tada, df with 'time_series__id' added from the database.
 
     def add(self, df: pd.DataFrame, type: Type | str | None = None) -> None:
@@ -142,9 +156,13 @@ def add(self, df: pd.DataFrame, type: Type | str | None = None) -> None:
                 - value
 
             Any combination of:
-                - step_year for ANNUAL data points
-                - step_year and step_category for CATEGORICAL data points
-                - step_datetime for DATETIME data points
+                - (``year`` or ``step_year``) for ANNUAL data points
+                - (``year`` or ``step_year``) and (``category`` or ``step_category``)
+                  for CATEGORICAL data points
+                - (``datetime`` or ``step_datetime``) for DATETIME data points
+                - ``time`` with integer and datetime values and optionally
+                  (``category`` or ``step_category``) for integer (year) rows.
+                  ``time`` will overwrite the ``year`` and ``datetime`` columns.
 
             You may optionally supply the type column for mixed data points:
                 - type
@@ -162,6 +180,8 @@ def add(self, df: pd.DataFrame, type: Type | str | None = None) -> None:
         """
 
         self._run.require_lock()
+        if "time" in df.columns:
+            df = self._split_time_col(df)
         df = self._rename_arg_cols(df)
         df["run__id"] = self._run.id
         df = self._get_or_create_ts(df)
@@ -221,6 +241,8 @@ def remove(self, df: pd.DataFrame, type: Type | str | None = None) -> None:
             ``with run.transact("message"):``.
         """
         self._run.require_lock()
+        if "time" in df.columns:
+            df = self._split_time_col(df)
         df = self._rename_arg_cols(df)
         df["run__id"] = self._run.id
         # NOTE: This creates ts and deletes them right after
@@ -230,7 +252,6 @@ def remove(self, df: pd.DataFrame, type: Type | str | None = None) -> None:
                 type = Type[type.upper()]
             df["type"] = type
 
-        df = df.drop(columns=["unit", "variable", "region"])
         self._backend.iamc.datapoints.bulk_delete(df)
 
     def tabulate(
diff --git a/tests/core/test_iamc_data.py b/tests/core/test_iamc_data.py
@@ -680,6 +680,32 @@ def expected_data(self) -> pd.DataFrame:
     def input_data(self, expected_data: pd.DataFrame) -> pd.DataFrame:
         return expected_data.copy()
 
+    def _canonical_sort_mixed_safe(self, df: pd.DataFrame) -> pd.DataFrame:
+        """Computes a sort index with categorical and object dtype special
+        case handling. Then uses the index to sort the original df."""
+        sorted_cols = df.columns.sort_values().to_list()
+        sort_df = df.copy()
+
+        for col in sorted_cols:
+            series = sort_df[col]
+
+            if isinstance(series.dtype, pd.CategoricalDtype):
+                # Avoid unordered categorical sort errors.
+                sort_df[col] = series.astype("string")
+                continue
+
+            if pd.api.types.is_object_dtype(series):
+                try:
+                    series.sort_values()
+                except TypeError:
+                    # Normalize only non-orderable mixed object columns.
+                    sort_df[col] = series.map(
+                        lambda v: "" if pd.isna(v) else f"{type(v).__name__}:{v}"
+                    )
+
+        order = sort_df.sort_values(by=sorted_cols).index
+        return df.loc[order].reset_index(drop=True)
+
     def test_iamc_data_input(
         self,
         run: ixmp4.Run,
@@ -691,8 +717,8 @@ def test_iamc_data_input(
 
         ret = run.iamc.tabulate()
         pdt.assert_frame_equal(
-            self.canonical_sort(expected_data),
-            self.canonical_sort(ret),
+            self._canonical_sort_mixed_safe(expected_data),
+            self._canonical_sort_mixed_safe(ret),
             check_like=True,
         )
 
@@ -702,10 +728,26 @@ def test_iamc_data_input(
         assert run.iamc.tabulate().empty
 
 
-class TestAnnualIamcInputData(IamcDataAnnual, IamcDataInputTest):
-    @pytest.fixture
-    def expected_data(self, test_data_add: pd.DataFrame) -> pd.DataFrame:
-        return test_data_add.copy()
+class TestAnnualIamcInputData(IamcDataInputTest):
+    @pytest.fixture(scope="class")
+    def expected_data(
+        self,
+        regions: list[ixmp4.Region],
+        units: list[ixmp4.Unit],
+    ) -> pd.DataFrame:
+        return pd.DataFrame(
+            [
+                ["Region 1", "Unit 1", "Variable 1", 2000, 1.1],
+                ["Region 1", "Unit 1", "Variable 1", 2010, 1.3],
+                ["Region 1", "Unit 2", "Variable 2", 2020, 1.5],
+                ["Region 1", "Unit 2", "Variable 2", 2030, 1.7],
+                ["Region 2", "Unit 1", "Variable 1", 2000, 2.1],
+                ["Region 2", "Unit 1", "Variable 1", 2010, 2.3],
+                ["Region 2", "Unit 2", "Variable 2", 2020, 2.5],
+                ["Region 2", "Unit 2", "Variable 2", 2030, 2.7],
+            ],
+            columns=["region", "unit", "variable", "year", "value"],
+        ).astype({"year": "Int64"})
 
     @pytest.fixture
     def input_data(self, expected_data: pd.DataFrame) -> pd.DataFrame:
@@ -722,15 +764,12 @@ class TestCategoricalIamcInputData(IamcDataInputTest):
     @pytest.fixture
     def expected_data(self) -> pd.DataFrame:
         return pd.DataFrame(
-            {
-                "region": ["Region 1", "Region 2"],
-                "variable": ["Variable 1", "Variable 2"],
-                "unit": ["Unit 1", "Unit 2"],
-                "year": pd.Series([2000, 2010], dtype="Int64"),
-                "subannual": ["Summer", "Winter"],
-                "value": [1.1, 2.3],
-            }
-        )
+            [
+                ["Region 1", "Variable 1", "Unit 1", 2000, "Summer", 1.1],
+                ["Region 2", "Variable 2", "Unit 2", 2010, "Winter", 2.3],
+            ],
+            columns=["region", "variable", "unit", "year", "subannual", "value"],
+        ).astype({"year": "Int64"})
 
     @pytest.fixture
     def input_data(self, expected_data: pd.DataFrame) -> pd.DataFrame:
@@ -748,33 +787,95 @@ class TestDatetimeIamcInputData(IamcDataInputTest):
     @pytest.fixture
     def expected_data(self) -> pd.DataFrame:
         return pd.DataFrame(
-            {
-                "region": ["Region 1", "Region 2"],
-                "variable": ["Variable 1", "Variable 2"],
-                "unit": ["Unit 1", "Unit 2"],
-                "time": pd.to_datetime(["2000-01-01 00:00:00", "2010-06-01 12:34:56"]),
-                "value": [1.1, 2.3],
-            }
+            [
+                [
+                    "Region 1",
+                    "Variable 1",
+                    "Unit 1",
+                    pd.Timestamp("2000-01-01 00:00:00"),
+                    1.1,
+                ],
+                [
+                    "Region 2",
+                    "Variable 2",
+                    "Unit 2",
+                    pd.Timestamp("2010-06-01 12:34:56"),
+                    2.3,
+                ],
+            ],
+            columns=["region", "variable", "unit", "time", "value"],
+        )
+
+
+class TestDatetimeAsStringIamcInputData(TestDatetimeIamcInputData):
+    @pytest.fixture
+    def input_data(self, expected_data: pd.DataFrame) -> pd.DataFrame:
+        input_df = expected_data.copy()
+        input_df["time"] = (
+            input_df["time"].dt.strftime("%Y-%m-%d %H:%M:%S").astype("string")
+        )
+        return input_df
+
+
+class TestSingleRowDatetimeIamcInputData(IamcDataInputTest):
+    @pytest.fixture
+    def expected_data(self) -> pd.DataFrame:
+        return pd.DataFrame(
+            [
+                [
+                    "Region 1",
+                    "Variable 1",
+                    "Unit 1",
+                    pd.Timestamp("2000-01-01 00:00:00"),
+                    1.1,
+                ],
+            ],
+            columns=["region", "variable", "unit", "time", "value"],
         )
 
     @pytest.fixture
     def input_data(self, expected_data: pd.DataFrame) -> pd.DataFrame:
         input_df = expected_data.copy()
-        input_df["region"] = input_df["region"].astype("category")
-        input_df["unit"] = input_df["unit"].astype("category")
-        input_df["variable"] = input_df["variable"].astype("category")
         input_df["time"] = (
             input_df["time"].dt.strftime("%Y-%m-%d %H:%M:%S").astype("string")
         )
-        input_df["value"] = input_df["value"].astype("float32")
         return input_df
 
 
-class TestObjectStringsIamcInputData(IamcDataAnnual, IamcDataInputTest):
+class TestMixedIamcInputData(IamcDataInputTest):
     @pytest.fixture
-    def expected_data(self, test_data_add: pd.DataFrame) -> pd.DataFrame:
-        return test_data_add.copy()
+    def expected_data(self) -> pd.DataFrame:
+        return pd.DataFrame(
+            [
+                # ANNUAL
+                ["Region 1", "Variable 1", "Unit 1", 2000, None, 0.1],
+                ["Region 2", "Variable 2", "Unit 2", 2010, None, 0.23],
+                # CATEGORICAL
+                ["Region 1", "Variable 1", "Unit 1", 2000, "Summer", 1.1],
+                ["Region 2", "Variable 2", "Unit 2", 2010, "Winter", 2.3],
+                # DATETIME
+                [
+                    "Region 1",
+                    "Variable 1",
+                    "Unit 1",
+                    pd.Timestamp("2000-01-01 00:00:00"),
+                    None,
+                    101.0,
+                ],
+                [
+                    "Region 2",
+                    "Variable 2",
+                    "Unit 2",
+                    pd.Timestamp("2010-06-01 12:34:56"),
+                    None,
+                    3.14,
+                ],
+            ],
+            columns=["region", "variable", "unit", "time", "subannual", "value"],
+        )
+
 
+class TestObjectStringsIamcInputData(TestAnnualIamcInputData):
     @pytest.fixture
     def input_data(self, expected_data: pd.DataFrame) -> pd.DataFrame:
         input_df = expected_data.copy()
