Improve IAMC endpoint query performance (#56)

* ad ondelete=CASCADE * adjust iamc filters * add count tests * fix linter errors * fix filter joins * remove ondelete and implement delete_orphans() * add deletion check to tests; remove ondelete and implement delete_orphans()
iiasa · Mar 6, 2024 · 6e04ed3 · 6e04ed3
1 parent c85d987
commit 6e04ed3
Show file tree

Hide file tree

Showing 14 changed files with 131 additions and 35 deletions.
diff --git a/ixmp4/data/db/filters/model.py b/ixmp4/data/db/filters/model.py
@@ -1,6 +1,6 @@
 from typing import ClassVar
 
-from ixmp4.db import filters
+from ixmp4.db import filters, utils
 
 from .. import Model, Run
 
@@ -12,4 +12,6 @@ class ModelFilter(filters.BaseFilter, metaclass=filters.FilterMeta):
     sqla_model: ClassVar[type] = Model
 
     def join(self, exc, **kwargs):
-        return exc.join(Model, Run.model)
+        if not utils.is_joined(exc, Model):
+            exc = exc.join(Model, Run.model)
+        return exc
diff --git a/ixmp4/data/db/filters/scenario.py b/ixmp4/data/db/filters/scenario.py
@@ -1,6 +1,6 @@
 from typing import ClassVar
 
-from ixmp4.db import filters
+from ixmp4.db import filters, utils
 
 from .. import Run, Scenario
 
@@ -12,4 +12,6 @@ class ScenarioFilter(filters.BaseFilter, metaclass=filters.FilterMeta):
     sqla_model: ClassVar[type] = Scenario
 
     def join(self, exc, **kwargs):
-        return exc.join(Scenario, Run.scenario)
+        if not utils.is_joined(exc, Scenario):
+            exc = exc.join(Scenario, Run.scenario)
+        return exc
diff --git a/ixmp4/data/db/iamc/datapoint/model.py b/ixmp4/data/db/iamc/datapoint/model.py
@@ -22,7 +22,10 @@ class DataPoint(base.BaseModel):
     @declared_attr
     def time_series__id(cls):
         return db.Column(
-            db.Integer, db.ForeignKey("iamc_timeseries.id"), nullable=False, index=True
+            db.Integer,
+            db.ForeignKey("iamc_timeseries.id"),
+            nullable=False,
+            index=True,
         )
 
     value = db.Column(db.Float)

diff --git a/ixmp4/data/db/iamc/datapoint/repository.py b/ixmp4/data/db/iamc/datapoint/repository.py
@@ -210,4 +210,17 @@ def bulk_update(self, df: DataFrame[UpdateDataPointFrameSchema]) -> None:
     @guard("edit")
     def bulk_delete(self, df: DataFrame[RemoveDataPointFrameSchema]) -> None:
         self.check_df_access(df)
-        return super().bulk_delete(df)
+        res = super().bulk_delete(df)
+        self.delete_orphans()
+        return res
+
+    def delete_orphans(self):
+        exc = db.delete(TimeSeries).where(
+            ~db.exists(
+                db.select(self.model_class.id).where(
+                    TimeSeries.id == self.model_class.time_series__id
+                )
+            )
+        )
+        self.session.execute(exc, execution_options={"synchronize_session": False})
+        self.session.commit()
diff --git a/ixmp4/data/db/iamc/variable/filter.py b/ixmp4/data/db/iamc/variable/filter.py
@@ -1,5 +1,4 @@
 from ixmp4.data.db import filters as base
-from ixmp4.data.db.iamc.datapoint import get_datapoint_model
 from ixmp4.data.db.iamc.timeseries import TimeSeries
 from ixmp4.db import filters, utils
 
@@ -22,8 +21,4 @@ def join(self, exc, session=None):
             exc = exc.join(
                 TimeSeries, onclause=TimeSeries.measurand__id == Measurand.id
             )
-
-        model = get_datapoint_model(session)
-        if not utils.is_joined(exc, model):
-            exc = exc.join(model, onclause=model.time_series__id == TimeSeries.id)
         return exc
diff --git a/ixmp4/data/db/meta/filter.py b/ixmp4/data/db/meta/filter.py
@@ -8,9 +8,8 @@
 class RunFilter(base.RunFilter, metaclass=filters.FilterMeta):
     def join(self, exc, **kwargs):
         if not utils.is_joined(exc, Run):
-            return exc.join(Run, onclause=RunMetaEntry.run__id == Run.id)
-        else:
-            return exc
+            exc = exc.join(Run, onclause=RunMetaEntry.run__id == Run.id)
+        return exc
 
 
 class RunMetaEntryFilter(base.RunMetaEntryFilter, metaclass=filters.FilterMeta):

diff --git a/ixmp4/data/db/model/filter.py b/ixmp4/data/db/model/filter.py
@@ -1,6 +1,5 @@
 from ixmp4 import db
 from ixmp4.data.db import filters as base
-from ixmp4.data.db.iamc.datapoint import get_datapoint_model
 from ixmp4.data.db.iamc.timeseries import TimeSeries
 from ixmp4.data.db.run.model import Run
 from ixmp4.db import filters, utils
@@ -16,9 +15,6 @@ def join_datapoints(self, exc: db.sql.Select, session=None):
         if not utils.is_joined(exc, TimeSeries):
             exc = exc.join(TimeSeries, onclause=TimeSeries.run__id == Run.id)
 
-        model = get_datapoint_model(session)
-        if not utils.is_joined(exc, model):
-            exc = exc.join(model, onclause=model.time_series__id == TimeSeries.id)
         return exc
 
 

diff --git a/ixmp4/data/db/optimization/indexset/filter.py b/ixmp4/data/db/optimization/indexset/filter.py
@@ -8,9 +8,8 @@
 class RunFilter(base.RunFilter, metaclass=filters.FilterMeta):
     def join(self, exc, **kwargs):
         if not utils.is_joined(exc, Run):
-            return exc.join(Run, onclause=IndexSet.run__id == Run.id)
-        else:
-            return exc
+            exc = exc.join(Run, onclause=IndexSet.run__id == Run.id)
+        return exc
 
 
 class OptimizationIndexSetFilter(

diff --git a/ixmp4/data/db/region/filter.py b/ixmp4/data/db/region/filter.py
@@ -1,6 +1,5 @@
 from ixmp4 import db
 from ixmp4.data.db import filters as base
-from ixmp4.data.db.iamc.datapoint import get_datapoint_model
 from ixmp4.data.db.iamc.timeseries import TimeSeries
 from ixmp4.db import filters, utils
 
@@ -11,10 +10,6 @@ class BaseIamcFilter(filters.BaseFilter, metaclass=filters.FilterMeta):
     def join_datapoints(self, exc: db.sql.Select, session=None):
         if not utils.is_joined(exc, TimeSeries):
             exc = exc.join(TimeSeries, onclause=TimeSeries.region__id == Region.id)
-
-        model = get_datapoint_model(session)
-        if not utils.is_joined(exc, model):
-            exc = exc.join(model, onclause=model.time_series__id == TimeSeries.id)
         return exc
 
 

diff --git a/ixmp4/data/db/scenario/filter.py b/ixmp4/data/db/scenario/filter.py
@@ -1,6 +1,5 @@
 from ixmp4 import db
 from ixmp4.data.db import filters as base
-from ixmp4.data.db.iamc.datapoint import get_datapoint_model
 from ixmp4.data.db.iamc.timeseries import TimeSeries
 from ixmp4.data.db.run.model import Run
 from ixmp4.db import filters, utils
@@ -16,9 +15,6 @@ def join_datapoints(self, exc: db.sql.Select, session=None):
         if not utils.is_joined(exc, TimeSeries):
             exc = exc.join(TimeSeries, onclause=TimeSeries.run__id == Run.id)
 
-        model = get_datapoint_model(session)
-        if not utils.is_joined(exc, model):
-            exc = exc.join(model, onclause=model.time_series__id == TimeSeries.id)
         return exc
 
 

diff --git a/ixmp4/data/db/unit/filter.py b/ixmp4/data/db/unit/filter.py
@@ -1,6 +1,5 @@
 from ixmp4 import db
 from ixmp4.data.db import filters as base
-from ixmp4.data.db.iamc.datapoint import get_datapoint_model
 from ixmp4.data.db.iamc.measurand import Measurand
 from ixmp4.data.db.iamc.timeseries import TimeSeries
 from ixmp4.db import filters, utils
@@ -17,10 +16,6 @@ def join_datapoints(self, exc, session=None):
             exc = exc.join(
                 TimeSeries, onclause=TimeSeries.measurand__id == Measurand.id
             )
-
-        model = get_datapoint_model(session)
-        if not utils.is_joined(exc, model):
-            exc = exc.join(model, onclause=model.time_series__id == TimeSeries.id)
         return exc
 
 

diff --git a/tests/core/test_iamc.py b/tests/core/test_iamc.py
@@ -147,6 +147,12 @@ def do_run_datapoints(test_mp, data, raw=True, _type=None):
         ret = ret.drop(columns=["id", "type"])
     assert_unordered_equality(remaining_data, ret, check_like=True)
 
+    ts_after_delete = test_mp.backend.iamc.timeseries.tabulate(join_parameters=True)
+    all_dp_after_delete = test_mp.backend.iamc.datapoints.tabulate()
+    assert set(ts_after_delete["id"].unique()) == set(
+        all_dp_after_delete["time_series__id"].unique()
+    )
+
     # == Partial Update / Partial Addition ==
     # Update all data values
     data["value"] = -9.9

diff --git a/tests/data/test_count.py b/tests/data/test_count.py
@@ -0,0 +1,81 @@
+from functools import reduce
+
+import pytest
+
+from ..utils import generated_db_platforms
+
+
+def deepgetattr(obj, attr):
+    return reduce(getattr, attr.split("."), obj)
+
+
+@generated_db_platforms
+@pytest.mark.parametrize(
+    "repo_name,filters",
+    [
+        [
+            "iamc.datapoints",
+            {
+                "model": {"name": "Model 0"},
+                "scenario": {"name": "Scenario 0"},
+                "run": {"default_only": False},
+            },
+        ],
+        [
+            "iamc.datapoints",
+            {
+                "scenario": {"name__like": "Scenario *"},
+                "run": {"default_only": False},
+            },
+        ],
+        [
+            "iamc.datapoints",
+            {
+                "model": {"name__like": "Model *"},
+                "unit": {"name__in": [f"Unit {i}" for i in range(10)]},
+                "variable": {"name__like": "Variable 1*"},
+                "region": {"name__in": [f"Region {i}" for i in range(10)]},
+                "run": {"default_only": False},
+            },
+        ],
+        [
+            "regions",
+            {
+                "name__like": "Region 1*",
+                "iamc": {
+                    "run": {"default_only": False},
+                },
+            },
+        ],
+        [
+            "models",
+            {
+                "name__like": "Model *",
+                "iamc": True,
+            },
+        ],
+        [
+            "scenarios",
+            {
+                "name__like": "Scenario *",
+                "iamc": {
+                    "variable": {"name__like": "Variable *"},
+                    "run": {"default_only": False},
+                },
+            },
+        ],
+        [
+            "units",
+            {
+                "iamc": {
+                    "variable": {"name__like": "Variable *"},
+                    "run": {"default_only": False},
+                },
+            },
+        ],
+    ],
+)
+def test_count(generated_mp, repo_name, filters, request):
+    generated_mp = request.getfixturevalue(generated_mp)
+    repo = deepgetattr(generated_mp.backend, repo_name)
+    assert len(repo.list(**filters)) == repo.count(**filters)
diff --git a/tests/utils.py b/tests/utils.py
@@ -69,6 +69,20 @@ def assert_unordered_equality(df1, df2, **kwargs):
     ],
 )
 
+generated_db_platforms = pytest.mark.parametrize(
+    "generated_mp",
+    [
+        "test_sqlite_mp_generated",
+        pytest.param(
+            "test_pgsql_mp_generated",
+            marks=pytest.mark.skipif(
+                SKIP_PGSQL_TESTS,
+                reason="Cannot connect to PostgreSQL database service, skipping test",
+            ),
+        ),
+    ],
+)
+
 
 generated_api_platforms = pytest.mark.parametrize(
     "generated_mp",