Next

Author: orazve

examples/kge-distmult.py

@@ -94,8 +94,13 @@ def read_data():
 
 
 def put_data_in_db(dataset):
-    for rel_split in tqdm(dataset, desc="Relationship"):
-        for rel_type in tqdm(dataset[rel_split], mininterval=1, leave=False):
+    res = gds.run_cypher("MATCH (m) RETURN count(m) as num_nodes")
+    if res['num_nodes'].values[0] > 0:
+        print("Data already in db, number of nodes: ", res['num_nodes'].values[0])
+        return
+    pbar = tqdm(desc='Putting data in db', total=sum([len(dataset[rel_split][rel_type]) for rel_split in dataset for rel_type in dataset[rel_split]]))
+    for rel_split in dataset:
+        for rel_type in dataset[rel_split]:
             edges = dataset[rel_split][rel_type]
 
             # MERGE (n)-[:{rel_type} {{text:l.rel_text}}]->(m)
@@ -109,6 +114,8 @@ def put_data_in_db(dataset):
                 """,
                 params={"ll": edges},
             )
+            pbar.update(len(edges))
+    pbar.close()
 
     for rel_split in dataset:
         res = gds.run_cypher(
@@ -120,7 +127,7 @@ def put_data_in_db(dataset):
         print(f"Number of relationships of type {rel_split} in db: ", res.numberOfRelationships)
 
 
-# put_data_in_db(dataset)
+put_data_in_db(dataset)
 
 ALL_RELS = dataset["TRAIN"].keys()
 gds.graph.drop("trainGraph", failIfMissing=False)
@@ -159,13 +166,18 @@ def inspect_graph(G):
 gds.set_compute_cluster_ip("localhost")
 
 kkge = gds.kge
+kmodel = gds.kge.model
+
+print(gds.debug.arrow())
 
 gds.kge.model.train(
     G_train,
-    scoring_function="distmult",
-    num_epochs=10,
-    embedding_dimension=100,
+    scoring_function="DistMult",
+    num_epochs=1,
+    embedding_dimension=10,
 )
+
+print('Finished training')
 #
 # node_projection = {"Entity": {"properties": "id"}}
 # relationship_projection = [

graphdatascience/graph_data_science.py

@@ -87,10 +87,12 @@ def __init__(
                 None if arrow is True else arrow,
             )
 
-        # if auth is not None:
-        #     with open(self._path("graphdatascience.resources.field-testing", "pub.pem"), "rb") as f:
-        #         pub_key = rsa.PublicKey.load_pkcs1(f.read())
-        #     self._encrypted_db_password = rsa.encrypt(auth[1].encode(), pub_key).hex()
+        if auth is not None:
+            with open(self._path("graphdatascience.resources.field-testing", "pub.pem"), "rb") as f:
+                pub_key = rsa.PublicKey.load_pkcs1(f.read())
+            self._encrypted_db_password = rsa.encrypt(auth[1].encode(), pub_key).hex()
+        # self._encrypted_db_password = None
+
 
         self._compute_cluster_ip = None
 
@@ -143,7 +145,6 @@ def fastpath(self) -> FastPathRunner:
 
     @property
     def kge(self) -> KgeRunner:
-        print("!kge")
         if not isinstance(self._query_runner, ArrowQueryRunner):
             raise ValueError("Running FastPath requires GDS with the Arrow server enabled")
         if self._compute_cluster_ip is None:
@@ -156,7 +157,7 @@ def kge(self) -> KgeRunner:
             self._server_version,
             self._compute_cluster_ip,
             self._encrypted_db_password,
-            None,
+            self._query_runner._gds_arrow_client._host + ":" + str(self._query_runner._gds_arrow_client._port),
         )
 
     def __getattr__(self, attr: str) -> IndirectCallBuilder:

graphdatascience/model/kge_runner.py

@@ -27,22 +27,21 @@ def __init__(
         encrypted_db_password: str,
         arrow_uri: str,
     ):
-        print("!init", flush=True)
         self._query_runner = query_runner
         self._namespace = namespace
         self._server_version = server_version
         self._compute_cluster_web_uri = f"http://{compute_cluster_ip}:5005"
-        self._compute_cluster_arrow_uri = f"grpc://{compute_cluster_ip}:8815"
+        # self._compute_cluster_arrow_uri = f"grpc://{compute_cluster_ip}:8491"
         self._compute_cluster_mlflow_uri = f"http://{compute_cluster_ip}:8080"
         self._encrypted_db_password = encrypted_db_password
         self._arrow_uri = arrow_uri
+        print("KgeRunner __dict__:")
+        print(self.__dict__)
 
-    @client_only_endpoint("gds.kge")
+    @property
     def model(self):
-        print("!model")
         return self
 
-    # @client_only_endpoint("gds.kge.model") and name is train
     # @compatible_with("stream", min_inclusive=ServerVersion(2, 5, 0))
     @client_only_endpoint("gds.kge.model")
     def train(
@@ -53,7 +52,6 @@ def train(
         embedding_dimension,
         mlflow_experiment_name: Optional[str] = None,
     ) -> Series:
-        print("!!!train")
         graph_config = {"name": G.name()}
 
         algo_config = {
@@ -85,6 +83,7 @@ def train(
         return Series({"status": "finished"})
 
     def _start_job(self, config: Dict[str, Any]) -> str:
+        print(config)
         res = requests.post(f"{self._compute_cluster_web_uri}/api/machine-learning/start", json=config)
         res.raise_for_status()
         job_id = res.json()["job_id"]

graphdatascience/query_runner/gds_arrow_client.py

@@ -92,6 +92,12 @@ def __init__(
         if tls_root_certs:
             client_options["tls_root_certs"] = tls_root_certs
 
+        print("location:")
+        print(location)
+        print("client_options:")
+        print(client_options)
+        print("auth:")
+        print(auth)
         self._flight_client = flight.FlightClient(location, **client_options)
 
     def connection_info(self) -> Tuple[str, int]:
@@ -128,6 +134,10 @@ def get_property(
             }
 
         ticket = flight.Ticket(json.dumps(payload).encode("utf-8"))
+        print("ticket:")
+        print(ticket)
+        print("_flight_client")
+        print(self._flight_client)
         get = self._flight_client.do_get(ticket)
         arrow_table = get.read_all()
 

graphdatascience/tests/integration/test_graph_construct.py

@@ -558,3 +558,36 @@ def test_graph_alpha_construct_backward_compat_with_arrow(gds: GraphDataScience)
 
     with pytest.warns(DeprecationWarning):
         gds.alpha.graph.construct("hello", nodes, relationships)
+
+@pytest.mark.enterprise
+@pytest.mark.compatible_with(min_inclusive=ServerVersion(2, 1, 0))
+def test_graph_alpha_construct_backward_compat_with_arrow(gds: GraphDataScience) -> None:
+    nodes = DataFrame({"nodeId": [0, 1, 2, 3]})
+    relationships = DataFrame({"sourceNodeId": [0, 1, 2, 3], "targetNodeId": [1, 2, 3, 0]})
+
+    with pytest.warns(DeprecationWarning):
+        gds.alpha.graph.construct("hello", nodes, relationships)
+
+
+@pytest.mark.compatible_with(min_inclusive=ServerVersion(2, 2, 0))
+def test_roundtrip_with_arrow(gds: GraphDataScience) -> None:
+    G, _ = gds.graph.project(GRAPH_NAME, {"Node": {"properties": ["x", "y"]}}, {"REL": {"properties": "relX"}})
+
+    rel_df = gds.graph.relationshipProperty.stream(G, "relX")
+    node_df = gds.graph.nodeProperty.stream(G, "x")
+
+    G_2 = gds.graph.construct("arrowGraph", node_df, rel_df)
+
+    res = gds.graph.list()
+    try:
+        assert set(res['graphName'].tolist()) == {'g', 'arrowGraph'}
+        assert G.node_count() == G_2.node_count()
+        assert G.relationship_count() == G_2.relationship_count()
+    finally:
+        G_2.drop()
+
+@pytest.mark.compatible_with(min_inclusive=ServerVersion(2, 2, 0))
+def test_drop_list_warning_reproduction(gds: GraphDataScience) -> None:
+    G, _ = gds.graph.project(GRAPH_NAME, {"Node": {"properties": ["x", "y"]}}, {"REL": {"properties": "relX"}})
+    res = gds.graph.list()
+    assert res['graphName'].tolist() == ['g']

graphdatascience/tests/integration/test_graph_ops.py

@@ -854,7 +854,7 @@ def test_graph_relationships_stream_without_arrow(gds_without_arrow: GraphDataSc
 
 @pytest.mark.compatible_with(min_inclusive=ServerVersion(2, 2, 0))
 def test_graph_relationships_stream_with_arrow(gds: GraphDataScience) -> None:
-    G, _ = gds.graph.project(GRAPH_NAME, "*", ["REL", "REL2"])
+    G, _ = gds.graph.project(GRAPH_NAME, "*", ["REL_0", "REL2"])
 
     if gds.server_version() >= ServerVersion(2, 5, 0):
         result = gds.graph.relationships.stream(G, ["REL", "REL2"])
@@ -1058,3 +1058,5 @@ def test_empty_relationships_stream(gds: GraphDataScience) -> None:
 
     result = gds.graph.relationships.stream(G, ["SIMILAR"])
     assert result.empty
+
+