change drug molecule map into a struct

Co-authored-by: project-defiant <szymonszyszkowski@gmail.com>

Author: remo87

src/main/scala/io/opentargets/etl/backend/Search.scala

@@ -447,6 +447,24 @@ object Transformers {
 
     }
 
+    /** Collect cross reference ids for the drugs
+      * @param crossReferences
+      *   crossReferences column from the drug dataframe
+      * @return
+      *   a column with an array of cross reference ids
+      */
+    def collectCrossReferenceIds(crossReferences: Column): Column =
+      sort_array(
+        array_distinct(
+          flatten(
+            transform(
+              crossReferences,
+              x => x.getField("ids")
+            )
+          )
+        )
+      ).alias("crossReferences")
+
     // uses target_ids, drug_id, target_labels, disease_id, disease_labels
     def setIdAndSelectFromDrugs(
         associatedDrugs: DataFrame,
@@ -502,6 +520,10 @@ object Transformers {
           when(col("disease_labels").isNull, Array.empty[String])
             .otherwise(col("disease_labels"))
         )
+        .withColumn(
+          "crossReferences",
+          collectCrossReferenceIds(col("crossReferences"))
+        )
       SearchIndex(
         id = col("id"),
         name = col("name"),
@@ -514,9 +536,7 @@ object Transformers {
           "array(name)",
           "array(id)",
           "childChemblIds",
-          "crossReferences.PubChem",
-          "crossReferences.drugbank",
-          "crossReferences.chEBI"
+          "crossReferences"
         ),
         prefixes = C.flattenCat(
           "synonyms",

src/main/scala/io/opentargets/etl/backend/drug/Drug.scala

@@ -70,7 +70,7 @@ object Drug extends Serializable with LazyLogging {
     )
 
     // We define a drug as having either a drugbank id, a mechanism of action, an indication, or if it is a chemical probe.
-    val isDrugMolecule: Column = array_contains(map_keys(col("crossReferences")), "drugbank") ||
+    val isDrugMolecule: Column = array_contains(col("crossReferences.source"), "drugbank") ||
       col("indications").isNotNull ||
       col("mechanismsOfAction").isNotNull ||
       col("chemicalProbeDrugId").isNotNull

src/main/scala/io/opentargets/etl/backend/drug/Molecule.scala

@@ -15,6 +15,7 @@ import org.apache.spark.sql.functions.{
   explode,
   lit,
   map_concat,
+  struct,
   typedLit,
   udf,
   upper,
@@ -164,7 +165,14 @@ object Molecule extends LazyLogging {
       .foldLeft(chemblCrossReferences)((agg, a) => mergeCrossReferenceMaps(agg, a))
       .filter(col(XREF_COLUMN_NAME).isNotNull)
       .withColumnRenamed(XREF_COLUMN_NAME, "crossReferences")
-    references
+
+    val transformedCrossReference = references
+      .select(col("id"), explode(col("crossReferences")))
+      .withColumnRenamed("key", "source")
+      .withColumnRenamed("value", "ids")
+      .groupBy("id")
+      .agg(collect_set(struct(col("source"), col("ids"))).as("crossReferences"))
+    transformedCrossReference
   }
 
   /** @param preProcessedMolecules

src/test/scala/io/opentargets/etl/backend/Drug/MoleculeTest.scala

@@ -202,10 +202,10 @@ class MoleculeTest extends EtlSparkUnitTest {
       .json(this.getClass.getResource("/sample_mol_after_preprocessing.json").getPath)
     // when
     val results = Molecule invokePrivate processMoleculeCrossReferences(sampleMolecule)
-    val xrefMap = results.head.getMap(1)
+    val xrefMap = results.head.getList(1)
     // then
     assertResult(4) {
-      xrefMap.keySet.size
+      xrefMap.size
     }
   }