Implement graph.filter for arrow endpoints

Author: DarthMax

graphdatascience/procedure_surface/api/catalog_endpoints.py

@@ -14,6 +14,42 @@
 class CatalogEndpoints(ABC):
     @abstractmethod
     def list(self, G: Optional[Union[Graph, str]] = None) -> List[GraphListResult]:
+        """List graphs in the graph catalog.
+
+        Args:
+            G (Optional[Union[Graph, str]], optional): Graph object or name to filter results.
+               If None, list all graphs. Defaults to None.
+
+        Returns:
+            List[GraphListResult]: List of graph metadata objects containing information like
+                                 graph name, node count, relationship count, etc.
+        """
+        pass
+
+    @abstractmethod
+    def filter(
+        self,
+        G: Graph,
+        graph_name: str,
+        node_filter: str,
+        relationship_filter: str,
+        concurrency: Optional[int] = None,
+        job_id: Optional[str] = None,
+    ) -> GraphFilterResult:
+        """Create a subgraph of a graph based on a filter expression.
+
+        Args:
+            G (Graph): Graph object to filter on
+            graph_name (str): Name of subgraph to create
+            node_filter (str): Filter expression for nodes
+            relationship_filter (str): Filter expression for relationships
+            concurrency (Optional[int], optional): Number of concurrent threads to use. Defaults to None.
+            job_id (Optional[str], optional): Unique identifier for the filtering job. Defaults to None.
+
+        Returns:
+            GraphFilterResult: Filter result containing information like
+                                graph name, node count, relationship count, etc.
+        """
         pass
 
 
@@ -40,3 +76,15 @@ def strip_timezone(cls, value: Any) -> Any:
         if isinstance(value, str):
             return re.sub(r"\[.*\]$", "", value)
         return value
+
+
+class GraphFilterResult(BaseModel):
+    model_config = ConfigDict(alias_generator=to_camel)
+
+    graph_name: str
+    from_graph_name: str
+    node_filter: str
+    relationship_filter: str
+    node_count: int
+    relationship_count: int
+    project_millis: int

graphdatascience/procedure_surface/arrow/catalog_arrow_endpoints.py

@@ -11,7 +11,12 @@
 from graphdatascience.arrow_client.authenticated_flight_client import AuthenticatedArrowClient
 from graphdatascience.arrow_client.v2.data_mapper_utils import deserialize
 from graphdatascience.arrow_client.v2.job_client import JobClient
-from graphdatascience.procedure_surface.api.catalog_endpoints import CatalogEndpoints, GraphListResult
+from graphdatascience.procedure_surface.api.catalog_endpoints import (
+    CatalogEndpoints,
+    GraphFilterResult,
+    GraphListResult,
+)
+from graphdatascience.procedure_surface.utils.config_converter import ConfigConverter
 from graphdatascience.query_runner.protocol.project_protocols import ProjectProtocol
 from graphdatascience.query_runner.termination_flag import TerminationFlag
 from graphdatascience.session.dbms.protocol_resolver import ProtocolVersionResolver
@@ -81,6 +86,28 @@ def project(
 
         return ProjectionResult(**JobClient.get_summary(self._arrow_client, job_id))
 
+    def filter(
+        self,
+        G: Graph,
+        graph_name: str,
+        node_filter: str,
+        relationship_filter: str,
+        concurrency: Optional[int] = None,
+        job_id: Optional[str] = None,
+    ) -> GraphFilterResult:
+        config = ConfigConverter.convert_to_gds_config(
+            fromGraphName=G.name(),
+            graphName=graph_name,
+            nodeFilter=node_filter,
+            relationshipFilter=relationship_filter,
+            concurrency=concurrency,
+            jobId=job_id,
+        )
+
+        job_id = JobClient.run_job_and_wait(self._arrow_client, "v2/graph.project.filter", config)
+
+        return GraphFilterResult(**JobClient.get_summary(self._arrow_client, job_id))
+
     def _arrow_config(self) -> dict[str, Any]:
         connection_info = self._arrow_client.connection_info()
 

graphdatascience/tests/integrationV2/procedure_surface/arrow/test_catalog_arrow_endpoints.py

@@ -13,9 +13,9 @@
 @pytest.fixture
 def sample_graph(arrow_client: AuthenticatedArrowClient) -> Generator[Graph, None, None]:
     gdl = """
-    (a: Node)
-    (b: Node)
-    (c: Node)
+    (a :Node:A)
+    (b :Node:A)
+    (c :Node:B)
     (a)-[:REL]->(c)
     """
 
@@ -76,3 +76,20 @@ def test_projection(arrow_client: AuthenticatedArrowClient, query_runner: QueryR
         assert len(endpoints.list("g")) == 1
     finally:
         arrow_client.do_action("v2/graph.drop", json.dumps({"graphName": "g"}).encode("utf-8"))
+
+
+def test_graph_filter(catalog_endpoints: CatalogArrowEndpoints, sample_graph: Graph) -> None:
+    try:
+        result = catalog_endpoints.filter(
+            sample_graph, graph_name="filtered", node_filter="n:A", relationship_filter="*"
+        )
+
+        assert result.node_count == 2
+        assert result.relationship_count == 0
+        assert result.from_graph_name == sample_graph.name()
+        assert result.graph_name == "filtered"
+        assert result.project_millis >= 0
+    finally:
+        catalog_endpoints._arrow_client.do_action(
+            "v2/graph.drop", json.dumps({"graphName": "filtered"}).encode("utf-8")
+        )