<commit_before>package spark.graph

import scala.collection.JavaConversions._
import scala.collection.mutable.ArrayBuffer

import com.esotericsoftware.kryo._

import it.unimi.dsi.fastutil.ints.IntArrayList

import spark.{ClosureCleaner, HashPartitioner, KryoRegistrator, SparkContext, RDD}
import spark.SparkContext._
import spark.storage.StorageLevel



class Vertex[@specialized(Char, Int, Boolean, Byte, Long, Float, Double) VD] {
  var id: Vid = _
  var data: VD = _

  def this(id: Int, data: VD) { this(); this.id = id; this.data = data; }
}


class Edge[@specialized(Char, Int, Boolean, Byte, Long, Float, Double) ED] {
  var src: Vid = _
  var dst: Vid = _
  var data: ED = _

  def this(src: Vid, dst: Vid, data: ED) {
    this(); this.src = src; this.dst = dst; this.data = data;
  }
}


class EdgeWithVertices[@specialized(Char, Int, Boolean, Byte, Long, Float, Double) VD,
                       @specialized(Char, Int, Boolean, Byte, Long, Float, Double) ED] {

  var src: Vertex[VD] = _
  var dst: Vertex[VD] = _
  var data: ED = _

  def otherVertex(vid: Vid): Vertex[VD] = if (src.id == vid) dst else src

  def vertex(vid: Vid): Vertex[VD] = if (src.id == vid) src else dst
}


/**
 * A Graph RDD that supports computation on graphs.
 */
class Graph[VD: Manifest, ED: Manifest](
  private val _vertices: RDD[Vertex[VD]],
  private val _edges: RDD[Edge[ED]]) {

  import Graph.EdgePartition

  var numEdgePartitions = 5
  var numVertexPartitions = 5

  private val eTable: RDD[(Pid, EdgePartition[ED])] = Graph.createETable(
    _edges, numEdgePartitions)

  private val vTable: RDD[(Vid, (VD, Array[Pid]))] = Graph.createVTable(
    _vertices, eTable, numVertexPartitions)

  def edges: RDD[Edge[ED]] = eTable.mapPartitions { iter => iter.next._2.iterator }

  def edgesWithVertices: RDD[EdgeWithVertices[VD, ED]] = new EdgeWithVerticesRDD(vTable, eTable)

}


object Graph {

  /**
   * A partition of edges. This is created so we can store edge data in columnar format so it is
   * more efficient to store the data in memory.
   */
  private[graph]
  class EdgePartition[@specialized(Char, Int, Boolean, Byte, Long, Float, Double) ED: Manifest] {
    val srcIds: IntArrayList = new IntArrayList
    val dstIds: IntArrayList = new IntArrayList
    // TODO: Specialize data.
    val data: ArrayBuffer[ED] = new ArrayBuffer[ED]

    /** Add a new edge to the partition. */
    def add(src: Vid, dst: Vid, d: ED) {
      srcIds.add(src)
      dstIds.add(dst)
      data += d
    }

    def trim() {<commit_msg>Added a vertices method to Graph.<commit_after>package spark.graph

import scala.collection.JavaConversions._
import scala.collection.mutable.ArrayBuffer

import com.esotericsoftware.kryo._

import it.unimi.dsi.fastutil.ints.IntArrayList

import spark.{ClosureCleaner, HashPartitioner, KryoRegistrator, SparkContext, RDD}
import spark.SparkContext._
import spark.storage.StorageLevel



class Vertex[@specialized(Char, Int, Boolean, Byte, Long, Float, Double) VD] {
  var id: Vid = _
  var data: VD = _

  def this(id: Int, data: VD) { this(); this.id = id; this.data = data; }
}


class Edge[@specialized(Char, Int, Boolean, Byte, Long, Float, Double) ED] {
  var src: Vid = _
  var dst: Vid = _
  var data: ED = _

  def this(src: Vid, dst: Vid, data: ED) {
    this(); this.src = src; this.dst = dst; this.data = data;
  }
}


class EdgeWithVertices[@specialized(Char, Int, Boolean, Byte, Long, Float, Double) VD,
                       @specialized(Char, Int, Boolean, Byte, Long, Float, Double) ED] {

  var src: Vertex[VD] = _
  var dst: Vertex[VD] = _
  var data: ED = _

  def otherVertex(vid: Vid): Vertex[VD] = if (src.id == vid) dst else src

  def vertex(vid: Vid): Vertex[VD] = if (src.id == vid) src else dst
}


/**
 * A Graph RDD that supports computation on graphs.
 */
class Graph[VD: Manifest, ED: Manifest](
  private val _vertices: RDD[Vertex[VD]],
  private val _edges: RDD[Edge[ED]]) {

  import Graph.EdgePartition

  var numEdgePartitions = 5
  var numVertexPartitions = 5

  private val eTable: RDD[(Pid, EdgePartition[ED])] = Graph.createETable(
    _edges, numEdgePartitions)

  private val vTable: RDD[(Vid, (VD, Array[Pid]))] = Graph.createVTable(
    _vertices, eTable, numVertexPartitions)

  def vertices: RDD[Vertex[VD]] = vTable.map { case(vid, (data, pids)) => new Vertex(vid, data) }

  def edges: RDD[Edge[ED]] = eTable.mapPartitions { iter => iter.next._2.iterator }

  def edgesWithVertices: RDD[EdgeWithVertices[VD, ED]] = new EdgeWithVerticesRDD(vTable, eTable)

}


object Graph {

  /**
   * A partition of edges. This is created so we can store edge data in columnar format so it is
   * more efficient to store the data in memory.
   */
  private[graph]
  class EdgePartition[@specialized(Char, Int, Boolean, Byte, Long, Float, Double) ED: Manifest] {
    val srcIds: IntArrayList = new IntArrayList
    val dstIds: IntArrayList = new IntArrayList
    // TODO: Specialize data.
    val data: ArrayBuffer[ED] = new ArrayBuffer[ED]

    /** Add a new edge to the partition. */
    def add(src: Vid, dst: Vid, d: ED) {
      srcIds.add(src)
      dstIds.add(dst)
      data += d
    }

    def trim() {
{'fraction_non_alphanumeric': 0.12318700444788241, 'fraction_numerical': 0.0011603171533552505, 'mean_word_length': 3.5973333333333333, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 6, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '3215762', 'n_tokens_mistral': 1807, 'n_tokens_neox': 1757, 'n_words': 537}