prepare_df fix

aphp · Dec 14, 2023 · ed2acb0 · ed2acb0
1 parent 3645178
commit ed2acb0
Showing 1 changed file with 7 additions and 5 deletions.
diff --git a/eds_scikit/biology/utils/prepare_df.py b/eds_scikit/biology/utils/prepare_df.py
@@ -6,7 +6,7 @@ def prepare_biology_relationship(
     data,
     source_terminologies,
     mapping,
-    concepts_sets
+    concept_codes=None
 ) -> pd.DataFrame:
     """Computes biology relationship table
 
@@ -34,7 +34,7 @@ def prepare_biology_relationship(
 
     """
 
-    check_tables(data=data, required_tables=["concept", "concept_relationship"])
+    #check_tables(data=data, required_tables=["concept", "concept_relationship"])
     concept_columns = [
         "concept_id",
         "concept_name",
@@ -61,6 +61,7 @@ def prepare_biology_relationship(
 
     concept = data.concept[concept_columns]
     concept_relationship = data.concept_relationship[concept_relationship_columns]
+    concept, concept_relationship = concept.to_pandas(), concept_relationship.to_pandas()
     concept_by_terminology = {}
     for terminology, regex in source_terminologies.items():
         concept_by_terminology[terminology] = (
@@ -92,13 +93,14 @@ def prepare_biology_relationship(
             relationship, on="{}_concept_id".format(source), how="left"
         )
 
-    if concepts_sets:
+    if concept_codes:
         code_colums = [column for column in biology_relationship.columns if "concept_code" in column]
-        concept_codes = [item for concept_set in concepts_sets for item in concept_set.concept_codes]
         isin_concept_set = biology_relationship[code_colums].isin(concept_codes).sum(axis=1) > 0
         biology_relationship = biology_relationship[isin_concept_set]
 
-    biology_relationship["concepts_set"] = "XXX"
+    #biology_relationship["concepts_set"] = "XXX"
     biology_relationship = biology_relationship.fillna("Unknown")
+
+    biology_relationship = to("koalas", biology_relationship)
 
     return biology_relationship