pyspark_graphframes_vertex_edge_dataframe_creation_and_queries.py

1from pyspark.sql import SparkSession
2from graphframes import GraphFrame
3
4# Create a SparkSession
5spark = SparkSession.builder \
6    .appName("GraphFramesQuickstart") \
7    .getOrCreate()
8
9# Create a Vertex DataFrame with unique ID column "id"
10v = spark.createDataFrame([
11  ("a", "Alice", 34),
12  ("b", "Bob", 36),
13  ("c", "Charlie", 30),
14], ["id", "name", "age"])
15
16# Create an Edge DataFrame with "src" and "dst" columns
17e = spark.createDataFrame([
18  ("a", "b", "friend"),
19  ("b", "c", "follow"),
20  ("c", "b", "follow"),
21], ["src", "dst", "relationship"])
22
23# Create a GraphFrame
24g = GraphFrame(v, e)
25
26# Query: Get vertices, edges, and the in-degree of each vertex
27g.vertices.show()
28g.edges.show()
29g.inDegrees.show()
30
31# Find the youngest user's age in the graph
32print(g.vertices.groupBy().min("age").collect())
33
34# Count the number of "follow" relationships in the graph
35print(g.edges.filter("relationship = 'follow'").count())