feature_engineering.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import findspark\n",
    "findspark.init()\n",
    "from pyspark import SparkConf\n",
    "from pyspark.sql import SparkSession\n",
    "conf = SparkConf().setAppName('featureEngineering').setMaster('local')\n",
    "spark = SparkSession.builder.config(conf=conf).getOrCreate()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Feature Engineering"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "from pyspark.ml import Pipeline\n",
    "from pyspark.ml.feature import OneHotEncoder, StringIndexer, QuantileDiscretizer, MinMaxScaler\n",
    "from pyspark.ml.linalg import VectorUDT, Vectors\n",
    "from pyspark.sql.functions import *\n",
    "from pyspark.sql.types import *\n",
    "from pyspark.sql import functions as F\n",
    "\n",
    "\n",
    "def oneHotEncoderExample(movieSamples):\n",
    "    # Add movieIdNumber column with intger type with the same value as movieId\n",
    "    samplesWithIdNumber = movieSamples.withColumn(\"movieIdNumber\", F.col(\"movieId\").cast(IntegerType()))\n",
    "    encoder = OneHotEncoder(inputCols=[\"movieIdNumber\"], outputCols=['movieIdVector'], dropLast=False)\n",
    "    oneHotEncoderSamples = encoder.fit(samplesWithIdNumber).transform(samplesWithIdNumber)\n",
    "    oneHotEncoderSamples.printSchema()\n",
    "    oneHotEncoderSamples.show(10)\n",
    "\n",
    "\n",
    "def array2vec(genreIndexes, indexSize):\n",
    "    genreIndexes.sort()\n",
    "    fill_list = [1.0 for _ in range(len(genreIndexes))]\n",
    "    return Vectors.sparse(indexSize, genreIndexes, fill_list)\n",
    "\n",
    "\n",
    "def multiHotEncoderExample(movieSamples):\n",
    "    # 把genre的多category explode\n",
    "    # movieId|               title|    genre|\n",
    "    #     +-------+--------------------+---------+\n",
    "    #     |      1|    Toy Story (1995)|Adventure|\n",
    "    #     |      1|    Toy Story (1995)|Animation|\n",
    "    #     |      1|    Toy Story (1995)| Children|\n",
    "    samplesWithGenre = movieSamples.select(\"movieId\", \"title\", explode(\n",
    "        split(F.col(\"genres\"), \"\\\\|\").cast(ArrayType(StringType()))).alias('genre'))\n",
    "    genreIndexer = StringIndexer(inputCol=\"genre\", outputCol=\"genreIndex\")\n",
    "    StringIndexerModel = genreIndexer.fit(samplesWithGenre)\n",
    "    # 把Adventure, Animation变成数值i.e. 5, 7, 9\n",
    "    genreIndexSamples = StringIndexerModel.transform(samplesWithGenre).withColumn(\"genreIndexInt\",\n",
    "                                                                                  F.col(\"genreIndex\").cast(IntegerType()))\n",
    "    indexSize = genreIndexSamples.agg(max(F.col(\"genreIndexInt\"))).head()[0] + 1\n",
    "    processedSamples = genreIndexSamples.groupBy('movieId').agg(\n",
    "        F.collect_list('genreIndexInt').alias('genreIndexes')).withColumn(\"indexSize\", F.lit(indexSize))\n",
    "    finalSample = processedSamples.withColumn(\"vector\",\n",
    "                                              udf(array2vec, VectorUDT())(F.col(\"genreIndexes\"), F.col(\"indexSize\")))\n",
    "    finalSample.printSchema()\n",
    "    finalSample.show(10)\n",
    "\n",
    "\n",
    "def ratingFeatures(ratingSamples):\n",
    "    ratingSamples.printSchema()\n",
    "    ratingSamples.show()\n",
    "    # calculate average movie rating score and rating count\n",
    "    movieFeatures = ratingSamples.groupBy('movieId').agg(F.count(F.lit(1)).alias('ratingCount'),\n",
    "                                                         F.avg(\"rating\").alias(\"avgRating\"),\n",
    "                                                         F.variance('rating').alias('ratingVar')) \\\n",
    "        .withColumn('avgRatingVec', udf(lambda x: Vectors.dense(x), VectorUDT())('avgRating'))\n",
    "    movieFeatures.show(10)\n",
    "    # bucketing\n",
    "    ratingCountDiscretizer = QuantileDiscretizer(numBuckets=100, inputCol=\"ratingCount\", outputCol=\"ratingCountBucket\")\n",
    "    # Normalization\n",
    "    ratingScaler = MinMaxScaler(inputCol=\"avgRatingVec\", outputCol=\"scaleAvgRating\")\n",
    "    pipelineStage = [ratingCountDiscretizer, ratingScaler]\n",
    "    featurePipeline = Pipeline(stages=pipelineStage)\n",
    "    movieProcessedFeatures = featurePipeline.fit(movieFeatures).transform(movieFeatures)\n",
    "    movieProcessedFeatures.show(10)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Raw Movie Samples:\n",
      "+-------+--------------------+--------------------+\n",
      "|movieId|               title|              genres|\n",
      "+-------+--------------------+--------------------+\n",
      "|      1|    Toy Story (1995)|Adventure|Animati...|\n",
      "|      2|      Jumanji (1995)|Adventure|Childre...|\n",
      "|      3|Grumpier Old Men ...|      Comedy|Romance|\n",
      "|      4|Waiting to Exhale...|Comedy|Drama|Romance|\n",
      "|      5|Father of the Bri...|              Comedy|\n",
      "|      6|         Heat (1995)|Action|Crime|Thri...|\n",
      "|      7|      Sabrina (1995)|      Comedy|Romance|\n",
      "|      8| Tom and Huck (1995)|  Adventure|Children|\n",
      "|      9| Sudden Death (1995)|              Action|\n",
      "|     10|    GoldenEye (1995)|Action|Adventure|...|\n",
      "+-------+--------------------+--------------------+\n",
      "only showing top 10 rows\n",
      "\n",
      "root\n",
      " |-- movieId: string (nullable = true)\n",
      " |-- title: string (nullable = true)\n",
      " |-- genres: string (nullable = true)\n",
      "\n",
      "OneHotEncoder Example:\n",
      "root\n",
      " |-- movieId: string (nullable = true)\n",
      " |-- title: string (nullable = true)\n",
      " |-- genres: string (nullable = true)\n",
      " |-- movieIdNumber: integer (nullable = true)\n",
      " |-- movieIdVector: vector (nullable = true)\n",
      "\n",
      "+-------+--------------------+--------------------+-------------+-----------------+\n",
      "|movieId|               title|              genres|movieIdNumber|    movieIdVector|\n",
      "+-------+--------------------+--------------------+-------------+-----------------+\n",
      "|      1|    Toy Story (1995)|Adventure|Animati...|            1| (1001,[1],[1.0])|\n",
      "|      2|      Jumanji (1995)|Adventure|Childre...|            2| (1001,[2],[1.0])|\n",
      "|      3|Grumpier Old Men ...|      Comedy|Romance|            3| (1001,[3],[1.0])|\n",
      "|      4|Waiting to Exhale...|Comedy|Drama|Romance|            4| (1001,[4],[1.0])|\n",
      "|      5|Father of the Bri...|              Comedy|            5| (1001,[5],[1.0])|\n",
      "|      6|         Heat (1995)|Action|Crime|Thri...|            6| (1001,[6],[1.0])|\n",
      "|      7|      Sabrina (1995)|      Comedy|Romance|            7| (1001,[7],[1.0])|\n",
      "|      8| Tom and Huck (1995)|  Adventure|Children|            8| (1001,[8],[1.0])|\n",
      "|      9| Sudden Death (1995)|              Action|            9| (1001,[9],[1.0])|\n",
      "|     10|    GoldenEye (1995)|Action|Adventure|...|           10|(1001,[10],[1.0])|\n",
      "+-------+--------------------+--------------------+-------------+-----------------+\n",
      "only showing top 10 rows\n",
      "\n",
      "MultiHotEncoder Example:\n",
      "root\n",
      " |-- movieId: string (nullable = true)\n",
      " |-- genreIndexes: array (nullable = true)\n",
      " |    |-- element: integer (containsNull = false)\n",
      " |-- indexSize: integer (nullable = false)\n",
      " |-- vector: vector (nullable = true)\n",
      "\n",
      "+-------+------------+---------+--------------------+\n",
      "|movieId|genreIndexes|indexSize|              vector|\n",
      "+-------+------------+---------+--------------------+\n",
      "|    296|[1, 5, 0, 3]|       19|(19,[0,1,3,5],[1....|\n",
      "|    467|         [1]|       19|      (19,[1],[1.0])|\n",
      "|    675|   [4, 0, 3]|       19|(19,[0,3,4],[1.0,...|\n",
      "|    691|      [1, 2]|       19|(19,[1,2],[1.0,1.0])|\n",
      "|    829| [1, 10, 14]|       19|(19,[1,10,14],[1....|\n",
      "|    125|         [1]|       19|      (19,[1],[1.0])|\n",
      "|    451|   [0, 8, 2]|       19|(19,[0,2,8],[1.0,...|\n",
      "|    800|  [0, 8, 16]|       19|(19,[0,8,16],[1.0...|\n",
      "|    853|         [0]|       19|      (19,[0],[1.0])|\n",
      "|    944|         [0]|       19|      (19,[0],[1.0])|\n",
      "+-------+------------+---------+--------------------+\n",
      "only showing top 10 rows\n",
      "\n",
      "Numerical features Example:\n",
      "root\n",
      " |-- userId: string (nullable = true)\n",
      " |-- movieId: string (nullable = true)\n",
      " |-- rating: string (nullable = true)\n",
      " |-- timestamp: string (nullable = true)\n",
      "\n",
      "+------+-------+------+----------+\n",
      "|userId|movieId|rating| timestamp|\n",
      "+------+-------+------+----------+\n",
      "|     1|      2|   3.5|1112486027|\n",
      "|     1|     29|   3.5|1112484676|\n",
      "|     1|     32|   3.5|1112484819|\n",
      "|     1|     47|   3.5|1112484727|\n",
      "|     1|     50|   3.5|1112484580|\n",
      "|     1|    112|   3.5|1094785740|\n",
      "|     1|    151|   4.0|1094785734|\n",
      "|     1|    223|   4.0|1112485573|\n",
      "|     1|    253|   4.0|1112484940|\n",
      "|     1|    260|   4.0|1112484826|\n",
      "|     1|    293|   4.0|1112484703|\n",
      "|     1|    296|   4.0|1112484767|\n",
      "|     1|    318|   4.0|1112484798|\n",
      "|     1|    337|   3.5|1094785709|\n",
      "|     1|    367|   3.5|1112485980|\n",
      "|     1|    541|   4.0|1112484603|\n",
      "|     1|    589|   3.5|1112485557|\n",
      "|     1|    593|   3.5|1112484661|\n",
      "|     1|    653|   3.0|1094785691|\n",
      "|     1|    919|   3.5|1094785621|\n",
      "+------+-------+------+----------+\n",
      "only showing top 20 rows\n",
      "\n",
      "+-------+-----------+------------------+------------------+--------------------+\n",
      "|movieId|ratingCount|         avgRating|         ratingVar|        avgRatingVec|\n",
      "+-------+-----------+------------------+------------------+--------------------+\n",
      "|    296|      14616| 4.165606185002737|0.9615737413069363| [4.165606185002737]|\n",
      "|    467|        174|3.4367816091954024|1.5075410271742742|[3.4367816091954024]|\n",
      "|    829|        402|2.6243781094527363|1.4982072182727266|[2.6243781094527363]|\n",
      "|    691|        254|3.1161417322834644|1.0842838691606236|[3.1161417322834644]|\n",
      "|    675|          6|2.3333333333333335|0.6666666666666667|[2.3333333333333335]|\n",
      "|    125|        788| 3.713197969543147|0.8598255922703314| [3.713197969543147]|\n",
      "|    800|       1609|4.0447482908638905|0.8325734596130598|[4.0447482908638905]|\n",
      "|    944|        259|3.8262548262548264|0.8534165394630507|[3.8262548262548264]|\n",
      "|    853|         20|               3.5| 1.526315789473684|               [3.5]|\n",
      "|    451|        159|  3.00314465408805|0.7800533397022527|  [3.00314465408805]|\n",
      "+-------+-----------+------------------+------------------+--------------------+\n",
      "only showing top 10 rows\n",
      "\n",
      "+-------+-----------+------------------+------------------+--------------------+-----------------+--------------------+\n",
      "|movieId|ratingCount|         avgRating|         ratingVar|        avgRatingVec|ratingCountBucket|      scaleAvgRating|\n",
      "+-------+-----------+------------------+------------------+--------------------+-----------------+--------------------+\n",
      "|    296|      14616| 4.165606185002737|0.9615737413069363| [4.165606185002737]|             57.0|[0.9170998054196596]|\n",
      "|    467|        174|3.4367816091954024|1.5075410271742742|[3.4367816091954024]|             21.0|[0.7059538707722662]|\n",
      "|    829|        402|2.6243781094527363|1.4982072182727266|[2.6243781094527363]|             32.0|[0.4705944962973248]|\n",
      "|    691|        254|3.1161417322834644|1.0842838691606236|[3.1161417322834644]|             26.0|[0.6130620985364005]|\n",
      "|    675|          6|2.3333333333333335|0.6666666666666667|[2.3333333333333335]|              3.0|[0.38627664627161...|\n",
      "|    125|        788| 3.713197969543147|0.8598255922703314| [3.713197969543147]|             40.0|[0.7860337592595664]|\n",
      "|    800|       1609|4.0447482908638905|0.8325734596130598|[4.0447482908638905]|             47.0|[0.8820863689021069]|\n",
      "|    944|        259|3.8262548262548264|0.8534165394630507|[3.8262548262548264]|             27.0|[0.8187871768460151]|\n",
      "|    853|         20|               3.5| 1.526315789473684|               [3.5]|              8.0|[0.7242687117592825]|\n",
      "|    451|        159|  3.00314465408805|0.7800533397022527|  [3.00314465408805]|             21.0|[0.5803259992335382]|\n",
      "+-------+-----------+------------------+------------------+--------------------+-----------------+--------------------+\n",
      "only showing top 10 rows\n",
      "\n"
     ]
    }
   ],
   "source": [
    "# Read data and do feature engineering\n",
    "file_path = '/Users/linchenxiao/Python_Code/Spark/SparrowRecsys'\n",
    "movieResourcesPath = file_path + \"/sampledata/movies.csv\"\n",
    "movieSamples = spark.read.format('csv').option('header', 'true').load(movieResourcesPath)\n",
    "print(\"Raw Movie Samples:\")\n",
    "movieSamples.show(10)\n",
    "movieSamples.printSchema()\n",
    "print(\"OneHotEncoder Example:\")\n",
    "oneHotEncoderExample(movieSamples)\n",
    "print(\"MultiHotEncoder Example:\")\n",
    "multiHotEncoderExample(movieSamples)\n",
    "print(\"Numerical features Example:\")\n",
    "ratingsResourcesPath = file_path + \"/sampledata/ratings.csv\"\n",
    "ratingSamples = spark.read.format('csv').option('header', 'true').load(ratingsResourcesPath)\n",
    "ratingFeatures(ratingSamples)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 如何使用Spark 生成Embedding"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "from pyspark.sql.functions import *\n",
    "from pyspark.sql.types import *\n",
    "from pyspark.ml.feature import BucketedRandomProjectionLSH\n",
    "from pyspark.mllib.feature import Word2Vec\n",
    "from pyspark.ml.linalg import Vectors\n",
    "import random\n",
    "from collections import defaultdict\n",
    "import numpy as np\n",
    "from pyspark.sql import functions as F\n",
    "\n",
    "\n",
    "class UdfFunction:\n",
    "    @staticmethod\n",
    "    def sortF(movie_list, timestamp_list):\n",
    "        \"\"\"\n",
    "        sort by time and return the corresponding movie sequence\n",
    "        eg:\n",
    "            input: movie_list:[1,2,3]\n",
    "                   timestamp_list:[1112486027,1212546032,1012486033]\n",
    "            return [3,1,2]\n",
    "        \"\"\"\n",
    "        pairs = []\n",
    "        for m, t in zip(movie_list, timestamp_list):\n",
    "            pairs.append((m, t))\n",
    "        # sort by time\n",
    "        pairs = sorted(pairs, key=lambda x: x[1])\n",
    "        return [x[0] for x in pairs]\n",
    "\n",
    "\n",
    "def processItemSequence(spark, rawSampleDataPath):\n",
    "    \"\"\"处理raw data生成sequence数据\"\"\"\n",
    "    # rating data\n",
    "    ratingSamples = spark.read.format(\"csv\").option(\"header\", \"true\").load(rawSampleDataPath)\n",
    "    #     +------+-------+------+----------+\n",
    "    #     |userId|movieId|rating| timestamp|\n",
    "    #     +------+-------+------+----------+\n",
    "    #     |     1|      2|   3.5|1112486027|\n",
    "    #     |     1|     29|   3.5|1112484676|\n",
    "    #     |     1|     32|   3.5|1112484819|\n",
    "    # ratingSamples.show(5)\n",
    "    # ratingSamples.printSchema()\n",
    "    # 自己定义的sortF函数\n",
    "    sortUdf = udf(UdfFunction.sortF, ArrayType(StringType()))\n",
    "    # 对每个user生成sequences\n",
    "    userSeq = ratingSamples \\\n",
    "        .where(F.col(\"rating\") >= 3.5) \\\n",
    "        .groupBy(\"userId\") \\\n",
    "        .agg(sortUdf(F.collect_list(\"movieId\"), F.collect_list(\"timestamp\")).alias('movieIds')) \\\n",
    "        .withColumn(\"movieIdStr\", array_join(F.col(\"movieIds\"), \" \")) # 把所有id连成一个String，方便后续word2vec模型处理\n",
    "    # userSeq.select(\"userId\", \"movieIdStr\").show(10, truncate = False)\n",
    "    # 把序列数据筛选出来，丢掉其他过程数据\n",
    "    # 最后输出 List[List[item]]\n",
    "    # ['858', '50', '593', '457'],\n",
    "    # ['1', '25', '32', '6', '608', '52', '58', '26', '30', '103', '582', '588'],\n",
    "    return userSeq.select('movieIdStr').rdd.map(lambda x: x[0].split(' '))\n",
    "\n",
    "\n",
    "def embeddingLSH(spark, movieEmbMap):\n",
    "    movieEmbSeq = []\n",
    "    for key, embedding_list in movieEmbMap.items():\n",
    "        embedding_list = [np.float64(embedding) for embedding in embedding_list]\n",
    "        movieEmbSeq.append((key, Vectors.dense(embedding_list)))\n",
    "    movieEmbDF = spark.createDataFrame(movieEmbSeq).toDF(\"movieId\", \"emb\")\n",
    "    bucketProjectionLSH = BucketedRandomProjectionLSH(inputCol=\"emb\", outputCol=\"bucketId\", bucketLength=0.1,\n",
    "                                                      numHashTables=3)\n",
    "    bucketModel = bucketProjectionLSH.fit(movieEmbDF)\n",
    "    embBucketResult = bucketModel.transform(movieEmbDF)\n",
    "    print(\"movieId, emb, bucketId schema:\")\n",
    "    embBucketResult.printSchema()\n",
    "    print(\"movieId, emb, bucketId data result:\")\n",
    "    embBucketResult.show(10, truncate=False)\n",
    "    print(\"Approximately searching for 5 nearest neighbors of the sample embedding:\")\n",
    "    sampleEmb = Vectors.dense(0.795, 0.583, 1.120, 0.850, 0.174, -0.839, -0.0633, 0.249, 0.673, -0.237)\n",
    "    bucketModel.approxNearestNeighbors(movieEmbDF, sampleEmb, 5).show(truncate=False)\n",
    "\n",
    "\n",
    "def trainItem2vec(spark, samples, embLength, embOutputPath, saveToRedis, redisKeyPrefix):\n",
    "    word2vec = Word2Vec().setVectorSize(embLength).setWindowSize(5).setNumIterations(10)\n",
    "    model = word2vec.fit(samples)\n",
    "    # 查找物品158最近的20个物品\n",
    "    synonyms = model.findSynonyms(\"158\", 20)\n",
    "    for synonym, cosineSimilarity in synonyms:\n",
    "        print(synonym, cosineSimilarity)\n",
    "    embOutputDir = '/'.join(embOutputPath.split('/')[:-1])\n",
    "    if not os.path.exists(embOutputDir):\n",
    "        os.makedirs(embOutputDir)\n",
    "    with open(embOutputPath, 'w') as f:\n",
    "        # getVectors(): {item1: embeddings1, item2: embeddings2}\n",
    "        for movie_id in model.getVectors():\n",
    "            vectors = \" \".join([str(emb) for emb in model.getVectors()[movie_id]])\n",
    "            f.write(movie_id + \":\" + vectors + \"\\n\")\n",
    "    embeddingLSH(spark, model.getVectors())\n",
    "    return model\n",
    "\n",
    "\n",
    "def generate_pair(x):\n",
    "    # eg:\n",
    "    # watch sequence:['858', '50', '593', '457']\n",
    "    # return:[['858', '50'],['50', '593'],['593', '457']]\n",
    "    pairSeq = []\n",
    "    previousItem = ''\n",
    "    for item in x:\n",
    "        if not previousItem:\n",
    "            previousItem = item\n",
    "        else:\n",
    "            pairSeq.append((previousItem, item))\n",
    "            previousItem = item\n",
    "    return pairSeq\n",
    "\n",
    "\n",
    "def generateTransitionMatrix(samples):\n",
    "    \"\"\"生成物品转移矩阵\"\"\"\n",
    "    # 通过flatMap将观影序列打碎成一个一个影片对，按照时间排列\n",
    "    pairSamples = samples.flatMap(lambda x: generate_pair(x))\n",
    "    # count每个影片对的数量，用作weight\n",
    "    # i.e. {('858', '50'): 465,('50', '593'): 395,('593', '457'): 419}\n",
    "    pairCountMap = pairSamples.countByValue()\n",
    "    # 一共有多少个pair\n",
    "    pairTotalCount = 0\n",
    "    # 二维dict表示从物品1到物品2的次数， 从物品858到物品50一共465次\n",
    "    # i.e. '858': {'50': 465, '527': 445, '356': 105, '750': 154, '593': 195}\n",
    "    transitionCountMatrix = defaultdict(dict)\n",
    "    # Item作为起点一共出现多少次 {'858': 7395,'50': 8769,'593': 11216}\n",
    "    itemCountMap = defaultdict(int)\n",
    "    for key, cnt in pairCountMap.items():\n",
    "        key1, key2 = key\n",
    "        transitionCountMatrix[key1][key2] = cnt\n",
    "        itemCountMap[key1] += cnt\n",
    "        pairTotalCount += cnt\n",
    "    # 二维dict表示从物品1到物品2的概率， 从物品858到物品50的概率是0.0628\n",
    "    # {'858': {'50': 0.06288032454361055,'527': 0.06017579445571332,'356': 0.014198782961460446}\n",
    "    transitionMatrix = defaultdict(dict)\n",
    "    # Item作为起点的比例是多少\n",
    "    itemDistribution = defaultdict(dict)\n",
    "    for key1, transitionMap in transitionCountMatrix.items():\n",
    "        for key2, cnt in transitionMap.items():\n",
    "            transitionMatrix[key1][key2] = transitionCountMatrix[key1][key2] / itemCountMap[key1]\n",
    "    for itemid, cnt in itemCountMap.items():\n",
    "        itemDistribution[itemid] = cnt / pairTotalCount\n",
    "    return transitionMatrix, itemDistribution\n",
    "\n",
    "\n",
    "def oneRandomWalk(transitionMatrix, itemDistribution, sampleLength):\n",
    "    sample = []\n",
    "    # pick the first element\n",
    "    # 根据物品出现次数分布随机选择一个起始物品\n",
    "    randomDouble = random.random()\n",
    "    firstItem = \"\"\n",
    "    accumulateProb = 0.0\n",
    "    for item, prob in itemDistribution.items():\n",
    "        accumulateProb += prob\n",
    "        if accumulateProb >= randomDouble:\n",
    "            firstItem = item\n",
    "            break\n",
    "    sample.append(firstItem)\n",
    "    curElement = firstItem\n",
    "    i = 1\n",
    "    # 不停的游走直到sampleLength\n",
    "    while i < sampleLength:\n",
    "        if (curElement not in itemDistribution) or (curElement not in transitionMatrix):\n",
    "            break\n",
    "        # 对curElement来说的游走概率\n",
    "        probDistribution = transitionMatrix[curElement]\n",
    "        randomDouble = random.random()\n",
    "        accumulateProb = 0.0\n",
    "        for item, prob in probDistribution.items():\n",
    "            accumulateProb += prob\n",
    "            if accumulateProb >= randomDouble:\n",
    "                curElement = item\n",
    "                break\n",
    "        sample.append(curElement)\n",
    "        i += 1\n",
    "    return sample\n",
    "\n",
    "\n",
    "def randomWalk(transitionMatrix, itemDistribution, sampleCount, sampleLength):\n",
    "    samples = []\n",
    "    for i in range(sampleCount):\n",
    "        samples.append(oneRandomWalk(transitionMatrix, itemDistribution, sampleLength))\n",
    "    return samples\n",
    "\n",
    "\n",
    "def graphEmb(samples, spark, embLength, embOutputFilename, saveToRedis, redisKeyPrefix):\n",
    "    #生成物品转移概率矩阵\n",
    "    transitionMatrix, itemDistribution = generateTransitionMatrix(samples)\n",
    "    sampleCount = 20000\n",
    "    sampleLength = 10\n",
    "    newSamples = randomWalk(transitionMatrix, itemDistribution, sampleCount, sampleLength)\n",
    "    rddSamples = spark.sparkContext.parallelize(newSamples)\n",
    "    trainItem2vec(spark, rddSamples, embLength, embOutputFilename, saveToRedis, redisKeyPrefix)\n",
    "\n",
    "\n",
    "def generateUserEmb(spark, rawSampleDataPath, model, embLength, embOutputPath, saveToRedis, redisKeyPrefix):\n",
    "    ratingSamples = spark.read.format(\"csv\").option(\"header\", \"true\").load(rawSampleDataPath)\n",
    "    Vectors_list = []\n",
    "    for key, value in model.getVectors().items():\n",
    "        Vectors_list.append((key, list(value)))\n",
    "    fields = [\n",
    "        StructField('movieId', StringType(), False),\n",
    "        StructField('emb', ArrayType(FloatType()), False)\n",
    "    ]\n",
    "    schema = StructType(fields)\n",
    "    Vectors_df = spark.createDataFrame(Vectors_list, schema=schema)\n",
    "    ratingSamples = ratingSamples.join(Vectors_df, on='movieId', how='inner')\n",
    "    # 对每个user来说，把所有看过电影的embeddings加起来（map以每一行为输入，输出是（userId, embedding)， reduceByKey 把embedding的每个diemension加起来)\n",
    "    # map的第一个参数就是key\n",
    "    # foo = sc.parallelize([(1, ('a','b')), (2, ('c','d')), (1, ('x','y'))])\n",
    "    # foo.map(lambda (x,y): (x, [y])).reduceByKey(lambda p,q: p+q).collect()\n",
    "    # [(1, [('a', 'b'), ('x', 'y')]), (2, [('c', 'd')])]\n",
    "    result = ratingSamples.select('userId', 'emb').rdd.map(lambda x: (x[0], x[1])) \\\n",
    "        .reduceByKey(lambda a, b: [a[i] + b[i] for i in range(len(a))]).collect()\n",
    "    with open(embOutputPath, 'w') as f:\n",
    "        for row in result:\n",
    "            vectors = \" \".join([str(emb) for emb in row[1]])\n",
    "            f.write(row[0] + \":\" + vectors + \"\\n\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "48 0.9580522179603577\n",
      "256 0.9551143050193787\n",
      "31 0.9187507629394531\n",
      "186 0.8950363397598267\n",
      "355 0.8806580901145935\n",
      "168 0.8764886260032654\n",
      "520 0.846748948097229\n",
      "252 0.8462055921554565\n",
      "276 0.842883288860321\n",
      "277 0.8394060730934143\n",
      "552 0.835608720779419\n",
      "432 0.8306246995925903\n",
      "44 0.8212148547172546\n",
      "333 0.800767719745636\n",
      "455 0.7899411916732788\n",
      "169 0.7763938903808594\n",
      "2 0.7761274576187134\n",
      "236 0.7647883892059326\n",
      "237 0.7625499963760376\n",
      "370 0.7538554668426514\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/Users/linchenxiao/.pyenv/versions/3.8.5/lib/python3.8/site-packages/pandas/compat/__init__.py:97: UserWarning: Could not import the lzma module. Your installed Python is incomplete. Attempting to use lzma compression will result in a RuntimeError.\n",
      "  warnings.warn(msg)\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "movieId, emb, bucketId schema:\n",
      "root\n",
      " |-- movieId: string (nullable = true)\n",
      " |-- emb: vector (nullable = true)\n",
      " |-- bucketId: array (nullable = true)\n",
      " |    |-- element: vector (containsNull = true)\n",
      "\n",
      "movieId, emb, bucketId data result:\n",
      "+-------+-----------------------------------------------------------------------------------------------------------------------+------------------------+\n",
      "|movieId|emb                                                                                                                    |bucketId                |\n",
      "+-------+-----------------------------------------------------------------------------------------------------------------------+------------------------+\n",
      "|710    |[-1.3764266,0.42805645,0.4078089,0.06357183,1.5329667,0.76017064,1.1100068,1.2193089,0.6592338,-1.0196304]             |[[22.0], [14.0], [18.0]]|\n",
      "|205    |[-0.030348117,0.37577343,-0.27565956,-0.21543641,0.81982124,-0.32763448,-0.44681698,0.64313394,0.46586126,-0.124947704]|[[2.0], [0.0], [3.0]]   |\n",
      "|45     |[0.18040743,0.4820949,-0.29451472,-0.17339988,0.49046096,-0.53235936,-0.66926533,0.44017255,-0.15411055,-0.013712002]  |[[-4.0], [-2.0], [2.0]] |\n",
      "|515    |[0.8509096,0.4795693,-0.11639882,-0.1505484,0.33250937,-0.6065106,-0.5707916,0.12644954,0.27913365,-0.14113115]        |[[-3.0], [0.0], [-3.0]] |\n",
      "|574    |[-0.33851513,0.3124887,-0.61625504,-0.05533133,0.56652004,-0.53284097,-0.2138791,0.8032379,0.16950141,-0.24174254]     |[[-1.0], [-2.0], [5.0]] |\n",
      "|858    |[-0.34093723,0.05677535,0.27658203,0.36985588,-0.24673973,-0.49019232,0.304704,-0.147099,0.10777886,-0.3371425]        |[[0.0], [-2.0], [0.0]]  |\n",
      "|619    |[-0.8324198,0.2705492,0.19470498,0.064283654,1.0003952,0.43707854,0.7038905,0.7485974,0.5390741,-0.67490506]           |[[13.0], [8.0], [11.0]] |\n",
      "|507    |[0.045675516,0.46570733,-0.4308784,-0.11986841,0.22520933,-0.1982101,0.2576734,0.31531927,-0.013505689,0.4613265]      |[[1.0], [-1.0], [2.0]]  |\n",
      "|113    |[-0.4647893,0.6728471,0.14973702,0.09040027,0.8007195,0.37355715,0.7152607,0.81182534,1.061501,-0.6431248]             |[[15.0], [10.0], [11.0]]|\n",
      "|34     |[-0.23965165,-0.34824345,-0.17979291,-0.44053462,-0.36419684,-0.24638802,-0.19067748,0.27056497,0.48327976,-0.18816309]|[[-4.0], [-6.0], [-3.0]]|\n",
      "+-------+-----------------------------------------------------------------------------------------------------------------------+------------------------+\n",
      "only showing top 10 rows\n",
      "\n",
      "Approximately searching for 5 nearest neighbors of the sample embedding:\n",
      "+-------+-------------------------------------------------------------------------------------------------------------------+-----------------------+------------------+\n",
      "|movieId|emb                                                                                                                |bucketId               |distCol           |\n",
      "+-------+-------------------------------------------------------------------------------------------------------------------+-----------------------+------------------+\n",
      "|41     |[-0.0025742408,0.55390096,0.051914547,0.49474746,0.52204704,-0.405865,-0.011662789,0.40014642,0.64678,-0.21186158] |[[4.0], [4.0], [5.0]]  |1.4965146395808553|\n",
      "|628    |[3.4593075E-4,0.7851311,0.518726,0.2162024,0.028984834,0.03238921,0.18920064,0.37531862,0.22485241,-0.1254263]     |[[7.0], [5.0], [6.0]]  |1.5839533976075804|\n",
      "|73     |[0.17417039,0.20182326,0.37762687,-0.25573578,0.5627994,-0.36293498,0.042584997,-0.02204218,0.51909876,-0.28036278]|[[4.0], [1.0], [-2.0]] |1.6710891317005365|\n",
      "|428    |[0.26189494,0.28420895,0.041300513,0.04024306,0.4413432,-0.6809347,0.11609199,0.46933374,0.053657662,0.069775194]  |[[0.0], [-1.0], [0.0]] |1.6875003412419964|\n",
      "|290    |[0.22048642,0.12811172,-0.21847986,0.28423643,0.5296796,-0.82837677,0.19064462,0.41539153,0.14734319,-0.015390343] |[[-2.0], [-2.0], [0.0]]|1.786848263781805 |\n",
      "+-------+-------------------------------------------------------------------------------------------------------------------+-----------------------+------------------+\n",
      "\n"
     ]
    }
   ],
   "source": [
    "# Change to your own filepath\n",
    "# Item2Vec\n",
    "file_path = '/Users/linchenxiao/Python_Code/Spark/SparrowRecsys'\n",
    "rawSampleDataPath = file_path + \"/sampledata/ratings.csv\"\n",
    "embLength = 10\n",
    "samples = processItemSequence(spark, rawSampleDataPath)\n",
    "model = trainItem2vec(spark, samples, embLength,\n",
    "                      embOutputPath=file_path[7:] + \"/modeldata2/item2vecEmb.csv\", saveToRedis=False,\n",
    "                      redisKeyPrefix=\"i2vEmb\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "355 0.9262107014656067\n",
      "552 0.9148286581039429\n",
      "540 0.8829371929168701\n",
      "256 0.8801469206809998\n",
      "257 0.8549433350563049\n",
      "60 0.8544407486915588\n",
      "577 0.8227375745773315\n",
      "277 0.8206712603569031\n",
      "31 0.8096217513084412\n",
      "520 0.7996008396148682\n",
      "276 0.7969487309455872\n",
      "48 0.795081615447998\n",
      "170 0.7835773825645447\n",
      "44 0.7819278240203857\n",
      "543 0.7695910334587097\n",
      "168 0.7638594508171082\n",
      "303 0.7603774070739746\n",
      "673 0.7584818601608276\n",
      "585 0.756323516368866\n",
      "419 0.7448134422302246\n",
      "movieId, emb, bucketId schema:\n",
      "root\n",
      " |-- movieId: string (nullable = true)\n",
      " |-- emb: vector (nullable = true)\n",
      " |-- bucketId: array (nullable = true)\n",
      " |    |-- element: vector (containsNull = true)\n",
      "\n",
      "movieId, emb, bucketId data result:\n",
      "+-------+---------------------------------------------------------------------------------------------------------------------+------------------------+\n",
      "|movieId|emb                                                                                                                  |bucketId                |\n",
      "+-------+---------------------------------------------------------------------------------------------------------------------+------------------------+\n",
      "|710    |[-0.37830722,0.20833682,-1.173933,0.8119011,0.76264876,-0.34173644,-0.032273203,-0.09984239,-1.2792163,-0.4804649]   |[[-7.0], [1.0], [6.0]]  |\n",
      "|205    |[-0.37038937,-0.3142308,-0.45939192,-0.0042979037,0.78656954,0.07423798,0.07393751,0.47588223,-0.15116447,-0.2611661]|[[0.0], [1.0], [3.0]]   |\n",
      "|45     |[-0.14799242,-0.085146025,-0.6453846,-0.31925118,0.3606454,0.037755728,-0.37077066,0.5372427,0.09998301,-0.056565717]|[[-3.0], [-2.0], [2.0]] |\n",
      "|515    |[-0.09304517,-0.374771,-0.52506846,-0.20027992,0.19847398,-0.19780791,-0.3654365,0.46015862,0.10549895,0.27023914]   |[[-5.0], [-5.0], [-1.0]]|\n",
      "|574    |[0.09345369,0.36273113,-0.5086564,-0.14573741,0.568588,0.40400565,-0.045167983,0.40989664,-0.4044854,-0.116133735]   |[[2.0], [4.0], [5.0]]   |\n",
      "|858    |[0.13481984,0.47227284,0.18306583,-0.12329879,0.007283765,-0.0019150455,0.34555617,0.07837073,0.041792817,0.46787563]|[[4.0], [1.0], [0.0]]   |\n",
      "|619    |[-0.61659074,0.4492626,-1.1492265,0.6867159,0.58156085,0.18676646,0.12614231,0.4715879,-0.92793,-0.06723567]         |[[-2.0], [3.0], [12.0]] |\n",
      "|507    |[-0.0676849,0.31852445,-0.06144982,-0.32029977,0.20954362,-0.04013603,-0.14846505,0.26755404,-0.15037872,-0.13548173]|[[1.0], [0.0], [2.0]]   |\n",
      "|113    |[-0.5871682,0.47120723,-0.50254214,0.75017315,0.89579195,0.039027177,0.06510872,0.028888738,-0.9532209,-0.4322837]   |[[1.0], [5.0], [10.0]]  |\n",
      "|34     |[0.08330008,0.1544467,-0.34220967,-0.2298119,0.44139984,-0.14771448,0.30452463,-0.27192995,0.39188313,0.056816924]   |[[2.0], [-1.0], [-3.0]] |\n",
      "+-------+---------------------------------------------------------------------------------------------------------------------+------------------------+\n",
      "only showing top 10 rows\n",
      "\n",
      "Approximately searching for 5 nearest neighbors of the sample embedding:\n",
      "+-------+----------------------------------------------------------------------------------------------------------------+----------------------+------------------+\n",
      "|movieId|emb                                                                                                             |bucketId              |distCol           |\n",
      "+-------+----------------------------------------------------------------------------------------------------------------+----------------------+------------------+\n",
      "|349    |[-0.3629742,0.3112525,0.5232163,0.13609006,0.47565883,-0.48908302,-0.107047364,-0.32795978,0.41968998,-0.432575]|[[4.0], [0.0], [0.0]] |1.7120809115516913|\n",
      "|165    |[-0.21248332,0.5663644,0.6283844,-0.0705075,0.47496602,-0.6264004,-0.22295822,-0.3036071,0.05160898,-0.51786196]|[[3.0], [0.0], [0.0]] |1.74240467991671  |\n",
      "|595    |[-0.41628656,0.3112009,0.11588539,0.43072867,0.14931744,-0.5770524,0.17977157,-0.20229392,0.603157,0.0794644]   |[[2.0], [-2.0], [0.0]]|1.7782373308347708|\n",
      "|231    |[-0.26435062,0.35859808,0.23310381,0.1347479,0.37020913,-0.5111924,0.37236208,-0.44786388,0.3110412,-0.47469538]|[[4.0], [0.0], [-1.0]]|1.8653793724782846|\n",
      "|898    |[0.6670893,0.8039998,-0.5181294,0.6888535,0.16336197,-0.20512353,-0.18808648,0.58270764,0.48033175,0.22049583]  |[[0.0], [5.0], [5.0]] |1.8841196020671545|\n",
      "+-------+----------------------------------------------------------------------------------------------------------------+----------------------+------------------+\n",
      "\n"
     ]
    }
   ],
   "source": [
    "# Graph Embedding\n",
    "graphEmb(samples, spark, embLength, embOutputFilename=file_path[7:] + \"/modeldata2/itemGraphEmb.csv\",\n",
    "         saveToRedis=True, redisKeyPrefix=\"graphEmb\")\n",
    "generateUserEmb(spark, rawSampleDataPath, model, embLength,\n",
    "                embOutputPath=file_path[7:] + \"/modeldata2/userEmb.csv\", saveToRedis=False,\n",
    "                redisKeyPrefix=\"uEmb\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Feature Engineering for model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pyspark.sql as sql\n",
    "from pyspark.sql.functions import *\n",
    "from pyspark.sql.types import *\n",
    "from collections import defaultdict\n",
    "from pyspark.sql import functions as F\n",
    "\n",
    "NUMBER_PRECISION = 2\n",
    "\n",
    "\n",
    "def addSampleLabel(ratingSamples):\n",
    "    # 大于3.5的是positive\n",
    "    ratingSamples.show(5, truncate=False)\n",
    "    ratingSamples.printSchema()\n",
    "    sampleCount = ratingSamples.count()\n",
    "    ratingSamples.groupBy('rating').count().orderBy('rating').withColumn('percentage',\n",
    "                                                                         F.col('count') / sampleCount).show()\n",
    "    ratingSamples = ratingSamples.withColumn('label', when(F.col('rating') >= 3.5, 1).otherwise(0))\n",
    "    return ratingSamples\n",
    "\n",
    "\n",
    "def extractReleaseYearUdf(title):\n",
    "    # add realease year\n",
    "    if not title or len(title.strip()) < 6:\n",
    "        return 1990\n",
    "    else:\n",
    "        yearStr = title.strip()[-5:-1]\n",
    "    return int(yearStr)\n",
    "\n",
    "\n",
    "def addMovieFeatures(movieSamples, ratingSamplesWithLabel):\n",
    "    # add movie basic features\n",
    "    samplesWithMovies1 = ratingSamplesWithLabel.join(movieSamples, on=['movieId'], how='left')\n",
    "    # add releaseYear,title. Title example: Toy Story (1995)\n",
    "    samplesWithMovies2 = samplesWithMovies1.withColumn('releaseYear',\n",
    "                                                       udf(extractReleaseYearUdf, IntegerType())('title')) \\\n",
    "        .withColumn('title', udf(lambda x: x.strip()[:-6].strip(), StringType())('title')) \\\n",
    "        .drop('title')\n",
    "    # split genres, 取三个genres\n",
    "    samplesWithMovies3 = samplesWithMovies2.withColumn('movieGenre1', split(F.col('genres'), \"\\\\|\")[0]) \\\n",
    "        .withColumn('movieGenre2', split(F.col('genres'), \"\\\\|\")[1]) \\\n",
    "        .withColumn('movieGenre3', split(F.col('genres'), \"\\\\|\")[2])\n",
    "    # add rating features movieRatingCount|movieAvgRating|movieRatingStddev\n",
    "    movieRatingFeatures = samplesWithMovies3.groupBy('movieId').agg(F.count(F.lit(1)).alias('movieRatingCount'),\n",
    "                                                                    format_number(F.avg(F.col('rating')),NUMBER_PRECISION).alias('movieAvgRating'),\n",
    "                                                                    F.stddev(F.col('rating')).alias('movieRatingStddev')).fillna(0) \\\n",
    "        .withColumn('movieRatingStddev', format_number(F.col('movieRatingStddev'), NUMBER_PRECISION))\n",
    "    # join movie rating features\n",
    "    samplesWithMovies4 = samplesWithMovies3.join(movieRatingFeatures, on=['movieId'], how='left')\n",
    "    samplesWithMovies4.printSchema()\n",
    "    samplesWithMovies4.show(5, truncate=False)\n",
    "    return samplesWithMovies4\n",
    "\n",
    "\n",
    "def extractGenres(genres_list):\n",
    "    '''\n",
    "    pass in a list which format like [\"Action|Adventure|Sci-Fi|Thriller\", \"Crime|Horror|Thriller\"]\n",
    "    count by each genre，return genre_list in reverse order\n",
    "    eg:\n",
    "    (('Thriller',2),('Action',1),('Sci-Fi',1),('Horror', 1), ('Adventure',1),('Crime',1))\n",
    "    return:['Thriller','Action','Sci-Fi','Horror','Adventure','Crime']\n",
    "    '''\n",
    "    genres_dict = defaultdict(int)\n",
    "    for genres in genres_list:\n",
    "        for genre in genres.split('|'):\n",
    "            genres_dict[genre] += 1\n",
    "    sortedGenres = sorted(genres_dict.items(), key=lambda x: x[1], reverse=True)\n",
    "    return [x[0] for x in sortedGenres]\n",
    "\n",
    "\n",
    "def addUserFeatures(samplesWithMovieFeatures):\n",
    "    extractGenresUdf = udf(extractGenres, ArrayType(StringType()))\n",
    "    # Window function防止未来信息\n",
    "    # userRatedMovie1, 对这个（用户，电影）来说，上一部看的电影\n",
    "    # userGenres用户喜欢的电影类型，是一个List. i.e. [crime, drama]\n",
    "    # userGenres1,2,3,4,5。没有那么多的时候用null补\n",
    "    samplesWithUserFeatures = samplesWithMovieFeatures \\\n",
    "        .withColumn('userPositiveHistory',\n",
    "                    F.collect_list(when(F.col('label') == 1, F.col('movieId')).otherwise(F.lit(None))).over(\n",
    "                        sql.Window.partitionBy(\"userId\").orderBy(F.col(\"timestamp\")).rowsBetween(-100, -1))) \\\n",
    "        .withColumn(\"userPositiveHistory\", reverse(F.col(\"userPositiveHistory\"))) \\\n",
    "        .withColumn('userRatedMovie1', F.col('userPositiveHistory')[0]) \\\n",
    "        .withColumn('userRatedMovie2', F.col('userPositiveHistory')[1]) \\\n",
    "        .withColumn('userRatedMovie3', F.col('userPositiveHistory')[2]) \\\n",
    "        .withColumn('userRatedMovie4', F.col('userPositiveHistory')[3]) \\\n",
    "        .withColumn('userRatedMovie5', F.col('userPositiveHistory')[4]) \\\n",
    "        .withColumn('userRatingCount',\n",
    "                    F.count(F.lit(1)).over(sql.Window.partitionBy('userId').orderBy('timestamp').rowsBetween(-100, -1))) \\\n",
    "        .withColumn('userAvgReleaseYear', F.avg(F.col('releaseYear')).over(\n",
    "        sql.Window.partitionBy('userId').orderBy('timestamp').rowsBetween(-100, -1)).cast(IntegerType())) \\\n",
    "        .withColumn('userReleaseYearStddev', F.stddev(F.col(\"releaseYear\")).over(\n",
    "        sql.Window.partitionBy('userId').orderBy('timestamp').rowsBetween(-100, -1))) \\\n",
    "        .withColumn(\"userAvgRating\", format_number(\n",
    "        F.avg(F.col(\"rating\")).over(sql.Window.partitionBy('userId').orderBy('timestamp').rowsBetween(-100, -1)),\n",
    "        NUMBER_PRECISION)) \\\n",
    "        .withColumn(\"userRatingStddev\", F.stddev(F.col(\"rating\")).over(\n",
    "        sql.Window.partitionBy('userId').orderBy('timestamp').rowsBetween(-100, -1))) \\\n",
    "        .withColumn(\"userGenres\", extractGenresUdf(\n",
    "        F.collect_list(when(F.col('label') == 1, F.col('genres')).otherwise(F.lit(None))).over(\n",
    "            sql.Window.partitionBy('userId').orderBy('timestamp').rowsBetween(-100, -1)))) \\\n",
    "        .withColumn(\"userRatingStddev\", format_number(F.col(\"userRatingStddev\"), NUMBER_PRECISION)) \\\n",
    "        .withColumn(\"userReleaseYearStddev\", format_number(F.col(\"userReleaseYearStddev\"), NUMBER_PRECISION)) \\\n",
    "        .withColumn(\"userGenre1\", F.col(\"userGenres\")[0]) \\\n",
    "        .withColumn(\"userGenre2\", F.col(\"userGenres\")[1]) \\\n",
    "        .withColumn(\"userGenre3\", F.col(\"userGenres\")[2]) \\\n",
    "        .withColumn(\"userGenre4\", F.col(\"userGenres\")[3]) \\\n",
    "        .withColumn(\"userGenre5\", F.col(\"userGenres\")[4]) \\\n",
    "        .drop(\"genres\", \"userGenres\", \"userPositiveHistory\") \\\n",
    "        .filter(F.col(\"userRatingCount\") > 1)\n",
    "    samplesWithUserFeatures.printSchema()\n",
    "    samplesWithUserFeatures.show(10)\n",
    "    samplesWithUserFeatures.filter(samplesWithMovieFeatures['userId'] == 1).orderBy(F.col('timestamp').asc()).show(\n",
    "        truncate=False)\n",
    "    return samplesWithUserFeatures\n",
    "\n",
    "\n",
    "def splitAndSaveTrainingTestSamples(samplesWithUserFeatures, file_path):\n",
    "    smallSamples = samplesWithUserFeatures.sample(0.1)\n",
    "    training, test = smallSamples.randomSplit((0.8, 0.2))\n",
    "    trainingSavePath = file_path + '/trainingSamples'\n",
    "    testSavePath = file_path + '/testSamples'\n",
    "    training.repartition(1).write.option(\"header\", \"true\").mode('overwrite') \\\n",
    "        .csv(trainingSavePath)\n",
    "    test.repartition(1).write.option(\"header\", \"true\").mode('overwrite') \\\n",
    "        .csv(testSavePath)\n",
    "\n",
    "\n",
    "def splitAndSaveTrainingTestSamplesByTimeStamp(samplesWithUserFeatures, file_path):\n",
    "    smallSamples = samplesWithUserFeatures.sample(0.1).withColumn(\"timestampLong\", F.col(\"timestamp\").cast(LongType()))\n",
    "    quantile = smallSamples.stat.approxQuantile(\"timestampLong\", [0.8], 0.05)\n",
    "    splitTimestamp = quantile[0]\n",
    "    training = smallSamples.where(F.col(\"timestampLong\") <= splitTimestamp).drop(\"timestampLong\")\n",
    "    test = smallSamples.where(F.col(\"timestampLong\") > splitTimestamp).drop(\"timestampLong\")\n",
    "    trainingSavePath = file_path + '/trainingSamples'\n",
    "    testSavePath = file_path + '/testSamples'\n",
    "    training.repartition(1).write.option(\"header\", \"true\").mode('overwrite') \\\n",
    "        .csv(trainingSavePath)\n",
    "    test.repartition(1).write.option(\"header\", \"true\").mode('overwrite') \\\n",
    "        .csv(testSavePath)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "+------+-------+------+----------+\n",
      "|userId|movieId|rating|timestamp |\n",
      "+------+-------+------+----------+\n",
      "|1     |2      |3.5   |1112486027|\n",
      "|1     |29     |3.5   |1112484676|\n",
      "|1     |32     |3.5   |1112484819|\n",
      "|1     |47     |3.5   |1112484727|\n",
      "|1     |50     |3.5   |1112484580|\n",
      "+------+-------+------+----------+\n",
      "only showing top 5 rows\n",
      "\n",
      "root\n",
      " |-- userId: string (nullable = true)\n",
      " |-- movieId: string (nullable = true)\n",
      " |-- rating: string (nullable = true)\n",
      " |-- timestamp: string (nullable = true)\n",
      "\n",
      "+------+------+--------------------+\n",
      "|rating| count|          percentage|\n",
      "+------+------+--------------------+\n",
      "|   0.5|  9788|0.008375561978987506|\n",
      "|   1.0| 45018| 0.03852176636392108|\n",
      "|   1.5| 11794|0.010092090108314123|\n",
      "|   2.0| 87084| 0.07451751526135553|\n",
      "|   2.5| 34269|0.029323879593167432|\n",
      "|   3.0|323616| 0.27691723185451783|\n",
      "|   3.5| 74376| 0.06364331811904114|\n",
      "|   4.0|324804|  0.2779337998593234|\n",
      "|   4.5| 53388| 0.04568395003414231|\n",
      "|   5.0|204501| 0.17499088682722966|\n",
      "+------+------+--------------------+\n",
      "\n",
      "+------+-------+------+----------+-----+\n",
      "|userId|movieId|rating|timestamp |label|\n",
      "+------+-------+------+----------+-----+\n",
      "|1     |2      |3.5   |1112486027|1    |\n",
      "|1     |29     |3.5   |1112484676|1    |\n",
      "|1     |32     |3.5   |1112484819|1    |\n",
      "|1     |47     |3.5   |1112484727|1    |\n",
      "|1     |50     |3.5   |1112484580|1    |\n",
      "|1     |112    |3.5   |1094785740|1    |\n",
      "|1     |151    |4.0   |1094785734|1    |\n",
      "|1     |223    |4.0   |1112485573|1    |\n",
      "|1     |253    |4.0   |1112484940|1    |\n",
      "|1     |260    |4.0   |1112484826|1    |\n",
      "+------+-------+------+----------+-----+\n",
      "only showing top 10 rows\n",
      "\n",
      "root\n",
      " |-- movieId: string (nullable = true)\n",
      " |-- userId: string (nullable = true)\n",
      " |-- rating: string (nullable = true)\n",
      " |-- timestamp: string (nullable = true)\n",
      " |-- label: integer (nullable = false)\n",
      " |-- genres: string (nullable = true)\n",
      " |-- releaseYear: integer (nullable = true)\n",
      " |-- movieGenre1: string (nullable = true)\n",
      " |-- movieGenre2: string (nullable = true)\n",
      " |-- movieGenre3: string (nullable = true)\n",
      " |-- movieRatingCount: long (nullable = true)\n",
      " |-- movieAvgRating: string (nullable = true)\n",
      " |-- movieRatingStddev: string (nullable = true)\n",
      "\n",
      "+-------+------+------+----------+-----+---------------------------+-----------+-----------+-----------+-----------+----------------+--------------+-----------------+\n",
      "|movieId|userId|rating|timestamp |label|genres                     |releaseYear|movieGenre1|movieGenre2|movieGenre3|movieRatingCount|movieAvgRating|movieRatingStddev|\n",
      "+-------+------+------+----------+-----+---------------------------+-----------+-----------+-----------+-----------+----------------+--------------+-----------------+\n",
      "|296    |1     |4.0   |1112484767|1    |Comedy|Crime|Drama|Thriller|1994       |Comedy     |Crime      |Drama      |14616           |4.17          |0.98             |\n",
      "|296    |8     |5.0   |833973081 |1    |Comedy|Crime|Drama|Thriller|1994       |Comedy     |Crime      |Drama      |14616           |4.17          |0.98             |\n",
      "|296    |11    |3.5   |1230858799|1    |Comedy|Crime|Drama|Thriller|1994       |Comedy     |Crime      |Drama      |14616           |4.17          |0.98             |\n",
      "|296    |13    |5.0   |849082366 |1    |Comedy|Crime|Drama|Thriller|1994       |Comedy     |Crime      |Drama      |14616           |4.17          |0.98             |\n",
      "|296    |15    |3.0   |840206642 |0    |Comedy|Crime|Drama|Thriller|1994       |Comedy     |Crime      |Drama      |14616           |4.17          |0.98             |\n",
      "+-------+------+------+----------+-----+---------------------------+-----------+-----------+-----------+-----------+----------------+--------------+-----------------+\n",
      "only showing top 5 rows\n",
      "\n",
      "root\n",
      " |-- movieId: string (nullable = true)\n",
      " |-- userId: string (nullable = true)\n",
      " |-- rating: string (nullable = true)\n",
      " |-- timestamp: string (nullable = true)\n",
      " |-- label: integer (nullable = false)\n",
      " |-- releaseYear: integer (nullable = true)\n",
      " |-- movieGenre1: string (nullable = true)\n",
      " |-- movieGenre2: string (nullable = true)\n",
      " |-- movieGenre3: string (nullable = true)\n",
      " |-- movieRatingCount: long (nullable = true)\n",
      " |-- movieAvgRating: string (nullable = true)\n",
      " |-- movieRatingStddev: string (nullable = true)\n",
      " |-- userRatedMovie1: string (nullable = true)\n",
      " |-- userRatedMovie2: string (nullable = true)\n",
      " |-- userRatedMovie3: string (nullable = true)\n",
      " |-- userRatedMovie4: string (nullable = true)\n",
      " |-- userRatedMovie5: string (nullable = true)\n",
      " |-- userRatingCount: long (nullable = false)\n",
      " |-- userAvgReleaseYear: integer (nullable = true)\n",
      " |-- userReleaseYearStddev: string (nullable = true)\n",
      " |-- userAvgRating: string (nullable = true)\n",
      " |-- userRatingStddev: string (nullable = true)\n",
      " |-- userGenre1: string (nullable = true)\n",
      " |-- userGenre2: string (nullable = true)\n",
      " |-- userGenre3: string (nullable = true)\n",
      " |-- userGenre4: string (nullable = true)\n",
      " |-- userGenre5: string (nullable = true)\n",
      "\n",
      "+-------+------+------+---------+-----+-----------+-----------+-----------+-----------+----------------+--------------+-----------------+---------------+---------------+---------------+---------------+---------------+---------------+------------------+---------------------+-------------+----------------+----------+----------+----------+----------+----------+\n",
      "|movieId|userId|rating|timestamp|label|releaseYear|movieGenre1|movieGenre2|movieGenre3|movieRatingCount|movieAvgRating|movieRatingStddev|userRatedMovie1|userRatedMovie2|userRatedMovie3|userRatedMovie4|userRatedMovie5|userRatingCount|userAvgReleaseYear|userReleaseYearStddev|userAvgRating|userRatingStddev|userGenre1|userGenre2|userGenre3|userGenre4|userGenre5|\n",
      "+-------+------+------+---------+-----+-----------+-----------+-----------+-----------+----------------+--------------+-----------------+---------------+---------------+---------------+---------------+---------------+---------------+------------------+---------------------+-------------+----------------+----------+----------+----------+----------+----------+\n",
      "|    514| 10096|   3.0|954365410|    0|       1994|     Comedy|       null|       null|            1038|          3.50|             0.86|            858|           null|           null|           null|           null|              2|              1982|                14.85|         3.50|            0.71|     Crime|     Drama|      null|      null|      null|\n",
      "|    608| 10096|   3.0|954365515|    0|       1996|     Comedy|      Crime|      Drama|            9505|          4.09|             0.93|            858|           null|           null|           null|           null|              3|              1986|                12.42|         3.33|            0.58|     Crime|     Drama|      null|      null|      null|\n",
      "|     50| 10096|   5.0|954365515|    1|       1995|      Crime|    Mystery|   Thriller|           10221|          4.35|             0.75|            858|           null|           null|           null|           null|              4|              1988|                11.24|         3.25|            0.50|     Crime|     Drama|      null|      null|      null|\n",
      "|    593| 10096|   4.0|954365552|    1|       1991|      Crime|     Horror|   Thriller|           13692|          4.18|             0.85|             50|            858|           null|           null|           null|              5|              1990|                10.12|         3.60|            0.89|     Crime|     Drama|   Mystery|  Thriller|      null|\n",
      "|     25| 10096|   2.0|954365571|    0|       1995|      Drama|    Romance|       null|            4684|          3.69|             1.04|            593|             50|            858|           null|           null|              6|              1990|                 9.06|         3.67|            0.82|     Crime|  Thriller|     Drama|   Mystery|    Horror|\n",
      "|    457| 10096|   5.0|954365571|    1|       1993|   Thriller|       null|       null|           10736|          3.97|             0.78|            593|             50|            858|           null|           null|              7|              1990|                 8.47|         3.43|            0.98|     Crime|  Thriller|     Drama|   Mystery|    Horror|\n",
      "|    541| 10096|   3.0|954365664|    0|       1982|     Action|     Sci-Fi|   Thriller|            6635|          4.14|             0.88|            457|            593|             50|            858|           null|              8|              1991|                 7.88|         3.62|            1.06|     Crime|  Thriller|     Drama|   Mystery|    Horror|\n",
      "|     32| 10351|   4.0|851791255|    1|       1995|    Mystery|     Sci-Fi|   Thriller|            9694|          3.89|             0.86|             25|           null|           null|           null|           null|              2|              1995|                 0.71|         3.50|            0.71|     Drama|   Romance|      null|      null|      null|\n",
      "|      1| 10351|   4.0|851791255|    1|       1995|  Adventure|  Animation|   Children|           10759|          3.91|             0.89|             32|             25|           null|           null|           null|              3|              1995|                 0.58|         3.67|            0.58|     Drama|   Romance|   Mystery|    Sci-Fi|  Thriller|\n",
      "|      6| 10351|   4.0|851791281|    1|       1995|     Action|      Crime|   Thriller|            5245|          3.84|             0.86|              1|             32|             25|           null|           null|              4|              1995|                 0.50|         3.75|            0.50|     Drama|   Romance|   Mystery|    Sci-Fi|  Thriller|\n",
      "+-------+------+------+---------+-----+-----------+-----------+-----------+-----------+----------------+--------------+-----------------+---------------+---------------+---------------+---------------+---------------+---------------+------------------+---------------------+-------------+----------------+----------+----------+----------+----------+----------+\n",
      "only showing top 10 rows\n",
      "\n",
      "+-------+------+------+----------+-----+-----------+-----------+-----------+-----------+----------------+--------------+-----------------+---------------+---------------+---------------+---------------+---------------+---------------+------------------+---------------------+-------------+----------------+----------+----------+----------+----------+----------+\n",
      "|movieId|userId|rating|timestamp |label|releaseYear|movieGenre1|movieGenre2|movieGenre3|movieRatingCount|movieAvgRating|movieRatingStddev|userRatedMovie1|userRatedMovie2|userRatedMovie3|userRatedMovie4|userRatedMovie5|userRatingCount|userAvgReleaseYear|userReleaseYearStddev|userAvgRating|userRatingStddev|userGenre1|userGenre2|userGenre3|userGenre4|userGenre5|\n",
      "+-------+------+------+----------+-----+-----------+-----------+-----------+-----------+----------------+--------------+-----------------+---------------+---------------+---------------+---------------+---------------+---------------+------------------+---------------------+-------------+----------------+----------+----------+----------+----------+----------+\n",
      "|653    |1     |3.0   |1094785691|0    |1996       |Action     |Adventure  |Fantasy    |3148            |3.22          |0.97             |919            |924            |null           |null           |null           |2              |1953              |20.51                |3.50         |0.00            |Adventure |Drama     |Sci-Fi    |Children  |Fantasy   |\n",
      "|337    |1     |3.5   |1094785709|1    |1993       |Drama      |null       |null       |3889            |3.76          |0.86             |919            |924            |null           |null           |null           |3              |1967              |28.50                |3.33         |0.29            |Adventure |Drama     |Sci-Fi    |Children  |Fantasy   |\n",
      "|151    |1     |4.0   |1094785734|1    |1995       |Action     |Drama      |Romance    |2767            |3.57          |0.92             |337            |919            |924            |null           |null           |4              |1974              |26.50                |3.38         |0.25            |Adventure |Drama     |Sci-Fi    |Children  |Fantasy   |\n",
      "|112    |1     |3.5   |1094785740|1    |1995       |Action     |Adventure  |Comedy     |2577            |3.41          |0.95             |151            |337            |919            |924            |null           |5              |1978              |24.79                |3.50         |0.35            |Drama     |Adventure |Sci-Fi    |Children  |Fantasy   |\n",
      "|50     |1     |3.5   |1112484580|1    |1995       |Crime      |Mystery    |Thriller   |10221           |4.35          |0.75             |112            |151            |337            |919            |924            |6              |1981              |23.21                |3.50         |0.32            |Adventure |Drama     |Action    |Sci-Fi    |Children  |\n",
      "|541    |1     |4.0   |1112484603|1    |1982       |Action     |Sci-Fi     |Thriller   |6635            |4.14          |0.88             |50             |112            |151            |337            |919            |7              |1983              |21.84                |3.50         |0.29            |Adventure |Drama     |Action    |Crime     |Sci-Fi    |\n",
      "|593    |1     |3.5   |1112484661|1    |1991       |Crime      |Horror     |Thriller   |13692           |4.18          |0.85             |541            |50             |112            |151            |337            |8              |1982              |20.22                |3.56         |0.32            |Adventure |Drama     |Action    |Sci-Fi    |Crime     |\n",
      "|29     |1     |3.5   |1112484676|1    |1995       |Adventure  |Drama      |Fantasy    |1859            |3.92          |0.97             |593            |541            |50             |112            |151            |9              |1983              |19.11                |3.56         |0.30            |Adventure |Drama     |Action    |Crime     |Thriller  |\n",
      "|293    |1     |4.0   |1112484703|1    |1994       |Action     |Crime      |Drama      |5587            |4.06          |0.81             |29             |593            |541            |50             |112            |10             |1984              |18.36                |3.55         |0.28            |Adventure |Drama     |Sci-Fi    |Action    |Crime     |\n",
      "|47     |1     |3.5   |1112484727|1    |1995       |Mystery    |Thriller   |null       |9335            |4.06          |0.87             |293            |29             |593            |541            |50             |11             |1985              |17.64                |3.59         |0.30            |Drama     |Adventure |Action    |Crime     |Thriller  |\n",
      "|296    |1     |4.0   |1112484767|1    |1994       |Comedy     |Crime      |Drama      |14616           |4.17          |0.98             |47             |293            |29             |593            |541            |12             |1986              |17.03                |3.58         |0.29            |Drama     |Thriller  |Adventure |Action    |Crime     |\n",
      "|318    |1     |4.0   |1112484798|1    |1994       |Crime      |Drama      |null       |13826           |4.45          |0.72             |296            |47             |293            |29             |593            |13             |1987              |16.43                |3.62         |0.30            |Drama     |Thriller  |Crime     |Adventure |Action    |\n",
      "|32     |1     |3.5   |1112484819|1    |1995       |Mystery    |Sci-Fi     |Thriller   |9694            |3.89          |0.86             |318            |296            |47             |293            |29             |14             |1987              |15.90                |3.64         |0.31            |Drama     |Crime     |Thriller  |Adventure |Action    |\n",
      "|260    |1     |4.0   |1112484826|1    |1977       |Action     |Adventure  |Sci-Fi     |11958           |4.19          |0.91             |32             |318            |296            |47             |293            |15             |1988              |15.44                |3.63         |0.30            |Drama     |Thriller  |Crime     |Adventure |Sci-Fi    |\n",
      "|253    |1     |4.0   |1112484940|1    |1994       |Drama      |Horror     |null       |6010            |3.50          |0.94             |260            |32             |318            |296            |47             |16             |1987              |15.17                |3.66         |0.30            |Drama     |Thriller  |Crime     |Adventure |Sci-Fi    |\n",
      "|589    |1     |3.5   |1112485557|1    |1991       |Action     |Sci-Fi     |null       |11483           |3.94          |0.90             |253            |260            |32             |318            |296            |17             |1987              |14.78                |3.68         |0.30            |Drama     |Thriller  |Crime     |Adventure |Sci-Fi    |\n",
      "|223    |1     |4.0   |1112485573|1    |1994       |Comedy     |null       |null       |5102            |3.86          |0.97             |589            |253            |260            |32             |318            |18             |1987              |14.35                |3.67         |0.30            |Drama     |Thriller  |Sci-Fi    |Action    |Crime     |\n",
      "|367    |1     |3.5   |1112485980|1    |1994       |Action     |Comedy     |Crime      |7577            |3.16          |1.00             |223            |589            |253            |260            |32             |19             |1988              |14.02                |3.68         |0.30            |Drama     |Thriller  |Sci-Fi    |Action    |Crime     |\n",
      "|2      |1     |3.5   |1112486027|1    |1995       |Adventure  |Children   |Fantasy    |4853            |3.21          |0.96             |367            |223            |589            |253            |260            |20             |1988              |13.71                |3.67         |0.29            |Drama     |Action    |Crime     |Thriller  |Sci-Fi    |\n",
      "+-------+------+------+----------+-----+-----------+-----------+-----------+-----------+----------------+--------------+-----------------+---------------+---------------+---------------+---------------+---------------+---------------+------------------+---------------------+-------------+----------------+----------+----------+----------+----------+----------+\n",
      "\n"
     ]
    }
   ],
   "source": [
    "file_path = '/Users/linchenxiao/Python_Code/Spark/SparrowRecsys'\n",
    "movieResourcesPath = file_path + \"/sampledata/movies.csv\"\n",
    "ratingsResourcesPath = file_path + \"/sampledata/ratings.csv\"\n",
    "movieSamples = spark.read.format('csv').option('header', 'true').load(movieResourcesPath)\n",
    "ratingSamples = spark.read.format('csv').option('header', 'true').load(ratingsResourcesPath)\n",
    "ratingSamplesWithLabel = addSampleLabel(ratingSamples)\n",
    "ratingSamplesWithLabel.show(10, truncate=False)\n",
    "samplesWithMovieFeatures = addMovieFeatures(movieSamples, ratingSamplesWithLabel)\n",
    "samplesWithUserFeatures = addUserFeatures(samplesWithMovieFeatures)\n",
    "# save samples as csv format\n",
    "splitAndSaveTrainingTestSamples(samplesWithUserFeatures, file_path + \"/sampledata\")\n",
    "# splitAndSaveTrainingTestSamplesByTimeStamp(samplesWithUserFeatures, file_path + \"/webroot/sampledata\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}