Speed up scenario filtered entity subquery

Seems better to compute and join element count earlier in the query rather than later
spine-tools · Feb 1, 2024 · 6cf762c · 6cf762c
1 parent 2da3e72
commit 6cf762c
Show file tree

Hide file tree

Showing 2 changed files with 21 additions and 6 deletions.
diff --git a/spinedb_api/filters/scenario_filter.py b/spinedb_api/filters/scenario_filter.py
@@ -258,6 +258,14 @@ def _make_scenario_filtered_entity_sq(db_map, state):
         Alias: a subquery for entity filtered by selected scenario
     """
     ext_entity_sq = _ext_entity_sq(db_map, state)
+    ext_entity_element_count_sq = (
+        db_map.query(
+            db_map.entity_element_sq.c.entity_id,
+            func.count(db_map.entity_element_sq.c.element_id).label("element_count"),
+        )
+        .group_by(db_map.entity_element_sq.c.entity_id)
+        .subquery()
+    )
     ext_entity_class_dimension_count_sq = (
         db_map.query(
             db_map.entity_class_dimension_sq.c.entity_class_id,
@@ -281,17 +289,21 @@ def _make_scenario_filtered_entity_sq(db_map, state):
                 and_(ext_entity_sq.c.active == None, ext_entity_sq.c.active_by_default == True),
             ),
         )
+        .outerjoin(
+            ext_entity_element_count_sq,
+            ext_entity_element_count_sq.c.entity_id == ext_entity_sq.c.id,
+        )
         .outerjoin(
             ext_entity_class_dimension_count_sq,
             ext_entity_class_dimension_count_sq.c.entity_class_id == ext_entity_sq.c.class_id,
         )
-        .outerjoin(db_map.entity_element_sq, ext_entity_sq.c.id == db_map.entity_element_sq.c.entity_id)
-        .group_by(ext_entity_sq.c.id)
-        .having(
+        .filter(
             or_(
-                ext_entity_class_dimension_count_sq.c.dimension_count == None,
-                ext_entity_class_dimension_count_sq.c.dimension_count
-                == func.count(db_map.entity_element_sq.c.element_id),
+                and_(
+                    ext_entity_element_count_sq.c.element_count == None,
+                    ext_entity_class_dimension_count_sq.c.dimension_count == None,
+                ),
+                ext_entity_element_count_sq.c.element_count == ext_entity_class_dimension_count_sq.c.dimension_count,
             )
         )
         .subquery()

diff --git a/spinedb_api/query.py b/spinedb_api/query.py
@@ -31,6 +31,9 @@ def __init__(self, bind, *entities):
         self._select = select(entities)
         self._from = None
 
+    def __str__(self):
+        return str(self._select)
+
     @property
     def column_descriptions(self):
         return [{"name": c.name} for c in self._select.columns]