diff --git a/psycop/common/feature_generation/loaders/raw/load_structured_sfi.py b/psycop/common/feature_generation/loaders/raw/load_structured_sfi.py index ff3fde427..5475d0a8e 100644 --- a/psycop/common/feature_generation/loaders/raw/load_structured_sfi.py +++ b/psycop/common/feature_generation/loaders/raw/load_structured_sfi.py @@ -2,7 +2,7 @@ from __future__ import annotations -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, Literal import polars as pl @@ -295,21 +295,34 @@ def smoking_categorical(mapping: dict[str, int] | None = None) -> pd.DataFrame: ).to_pandas() -def systolic_blood_pressure() -> pd.DataFrame: - df = pl.from_pandas( - sql_load( - query="SELECT * FROM [fct].[FOR_SFI_Blodtyk_Puls_psyk_somatik_inkl_2021]", - ), +def _get_blood_pressure_pulse( + subtype: Literal["Systolisk", "Diastolisk", "Pulsslag / min"] +) -> pl.LazyFrame: + df = ( + pl.from_pandas( + sql_load( + query="SELECT * FROM [fct].[FOR_SFI_Blodtyk_Puls_psyk_somatik_inkl_2021]", + ), + ) + .lazy() + .rename( + {"datotid_senest_aendret_i_sfien": "timestamp", "numelementvaerdi": "value"} + ) ) - df_pl_subset = df.select( - [ - "dw_ek_borger", - "datotid_senest_aendret_i_sfien", - "numelementvaerdi", - ], - ).filter(pl.col("numelementvaerdi").is_not_null()) + subset = df.select(["dw_ek_borger", "timestamp", "value", "elementledetekst"]) + return df.filter(pl.col("elementledetekst") == pl.lit(subtype)).select( + ["dw_ek_borger", "timestamp", "value"] + ) - return df_pl_subset.rename( - {"datotid_senest_aendret_i_sfien": "timestamp", "numelementvaerdi": "value"}, - ).to_pandas() + +def systolic_blood_pressure() -> pd.DataFrame: + df = _get_blood_pressure_pulse(subtype="Systolisk") + + return df.collect().to_pandas() + + +def diastolic_blood_pressure() -> pd.DataFrame: + df = _get_blood_pressure_pulse(subtype="Diastolisk") + + return df.collect().to_pandas()