Graph Embeddings at scale with Spark and GraphFrames

I'm the currently most active maintainer of the GraphFrames project and a big enthusiast of Apache Spark and Graph analysis. I'm not affiliated to any company selling GraphFrames or Apache Spark itself and my work on the project is based on a pure enthusiasm and passion to Open Source. I can estimate my knowledge of Apache Spark as medium: I know the API and some details of internal implementation and how Spark works from the top-level of view. But I'm not a contributor of Spark or engineer who is working on Spark internals. I do not have any CS-degree or DS-degree and my knowledge of the Data Science field as well as Graph Data Science are limited to what I read and try. And this post will be as honest as possible, and if I'm not sure in something or don't know something I will tell about it explicitly. No LLM was used to generate the text or any of its parts. However, because English is not my native language, I use DeepL Write from time to time.

GraphFrames is an open source extension for the Apache Spark that aims to bring easy to use graph algorithms to the Spark ecosystem. While there are a lot of single-node libraries for graph processing and graph machine learning, they cannot be scaled to billions of nodes and edges. At the same time, there are a lot of specialized solutions for distributed graph processing, but they required a separate infrastructure and are hard to maintain. GraphFrames aims to somewhere in the middle. If you already have a huge graphs, like user-item interactions, financial transactions or identity graphs, but still not ready to heavily invest in the specialized infrastructure, you can re-use existing Spark with GraphFrames. Spark today is a de-facto standard for big-data industry, so with GraphFrames teams can easily add some big-data graph features to their pipelines on top of existing processing. GraphFrames is a pure open-source project, maintained by the group of independent contributors and it is not backed by any vendors as well GraphFrames does not have any kind of paid version or enterprise support.

I will not write again about basic graph concepts and how the graph can be processed in a Map-Reduce ways. You can read about it in my another blog post. For this post it is enough to understand that Graph in GraphFrames is modeled as two relations: edges and vertices. Because code snippets will be in Scala and Spark, some basic knowledge of this two is required.

An obvious question is why do we even need this embeddings? So, let's briefly view what can we do with embeddings.

• Node classification. One of the most-common case. We have some labels, like who is who in our social network. Who is fraud-account or bot, who is not. Or, for example, if we have an electricity grid network, what is the most risky node? And a lot of other examples;
• Link Prediction. If we assume that our graph is not "complete" and there are some additional connections we are not aware about, we can try to predict the probability of edges between nodes. It naturally leads us to the RecSys story. If we are "building Netflix" and want to recommend movies to people, we can represent the problem as a bipartite graph where user is connected to the movie if they liked / watched it. And our goal is to recommend to the user a new movie or to predict edges in this bipartite graph in other words;
• Graph Clustering. When we have numeric representation of vertices, we can apply any of known clustering algorithm on top of our embeddings to get clusters;
• De-duplication and similarity. One of the common graph problems. We have a lot of nodes and we are knowing that some of them are actually the same entity. Having compact numeric representation of vertices allows us to use nearest neighbors and similarity search to identify them.

And a lot of other problems that can be solved with numeric representation of vertices.

If we can define anyhow the similarity between vertices of the graph, we can use the well known encoder-decoder approach.

In this case, we can sample pairs of very similar vertices and sample pairs of not similar vertices (so called negative-sampling) and try to maximize something like a cosine distance in embeddings space between vertex representations. It may even work without negative sampling. The main problem here is to define a similarity metric. What should it be? Jaccard score? Links? The same cluster?

While the idea looks interesting, I don't see a lot of real-world applications of it.

Each graph can be represented as a matrix actually. Let's imagine we have the graph with vertices 0, 1, 2, 3, 4. And we have the following edges: 1 -> 2, 1 -> 3, 2 -> 4, 3 -> 0, 0 -> 4. We can build a so-called adjacency matrix:

In this matrix element \( e_{ij}\) is equal to one if there is an edge between vertices \(i\) and \(j\) and zero otherwise.

While an adjacency matrix is already a representation of vertices by vectors, the size of such a matrix is growing very fast as well as adding new vertex will change this matrix. The second problem cannot be resolved easily, but the size can be addressed by matrix decomposition approach. The most used is Singular Value Decomposition (SVD). I won't stopping on it, because the problem of changing this matrix if a vertex is added cannot be resolved anyway. As well as SVD is very expensive from the computation point of view and such an approach barely will work on real world graphs with hundreds of millions of vertices.

Graphlets is a topic, but long story short, we can consider subgraphs of small size and there are not to many of them actually.

If we enumerate all the graphlets as index and use an amount of graphlets of each type as a values, we can construct a vector. This vector is called a Graphlet Degree Vector (GDV) and is a kind of vertex representation. While the are efficient algorithms to (approximately) enumerate graphlets, they are not easy to implement and GDV itself encodes very specific information. I have a plan to implement distributed GDV in GraphFrames one day, but at the moment the project needs a simpler and more generic approach.

The most advanced methods of vertex representation learning is Graph Neural Networks (GNN) and Graph Convolutional Neural Networks (GCNN). The idea is quite a simple (but only idea, implementation, and tuning is very hard). We can take information from neighborhood of each vertex and apply a convolution operation. Each level of convolution (1-st neighbors, 2-d neighbors, etc.) have each own matrix of numerical weights (that should be learned) and we do matrix multiplication to compute the final representation.

While it is the most powerful method, it is not very easy to implement, very hard to tune and even more hard to scale up to graphs for that GraphFrames is supposed to be used. I'm still going to look into it after implementing basic embeddings, but it is definitely not the first step.

Now it's time to implement it with Spark and GraphFrames. Graph in GraphFrames is modeled by two Spark's DataFrame objects, one for edges, one for vertices. As well as because it is Spark we should keep in mind the following limitations:

• Avoiding shuffling of long arrays and generating heavy rows. In our DuckDB toy example, each iteration is join that on Spark and GraphFrames will lead to shuffle;
• Distribute the work equally across the cluster. If one vertex has 10 possible candidates for the next walk, but another has 1000 it will create a skew;
• Avoid long chaining transformations without writing or checkpointing: growing lineage and long plan optimization time is a curse of Spark;

After thinking a lot, I stopped with the following concept.

The idea is to split algorithm to batches and limit the sequence length that is kept in memory. On each batch we read the "end" of the walk from the previous batch (or initialize a new if it is the first iteration). For 1st order Random Walks it is enough to read only the currently visiting vertex, for the no-backtracking we need also a previously visited one, etc.

Inside each batch we are sampling at most K neighbors per vertex that are possible candidates for the next walk. Sampling is very important here, otherwise on power-law real-world graphs we can easily face huge skew due to fact that 1% of vertices can handle half of the whole edges of the graph.

When we sampled neighbors and have a prepared walks, we can do joins in the same way like in the above DuckDB toy example.

  package org.apache.spark.sql.graphframes.expressions

  import org.apache.spark.sql.Encoder
  import org.apache.spark.sql.Encoders
  import org.apache.spark.sql.catalyst.encoders.ExpressionEncoder
  import org.apache.spark.sql.expressions.Aggregator

  import scala.reflect.runtime.universe.TypeTag

  import collection.mutable.ArrayBuffer

  case class Reservoir[T](seq: ArrayBuffer[T], elements: Int) extends Serializable

  case class ReservoirSamplingAgg[T: TypeTag](size: Int)
     extends Aggregator[T, Reservoir[T], Seq[T]]
      with Serializable {}

  override def zero: Reservoir[T] = Reservoir[T](ArrayBuffer.empty, 0)

  override def reduce(b: Reservoir[T], a: T): Reservoir[T] = {
    if (b.seq.size < size) {
      Reservoir(b.seq += a, b.elements + 1)
    } else {
      val j = java.util.concurrent.ThreadLocalRandom.current().nextInt(b.elements + 1)
      if (j < size) {
        b.seq(j) = a
      }
      Reservoir(b.seq, b.elements + 1)
    }
  }

  private def mergeFull(left: Reservoir[T], right: Reservoir[T]): Reservoir[T] = {
    val total_cnt = left.elements + right.elements
    val rng = java.util.concurrent.ThreadLocalRandom.current()
    val pLeft = left.elements.toDouble / total_cnt.toDouble

    var newSeq = ArrayBuffer.empty[T]
    val leftCloned = left.seq.clone()
    val rightCloned = right.seq.clone()
    for (_ <- (1 to size)) {
      if (rng.nextDouble() <= pLeft) {
        newSeq = newSeq += leftCloned.remove(rng.nextInt(leftCloned.size))
      } else {
        newSeq = newSeq += rightCloned.remove(rng.nextInt(rightCloned.size))
      }
    }

    Reservoir(newSeq, total_cnt)
  }

  private def mergeTwoPartial(left: Reservoir[T], right: Reservoir[T]): Reservoir[T] = {
    val total_cnt = left.elements + right.elements
    val rng = java.util.concurrent.ThreadLocalRandom.current()
    if (total_cnt <= size) {
      Reservoir(left.seq ++ right.seq, total_cnt)
    } else {
      val currElements = left.seq ++ right.seq.slice(0, size - left.elements)
      var currSize = size + 1

      for (i <- ((size - left.elements) to right.elements)) {
        val j = rng.nextInt(currSize)
        if (j < size) {
          currElements(j) = right.seq(i)
        }
        currSize += 1
      }

      Reservoir(currElements, currSize)
    }
  }

  private def mergePartialRight(left: Reservoir[T], right: Reservoir[T]): Reservoir[T] = {
    val total_cnt = left.elements + right.elements
    val pLeft = left.elements.toDouble / total_cnt.toDouble
    val currElements = ArrayBuffer.empty[T]
    val rng = java.util.concurrent.ThreadLocalRandom.current()

    val clonedLeft = left.seq.clone()
    val clonedRight = right.seq.clone()
    for (_ <- (1 to size)) {
      if ((clonedRight.isEmpty) || (rng.nextDouble() <= pLeft)) {
        val idx = rng.nextInt(clonedLeft.size)
        currElements += clonedLeft.remove(idx)
      } else {
        val idx = rng.nextInt(clonedRight.size)
        currElements += clonedRight.remove(idx)
      }
    }

    Reservoir(currElements, total_cnt)
  }

  override def merge(b1: Reservoir[T], b2: Reservoir[T]): Reservoir[T] = {
    val (left, right) = if (b1.seq.size > b2.seq.size) {
      (b1, b2)
    } else {
      (b2, b1)
    }

    if (left.elements < size) {
      mergeTwoPartial(left, right)
    } else if (right.elements < size) {
      mergePartialRight(left, right)
    } else {
      mergeFull(left, right)
    }
  }

  override def finish(reduction: Reservoir[T]): Seq[T] = reduction.seq.toSeq

  override def bufferEncoder: Encoder[Reservoir[T]] = Encoders.product

  override def outputEncoder: Encoder[Seq[T]] = ExpressionEncoder[Seq[T]]()

  package org.graphframes.rw

  import org.apache.spark.sql.DataFrame
  import org.apache.spark.sql.Encoders
  import org.apache.spark.sql.functions.array_union
  import org.apache.spark.sql.functions.col
  import org.apache.spark.sql.functions.udaf
  import org.apache.spark.sql.graphframes.expressions.ReservoirSamplingAgg
  import org.apache.spark.sql.types.ByteType
  import org.apache.spark.sql.types.IntegerType
  import org.apache.spark.sql.types.LongType
  import org.apache.spark.sql.types.ShortType
  import org.apache.spark.sql.types.StringType
  import org.graphframes.GraphFrame
  import org.graphframes.GraphFramesUnsupportedVertexTypeException
  import org.graphframes.Logging
  import org.graphframes.WithIntermediateStorageLevel

  import scala.util.Random

  trait RandomWalkBase extends Serializable with Logging with WithIntermediateStorageLevel {
    protected var maxNbrs: Int = 50
    protected var graph: GraphFrame = null
    protected var numWalksPerNode: Int = 5
    protected var batchSize: Int = 10
    protected var numBatches: Int = 5
    protected var useEdgeDirection: Boolean = false
    protected var globalSeed: Long = 42L
    protected var temporaryPrefix: Option[String] = None
    protected var runID: String = ""
  }

  protected def prepareGraph(): GraphFrame = {
    val preAggs = if (useEdgeDirection) {
      graph.edges
        .select(col(GraphFrame.SRC), col(GraphFrame.DST))
        .groupBy(col(GraphFrame.SRC).alias(GraphFrame.ID))
    } else {
      graph.edges
        .select(GraphFrame.SRC, GraphFrame.DST)
        .union(graph.edges.select(GraphFrame.DST, GraphFrame.SRC))
        .distinct()
        .groupBy(col(GraphFrame.SRC).alias(GraphFrame.ID))
    }

    val vertices = graph.vertices.schema(GraphFrame.ID).dataType match {
      case StringType =>
        preAggs.agg(
          udaf(ReservoirSamplingAgg[java.lang.String](maxNbrs), Encoders.STRING)
            .apply(col(GraphFrame.DST))
            .alias(RandomWalkBase.nbrsColName))
      case ShortType =>
        preAggs.agg(
          udaf(ReservoirSamplingAgg[java.lang.Short](maxNbrs), Encoders.SHORT)
            .apply(col(GraphFrame.DST))
            .alias(RandomWalkBase.nbrsColName))
      case ByteType =>
        preAggs.agg(
          udaf(ReservoirSamplingAgg[java.lang.Byte](maxNbrs), Encoders.BYTE)
            .apply(col(GraphFrame.DST))
            .alias(RandomWalkBase.nbrsColName))
      case IntegerType =>
        preAggs.agg(
          udaf(ReservoirSamplingAgg[java.lang.Integer](maxNbrs), Encoders.INT)
            .apply(col(GraphFrame.DST))
            .alias(RandomWalkBase.nbrsColName))
      case LongType =>
        preAggs.agg(
          udaf(ReservoirSamplingAgg[java.lang.Long](maxNbrs), Encoders.LONG)
            .apply(col(GraphFrame.DST))
            .alias(RandomWalkBase.nbrsColName))
      case _ => throw new GraphFramesUnsupportedVertexTypeException("unsupported vertex type")
    }

    val edges = graph.edges
    
    GraphFrame(vertices, edges)
  }

  protected def runIter(
      graph: GraphFrame,
      prevIterationDF: Option[DataFrame],
      iterSeed: Long): DataFrame

  def run(): DataFrame = {
    if (graph == null) {
      throw new IllegalArgumentException("Graph is not set")
    }
    if (temporaryPrefix.isEmpty) {
      throw new IllegalArgumentException("Temporary prefix is required for random walks.")
    }
    runID = java.util.UUID.randomUUID().toString
    logInfo(s"Starting random walk with runID: $runID")
    val iterationsRng = new Random()
    iterationsRng.setSeed(globalSeed)
    val spark = graph.vertices.sparkSession

    for (i <- 1 to numBatches) {
      logInfo(s"Starting batch $i of $numBatches")
      val iterSeed = iterationsRng.nextLong()
      val preparedGraph = prepareGraph()
      val prevIterationDF = if (i == 1) { None }
      else {
        Some(spark.read.parquet(iterationTmpPath(i - 1)))
      }
      val iterationResult: DataFrame = runIter(preparedGraph, prevIterationDF, iterSeed)
      iterationResult.write.parquet(iterationTmpPath(i))
    }

    logInfo("Finished all batches, merging results.")
    var result = spark.read.parquet(iterationTmpPath(1))

    for (i <- 2 to numBatches) {
      val tmpDF = spark.read
        .parquet(iterationTmpPath(i))
        .withColumnRenamed(RandomWalkBase.rwColName, "toMerge")
      result = result
        .join(tmpDF, Seq(RandomWalkBase.walkIdCol))
        .select(
          col(RandomWalkBase.walkIdCol),
          array_union(col(RandomWalkBase.rwColName), col("toMerge"))
            .alias(RandomWalkBase.rwColName))
    }
    result = result.persist(intermediateStorageLevel)

    val cnt = result.count()
    resultIsPersistent()
    logInfo(s"$cnt random walks are returned")
    result
  }

The first and the only present at the moment implementation is a simple Random Walks with Restart. While it is still a 1st order algorithm, based on the papers and reviews I read it is one of the best in terms of balance between scalability and quality. Compared to the RandomWalkBase trait it will have one additional hyperparameter: the restart probability.

  package org.graphframes.rw

  import org.apache.spark.sql.DataFrame
  import org.apache.spark.sql.functions.*
  import org.graphframes.GraphFrame

  class RandomWalkWithRestart extends RandomWalkBase {
      private var restartProbability: Double = 0.1
  }

And finally the concrete implementation.

  override protected def runIter(
      graph: GraphFrame,
      prevIterationDF: Option[DataFrame],
      iterSeed: Long): DataFrame = {
    val neighbors = graph.vertices.select(col(GraphFrame.ID), col(RandomWalkBase.nbrsColName))
    var walks = if (prevIterationDF.isEmpty) {
      graph.vertices.select(
        col(GraphFrame.ID).alias("startingNode"),
        col(GraphFrame.ID).alias(RandomWalkBase.currVisitingVertexColName),
        explode(
          when(
            array_size(col(RandomWalkBase.nbrsColName)) > lit(0),
            array((0 until numWalksPerNode).map(_ => uuid()): _*)).otherwise(array()))
          .alias(RandomWalkBase.walkIdCol),
        array(col(GraphFrame.ID)).alias(RandomWalkBase.rwColName))
    } else {
      prevIterationDF.get.select(
        col("startingNode"),
        col(RandomWalkBase.currVisitingVertexColName),
        col(RandomWalkBase.walkIdCol),
        array(col(RandomWalkBase.currVisitingVertexColName)).alias(RandomWalkBase.rwColName))
    }

    for (_ <- (0 until batchSize)) {
      walks = walks
        .join(
          neighbors,
          col(GraphFrame.ID) === col(RandomWalkBase.currVisitingVertexColName),
          "left")
        .withColumn("doRestart", rand() <= lit(restartProbability))
        .withColumn(
          "nextNode",
          when(col("doRestart"), col("startingNode")).otherwise(
            element_at(shuffle(col(RandomWalkBase.nbrsColName)), 1)))
        .select(
          col(RandomWalkBase.walkIdCol),
          col("startingNode"),
          col("nextNode").alias(RandomWalkBase.currVisitingVertexColName),
          array_append(
            col(RandomWalkBase.rwColName),
            col(RandomWalkBase.currVisitingVertexColName)).alias(RandomWalkBase.rwColName))
    }

    walks
  }

Here I'm using Spark's uuid to generate walks ID on the first iteration (the case when prevIterationDF is empty). Unfortunately spark does not have a built-in choice function to choose the random element from the array, so I just used shuffle + taking the first element. In the case of the restart we are just moving walker to the starting node (and that is reason why we need to keep it).

But what's next? The goal was to get representations, not just sequences. How to go from sequence to embeddings? And the simplest answer is Word2vec. It is a way to learn the representations by optimizing the similarity of often co-occurred items in sequences. It takes the sequences of elements (words, items, vertices, etc.) and train the vocabulary as the Map[String, Array[Double]].

There are a lot of educational materials about Word2vec model and modifications, so I won't stop on the details. One important thing to mention is that word2vec required "learning" of the vertex representations. Let's see how the distributed word2vec works in Apache Spark MLLib (there is also a word2vec in the Apache Spark ML, but it is just a wrapper). It creates the global vocabulary of all the words (or vertices in our case) and their initial vectors. It is done on the Spark's driver.

  val words = dataset.flatMap(x => x)
  vocab = words.map(w => (w, 1))
    .reduceByKey(_ + _)
    .filter(_._2 >= minCount)
    .map(x => VocabWord(
      x._1,
      x._2,
      new Array[Int](MAX_CODE_LENGTH),
      new Array[Int](MAX_CODE_LENGTH),
      0))
    .collect()
    .sortBy(_.cn)(Ordering[Long].reverse)

After that (and some additional magic that I will skip) the vocabulary is broadcasted from driver to workers that process each own subset of sequences. On the reduce step workers send their updates for the vocabulary to driver that do all-reduce operation. And this steps are repeated until convergence (or until the numIterations reached).

During my research, I checked existing implementations of random-walk based embeddings on Spark. All of them are generating sequences from Random Walks in one or another way and at the end just put these sequences to Spark's Word2vec:

• Mercury Graph
• Stellar Random Walk
• Graph Embedding Spark

For example, that is the code from the Mercury Graph project.

  w2v = Word2Vec(
      vectorSize=self.dimension,
      maxIter=self.w2v_max_iter,
      numPartitions=self.w2v_num_partitions,
      stepSize=self.w2v_step_size,
      inputCol="random_walks",
      outputCol="model",
      minCount=self.w2v_min_count,
  )

  self.node2vec_ = w2v.fit(self.paths_)

I cannot try to blame anyone for using such an approach because it works. The main issue is scalability of it to the real-world graphs with billions of vertices… The story is that Word2vec itself as well as its implementation in Apache Spark were designed to solve a different problem. The algorithm was created for Natural Language Processing (NLP). And the story is there is a natural limitations for the problems size. An amount of words in languages are barely hit a million of unique items. The Oxford English Dictionary provides even a smaller estimation, like \(\simeq 500000\) of unique words. But in the graph worlds, 500000 of vertices is a small-medium sized graph. For example, we are running benchmarks of GraphFrames inside GitHub action agents on the graph with 2M of vertices that is already much more than amount of possible words in English.

While it is really hard to imagine a Spark driver that can handle a hash map of the size of ten millions of keys and Array[Double] of the length 100 as values, there is a more explicit limitation. Let's examine the code of Spark's Word2vec again and take a look on this line:

  if (vocabSize.toLong * vectorSize >= Int.MaxValue) {
    throw new RuntimeException("Please increase minCount or decrease vectorSize in Word2Vec" +
      " to avoid an OOM. You are highly recommended to make your vocabSize*vectorSize, " +
      "which is " + vocabSize + "*" + vectorSize + " for now, less than `Int.MaxValue`.")
  }

In other words, the size of vocabulary multiplied by the size of embeddings should be less than Int.MaxValue. It leads to the following picture where on the X-axis is an amount of the vertices in the graph and on the Y-axis is the maximal possible size of embeddings.

For example, the default size of embeddings in Spark is 100 that means you can process graphs up to \(\simeq 21.6 \cdot 10^6\) vertices. If your graph has 1B of vertices, the best you can learn is embeddings of the size 2… Is 2 float numbers enough to encode the information about 1B of vertices? I barely think so. Moreover, if you are working with a graph with a couple of millions of vertices you barely need GraphFrames, Spark and distributed graph processing, because such a size is much easier to process on a single node.

While it is possible to left Word2vec as an option for users of GraphFrames, it is obvious that the library requires more scalable solutions. I read more than 15 different papers with alternative approaches and my group them in the following way.

• Parameter Server based solutions
• Hashing trick based solutions

Parameter Server Architecture

Parameter Server (PS) architecture is a way to train models like Word2vec without limitations. Instead of using Spark driver for keeping the whole vocabulary and do all-reduce, we are creating a separate Parameter Server that is a Key-Value store in the simplest scenario. With a PS we do not need to materialize the huge vocabulary, instead we can store it offloaded and send to workers only what is actually required.

While such an approach is the industry standard, it is not a fit for GraphFrames that aims to bring Graph Machine Learning on top of existing infrastructure. If the strong engineering teem can maintain a Parameter Server they just do not need GraphFrames. And I cannot realize how I can add a Parameter Server "without server" on top of existing Spark Clusters.

Hashing Trick

Another possible direction is to try to reduce the vocabulary size to fit it into the driver. Hashing Trick is an attempt to make elements in vocabulary sharable for different elements of sequences. The key work is Tito Svenstrup, Dan, Jonas Hansen, and Ole Winther. "Hash embeddings for efficient word representations." Advances in neural information processing systems 30 (2017). It suggest the idea of applying multiple hash function to each element to map it into few buckets and the final representation of element in that case is a weighted sum of buckets embeddings.

This is already feasible to implement on Spark beside the fact it would require to write Word2vec from scratch with a lot of complications. Maybe one day there will be a true hashing-trick in GraphFrames, at the moment I'm not ready to do it. My knowledge of linear algebra and hard math is quite a limited.

But what we can learn from the hashing trick is the direction. And reading one paper after another I finally found one that looks like a perfect balance between scalability and feasibility of the implementation. It is a work Tito Svenstrup, Dan, Jonas Hansen, and Ole Winther. "Hash embeddings for efficient word representations." Advances in neural information processing systems 30 (2017). The work itself is based on the well known theory of Random Projections that states, that random projections preserve the distances.

Algorithm itself is very-very simple and works in linear time (actually just in a single pass over data). As well the approach does not require any kind of global shared vocabulary: each sequence of the dataset can be processed independently and final vectors of elements can be done in a single all-reduce step.

  package org.graphframes.embeddings

  import org.apache.spark.ml.linalg.SQLDataTypes.VectorType
  import org.apache.spark.ml.linalg.Vectors
  import org.apache.spark.ml.stat.Summarizer
  import org.apache.spark.rdd.RDD
  import org.apache.spark.sql.DataFrame
  import org.apache.spark.sql.Row
  import org.apache.spark.sql.functions.col
  import org.apache.spark.sql.types.ArrayType
  import org.apache.spark.sql.types.ByteType
  import org.apache.spark.sql.types.IntegerType
  import org.apache.spark.sql.types.LongType
  import org.apache.spark.sql.types.ShortType
  import org.apache.spark.sql.types.StringType
  import org.apache.spark.sql.types.StructField
  import org.apache.spark.sql.types.StructType
  import org.apache.spark.unsafe.hash.Murmur3_x86_32.*
  import org.apache.spark.unsafe.types.UTF8String
  import org.graphframes.GraphFramesUnsupportedVertexTypeException
  import org.graphframes.rw.RandomWalkBase

  import scala.annotation.nowarn
  import scala.jdk.CollectionConverters.*
  import scala.reflect.ClassTag

  class Hash2Vec extends Serializable {
    private var contextSize: Int = 5
    private var numPartitions: Int = 5
    private var embeddingsDim: Int = 256
    private var sequenceCol: String = RandomWalkBase.rwColName
    private var decayFunction: String = "gaussian"
    private var gaussianSigma: Double = 1.0
    private var hashingSeed: Int = 42
    private var signHashingSeed: Int = 18

  private def nonNegativeMod(x: Int, mod: Int): Int = {
    val rawMod = x % mod
    rawMod + (if (rawMod < 0) mod else 0)
  }

  private def hashFunc(term: Any, seed: Int): Int = {
    term match {
      case null => seed
      case b: Boolean => hashInt(if (b) 1 else 0, seed)
      case b: Byte => hashInt(b.toInt, seed)
      case s: Short => hashInt(s.toInt, seed)
      case i: Int => hashInt(i, seed)
      case l: Long => hashLong(l, seed)
      case f: Float => hashInt(java.lang.Float.floatToIntBits(f), seed)
      case d: Double => hashLong(java.lang.Double.doubleToLongBits(d), seed)
      case s: String =>
        val utf8 = UTF8String.fromString(s)
        hashUnsafeBytes(utf8.getBaseObject, utf8.getBaseOffset, utf8.numBytes(), seed)
      case _ =>
        throw new GraphFramesUnsupportedVertexTypeException(
          "Hashing2vec with murmur3 algorithm does not " +
            s"support type ${term.getClass.getCanonicalName} of input data.")
    }
  }

  private def decayGaussian(d: Int, sigma: Double): Double = {
    math.exp(-(d * d) / (sigma * sigma))
  }

  private def processPartition[T](iter: Iterator[Seq[T]]): Iterator[(T, Array[Double])] = {
    val localVocab = new java.util.concurrent.ConcurrentHashMap[T, Array[Double]]()

    for (seq <- iter) {
      val currentSeqSize = seq.length
      for (idx <- (0 until currentSeqSize)) {
        val currentWord = seq(idx)
        if (!localVocab.containsKey(currentWord)) {
          localVocab.put(currentWord, Array.fill(embeddingsDim)(0.0))
        }
        val context = ((idx - contextSize) to (idx + contextSize)).filter(i =>
          (i >= 0) && (i < currentSeqSize) && (i != idx))
        for (cIdx <- context) {
          val word = seq(cIdx)
          val weight = weightFunction(math.abs(cIdx - idx))
          val sign = 2.0 * signHash(word) - 1.0
          val embeddingIdx = valueHash(word)

          val currentEmbedding = localVocab.get(currentWord)
          currentEmbedding(embeddingIdx) += sign * weight
        }
      }
    }

    localVocab
      .entrySet()
      .asScala
      .map(entry => (entry.getKey(), entry.getValue()))
      .iterator
  }

  private def runTyped[T: ClassTag](data: DataFrame): RDD[(T, Array[Double])] = {
    data
      .select(col(sequenceCol))
      .rdd
      .map(_.getAs[Seq[T]](0))
      .repartition(numPartitions)
      .mapPartitions(processPartition[T])
  }

  def run(data: DataFrame): DataFrame = {
    val spark = data.sparkSession
    require(data.schema(sequenceCol).dataType.isInstanceOf[ArrayType], "sequence should be array")
    val elDataType = data.schema(sequenceCol).dataType.asInstanceOf[ArrayType].elementType

    weightFunction = decayFunction match {
      case "gaussian" => (d: Int) => decayGaussian(d, gaussianSigma)
      case "constant" => (_: Int) => 1.0
      case _ => throw new RuntimeException(s"unsupported decay functions $decayFunction")
    }

    valueHash = (el: Any) => nonNegativeMod(hashFunc(el, hashingSeed), embeddingsDim)
    signHash = (el: Any) => nonNegativeMod(hashFunc(el, signHashingSeed), 2)

    val (rowRDD, schema) = elDataType match {
      case _: StringType =>
        (
          runTyped[String](data).map(f => Row(f._1, Vectors.dense(f._2))),
          StructType(Seq(StructField("id", StringType), StructField("vector", VectorType))))
      case _: ByteType =>
        (
          runTyped[Byte](data).map(f => Row(f._1, Vectors.dense(f._2))),
          StructType(Seq(StructField("id", ByteType), StructField("vector", VectorType))))
      case _: ShortType =>
        (
          runTyped[Short](data).map(f => Row(f._1, Vectors.dense(f._2))),
          StructType(Seq(StructField("id", ShortType), StructField("vector", VectorType))))
      case _: IntegerType =>
        (
          runTyped[Int](data).map(f => Row(f._1, Vectors.dense(f._2))),
          StructType(Seq(StructField("id", IntegerType), StructField("vector", VectorType))))
      case _: LongType =>
        (
          runTyped[Long](data).map(f => Row(f._1, Vectors.dense(f._2))),
          StructType(Seq(StructField("id", LongType), StructField("vector", VectorType))))
      case _ =>
        throw new GraphFramesUnsupportedVertexTypeException(
          s"Hash2vec supports only string or numeric types of elements but gor ${elDataType.toString()}")
    }

    spark.createDataFrame(rowRDD, schema).groupBy("id").agg(Summarizer.sum(col("vector")).alias("embedding"))
  }

But generating cool pictures in Gephi is barely a goal of anyone who is going to use GraphFrames. The fair approach is to check how much predictive power on ML tasks. I took a few graphs from the Karateclub Project. This graphs are small but they a) fit into my laptop b) have classification targets. For all of them I generated embeddings with both, Hash2vec and Word2vec. To check the power of embeddings instead of power of ML classifier I used a K-Nearest-Neighbor algorithm with cosine similarity functions. It just compute the distances between elements, take K nearest neighbors of each and use an average of their targets as a prediction. For evaluation I used Area Under the Receiver Operating Characteristic Curve (ROC-AUC) metric that can be explained in the following way. If we range all our predictions by the predicted score, the ROC-AUS shows the probability that a random element with label 1 will have a score bigger than a random element with label 0. ROC-AUC was computed on cross-validation. That are results:

As you can see while Hash2vec provides some prediction power, Word2vec outperforms it significantly. I'm not sure may it be the problem with huge dispersion of L2 norms of Hash2vec vectors because I did not apply normalization. Just for understanding, the norm of Hash2vec vectors is proportional to the vertex degree while norm of the Word2vec vectors should distributed equally. I will discuss this in the next section, but as you can see, Hash2vec already has a predictive power, even without normalization!

While there are a lot of space for improvements I need to stop here. The Pull Request is already 1000 lines of the Scala code. I think having almost free to compute graph embeddings from Random Walks is an excellent starting point for the project.

I need also to write tests for corner cases, write the documentation. And the most boring but required part is to write PySpark bindings for this new API because Scala is nice for the low-level Spark, but people and especially Data Scientists prefer to work with Python. It would be natural to propose Hash2vec as an addition for SparkML. But I'm not sure that maintainer of the project are ready to accept anything new to the SparkML module. I'm not ready to jump into it without a shepherd. So, if by any chance you know any PMC of the project who can help to promote it, I'm willing to rewrite it according Spark standards and push.

As I already mentioned, Hash2vec vectors are sparse and have L2 norm proportional to vertex degree. From one point of view, encoding the degree is important. From another, normalization may increase the prediction power. I'm going to at least add this option in the future. This is question that requires hard math knowledge, so I decided to postpone it.

I read recently two papers (The Unreasonable Effectiveness of Randomized Representations in Online Continual Graph Learning, On the Effectiveness of Random Weights in Graph Neural Networks) that give me the following idea. The main problem of Hash2vec embeddings is that they are very "noisy". But what if we smooth them by applying a GraphSAGE like convolution on top with fixed random weights? Because it does not require any kind of training it may be done in Spark with a pure BLAS operations. Can't say it will be easy, but it sounds feasible at least. In my mind, the convolution layer with random weights should add some local information to embeddings. So, we will have the best of two worlds: deep structure information about the graph from Random Walks and local information from graph convolutions… Definitely worth trying I think.