From 3f6ac20f392aebe0f524cae20a485d1a98e8e2c6 Mon Sep 17 00:00:00 2001
From: rjawesome <rohanj2006@gmail.com>
Date: Fri, 4 Oct 2024 16:29:16 -0700
Subject: [PATCH] store edge attributes as arrays, convert to set later if
 needed

---
 src/graph/kg_edge.ts         |  8 +++-----
 src/graph/knowledge_graph.ts | 18 +++++++++++++-----
 src/index.ts                 |  2 +-
 3 files changed, 17 insertions(+), 11 deletions(-)

diff --git a/src/graph/kg_edge.ts b/src/graph/kg_edge.ts
index 797c82f1..14ec7c40 100644
--- a/src/graph/kg_edge.ts
+++ b/src/graph/kg_edge.ts
@@ -29,7 +29,7 @@ export default class KGEdge {
     [qualifier_type_id: string]: string | string[];
   };
   attributes: {
-    [attribute_type_id: string]: Set<string> | TrapiAttribute[];
+    [attribute_type_id: string]: string[] | TrapiAttribute[];
     'edge-attributes'?: TrapiAttribute[];
   };
   constructor(id: string, info: KGEdgeInfo) {
@@ -125,13 +125,11 @@ export default class KGEdge {
     }
 
     if (!(name in this.attributes)) {
-      this.attributes[name] = new Set();
+      this.attributes[name] = [];
     }
     if (!Array.isArray(value)) {
       value = [value];
     }
-    (value as string[]).map((item) => {
-      (this.attributes[name] as Set<string>).add(item);
-    });
+    (this.attributes[name] as string[]).push(...(value as string[]));
   }
 }
diff --git a/src/graph/knowledge_graph.ts b/src/graph/knowledge_graph.ts
index 9aa4b2da..c5f2b870 100644
--- a/src/graph/knowledge_graph.ts
+++ b/src/graph/knowledge_graph.ts
@@ -17,7 +17,8 @@ import { toArray, Telemetry } from '@biothings-explorer/utils';
 
 const debug = Debug('bte:biothings-explorer-trapi:KnowledgeGraph');
 
-const NON_ARRAY_ATTRIBUTES = ['biolink:knowledge_level', 'biolink:agent_type', 'biolink:evidence_count'];
+const NON_ARRAY_ATTRIBUTES = ['biolink:knowledge_level', 'biolink:agent_type'];
+const SUM_ATTRIBUTES = ['biolink:evidence_count'];
 
 interface SpecialAttributeHandlers {
   [attribute_type_id: string]: (value: Set<string | number>, kgEdge: KGEdge) => TrapiAttribute['value'];
@@ -149,12 +150,19 @@ export default class KnowledgeGraph {
     Object.entries(kgEdge.attributes).forEach(([key, value]) => {
       if (key === 'edge-attributes') return;
 
-      let formatted_value: TrapiAttribute['value'] = NON_ARRAY_ATTRIBUTES.includes(key)
-        ? Array.from(value as Set<string>).reduce((acc, val) => acc + val)
-        : Array.from(value as Set<string>);
+      let formatted_value: TrapiAttribute['value'];
+      if (SUM_ATTRIBUTES.includes(key)) {
+        // for sums we don't want to remove duplicates
+        formatted_value = (value as string[]).reduce((acc, val) => acc + val);
+      } else if (NON_ARRAY_ATTRIBUTES.includes(key)) {
+        // for non array attributes we want to remove duplicates (ie. same string for knowledge_level multiple times)
+        formatted_value = Array.from(new Set(value as string[])).reduce((acc, val) => acc + val);
+      } else {
+        formatted_value = Array.from(new Set(value as string[]));
+      }
 
       if (key in SPECIAL_ATTRIBUTE_HANDLERS) {
-        formatted_value = SPECIAL_ATTRIBUTE_HANDLERS[key](value as Set<string | number>, kgEdge);
+        formatted_value = SPECIAL_ATTRIBUTE_HANDLERS[key](new Set(value as string[]), kgEdge);
       }
 
       attributes.push({
diff --git a/src/index.ts b/src/index.ts
index 9ce86640..e933d563 100644
--- a/src/index.ts
+++ b/src/index.ts
@@ -278,7 +278,7 @@ export default class TRAPIQueryHandler {
           ]);
           this.bteGraph.edges[boundEdgeID] = boundEdge;
         } else {
-          (this.bteGraph.edges[boundEdgeID].attributes['biolink:support_graphs'] as Set<string>).add(supportGraphID);
+          this.bteGraph.edges[boundEdgeID].addAdditionalAttributes('biolink:support_graphs', supportGraphID);
         }
         if (!edgesToRebind[edgeID]) edgesToRebind[edgeID] = {};
         if (!edgesToRebind[edgeID][subject]) edgesToRebind[edgeID][subject] = {};