feat: allow saving additional interval attributes in result_interval table

Author: scymtym

execution_engine/omop/db/celida/tables.py

@@ -169,7 +169,9 @@ class ResultInterval(Base):  # noqa: D101
     interval_start: Mapped[datetime]
     interval_end: Mapped[datetime]
     interval_type = mapped_column(IntervalTypeEnum)
-
+    interval_ratio: Mapped[float] = mapped_column(
+        nullable=True
+    )
     execution_run: Mapped["ExecutionRun"] = relationship(
         primaryjoin="ResultInterval.run_id == ExecutionRun.run_id",
     )

execution_engine/omop/db/celida/views.py

@@ -199,6 +199,7 @@ def interval_result_view() -> Select:
             rri.c.interval_type,
             rri.c.interval_start,
             rri.c.interval_end,
+            rri.c.interval_ratio,
         )
         .select_from(rri)
         .outerjoin(pip, (rri.c.pi_pair_id == pip.c.pi_pair_id))

execution_engine/task/process/__init__.py

@@ -1,3 +1,4 @@
+import copy
 import importlib
 import os
 import sys
@@ -44,4 +45,21 @@ def get_processing_module(
 AnyInterval = Interval | IntervalWithCount | IntervalWithTypeCounts
 GeneralizedInterval = None | AnyInterval
 
-TInterval = TypeVar('TInterval', bound = AnyInterval)
+TInterval = TypeVar("TInterval", bound=AnyInterval)
+
+def interval_like(interval: TInterval, start: int, end: int) -> TInterval:
+    """
+    Return a copy of the given interval with its lower and upper bounds replaced.
+
+    Args:
+        interval (I): The interval to copy. Must be one of Interval, IntervalWithCount, or IntervalWithTypeCounts.
+        start (datetime): The new lower bound.
+        end (datetime): The new upper bound.
+
+    Returns:
+        I: A copy of the interval with updated lower and upper bounds.
+    """
+
+    return copy.copy(interval)._replace(
+        lower=start, upper=end
+    )  # type: ignore[return-value]m

execution_engine/task/process/rectangle.py

@@ -2,7 +2,7 @@
 import importlib
 import logging
 import os
-from typing import Callable, cast
+from typing import Callable, cast, List
 
 import numpy as np
 import pendulum
@@ -12,7 +12,7 @@
 from execution_engine.util.interval import IntervalType, interval_datetime
 from execution_engine.util.types import TimeRange
 
-from . import Interval, IntervalWithCount
+from . import Interval, IntervalWithCount, AnyInterval, GeneralizedInterval, interval_like
 
 PROCESS_RECTANGLE_VERSION = os.getenv("PROCESS_RECTANGLE_VERSION", "auto")
 
@@ -566,17 +566,11 @@ def mask_intervals(
         for person_id, intervals in mask.items()
     }
 
-    result = {}
-    for person_id in data:
-        # intersect every interval in data with every interval in mask
-        person_result = _impl.intersect_interval_lists(
-            data[person_id], person_mask[person_id]
-        )
-        if not person_result:
-            continue
-
-        result[person_id] = person_result
-
+    def intersection_interval(start: int, end: int, intervals: List[GeneralizedInterval]) -> GeneralizedInterval:
+        left_interval, right_interval = intervals
+        if left_interval is not None and right_interval is not None:
+            return interval_like(right_interval, start, end)
+    result = find_rectangles([person_mask, data], intersection_interval)
     return result
 
 

execution_engine/task/task.py

@@ -578,6 +578,18 @@ def store_result_in_db(
             run_id=bind_params["run_id"],
             cohort_category=self.category,
         )
+        def interval_data(interval):
+            data = dict(
+                interval_start=interval.lower,
+                interval_end=interval.upper,
+                interval_type=interval.type,
+            )
+            if isinstance(interval, Interval):
+                data["interval_ratio"] = None
+            else:
+                assert isinstance(interval, IntervalWithCount)
+                data["interval_ratio"] = interval.count
+            return data
 
         try:
             with get_engine().begin() as conn:
@@ -586,9 +598,7 @@ def store_result_in_db(
                     [
                         {
                             "person_id": person_id,
-                            "interval_start": normalized_interval.lower,
-                            "interval_end": normalized_interval.upper,
-                            "interval_type": normalized_interval.type,
+                            **interval_data(normalized_interval),
                             **params,
                         }
                         for person_id, intervals in result.items()

execution_engine/util/types.py

@@ -13,8 +13,9 @@
 )
 from execution_engine.util.value import ValueNumber, ValueNumeric
 from execution_engine.util.value.time import ValueCount, ValueDuration, ValuePeriod
+from execution_engine.task.process import AnyInterval
 
-PersonIntervals = dict[int, Any]
+PersonIntervals = dict[int, AnyInterval]
 
 
 class TimeRange(BaseModel):