package com.gmei
import java.io.Serializable
import scala.util.Try
import org.apache.spark.storage.StorageLevel
import org.apache.spark.mllib.linalg.Vectors
import org.apache.spark.mllib.linalg.distributed.{IndexedRow, IndexedRowMatrix}
import org.apache.spark.sql.{SaveMode, TiContext}
import org.apache.log4j.{Level, Logger}
import scopt.OptionParser
import com.gmei.lib.AbstractParams
import com.soundcloud.lsh.Lsh
object Main {
Logger.getLogger("org.apache.spark").setLevel(Level.WARN)
Logger.getLogger("org.apache.eclipse.jetty.server").setLevel(Level.OFF)
case class Params(iter: Int = 10,
lr: Double = 0.025,
numPartition: Int = 10,
dim: Int = 128,
window: Int = 10,
walkLength: Int = 80,
numWalks: Int = 10,
p: Double = 1.0,
q: Double = 1.0,
weighted: Boolean = true,
directed: Boolean = false,
degree: Int = 30,
indexed: Boolean = false,
env: String = ENV.DEV,
nodePath: String = null
) extends AbstractParams[Params] with Serializable
val defaultParams = Params()
val parser = new OptionParser[Params]("Node2Vec_Spark") {
head("Main")
opt[Int]("walkLength")
.text(s"walkLength: ${defaultParams.walkLength}")
.action((x, c) => c.copy(walkLength = x))
opt[Int]("numWalks")
.text(s"numWalks: ${defaultParams.numWalks}")
.action((x, c) => c.copy(numWalks = x))
opt[Double]("p")
.text(s"return parameter p: ${defaultParams.p}")
.action((x, c) => c.copy(p = x))
opt[Double]("q")
.text(s"in-out parameter q: ${defaultParams.q}")
.action((x, c) => c.copy(q = x))
opt[Boolean]("weighted")
.text(s"weighted: ${defaultParams.weighted}")
.action((x, c) => c.copy(weighted = x))
opt[Boolean]("directed")
.text(s"directed: ${defaultParams.directed}")
.action((x, c) => c.copy(directed = x))
opt[Int]("degree")
.text(s"degree: ${defaultParams.degree}")
.action((x, c) => c.copy(degree = x))
opt[Boolean]("indexed")
.text(s"Whether nodes are indexed or not: ${defaultParams.indexed}")
.action((x, c) => c.copy(indexed = x))
opt[String]("env")
.text(s"the databases environment you used")
.action((x,c) => c.copy(env = x))
opt[String]("nodePath")
.text("Input node2index file path: empty")
.action((x, c) => c.copy(nodePath = x))
note(
"""
|For example, the following command runs this app on a tidb dataset:
|
| spark-submit --class com.gmei.Main ./target/scala-2.11/Node2vec-assembly-0.2.jar \
""".stripMargin +
s"| --env ${defaultParams.env}"
)
}
def main(args: Array[String]):Unit = {
parser.parse(args, defaultParams).map { param =>
//1. get the input and node2vec
GmeiConfig.setup(param)
val spark_env = GmeiConfig.getSparkSession()
val context = spark_env._1
val sc = spark_env._2
val ti = new TiContext(sc)
ti.tidbMapTable(dbName = GmeiConfig.config.getString("tidb.database"),tableName = "nd_data_meigou_cid")
ti.tidbMapTable(dbName = GmeiConfig.config.getString("tidb.database"),tableName = "data_feed_click")
val date8 = GmeiConfig.getMinusNDate(30)
val tidb_inupt = sc.sql(
s"""
|SELECT
| service_id,cid
|FROM nd_data_meigou_cid
|where stat_date > '${date8}'
""".stripMargin
)
Node2vec.setup(context, param)
Node2vec.load(tidb_inupt)
.initTransitionProb()
.randomWalk()
.embedding()
val node2vector = context.parallelize(Word2vec.getVectors.toList)
.map { case (nodeId, vector) =>
(nodeId.toLong, vector.map(x => x.toDouble))
}
val id2Node = Node2vec.node2id.map{ case (strNode, index) =>
(index, strNode)
}
val node2vec_2 = node2vector.join(id2Node)
.map { case (nodeId, (vector, name)) => (name,vector) }
.repartition(200)
//2. compute similar cid and then take top k
val storageLevel = StorageLevel.MEMORY_AND_DISK
val indexed = node2vec_2.zipWithIndex.persist(storageLevel)
// create indexed row matrix where every row represents one word
val rows = indexed.map {
case ((word, features), index) =>
IndexedRow(index, Vectors.dense(features))
}
// store index for later re-mapping (index to word)
val index = indexed.map {
case ((word, features), index) =>
(index, word)
}.persist(storageLevel)
// create an input matrix from all rows and run lsh on it
val matrix = new IndexedRowMatrix(rows)
val lsh = new Lsh(
minCosineSimilarity = 0.5,
dimensions = 20,
numNeighbours = 200,
numPermutations = 10,
partitions = 200,
storageLevel = storageLevel
)
val similarityMatrix = lsh.join(matrix)
import sc.implicits._
// remap both ids back to words
val remapFirst = similarityMatrix.entries.keyBy(_.i).join(index).values
val remapSecond = remapFirst.keyBy { case (entry, word1) => entry.j }.join(index).values.map {
case ((entry, word1), word2) =>
(word1, word2, entry.value)
}
remapSecond.take(20).foreach(println)
val score_result = remapSecond.toDF("cid1","cid2","score")
GmeiConfig.writeToJDBCTable(score_result, table="nd_cid_pairs_cosine_distince", SaveMode.Overwrite)
// group by neighbours to get a list of similar words and then take top k
val result = remapSecond.filter(_._1.startsWith("diary")).groupBy(_._1).map {
case (word1, similarWords) =>
// sort by score desc. and take top 20 entries
val similar = Try(similarWords.toSeq.sortBy(-1 * _._3).filter(_._2.startsWith("diary")).take(20).map(_._2).mkString(",")).getOrElse(null)
(word1,s"$similar")
}.filter(_._2.split(",").length > 4)
result.take(20).foreach(println)
val similar_result = result.toDF("cid","similarity_cid")
GmeiConfig.writeToJDBCTable(similar_result, table="nd_cid_similarity_matrix", SaveMode.Overwrite)
//3. cids queue map to device_id
ti.tidbMapTable(dbName = GmeiConfig.config.getString("tidb.database"),tableName = "nd_cid_similarity_matrix")
val device_id = sc.sql(
s"""
|select a.device_id device_id,a.city_id city_id ,b.similarity_cid similarity_cid from
|(select device_id,first(city_id) as city_id,first(cid) as cid from data_feed_click
|where cid in (select cid from nd_cid_similarity_matrix)
|group by device_id) a left join
|nd_cid_similarity_matrix b
|on a.cid = b.cid
|where b.similarity_cid is not null
""".stripMargin
).na.fill(Map("city_id"->"beijing"))
device_id.show()
val device_queue = device_id.rdd.map {item =>
val parts = (item.getAs[String](fieldName = "device_id"),item.getAs[String](fieldName = "city_id"),item.getAs[String](fieldName = "similarity_cid"))
Try {
(parts._1,Try(parts._2.toString.replace("worldwide","beijing")).getOrElse(null),Try(parts._3.toString.replace("diary|","")).getOrElse(null))
}.getOrElse(null)
}.filter(_!=null).toDF("device_id","city_id","similarity_cid")
device_queue.take(20).foreach(println)
GmeiConfig.writeToJDBCTable(device_queue, table="nd_device_cid_similarity_matrix", SaveMode.Overwrite)
sc.stop()
}
} getOrElse {
sys.exit(1)
}
}