# files containing vertices and edges, these are donwloaded from SNAP website, cleaned, and formatted for our use
# files can be accessed from https://community.cloud.databricks.com/files/tables/fovepx7h1479410900486/vertices.csv
# create new RDDs from each file
e = sc.textFile("/FileStore/tables/fovepx7h1479410900486/edges.csv")
v = sc.textFile("/FileStore/tables/fovepx7h1479410900486/vertices.csv")
 
# let's view what's in the vertices RDD
v.collect()

# helper function to convert from string value to integer
def convertToInt(s):
  if s == None or s == 'None':
    return None
  else:
    return int(s)
 
# the uploaded files have header and data is stored as string
# we need to remove the header and convert the strings to integers; None fields will be translated to null
eheader = e.first() #extract edge header
edges = e.filter(lambda row:row != eheader and row != "").map(lambda line:line.split(",")).map(lambda line:(convertToInt(line[0]), convertToInt(line[1])))
 
vheader = v.first() #extract vertex header
vertices = v.filter(lambda row:row != vheader).map(lambda line:line.split(",")).map(lambda line:(convertToInt(line[0]), convertToInt(line[1]), convertToInt(line[2]), convertToInt(line[3]), convertToInt(line[4]), convertToInt(line[5])))
 
print("Vertex data: %s" % vheader)
 
vertices.collect()

# Next we will convert our RDDs to DataFrames
# We will first create the schema and then create the dataframe
from pyspark.sql.types import *
 
# GraphFrames will expect to have our key named as 'id'
vertexSchema = StructType([StructField("id", IntegerType(), False),
                           StructField("birthday", IntegerType(), True),
                           StructField("hometown_id", IntegerType(), True),
                           StructField("work_employer_id", IntegerType(), True),
                           StructField("education_school_id", IntegerType(), True),
                           StructField("education_year_id", IntegerType(), True)])
 
# Create a data frame
vdf = sqlContext.createDataFrame(vertices, vertexSchema)
 
# We will do the same for edges.
# Note that later on GraphFrames require 'src' and 'dst' headers
edgeSchema = StructType([StructField("src", IntegerType(), False),
                           StructField("dst", IntegerType(), False)])
 
# Create a data frame
edf = sqlContext.createDataFrame(edges, edgeSchema)
 
# Let's look at the contents of the dataframe
display(vdf)

# Once we have our two DataFrames, we will now create the GraphFrame
from graphframes import *
 
g = GraphFrame(vdf, edf)
 
display(g.edges)

display(g.outDegrees)

# find all connected vertices with the same birthday identifier
print "same birthdays"
res = g.find("(a)-[]->(b)") \
         .filter("a.birthday = b.birthday")
print "count: %d" % res.count()
display(res.select("a.id", "b.id", "b.birthday"))

from pyspark.sql.functions import *
 
# We need to do some heading manipulation so that the two columns named 'id' are changed to 'source' and 'target'
# clean filesystem from previous runs
dbutils.fs.rm("/FileStore/tables/fovepx7h1479410900486/birthdays.csv", True)
 
# we need Spark to run this job on one node so that only one file is generated.  Thus, we use the coalesce() command
res.select(col("a.id").alias("source"), col("b.id").alias("target")).coalesce(1).write.format("com.databricks.spark.csv").save("/FileStore/tables/fovepx7h1479410900486/birthdays.csv")
 
# Lists the file generated
dbutils.fs.ls("/FileStore/tables/fovepx7h1479410900486/birthdays.csv")

# find "friends of friends" who are not connected to us, but graduated the same
# year from the same school
print "same class"
res = g.find("(a)-[]->(b); (b)-[]->(c); !(a)-[]->(c)") \
         .filter("a.education_school_id = c.education_school_id and " \
                 "a.education_year_id = c.education_year_id")
res = res.filter("a.id != c.id").select(col("a.id").alias("source"), "a.education_school_id", "a.education_year_id",
        col("c.id").alias("target"), "c.education_school_id", "c.education_year_id") 
 
# Run PageRank on our graph
pgr = g.pageRank(resetProbability=0.15, tol=0.01).vertices.sort(
    	    'pagerank', ascending=False)
display(pgr.select("id", "pagerank"))