<commit_before>package com.databricks.dfgraph

import org.apache.spark.graphx.{Edge, Graph, EdgeRDD, VertexRDD}
import org.apache.spark.sql.{Row, SQLContext, DataFrame}

import scala.reflect.ClassTag
import scala.reflect.runtime.universe.TypeTag

/**
 * Wraps a GraphX graph inside a Spark Dataframe.
 * @tparam VD vertex attribute type
 */
class DFGraph[VD: TypeTag:ClassTag, ED: TypeTag:ClassTag](
    private val vertexDF: DataFrame,
    private val edgeDF: DataFrame) extends Serializable {
  import DFGraph._
  private val _sqlContext: SQLContext = vertexDF.sqlContext

  def toGraph(): Graph[VD, ED] = {
    val verts = VertexRDD(vertexDF.select(VERTEX_ID_COLNAME, VERTEX_ATTR_COLNAME).map {
      case Row(id: Long, attr: VD) => (id, attr)
    })
    edgeDF.select(EDGE_SRC_ID_COLNAME, EDGE_TGT_ID_COLNAME, EDGE_ATTR_COLNAME).printSchema()
    edgeDF.select(EDGE_SRC_ID_COLNAME, EDGE_TGT_ID_COLNAME, EDGE_ATTR_COLNAME).show()
    val edges = EdgeRDD.fromEdges[ED, VD](
      edgeDF.select(EDGE_SRC_ID_COLNAME, EDGE_TGT_ID_COLNAME, EDGE_ATTR_COLNAME).map {
        // TODO: fix matcherror
        case Row(src: Long, dst: Long, attr: ED) => Edge(src, dst, attr)
      }
    )
    Graph(verts, edges)
  }
}

private object DFGraph {
  val VERTEX_ID_COLNAME = "vtxID"
  val VERTEX_ATTR_COLNAME = "vtxAttr"
  val EDGE_SRC_ID_COLNAME = "edgeSrcId"
  val EDGE_TGT_ID_COLNAME = "edgeTgtId"
  val EDGE_ATTR_COLNAME = "edgeAttr"<commit_msg>Fixes tests with ugly type casting<commit_after>package com.databricks.dfgraph

import org.apache.spark.graphx.{Edge, Graph, EdgeRDD, VertexRDD}
import org.apache.spark.sql.{Row, SQLContext, DataFrame}

import scala.reflect.ClassTag
import scala.reflect.runtime.universe.TypeTag

/**
 * Wraps a GraphX graph inside a Spark Dataframe.
 * @tparam VD vertex attribute type
 */
class DFGraph[VD: TypeTag:ClassTag, ED: TypeTag:ClassTag](
    private val vertexDF: DataFrame,
    private val edgeDF: DataFrame) extends Serializable {
  import DFGraph._
  private val _sqlContext: SQLContext = vertexDF.sqlContext

  def toGraph(): Graph[VD, ED] = {
    val verts = VertexRDD(vertexDF.select(VERTEX_ID_COLNAME, VERTEX_ATTR_COLNAME).map {
      case Row(id: Long, attr) => (id, attr.asInstanceOf[VD])
    })

    edgeDF.select(EDGE_SRC_ID_COLNAME, EDGE_TGT_ID_COLNAME, EDGE_ATTR_COLNAME).printSchema()
    val edges = EdgeRDD.fromEdges[ED, VD](
      edgeDF.select(EDGE_SRC_ID_COLNAME, EDGE_TGT_ID_COLNAME, EDGE_ATTR_COLNAME).map {
        case Row(src: Long, dst: Long, attr) => Edge(src, dst, attr.asInstanceOf[ED])
      }
    )
    Graph(verts, edges)
  }
}

private object DFGraph {
  val VERTEX_ID_COLNAME = "vtxID"
  val VERTEX_ATTR_COLNAME = "vtxAttr"
  val EDGE_SRC_ID_COLNAME = "edgeSrcId"
  val EDGE_TGT_ID_COLNAME = "edgeTgtId"
  val EDGE_ATTR_COLNAME = "edgeAttr"
{'fraction_non_alphanumeric': 0.10441910192444762, 'fraction_numerical': 0.0, 'mean_word_length': 4.482421875, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 7, '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': '3396072', 'n_tokens_mistral': 1055, 'n_tokens_neox': 1036, 'n_words': 218}