Repository: iossifovlab/gpf

Issue:
[BUG] Pheno browser download measures uses too much memory
**What happened**:
Pheno browser download is complex and collects a bunch of data in memory. This is tedious and slow
**What you expected to happen**:
The request takes too long and should be more optimized


Discussion:

Patch:
diff --git a/dae/dae/pheno/pheno_db.py b/dae/dae/pheno/pheno_db.py
index fa7eacc0a..5ed234766 100644
--- a/dae/dae/pheno/pheno_db.py
+++ b/dae/dae/pheno/pheno_db.py
@@ -9,12 +9,9 @@ from abc import ABC, abstractmethod
 from collections import defaultdict
 from itertools import chain
 
-import csv
-from io import StringIO
-
 import pandas as pd
-from sqlalchemy.sql import select, text
-from sqlalchemy import not_, desc
+from sqlalchemy.sql import select, text, Select
+from sqlalchemy import not_, desc, Column
 
 from dae.pedigrees.family import Person
 from dae.pedigrees.families_data import FamiliesData
@@ -280,7 +277,7 @@ class PhenotypeData(ABC):
     def get_measures(
             self,
             instrument_name: Optional[str] = None,
-            measure_type: Optional[str] = None) -> Dict[str, Measure]:
+            measure_type: Optional[MeasureType] = None) -> Dict[str, Measure]:
         """
         Return a dictionary of measures objects.
 
@@ -304,7 +301,7 @@ class PhenotypeData(ABC):
         if measure_type is not None:
             measure_type = MeasureType.from_str(measure_type)
 
-        for instrument_name, instrument in instruments.items():
+        for _, instrument in instruments.items():
             for measure in instrument.measures.values():
                 if measure_type is not None and \
                         measure.measure_type != measure_type:
@@ -493,7 +490,7 @@ class PhenotypeData(ABC):
             value_df.set_index("person_id"),
             on="person_id",
             how="right",
-            rsuffix="_val")  # type: ignore
+            rsuffix="_val")
         df = df.set_index("person_id")
         df = df.reset_index()
 
@@ -544,17 +541,17 @@ class PhenotypeData(ABC):
         return self.get_values(measure_ids, person_ids, family_ids, role)
 
     @abstractmethod
-    def get_values_streaming_csv(
+    def get_people_measure_values(
         self,
         measure_ids: list[str],
         person_ids: Optional[list[str]] = None,
         family_ids: Optional[list[str]] = None,
         roles: Optional[list[str]] = None,
-    ) -> Generator[str, None, None]:
+    ) -> Generator[dict[str, Any], None, None]:
         """
-        Collect and format the values of the given measures in CSV format.
+        Collect and format the values of the given measures in dict format.
 
-        Yields lines.
+        Yields a dict representing every row.
 
         `measure_ids` -- list of measure ids which values should be returned.
 
@@ -857,7 +854,7 @@ class PhenotypeStudy(PhenotypeData):
         """
         assert isinstance(measure_ids, list)
         assert len(measure_ids) >= 1
-        assert all([self.has_measure(m) for m in measure_ids])
+        assert all(self.has_measure(m) for m in measure_ids)
 
         dfs = [
             self.get_measure_values_df(
@@ -877,88 +874,133 @@ class PhenotypeStudy(PhenotypeData):
 
         return res_df
 
-    def get_values_streaming_csv(
+    def _build_measures_subquery(
         self,
+        measure_id_map: dict[str, str],
         measure_ids: list[str],
         person_ids: Optional[list[str]] = None,
         family_ids: Optional[list[str]] = None,
         roles: Optional[list[str]] = None,
-    ) -> Generator[str, None, None]:
-        assert isinstance(measure_ids, list)
-        assert len(measure_ids) >= 1
-        assert all(self.has_measure(m) for m in measure_ids)
-
-        columns = [
-            self.db.measure.c.measure_id,
-            self.db.person.c.person_id,
+    ) -> Select:
+        select_columns = [
+            self.db.person.c.person_id
         ]
-        for measure_id in measure_ids:
-            assert measure_id in self.measures, measure_id
-            measure = self.measures[measure_id]
+        for m_id in measure_ids:
+            measure = self.measures[m_id]
             measure_type = measure.measure_type
             if measure_type is None:
                 raise ValueError(
                     f"bad measure: {measure.measure_id}; unknown value type"
                 )
-        value_tables = [
-            self.db.get_value_table(MeasureType.categorical),
-            self.db.get_value_table(MeasureType.continuous),
-            self.db.get_value_table(MeasureType.ordinal),
-            self.db.get_value_table(MeasureType.raw)
-        ]
-
-        header = ["person_id"] + measure_ids
-        field_indexes = {field: index for index, field in enumerate(header)}
-
-        buffer = StringIO()
-        writer = csv.writer(buffer, delimiter=",")
-        writer.writerow(header)
-        yield buffer.getvalue()
-        buffer.seek(0)
-        buffer.truncate(0)
+            select_columns.append(cast(Column[Any], text(
+                f"\"{m_id}_value\".value AS '{m_id}'"
+            )))
+        query = select(*select_columns).select_from(
+            self.db.person.join(self.db.family)
+        )
 
-        output = {}
+        for m_id in measure_ids:
+            db_id = measure_id_map[m_id]
+            measure = self.measures[m_id]
+            measure_type = self.measures[m_id].measure_type
+            measure_table = self.db.get_value_table(measure_type)
+            table_alias = f"{m_id}_value"
+            query = query.join(
+                text(f'{measure_table.name} as "{table_alias}"'),
+                text(
+                    f'"{table_alias}".person_id = person.id AND '
+                    f'"{table_alias}".measure_id = {db_id}'
+                ),
+                isouter=True
+            )
 
-        for table in value_tables:
-            query = select(*columns, table.c.value)
+        if roles is not None:
+            query = query.where(self.db.person.c.role.in_(roles))
+        if person_ids is not None:
+            query = query.where(self.db.person.c.person_id.in_(person_ids))
+        if family_ids is not None:
+            query = query.where(self.db.family.c.family_id.in_(family_ids))
 
-            join = (
-                self.db.family.join(self.db.person)
-                .join(table, isouter=True)
-                .join(self.db.measure)
-            )
+        query = query.order_by(desc(self.db.person.c.person_id))
+        return query
+
+    def _split_measures_into_groups(
+        self, measure_ids: list[str], group_size: int = 60
+    ) -> list[list[str]]:
+        groups_count = int(len(measure_ids) / group_size) + 1
+        if (groups_count) == 1:
+            return [measure_ids]
+        measure_groups = []
+        for i in range(groups_count):
+            begin = i * group_size
+            end = (i + 1) * group_size
+            group = measure_ids[begin:end]
+            if len(group) > 0:
+                measure_groups.append(group)
+        return measure_groups
+
+    def get_people_measure_values(
+        self,
+        measure_ids: list[str],
+        person_ids: Optional[list[str]] = None,
+        family_ids: Optional[list[str]] = None,
+        roles: Optional[list[str]] = None,
+    ) -> Generator[dict[str, Any], None, None]:
+        assert isinstance(measure_ids, list)
+        assert len(measure_ids) >= 1
+        assert all(self.has_measure(m) for m in measure_ids)
 
-            query = query.select_from(join).where(
-                self.db.measure.c.measure_id.in_(measure_ids))
-
-            if roles is not None:
-                query = query.where(self.db.person.c.role.in_(roles))
-            if person_ids is not None:
-                query = query.where(self.db.person.c.person_id.in_(person_ids))
-            if family_ids is not None:
-                query = query.where(self.db.family.c.family_id.in_(family_ids))
-
-            query = query.order_by(desc(self.db.person.c.person_id))
-
-            with self.db.pheno_engine.connect() as connection:
-                results = connection.execute(query)
-                for row in results:
-                    row = row._mapping  # pylint: disable=protected-access
-                    person_id = row["person_id"]
-                    if person_id not in output:
-                        output[person_id] = ["-"] * len(header)
-                        output[person_id][field_indexes["person_id"]] = \
-                            person_id
-                    output[person_id][field_indexes[row["measure_id"]]] = \
-                        row["value"]
-
-        for value in output.values():
-            writer.writerow(value)
-            yield buffer.getvalue()
-            buffer.seek(0)
-            buffer.truncate(0)
-        buffer.close()
-        del output
+        measure_id_map = {m_id: None for m_id in measure_ids}
+        with self.db.pheno_engine.connect() as connection:
+            query = select(
+                self.db.measure.c.id, self.db.measure.c.measure_id
+            ).where(self.db.measure.c.measure_id.in_(measure_ids))
+            results = connection.execute(query)
+            for result_row in results:
+                measure_id_map[result_row.measure_id] = result_row.id
+
+        measure_groups = self._split_measures_into_groups(measure_ids)
+
+        queries = []
+        for group in measure_groups:
+            queries.append(self._build_measures_subquery(
+                cast(dict[str, str], measure_id_map),
+                group,
+                person_ids,
+                family_ids,
+                roles
+            ))
+
+        with self.db.pheno_engine.connect() as connection:
+            query_results = []
+            for query in queries:
+                query_results.append(connection.execute(query))
+            for result_row in query_results[0]:
+                row = result_row._mapping  # pylint: disable=protected-access
+                output = {**row}
+                skip = True
+                for measure_id in measure_groups[0]:
+                    if row[measure_id] is not None:
+                        skip = False
+                if len(query_results) > 1:
+                    for i in range(1, len(measure_groups)):
+                        # pylint: disable=protected-access
+                        try:
+                            row = next(query_results[i])._mapping
+                            assert row["person_id"] == output["person_id"]
+                        except StopIteration:
+                            logger.error(
+                                "Subquery %s has different length",
+                                i
+                            )
+                            continue
+                        output.update({**row})
+                        for measure_id in measure_groups[i]:
+                            if row[measure_id] is not None:
+                                skip = False
+                if skip:
+                    continue
+                yield output
 
     def get_regressions(self) -> dict[str, Any]:
         return cast(dict[str, Any], self.db.regression_display_names_with_ids)
@@ -1159,7 +1201,7 @@ class PhenotypeGroup(PhenotypeData):
             roles: Optional[Iterable[Role]] = None,
             default_filter: str = "apply") -> pd.DataFrame:
 
-        assert all([self.has_measure(mid) for mid in measure_ids]), measure_ids
+        assert all(self.has_measure(mid) for mid in measure_ids), measure_ids
 
         dfs = []
         for pheno in self.phenotype_data:
@@ -1218,13 +1260,13 @@ class PhenotypeGroup(PhenotypeData):
         measures = chain(*generators)
         yield from measures
 
-    def get_values_streaming_csv(
+    def get_people_measure_values(
         self,
         measure_ids: list[str],
         person_ids: Optional[list[str]] = None,
         family_ids: Optional[list[str]] = None,
         roles: Optional[list[str]] = None,
-    ) -> Generator[str, None, None]:
+    ) -> Generator[dict[str, Any], None, None]:
         raise NotImplementedError()
 
 
diff --git a/dev-environment.yml b/dev-environment.yml
index f918ae326..2d414d93d 100644
--- a/dev-environment.yml
+++ b/dev-environment.yml
@@ -32,7 +32,6 @@ dependencies:
   - django-stubs=4.2
   - types-pyyaml=6.0.12.2
   - types-setuptools=67
-  - sqlalchemy-stubs=0.4
   - types-toml=0.10.3
   - types-python-dateutil=2.8.9
   - types-requests=2.27.30
diff --git a/wdae/wdae/pheno_browser_api/views.py b/wdae/wdae/pheno_browser_api/views.py
index 82fe2b582..e85934c0c 100644
--- a/wdae/wdae/pheno_browser_api/views.py
+++ b/wdae/wdae/pheno_browser_api/views.py
@@ -1,5 +1,8 @@
 import logging
 
+import csv
+from io import StringIO
+
 from rest_framework.response import Response
 from rest_framework import status
 from django.http.response import StreamingHttpResponse
@@ -126,6 +129,30 @@ class PhenoMeasuresView(PhenoBrowserBaseView):
 
 
 class PhenoMeasuresDownload(QueryDatasetView):
+
+    def csv_value_iterator(self, dataset, measure_ids):
+        header = ["person_id"] + measure_ids
+        buffer = StringIO()
+        writer = csv.writer(buffer, delimiter=",")
+        writer.writerow(header)
+        yield buffer.getvalue()
+        buffer.seek(0)
+        buffer.truncate(0)
+
+        values_iterator = dataset.phenotype_data.get_people_measure_values(
+            measure_ids)
+
+        for values_dict in values_iterator:
+            output = []
+            for col in header:
+                output.append(values_dict[col])
+            writer.writerow(output)
+            yield buffer.getvalue()
+            buffer.seek(0)
+            buffer.truncate(0)
+
+        buffer.close()
+
     def post(self, request):
         data = request.data
         if "dataset_id" not in data:
@@ -153,8 +180,13 @@ class PhenoMeasuresDownload(QueryDatasetView):
             if not set(measure_ids).issubset(set(instrument_measures)):
                 return Response(status=status.HTTP_400_BAD_REQUEST)
 
-        values_iterator = dataset.phenotype_data.get_values_streaming_csv(
-            measure_ids)
+        if dataset.is_remote:
+            values_iterator = dataset.phenotype_data.get_people_measure_values(
+                measure_ids
+            )
+        else:
+            values_iterator = self.csv_value_iterator(dataset, measure_ids)
+
         response = StreamingHttpResponse(
             values_iterator, content_type="text/csv")
 
diff --git a/wdae/wdae/remote/remote_phenotype_data.py b/wdae/wdae/remote/remote_phenotype_data.py
index 233ec3817..085f81b7a 100644
--- a/wdae/wdae/remote/remote_phenotype_data.py
+++ b/wdae/wdae/remote/remote_phenotype_data.py
@@ -201,7 +201,7 @@ class RemotePhenotypeData(PhenotypeData):
 
         return instrument_values
 
-    def get_values_streaming_csv(
+    def get_people_measure_values(
         self,
         measure_ids,
         person_ids=None,


Test patch:
diff --git a/dae/dae/pheno/tests/test_pheno_db.py b/dae/dae/pheno/tests/test_pheno_db.py
index a851b230e..56824634e 100644
--- a/dae/dae/pheno/tests/test_pheno_db.py
+++ b/dae/dae/pheno/tests/test_pheno_db.py
@@ -31,6 +31,13 @@ def str_check(str_, expected_count, expected_cols):
     assert len(values_lines) == expected_count
 
 
+def dict_list_check(dict_list, expected_count, expected_cols):
+    assert isinstance(dict_list, list)
+    for dict_ in dict_list:
+        assert set(dict_.keys()) == set(expected_cols)
+    assert len(dict_list) == expected_count
+
+
 def dict_check_measure(dict_, expected_count, *args):
     assert isinstance(dict_, dict)
     for _id, measure in dict_.items():
@@ -66,31 +73,96 @@ def test_get_values(fake_phenotype_data, query_cols, get, check):
 
 
 @pytest.mark.parametrize("query_cols", [(["i1.m1"]), (["i1.m1", "i1.m2"])])
-def test_get_values_streaming_csv(fake_phenotype_data, query_cols):
-    result_it = fake_phenotype_data.get_values_streaming_csv(query_cols)
-    result = "".join(list(result_it))
-    str_check(result, 195, query_cols)
+def test_get_people_measure_values(fake_phenotype_data, query_cols):
+    result_it = fake_phenotype_data.get_people_measure_values(query_cols)
+    result = list(result_it)
+    dict_list_check(result, 195, ["person_id"] + query_cols)
 
-    result_it = fake_phenotype_data.get_values_streaming_csv(
+    result_it = fake_phenotype_data.get_people_measure_values(
         query_cols, ["f20.p1"])
-    result = "".join(list(result_it))
-    str_check(result, 1, query_cols)
+    result = list(result_it)
+    dict_list_check(result, 1, ["person_id"] + query_cols)
 
-    result_it = fake_phenotype_data.get_values_streaming_csv(
+    result_it = fake_phenotype_data.get_people_measure_values(
         query_cols, ["f20.p1", "f21.p1"])
-    result = "".join(list(result_it))
-    str_check(result, 2, query_cols)
+    result = list(result_it)
+    dict_list_check(result, 2, ["person_id"] + query_cols)
 
-    result_it = fake_phenotype_data.get_values_streaming_csv(
+    result_it = fake_phenotype_data.get_people_measure_values(
         query_cols, roles=["prb"])
-    result = "".join(list(result_it))
-    str_check(result, 39, query_cols)
+    result = list(result_it)
+    dict_list_check(result, 39, ["person_id"] + query_cols)
 
 
-def test_get_values_streaming_csv_correct_values(fake_phenotype_data):
-    result_list = list(fake_phenotype_data.get_values_streaming_csv(
+def test_get_people_measure_values_correct_values(fake_phenotype_data):
+    result_list = list(fake_phenotype_data.get_people_measure_values(
         ["i1.m1", "i1.m2"], roles=["prb"]))
-    assert result_list[-1] == "f1.p1,34.76285793898369,48.44644402952317\r\n"
+    assert result_list[-1] == {
+        "person_id": "f1.p1",
+        "i1.m1": 34.76285793898369,
+        "i1.m2": 48.44644402952317
+    }
+
+
+def test_split_into_groups(fake_phenotype_data):
+    measures = [f"measure_{i}" for i in range(1, 101)]
+    groups = fake_phenotype_data._split_measures_into_groups(measures)
+    assert len(groups) == 2
+    assert len(groups[0]) == 60
+    assert groups[0][0] == "measure_1"
+    assert groups[0][-1] == "measure_60"
+    assert len(groups[1]) == 40
+    assert groups[1][0] == "measure_61"
+    assert groups[1][-1] == "measure_100"
+
+    groups = fake_phenotype_data._split_measures_into_groups(
+        measures, group_size=25
+    )
+    assert len(groups) == 4
+    assert len(groups[0]) == 25
+    assert groups[0][0] == "measure_1"
+    assert groups[0][-1] == "measure_25"
+    assert len(groups[1]) == 25
+    assert groups[1][0] == "measure_26"
+    assert groups[1][-1] == "measure_50"
+    assert len(groups[2]) == 25
+    assert groups[2][0] == "measure_51"
+    assert groups[2][-1] == "measure_75"
+    assert len(groups[3]) == 25
+    assert groups[3][0] == "measure_76"
+    assert groups[3][-1] == "measure_100"
+
+
+def test_subquery_generation(fake_phenotype_data):
+    fake_measures_ids = {
+        "i1.m1": "1",
+        "i1.m2": "2"
+    }
+    query = str(fake_phenotype_data._build_measures_subquery(
+        fake_measures_ids,
+        fake_measures_ids.keys(),
+        person_ids=["person1", "person2"],
+        family_ids=["family1"],
+        roles=["unaffected"]
+    ))
+
+    print(query)
+    expected = (
+        "SELECT person.person_id, \"i1.m1_value\".value AS 'i1.m1', "
+        "\"i1.m2_value\".value AS 'i1.m2' \n"
+        "FROM person JOIN family ON family.id = person.family_id "
+        'LEFT OUTER JOIN value_continuous as "i1.m1_value" '
+        'ON "i1.m1_value".person_id = person.id '
+        'AND "i1.m1_value".measure_id = 1 '
+        'LEFT OUTER JOIN value_continuous as "i1.m2_value" '
+        'ON "i1.m2_value".person_id = person.id '
+        'AND "i1.m2_value".measure_id = 2 \n'
+        "WHERE person.role IN (__[POSTCOMPILE_role_1]) "
+        "AND person.person_id IN (__[POSTCOMPILE_person_id_1]) "
+        "AND family.family_id IN (__[POSTCOMPILE_family_id_1]) "
+        "ORDER BY person.person_id DESC"
+    )
+    assert query == expected
 
 
 @pytest.mark.parametrize(