Converting Azure Cosmos DB graph database to python graph object

StatStein 0
Viewed 32

I need to run some link prediction models on a graph database using the networkx package in python. The graph database is currently stored in Azure Cosmos DB. I have been able to connect to the database using gremlin-python package, but I cannot extract the data in any form that can be converted to a graph object for networkx. I need the vertices and all of their attributes as well as the edges connecting them.

Here’s what I have so far, which just returns an empty dataframe. The vertex information will appear in the terminal, but I cannot save it as an object. How can I extract the data from Cosmos DB into a format that can be converted to a graph object in python? Thanks.

Code:

from gremlin_python.driver import client, serializer,protocol
import pandas as pd

def get_vertices(client, vertices):
   callback = client.submitAsync(vertices)
   for result in callback.result():
       print("\t{0}".format(result))
 
   return(pd.DataFrame(callback.result()))

   if __name__ == "__main__":

       client = client.Client('<endpoint>','g',
       username=="<username>”
       password="<password>”
 
       message_serializer=serializer.GraphSONSerializersV2d0())

   vertices = "g.V().valueMap().by(unfold()).toList()"

   Vdata = get_vertices(client, vertices)
   print(Vdata.head())
gremlinpython graph-databases networkx azure-cosmosdb python
1 Answers

I think the specific issue in your above code is how you are accessing the result. You are using callback.result() where you might need to be doing something like callback.result().all().result(). callback is a future and you are correctly getting the result of the future (.result()) but that returns a ResultSet so you need the additional all().result() to access what you need.

Azure has some samples and you can see how they access the return results here: https://github.com/Azure-Samples/azure-cosmos-db-graph-python-getting-started/blob/master/connect.py

I did this a while back as I wanted to visualise data quickly in Python from CosmosDB. I didn't use Pandas, I directly put the data in a NetworkX graph. I am not sure if this is the most efficient way of going about it but it worked:

  1. Create your NetworkX graph G = nx.DiGraph()
  2. Execute query to get vertices you want from your graph client.submit(<YOUR GREMLIN VERTEX QUERY e.g. g.V()>)
  3. Iterate over returned result, for each vertex add a node into the NetworkX graph G.add_node()
    vertices_response = client.submit(<YOUR GREMLIN VERTEX QUERY e.g. g.V()>)
    vertices = vertices_response.all().result()
    for vertex in vertices:
        G.add_node(< Up to you how you transform your data to be represented here. You could explicitly state what fields you need e.g. vertex["id"] >)
  1. Repeat for edges:
    edges_response = client.submit(<YOUR GREMLIN VERTEX QUERY e.g. g.E()>)
    edges = edges_response.all().result()

    for edge in edges:
        logging.debug(edge)
        G.add_edge(edge["inV"], edge["outV"])

Note: Mine is slightly different using submit instead of AsyncSubmit but you should be able to get the rough idea

dominicst0
Related