Add faceting/filtering to association endpoints (#786)

Updates association endpoints to allow facet field, facet query and
(open) filter query params

backend/src/monarch_py/api/association.py

@@ -29,6 +29,9 @@ async def _get_associations(
     object_taxon: Union[List[str], None] = Query(default=None),
     entity: Union[List[str], None] = Query(default=None),
     direct: bool = Query(default=False),
+    facet_fields: List[str] = Query(default_factory=list),
+    facet_queries: List[str] = Query(default_factory=list),
+    filter_queries: List[str] = Query(default_factory=list),
     compact: bool = Query(default=False),
     pagination: PaginationParams = Depends(),
     format: OutputFormat = Query(
@@ -52,6 +55,9 @@ async def _get_associations(
         object_namespace=object_namespace,
         direct=direct,
         compact=compact,
+        facet_fields=facet_fields,
+        facet_queries=facet_queries,
+        filter_queries=filter_queries,
         offset=pagination.offset,
         limit=pagination.limit,
     )

backend/src/monarch_py/api/entity.py

@@ -86,6 +86,17 @@ def _association_table(
         title="Only return direct associations",
         examples=[True, False],
     ),
+    facet_fields: List[str] = Query(
+        default=None,
+        title="Facet fields to include in the response",
+        examples=["subject", "subject_taxon", "predicate"],
+    ),
+    facet_queries: List[str] = Query(
+        default=None, title="Facet queries to include in the response", examples=['subject_category:"biolink:Gene"']
+    ),
+    filter_queries: List[str] = Query(
+        default=None, title="Filter queries to limit the response", examples=['subject_category:"biolink:Gene"']
+    ),
 ) -> Union[AssociationTableResults, str]:
     """
     Retrieves association table data for a given entity and association type
@@ -105,6 +116,9 @@ def _association_table(
         q=query,
         traverse_orthologs=traverse_orthologs,
         direct=direct,
+        facet_fields=facet_fields,
+        facet_queries=facet_queries,
+        filter_queries=filter_queries,
         sort=sort,
         offset=pagination.offset,
         limit=pagination.limit,

backend/src/monarch_py/datamodels/model.py

@@ -2205,6 +2205,36 @@ class AssociationResults(Results):
             }
         },
     )
+    facet_fields: Optional[List[FacetField]] = Field(
+        None,
+        description="""Collection of facet field responses with the field values and counts""",
+        json_schema_extra={
+            "linkml_meta": {
+                "alias": "facet_fields",
+                "domain_of": [
+                    "AssociationResults",
+                    "CompactAssociationResults",
+                    "AssociationTableResults",
+                    "SearchResults",
+                ],
+            }
+        },
+    )
+    facet_queries: Optional[List[FacetValue]] = Field(
+        None,
+        description="""Collection of facet query responses with the query string values and counts""",
+        json_schema_extra={
+            "linkml_meta": {
+                "alias": "facet_queries",
+                "domain_of": [
+                    "AssociationResults",
+                    "CompactAssociationResults",
+                    "AssociationTableResults",
+                    "SearchResults",
+                ],
+            }
+        },
+    )
     limit: int = Field(
         ...,
         description="""number of items to return in a response""",
@@ -2250,6 +2280,36 @@ class CompactAssociationResults(Results):
             }
         },
     )
+    facet_fields: Optional[List[FacetField]] = Field(
+        None,
+        description="""Collection of facet field responses with the field values and counts""",
+        json_schema_extra={
+            "linkml_meta": {
+                "alias": "facet_fields",
+                "domain_of": [
+                    "AssociationResults",
+                    "CompactAssociationResults",
+                    "AssociationTableResults",
+                    "SearchResults",
+                ],
+            }
+        },
+    )
+    facet_queries: Optional[List[FacetValue]] = Field(
+        None,
+        description="""Collection of facet query responses with the query string values and counts""",
+        json_schema_extra={
+            "linkml_meta": {
+                "alias": "facet_queries",
+                "domain_of": [
+                    "AssociationResults",
+                    "CompactAssociationResults",
+                    "AssociationTableResults",
+                    "SearchResults",
+                ],
+            }
+        },
+    )
     limit: int = Field(
         ...,
         description="""number of items to return in a response""",
@@ -2295,6 +2355,36 @@ class AssociationTableResults(Results):
             }
         },
     )
+    facet_fields: Optional[List[FacetField]] = Field(
+        None,
+        description="""Collection of facet field responses with the field values and counts""",
+        json_schema_extra={
+            "linkml_meta": {
+                "alias": "facet_fields",
+                "domain_of": [
+                    "AssociationResults",
+                    "CompactAssociationResults",
+                    "AssociationTableResults",
+                    "SearchResults",
+                ],
+            }
+        },
+    )
+    facet_queries: Optional[List[FacetValue]] = Field(
+        None,
+        description="""Collection of facet query responses with the query string values and counts""",
+        json_schema_extra={
+            "linkml_meta": {
+                "alias": "facet_queries",
+                "domain_of": [
+                    "AssociationResults",
+                    "CompactAssociationResults",
+                    "AssociationTableResults",
+                    "SearchResults",
+                ],
+            }
+        },
+    )
     limit: int = Field(
         ...,
         description="""number of items to return in a response""",
@@ -2673,12 +2763,32 @@ class SearchResults(Results):
     facet_fields: Optional[List[FacetField]] = Field(
         None,
         description="""Collection of facet field responses with the field values and counts""",
-        json_schema_extra={"linkml_meta": {"alias": "facet_fields", "domain_of": ["SearchResults"]}},
+        json_schema_extra={
+            "linkml_meta": {
+                "alias": "facet_fields",
+                "domain_of": [
+                    "AssociationResults",
+                    "CompactAssociationResults",
+                    "AssociationTableResults",
+                    "SearchResults",
+                ],
+            }
+        },
     )
     facet_queries: Optional[List[FacetValue]] = Field(
         None,
         description="""Collection of facet query responses with the query string values and counts""",
-        json_schema_extra={"linkml_meta": {"alias": "facet_queries", "domain_of": ["SearchResults"]}},
+        json_schema_extra={
+            "linkml_meta": {
+                "alias": "facet_queries",
+                "domain_of": [
+                    "AssociationResults",
+                    "CompactAssociationResults",
+                    "AssociationTableResults",
+                    "SearchResults",
+                ],
+            }
+        },
     )
     limit: int = Field(
         ...,

backend/src/monarch_py/datamodels/model.yaml

@@ -125,6 +125,8 @@ classes:
     is_a: Results
     slots:
       - items
+      - facet_fields
+      - facet_queries
     slot_usage:
       items:
         range: Association
@@ -141,13 +143,17 @@ classes:
     is_a: Results
     slots:
       - items
+      - facet_fields
+      - facet_queries
     slot_usage:
       items:
         range: CompactAssociation
   AssociationTableResults:
     is_a: Results
     slots:
       - items
+      - facet_fields
+      - facet_queries
     slot_usage:
       items:
         range: DirectionalAssociation

backend/src/monarch_py/datamodels/solr.py

@@ -44,6 +44,7 @@ class SolrQuery(BaseModel):
     facet_fields: Optional[List[str]] = Field(default_factory=list)
     facet_queries: Optional[List[str]] = Field(default_factory=list)
     filter_queries: Optional[List[str]] = Field(default_factory=list)
+    facet_mincount: int = 1
     query_fields: Optional[str] = None
     def_type: str = "edismax"
     q_op: str = "AND"  # See SOLR-8812, need this plus mm=100% to allow boolean operators in queries
@@ -85,6 +86,8 @@ def _solrize(self, value):
             return "facet.query"
         elif value == "filter_queries":
             return "fq"
+        elif value == "facet_mincount":
+            return "facet.mincount"
         elif value == "query_fields":
             return "qf"
         elif value == "def_type":

backend/src/monarch_py/implementations/solr/solr_implementation.py

@@ -123,7 +123,7 @@ def get_entity(self, id: str, extra: bool) -> Optional[Union[Node, Entity]]:
                 ).items
             ]
         node: Node = Node(
-            **entity.dict(),
+            **entity.model_dump(),
             node_hierarchy=self._get_node_hierarchy(entity),
             association_counts=self.get_association_counts(id).items,
             external_links=get_links_for_field(entity.xref) if entity.xref else [],
@@ -244,6 +244,9 @@ def get_associations(
         entity: Optional[List[str]] = None,
         direct: bool = False,
         q: Optional[str] = None,
+        facet_fields: Optional[List[str]] = None,
+        facet_queries: Optional[List[str]] = None,
+        filter_queries: Optional[List[str]] = None,
         compact: bool = False,
         offset: int = 0,
         limit: int = 20,
@@ -259,6 +262,9 @@ def get_associations(
             object_closure: Filter to only associations with the specified term ID as an ancestor of the object. Defaults to None
             entity: Filter to only associations where the specified entities are the subject or the object. Defaults to None.
             q: Query string to search within matches. Defaults to None.
+            facet_fields: List of fields to include facet counts for. Defaults to None.
+            facet_queries: List of queries to include facet counts for. Defaults to None.
+            filter_queries: List of queries to filter results by. Defaults to None.
             compact: Return compact results with fewer fields. Defaults to False.
             offset: Result offset, for pagination. Defaults to 0.
             limit: Limit results to specified number. Defaults to 20.
@@ -283,6 +289,9 @@ def get_associations(
             object_namespace=[object_namespace] if isinstance(object_namespace, str) else object_namespace,
             direct=direct,
             q=q,
+            facet_fields=facet_fields,
+            facet_queries=facet_queries,
+            filter_queries=filter_queries,
             offset=offset,
             limit=limit,
         )
@@ -430,6 +439,7 @@ def get_association_facets(
         entity: Optional[List[str]] = None,
         facet_fields: Optional[List[str]] = None,
         facet_queries: Optional[List[str]] = None,
+        filter_queries: Optional[List[str]] = None,
     ) -> SearchResults:
         solr = SolrService(base_url=self.base_url, core=core.ASSOCIATION)
 
@@ -445,6 +455,7 @@ def get_association_facets(
             limit=0,
             facet_fields=facet_fields,
             facet_queries=facet_queries,
+            filter_queries=filter_queries,
         )
         query_result = solr.query(query)
         return SearchResults(
@@ -467,6 +478,9 @@ def get_association_table(
         traverse_orthologs: bool = False,
         direct: bool = False,
         q: Optional[str] = None,
+        facet_fields: Optional[List[str]] = None,
+        facet_queries: Optional[List[str]] = None,
+        filter_queries: Optional[List[str]] = None,
         sort: Optional[List[str]] = None,
         offset: int = 0,
         limit: int = 5,
@@ -485,6 +499,9 @@ def get_association_table(
             category=category.value,
             direct=direct,
             q=q,
+            facet_fields=facet_fields,
+            facet_queries=facet_queries,
+            filter_queries=filter_queries,
             sort=sort,
             offset=offset,
             limit=limit,

backend/src/monarch_py/implementations/solr/solr_parsers.py

@@ -54,7 +54,14 @@ def parse_associations(
             )
             for doc in query_result.response.docs
         ]
-        return CompactAssociationResults(items=associations, limit=limit, offset=offset, total=total)
+        return CompactAssociationResults(
+            items=associations,
+            limit=limit,
+            offset=offset,
+            total=total,
+            facet_fields=convert_facet_fields(query_result.facet_counts.facet_fields),
+            facet_queries=convert_facet_queries(query_result.facet_counts.facet_queries),
+        )
     else:
         for doc in query_result.response.docs:
             try:
@@ -72,7 +79,14 @@ def parse_associations(
                 get_links_for_field(association.publications) if association.publications else []
             )
             associations.append(association)
-        return AssociationResults(items=associations, limit=limit, offset=offset, total=total)
+        return AssociationResults(
+            items=associations,
+            limit=limit,
+            offset=offset,
+            total=total,
+            facet_fields=convert_facet_fields(query_result.facet_counts.facet_fields),
+            facet_queries=convert_facet_queries(query_result.facet_counts.facet_queries),
+        )
 
 
 def parse_association_counts(query_result: SolrQueryResult, entity: str) -> AssociationCountList:
@@ -140,7 +154,14 @@ def parse_association_table(
         except ValidationError:
             logger.error(f"Validation error for {doc}")
             raise
-    results = AssociationTableResults(items=associations, limit=limit, offset=offset, total=total)
+    results = AssociationTableResults(
+        items=associations,
+        limit=limit,
+        offset=offset,
+        total=total,
+        facet_fields=convert_facet_fields(query_result.facet_counts.facet_fields),
+        facet_queries=convert_facet_queries(query_result.facet_counts.facet_queries),
+    )
     for i in zip(results.items, query_result.response.docs):
         assert i[0].subject == i[1]["subject"]
     return results