sc.version // link the right library depending on Spark version of the cluster that's running
// spark version 2.3.0 works with graphframes:graphframes:0.7.0-spark2.3-s_2.11
// spark version 3.0.1 works with graphframes:graphframes:0.8.1-spark3.0-s_2.12

res1: String = 3.2.1

// we first need to install the library - graphframes as a Spark package - and attach it to our cluster - see note two cells above!
import org.apache.spark.sql._
import org.apache.spark.sql.functions._

import org.graphframes._

import org.apache.spark.sql._
import org.apache.spark.sql.functions._
import org.graphframes._

// Vertex DataFrame
val v = sqlContext.createDataFrame(List(
  ("a", "Alice", 34),
  ("b", "Bob", 36),
  ("c", "Charlie", 30),
  ("d", "David", 29),
  ("e", "Esther", 32),
  ("f", "Fanny", 36),
  ("g", "Gabby", 60)
)).toDF("id", "name", "age")

// Edge DataFrame
val e = sqlContext.createDataFrame(List(
  ("a", "b", "friend"),
  ("b", "c", "follow"),
  ("c", "b", "follow"),
  ("f", "c", "follow"),
  ("e", "f", "follow"),
  ("e", "d", "friend"),
  ("d", "a", "friend"),
  ("a", "e", "friend")
)).toDF("src", "dst", "relationship")

v: org.apache.spark.sql.DataFrame = [id: string, name: string ... 1 more field]
e: org.apache.spark.sql.DataFrame = [src: string, dst: string ... 1 more field]

val g = GraphFrame(v, e)

g: org.graphframes.GraphFrame = GraphFrame(v:[id: string, name: string ... 1 more field], e:[src: string, dst: string ... 1 more field])

d3.graphs.help()

import org.apache.spark.sql.functions.lit // import the lit function in sql
val gE= g.edges.select($"src", $"dst".as("dest"), lit(1L).as("count")) // for us the column count is just an edge incidence

import org.apache.spark.sql.functions.lit
gE: org.apache.spark.sql.DataFrame = [src: string, dest: string ... 1 more field]

display(gE)

d3.graphs.force(
  height = 500,
  width = 500,
  clicks = gE.as[d3.Edge])

// This example graph also comes with the GraphFrames package.
val g0 = examples.Graphs.friends

g0: org.graphframes.GraphFrame = GraphFrame(v:[id: string, name: string ... 1 more field], e:[src: string, dst: string ... 1 more field])

d3.graphs.force( // let us see g0 now in one cell
  height = 500,
  width = 500,
  clicks = g0.edges.select($"src", $"dst".as("dest"), lit(1L).as("count")).as[d3.Edge])

display(g.vertices)

display(g0.vertices) // this is the same query on the graph loaded as an example from GraphFrame package

display(g.edges)

display(g.inDegrees)

display(g.outDegrees)

display(g.degrees)

val youngest = g.vertices.groupBy().min("age")
display(youngest)

val numFollows = g.edges.filter("relationship = 'follow'").count()

numFollows: Long = 4

// Search for pairs of vertices with edges in both directions between them, i.e., find undirected or bidirected edges.
val motifs = g.find("(a)-[e1]->(b); (b)-[e2]->(a)")
display(motifs)

val filtered = motifs.filter("b.age > 30")
display(filtered)

//Search for all "directed triangles" or triplets of vertices: a,b,c with edges: a->b, b->c and c->a
//uncomment the next 2 lines and replace the "XXX" below
//val motifs3 = g.find("(a)-[e1]->(b); (b)-[e2]->(c); (c)-[e3]->(XXX)")
//display(motifs3)

// Find chains of 4 vertices.
val chain4 = g.find("(a)-[ab]->(b); (b)-[bc]->(c); (c)-[cd]->(d)")

// Query on sequence, with state (cnt)
//  (a) Define method for updating state given the next element of the motif.
def sumFriends(cnt: Column, relationship: Column): Column = {
  when(relationship === "friend", cnt + 1).otherwise(cnt)
}
//  (b) Use sequence operation to apply method to sequence of elements in motif.
//      In this case, the elements are the 3 edges.
val condition = Seq("ab", "bc", "cd").
  foldLeft(lit(0))((cnt, e) => sumFriends(cnt, col(e)("relationship")))
//  (c) Apply filter to DataFrame.
val chainWith2Friends2 = chain4.where(condition >= 2)
display(chainWith2Friends2)

chain4

res22: org.apache.spark.sql.DataFrame = [a: struct<id: string, name: string ... 1 more field>, ab: struct<src: string, dst: string ... 1 more field> ... 5 more fields]

chain4.printSchema

root
 |-- a: struct (nullable = false)
 |    |-- id: string (nullable = true)
 |    |-- name: string (nullable = true)
 |    |-- age: integer (nullable = false)
 |-- ab: struct (nullable = false)
 |    |-- src: string (nullable = true)
 |    |-- dst: string (nullable = true)
 |    |-- relationship: string (nullable = true)
 |-- b: struct (nullable = false)
 |    |-- id: string (nullable = true)
 |    |-- name: string (nullable = true)
 |    |-- age: integer (nullable = false)
 |-- bc: struct (nullable = false)
 |    |-- src: string (nullable = true)
 |    |-- dst: string (nullable = true)
 |    |-- relationship: string (nullable = true)
 |-- c: struct (nullable = false)
 |    |-- id: string (nullable = true)
 |    |-- name: string (nullable = true)
 |    |-- age: integer (nullable = false)
 |-- cd: struct (nullable = false)
 |    |-- src: string (nullable = true)
 |    |-- dst: string (nullable = true)
 |    |-- relationship: string (nullable = true)
 |-- d: struct (nullable = false)
 |    |-- id: string (nullable = true)
 |    |-- name: string (nullable = true)
 |    |-- age: integer (nullable = false)

// Select subgraph of users older than 30, and edges of type "friend"
val v2 = g.vertices.filter("age > 30")
val e2 = g.edges.filter("relationship = 'friend'")
val g2 = GraphFrame(v2, e2)

v2: org.apache.spark.sql.Dataset[org.apache.spark.sql.Row] = [id: string, name: string ... 1 more field]
e2: org.apache.spark.sql.Dataset[org.apache.spark.sql.Row] = [src: string, dst: string ... 1 more field]
g2: org.graphframes.GraphFrame = GraphFrame(v:[id: string, name: string ... 1 more field], e:[src: string, dst: string ... 1 more field])

display(g2.vertices)

display(g2.edges)

d3.graphs.force( // let us see g2 now in one cell
  height = 500,
  width = 500,
  clicks = g2.edges.select($"src", $"dst".as("dest"), lit(1L).as("count")).as[d3.Edge])

// Select subgraph based on edges "e" of type "follow"
// pointing from a younger user "a" to an older user "b".
val paths = g.find("(a)-[e]->(b)")
  .filter("e.relationship = 'follow'")
  .filter("a.age < b.age")
// "paths" contains vertex info. Extract the edges.
val e2 = paths.select("e.src", "e.dst", "e.relationship")
// In Spark 1.5+, the user may simplify this call:
//  val e2 = paths.select("e.*")

// Construct the subgraph
val g2 = GraphFrame(g.vertices, e2)

paths: org.apache.spark.sql.Dataset[org.apache.spark.sql.Row] = [a: struct<id: string, name: string ... 1 more field>, e: struct<src: string, dst: string ... 1 more field> ... 1 more field]
e2: org.apache.spark.sql.DataFrame = [src: string, dst: string ... 1 more field]
g2: org.graphframes.GraphFrame = GraphFrame(v:[id: string, name: string ... 1 more field], e:[src: string, dst: string ... 1 more field])

display(g2.vertices)

display(g2.edges)

// Search from "Esther" for users of age <= 32.
val paths: DataFrame = g.bfs.fromExpr("name = 'Esther'").toExpr("age < 32").run()
display(paths)

val paths: DataFrame = g.bfs.fromExpr("name = 'Esther' OR name = 'Bob'").toExpr("age < 32").run()
display(paths)

val filteredPaths = g.bfs.fromExpr("name = 'Esther'").toExpr("age < 32")
  .edgeFilter("relationship != 'friend'")
  .maxPathLength(3)
  .run()
display(filteredPaths)

sc.setCheckpointDir("/_checkpoint") // just a directory in distributed file system
val result = g.connectedComponents.run() 
display(result)

val result = g.stronglyConnectedComponents.maxIter(10).run()
display(result.orderBy("component"))

val result = g.labelPropagation.maxIter(5).run()
display(result.orderBy("label"))

// Run PageRank until convergence to tolerance "tol".
val results = g.pageRank.resetProbability(0.15).tol(0.01).run()
display(results.vertices)

display(results.edges)

// Run PageRank for a fixed number of iterations.
val results2 = g.pageRank.resetProbability(0.15).maxIter(10).run()
display(results2.vertices)

// Run PageRank personalized for vertex "a"
val results3 = g.pageRank.resetProbability(0.15).maxIter(10).sourceId("a").run()
display(results3.vertices)

val paths = g.shortestPaths.landmarks(Seq("a", "d")).run()
display(paths)

g.edges.show()

+---+---+------------+
|src|dst|relationship|
+---+---+------------+
|  a|  b|      friend|
|  b|  c|      follow|
|  c|  b|      follow|
|  f|  c|      follow|
|  e|  f|      follow|
|  e|  d|      friend|
|  d|  a|      friend|
|  a|  e|      friend|
+---+---+------------+

val results = g.triangleCount.run()
display(results)

import org.graphframes.{examples,GraphFrame}
import org.graphframes.lib.AggregateMessages
val g: GraphFrame = examples.Graphs.friends  // get example graph

// We will use AggregateMessages utilities later, so name it "AM" for short.
val AM = AggregateMessages

// For each user, sum the ages of the adjacent users.
val msgToSrc = AM.dst("age")
val msgToDst = AM.src("age")
val agg = { g.aggregateMessages
  .sendToSrc(msgToSrc)  // send destination user's age to source
  .sendToDst(msgToDst)  // send source user's age to destination
  .agg(sum(AM.msg).as("summedAges")) } // sum up ages, stored in AM.msg column
agg.show()

+---+----------+
| id|summedAges|
+---+----------+
|  a|        97|
|  c|       108|
|  e|        99|
|  d|        66|
|  b|        94|
|  f|        62|
+---+----------+

import org.graphframes.{examples, GraphFrame}
import org.graphframes.lib.AggregateMessages
g: org.graphframes.GraphFrame = GraphFrame(v:[id: string, name: string ... 1 more field], e:[src: string, dst: string ... 1 more field])
AM: org.graphframes.lib.AggregateMessages.type = org.graphframes.lib.AggregateMessages$@706833c8
msgToSrc: org.apache.spark.sql.Column = dst[age]
msgToDst: org.apache.spark.sql.Column = src[age]
agg: org.apache.spark.sql.DataFrame = [id: string, summedAges: bigint]

/*
 * Licensed to the Apache Software Foundation (ASF) under one or more
 * contributor license agreements.  See the NOTICE file distributed with
 * this work for additional information regarding copyright ownership.
 * The ASF licenses this file to You under the Apache License, Version 2.0
 * (the "License"); you may not use this file except in compliance with
 * the License.  You may obtain a copy of the License at
 *
 *    http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

package org.graphframes.examples

import org.apache.spark.{SparkConf, SparkContext}
import org.apache.spark.graphx.{Graph, VertexRDD, Edge => GXEdge}
import org.apache.spark.sql.{Column, Row, SparkSession, SQLContext}
import org.apache.spark.sql.functions.{col, lit, sum, udf, when}

import org.graphframes.GraphFrame
import org.graphframes.examples.Graphs.gridIsingModel
import org.graphframes.lib.AggregateMessages


/**
 * Example code for Belief Propagation (BP)
 *
 * This provides a template for building customized BP algorithms for different types of
 * graphical models.
 *
 * This example:
 *  - Ising model on a grid
 *  - Parallel Belief Propagation using colored fields
 *
 * Ising models are probabilistic graphical models over binary variables x,,i,,.
 * Each binary variable x,,i,, corresponds to one vertex, and it may take values -1 or +1.
 * The probability distribution P(X) (over all x,,i,,) is parameterized by vertex factors a,,i,,
 * and edge factors b,,ij,,:
 * {{{
 *  P(X) = (1/Z) * exp[ \sum_i a_i x_i + \sum_{ij} b_{ij} x_i x_j ]
 * }}}
 * where Z is the normalization constant (partition function).
 * See [[https://en.wikipedia.org/wiki/Ising_model Wikipedia]] for more information on Ising models.
 *
 * Belief Propagation (BP) provides marginal probabilities of the values of the variables x,,i,,,
 * i.e., P(x,,i,,) for each i.  This allows a user to understand likely values of variables.
 * See [[https://en.wikipedia.org/wiki/Belief_propagation Wikipedia]] for more information on BP.
 *
 * We use a batch synchronous BP algorithm, where batches of vertices are updated synchronously.
 * We follow the mean field update algorithm in Slide 13 of the
 * [[http://www.eecs.berkeley.edu/~wainwrig/Talks/A_GraphModel_Tutorial  talk slides]] from:
 *  Wainwright. "Graphical models, message-passing algorithms, and convex optimization."
 *
 * The batches are chosen according to a coloring.  For background on graph colorings for inference,
 * see for example:
 *  Gonzalez et al. "Parallel Gibbs Sampling: From Colored Fields to Thin Junction Trees."
 *  AISTATS, 2011.
 *
 * The BP algorithm works by:
 *  - Coloring the graph by assigning a color to each vertex such that no neighboring vertices
 *    share the same color.
 *  - In each step of BP, update all vertices of a single color.  Alternate colors.
 */
object BeliefPropagation {

  def main(args: Array[String]): Unit = {
    val spark = SparkSession
      .builder()
      .appName("BeliefPropagation example")
      .getOrCreate()

    val sql = spark.sqlContext

    // Create graphical model g of size 3 x 3.
    val g = gridIsingModel(sql, 3)

    println("Original Ising model:")
    g.vertices.show()
    g.edges.show()

    // Run BP for 5 iterations.
    val numIter = 5
    val results = runBPwithGraphX(g, numIter)

    // Display beliefs.
    val beliefs = results.vertices.select("id", "belief")
    println(s"Done with BP. Final beliefs after $numIter iterations:")
    beliefs.show()

    spark.stop()
  }

  /**
   * Given a GraphFrame, choose colors for each vertex.  No neighboring vertices will share the
   * same color.  The number of colors is minimized.
   *
   * This is written specifically for grid graphs. For non-grid graphs, it should be generalized,
   * such as by using a greedy coloring scheme.
   *
   * @param g  Grid graph generated by [[org.graphframes.examples.Graphs.gridIsingModel()]]
   * @return  Same graph, but with a new vertex column "color" of type Int (0 or 1)
   */
  private def colorGraph(g: GraphFrame): GraphFrame = {
    val colorUDF = udf { (i: Int, j: Int) => (i + j) % 2 }
    val v = g.vertices.withColumn("color", colorUDF(col("i"), col("j")))
    GraphFrame(v, g.edges)
  }

  /**
   * Run Belief Propagation.
   *
   * This implementation of BP shows how to use GraphX's aggregateMessages method.
   * It is simple to convert to and from GraphX format.  This method does the following:
   *  - Color GraphFrame vertices for BP scheduling.
   *  - Convert GraphFrame to GraphX format.
   *  - Run BP using GraphX's aggregateMessages API.
   *  - Augment the original GraphFrame with the BP results (vertex beliefs).
   *
   * @param g  Graphical model created by `org.graphframes.examples.Graphs.gridIsingModel()`
   * @param numIter  Number of iterations of BP to run.  One iteration includes updating each
   *                 vertex's belief once.
   * @return  Same graphical model, but with [[GraphFrame.vertices]] augmented with a new column
   *          "belief" containing P(x,,i,, = +1), the marginal probability of vertex i taking
   *          value +1 instead of -1.
   */
  def runBPwithGraphX(g: GraphFrame, numIter: Int): GraphFrame = {
    // Choose colors for vertices for BP scheduling.
    val colorG = colorGraph(g)
    val numColors: Int = colorG.vertices.select("color").distinct.count().toInt

    // Convert GraphFrame to GraphX, and initialize beliefs.
    val gx0 = colorG.toGraphX
    // Schema maps for extracting attributes
    val vColsMap = colorG.vertexColumnMap
    val eColsMap = colorG.edgeColumnMap
    // Convert vertex attributes to nice case classes.
    val gx1: Graph[VertexAttr, Row] = gx0.mapVertices { case (_, attr) =>
      // Initialize belief at 0.0
      VertexAttr(attr.getDouble(vColsMap("a")), 0.0, attr.getInt(vColsMap("color")))
    }
    // Convert edge attributes to nice case classes.
    val extractEdgeAttr: (GXEdge[Row] => EdgeAttr) = { e =>
      EdgeAttr(e.attr.getDouble(eColsMap("b")))
    }
    var gx: Graph[VertexAttr, EdgeAttr] = gx1.mapEdges(extractEdgeAttr)

    // Run BP for numIter iterations.
    for (iter <- Range(0, numIter)) {
      // For each color, have that color receive messages from neighbors.
      for (color <- Range(0, numColors)) {
        // Send messages to vertices of the current color.
        val msgs: VertexRDD[Double] = gx.aggregateMessages(
          ctx =>
            // Can send to source or destination since edges are treated as undirected.
            if (ctx.dstAttr.color == color) {
              val msg = ctx.attr.b * ctx.srcAttr.belief
              // Only send message if non-zero.
              if (msg != 0) ctx.sendToDst(msg)
            } else if (ctx.srcAttr.color == color) {
              val msg = ctx.attr.b * ctx.dstAttr.belief
              // Only send message if non-zero.
              if (msg != 0) ctx.sendToSrc(msg)
            },
          _ + _)
        // Receive messages, and update beliefs for vertices of the current color.
        gx = gx.outerJoinVertices(msgs) {
          case (vID, vAttr, optMsg) =>
            if (vAttr.color == color) {
              val x = vAttr.a + optMsg.getOrElse(0.0)
              val newBelief = math.exp(-log1pExp(-x))
              VertexAttr(vAttr.a, newBelief, color)
            } else {
              vAttr
            }
        }
      }
    }

    // Convert back to GraphFrame with a new column "belief" for vertices DataFrame.
    val gxFinal: Graph[Double, Unit] = gx.mapVertices((_, attr) => attr.belief).mapEdges(_ => ())
    GraphFrame.fromGraphX(colorG, gxFinal, vertexNames = Seq("belief"))
  }

  case class VertexAttr(a: Double, belief: Double, color: Int)

  case class EdgeAttr(b: Double)

  /**
   * Run Belief Propagation.
   *
   * This implementation of BP shows how to use GraphFrame's aggregateMessages method.
   *  - Color GraphFrame vertices for BP scheduling.
   *  - Run BP using GraphFrame's aggregateMessages API.
   *  - Augment the original GraphFrame with the BP results (vertex beliefs).
   *
   * @param g  Graphical model created by `org.graphframes.examples.Graphs.gridIsingModel()`
   * @param numIter  Number of iterations of BP to run.  One iteration includes updating each
   *                 vertex's belief once.
   * @return  Same graphical model, but with [[GraphFrame.vertices]] augmented with a new column
   *          "belief" containing P(x,,i,, = +1), the marginal probability of vertex i taking
   *          value +1 instead of -1.
   */
  def runBPwithGraphFrames(g: GraphFrame, numIter: Int): GraphFrame = {
    // Choose colors for vertices for BP scheduling.
    val colorG = colorGraph(g)
    val numColors: Int = colorG.vertices.select("color").distinct.count().toInt

    // TODO: Handle vertices without any edges.

    // Initialize vertex beliefs at 0.0.
    var gx = GraphFrame(colorG.vertices.withColumn("belief", lit(0.0)), colorG.edges)

    // Run BP for numIter iterations.
    for (iter <- Range(0, numIter)) {
      // For each color, have that color receive messages from neighbors.
      for (color <- Range(0, numColors)) {
        // Define "AM" for shorthand for referring to the src, dst, edge, and msg fields.
        // (See usage below.)
        val AM = AggregateMessages
        // Send messages to vertices of the current color.
        // We may send to source or destination since edges are treated as undirected.
        val msgForSrc: Column = when(AM.src("color") === color, AM.edge("b") * AM.dst("belief"))
        val msgForDst: Column = when(AM.dst("color") === color, AM.edge("b") * AM.src("belief"))
        val logistic = udf { (x: Double) => math.exp(-log1pExp(-x)) }
        val aggregates = gx.aggregateMessages
          .sendToSrc(msgForSrc)
          .sendToDst(msgForDst)
          .agg(sum(AM.msg).as("aggMess"))
        val v = gx.vertices
        // Receive messages, and update beliefs for vertices of the current color.
        val newBeliefCol = when(v("color") === color && aggregates("aggMess").isNotNull,
          logistic(aggregates("aggMess") + v("a")))
          .otherwise(v("belief"))  // keep old beliefs for other colors
        val newVertices = v
          .join(aggregates, v("id") === aggregates("id"), "left_outer")  // join messages, vertices
          .drop(aggregates("id"))  // drop duplicate ID column (from outer join)
          .withColumn("newBelief", newBeliefCol)  // compute new beliefs
          .drop("aggMess")  // drop messages
          .drop("belief")  // drop old beliefs
          .withColumnRenamed("newBelief", "belief")
        // Cache new vertices using workaround for SPARK-13346
        val cachedNewVertices = AM.getCachedDataFrame(newVertices)
        gx = GraphFrame(cachedNewVertices, gx.edges)
      }
    }

    // Drop the "color" column from vertices
    GraphFrame(gx.vertices.drop("color"), gx.edges)
  }

  /** More numerically stable `log(1 + exp(x))` */
  private def log1pExp(x: Double): Double = {
    if (x > 0) {
      x + math.log1p(math.exp(-x))
    } else {
      math.log1p(math.exp(x))
    }
  }
}