Handle is_set

strider/node_sets.py

@@ -0,0 +1,61 @@
+"""Node sets."""
+from collections import defaultdict
+
+
+def collapse_sets(message: dict) -> None:
+    """Collase results according to is_set qnode notations."""
+    unique_qnodes = {
+        qnode_id
+        for qnode_id, qnode in message["query_graph"]["nodes"].items()
+        if not qnode.get("is_set", False)
+    }
+    if len(unique_qnodes) == len(message["query_graph"]["nodes"]):
+        return
+    unique_qedges = {
+        qedge_id
+        for qedge_id, qedge in message["query_graph"]["edges"].items()
+        if (
+            qedge["subject"] in unique_qnodes
+            and qedge["object"] in unique_qnodes
+        )
+    }
+    result_buckets = defaultdict(lambda: {
+        "node_bindings": defaultdict(set),
+        "edge_bindings": defaultdict(set),
+    })
+    for result in message["results"]:
+        bucket_key = tuple([
+            binding["id"]
+            for qnode_id in unique_qnodes
+            for binding in result["node_bindings"][qnode_id]
+        ] + [
+            binding["id"]
+            for qedge_id in unique_qedges
+            for binding in result["edge_bindings"][qedge_id]
+        ])
+        for qnode_id in message["query_graph"]["nodes"]:
+            result_buckets[bucket_key]["node_bindings"][qnode_id] |= {
+                binding["id"]
+                for binding in result["node_bindings"][qnode_id]
+            }
+        for qedge_id in message["query_graph"]["edges"]:
+            result_buckets[bucket_key]["edge_bindings"][qedge_id] |= {
+                binding["id"]
+                for binding in result["edge_bindings"][qedge_id]
+            }
+    for result in result_buckets.values():
+        result["node_bindings"] = {
+            qnode_id: [
+                {"id": binding}
+                for binding in bindings
+            ]
+            for qnode_id, bindings in result["node_bindings"].items()
+        }
+        result["edge_bindings"] = {
+            qedge_id: [
+                {"id": binding}
+                for binding in bindings
+            ]
+            for qedge_id, bindings in result["edge_bindings"].items()
+        }
+    message["results"] = list(result_buckets.values())

strider/server.py

@@ -21,6 +21,7 @@
 from reasoner_pydantic import Query, AsyncQuery, Message, Response as ReasonerResponse
 
 from .fetcher import Binder
+from .node_sets import collapse_sets
 from .query_planner import NoAnswersError, generate_plan
 from .scoring import score_graph
 from .storage import RedisGraph, RedisList, get_client as get_redis_client
@@ -41,7 +42,7 @@
     title="Strider",
     description=DESCRIPTION,
     docs_url=None,
-    version="3.12.0",
+    version="3.13.0",
     terms_of_service=(
         "http://robokop.renci.org:7055/tos"
         "?service_long=Strider"
@@ -243,13 +244,15 @@ async def lookup(
                 kgraph,
             ])
             results.append(result)
+    message = {
+        "query_graph": qgraph,
+        "knowledge_graph": kgraph,
+        "results": results,
+    }
+    collapse_sets(message)
     logs = list(RedisList(f"{qid}:log", redis_client).get())
     return {
-        "message": {
-            "query_graph": qgraph,
-            "knowledge_graph": kgraph,
-            "results": results,
-        },
+        "message": message,
         "logs": logs,
     }
 

tests/test_node_sets.py

@@ -0,0 +1,54 @@
+"""Test node set handling."""
+from strider.node_sets import collapse_sets
+
+
+def test_node_sets():
+    """Test collapsing one edge of a two-hop query."""
+    qgraph = {
+        "nodes": {
+            "n0": {},
+            "n1": {"is_set": True}
+        },
+        "edges": {
+            "e01": {
+                "subject": "n0",
+                "object": "n1",
+            },
+        },
+    }
+    results = [
+        {
+            "node_bindings": {
+                "n0": [{"id": "a0"}],
+                "n1": [{"id": "b0"}],
+            },
+            "edge_bindings": {
+                "e01": [{"id": "c0"}],
+            },
+        },
+        {
+            "node_bindings": {
+                "n0": [{"id": "a1"}],
+                "n1": [{"id": "b0"}],
+            },
+            "edge_bindings": {
+                "e01": [{"id": "c1"}],
+            },
+        },
+        {
+            "node_bindings": {
+                "n0": [{"id": "a0"}],
+                "n1": [{"id": "b1"}],
+            },
+            "edge_bindings": {
+                "e01": [{"id": "c2"}],
+            },
+        },
+    ]
+    message = {
+        "query_graph": qgraph,
+        "results": results,
+    }
+    collapse_sets(message)
+    assert len(message["results"]) == 2
+    assert len(message["results"][0]["node_bindings"]["n1"]) == 2

tests/test_query.py

@@ -377,7 +377,8 @@ async def test_trivial_unbatching(redis):
         },
         output["message"]
     )
-
+
+
 @pytest.mark.asyncio
 @with_translator_overlay(
     settings.kpregistry_url,
@@ -464,3 +465,57 @@ async def test_gene_protein_conflation(redis):
         },
         output["message"]
     )
+
+
+@pytest.mark.asyncio
+@with_translator_overlay(
+    settings.kpregistry_url,
+    settings.normalizer_url,
+    {
+        "ctd":
+        """
+            CHEBI:6801(( category biolink:ChemicalSubstance ))
+            MONDO:0005148(( category biolink:Disease ))
+            CHEBI:6801-- predicate biolink:treats -->MONDO:0005148
+            MONDO:0003757(( category biolink:Disease ))
+            CHEBI:6801-- predicate biolink:treats -->MONDO:0003757
+        """,
+        "pharos":
+        """
+            MONDO:0003757(( category biolink:Disease ))
+            MONDO:0005148(( category biolink:Disease ))
+            HP:XXX(( category biolink:PhenotypicFeature ))
+            HP:YYY(( category biolink:PhenotypicFeature ))
+            MONDO:0003757-- predicate biolink:has_phenotype -->HP:XXX
+            MONDO:0003757-- predicate biolink:has_phenotype -->HP:YYY
+            MONDO:0005148-- predicate biolink:has_phenotype -->HP:YYY
+        """,
+    }
+)
+async def test_node_set(redis):
+    """Test that is_set is handled correctly."""
+    QGRAPH = query_graph_from_string(
+        """
+        n0(( ids[] CHEBI:6801 ))
+        n0(( categories[] biolink:ChemicalSubstance ))
+        n1(( categories[] biolink:Disease ))
+        n2(( categories[] biolink:PhenotypicFeature ))
+        n0-- biolink:treats -->n1
+        n1-- biolink:has_phenotype -->n2
+        """
+    )
+    QGRAPH["nodes"]["n1"]["is_set"] = True
+
+    # Create query
+    q = {
+        "message" : {"query_graph" : QGRAPH},
+        "log_level" : "WARNING",
+    }
+
+    # Run
+    output = await lookup(q, redis)
+    assert len(output["message"]["results"]) == 2
+    assert {
+        len(result["node_bindings"]["n1"])
+        for result in output["message"]["results"]
+    } == {1, 2}