Week4-Spark-SQL-1-Couting-mutual-friends.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "#execfile(os.path.join(os.environ[\"SPARK_HOME\"], 'python/pyspark/shell.py'))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "from pyspark.sql import SparkSession\n",
    "sparkSession = SparkSession.builder.enableHiveSupport().master(\"local\").getOrCreate()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "graphPath = \"/data/graphDFSample\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "from pyspark.sql.functions import udf\n",
    "import pyspark.sql.types as t\n",
    "\n",
    "def emit_all_pairs(friends_array):\n",
    "    ret = []\n",
    "    friends_len = len(friends_array)\n",
    "    for i in range(friends_len):\n",
    "        for j in range(i+1,friends_len):\n",
    "            ret.append((friends_array[i], friends_array[j]))\n",
    "    return ret\n",
    "\n",
    "pair_schema = t.StructType([\n",
    "    t.StructField(\"u1\", t.IntegerType(), False),\n",
    "    t.StructField(\"u2\", t.IntegerType(), False)\n",
    "])\n",
    "\n",
    "emit_all_pairs_udf = udf(emit_all_pairs, t.ArrayType(pair_schema))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "from pyspark.sql.functions import explode, collect_list, size, col, row_number, sort_array, count, desc, size\n",
    "from pyspark.sql import Window\n",
    "\n",
    "reversedGraph = sparkSession.read.parquet(graphPath) \\\n",
    "    .withColumn(\"friend\", explode('friends')) \\\n",
    "    .groupBy(\"friend\") \\\n",
    "    .agg(collect_list(\"user\").alias(\"users\")) \\\n",
    "    .withColumn(\"users\", sort_array('users')) \\\n",
    "    .where(size(\"users\") > 1)\n",
    "                \n",
    "    \n",
    "pairs = reversedGraph.withColumn(\"pairs\", emit_all_pairs_udf(\"users\")) \\\n",
    "    .withColumn(\"pair\", explode(\"pairs\")) \\\n",
    "    .groupBy(\"pair\") \\\n",
    "    .agg(count(\"pair\").alias(\"pair_count\"))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "results = pairs.select(col(\"pair_count\"), \"pair.*\") \\\n",
    "           .orderBy(desc(\"pair_count\"), desc(\"u1\"), desc(\"u2\")) \\\n",
    "           .limit(49) \\\n",
    "          .collect()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "3206 27967558 42973992\n",
      "3130 20158643 42973992\n",
      "3066 22582764 42973992\n",
      "3044 21864412 51640390\n",
      "3021 17139850 51640390\n",
      "3010 14985079 51640390\n",
      "2970 17139850 21864412\n",
      "2913 20158643 27967558\n",
      "2903 22280814 51151280\n",
      "2870 23848749 51640390\n",
      "2855 20158643 22582764\n",
      "2849 20158643 44996025\n",
      "2846 22280814 42973992\n",
      "2784 21864412 23848749\n",
      "2779 31964081 51640390\n",
      "2776 39205988 51640390\n",
      "2754 17139850 23848749\n",
      "2749 22582764 27967558\n",
      "2728 50561859 51640390\n",
      "2724 15485897 51640390\n",
      "2700 28135661 42973992\n",
      "2655 22280814 27967558\n",
      "2653 42973992 43548989\n",
      "2639 26755857 51640390\n",
      "2621 14635589 51640390\n",
      "2608 15485897 17139850\n",
      "2606 17139850 26755857\n",
      "2601 21864412 39205988\n",
      "2600 8406745 51640390\n",
      "2599 37735419 51640390\n",
      "2597 20158643 28135661\n",
      "2585 40003405 42973992\n",
      "2585 21864412 31964081\n",
      "2581 27967558 43548989\n",
      "2579 23848749 31964081\n",
      "2578 27967558 28135661\n",
      "2578 15485897 21864412\n",
      "2577 42973992 64755069\n",
      "2574 51151280 57077210\n",
      "2573 20158643 43548989\n",
      "2566 21864412 26755857\n",
      "2564 22280814 64755069\n",
      "2561 42973992 44996025\n",
      "2556 17139850 39205988\n",
      "2543 23848749 39205988\n",
      "2521 17139850 31964081\n",
      "2515 27967558 44996025\n",
      "2506 41629539 51640390\n",
      "2505 51151280 64755069\n"
     ]
    }
   ],
   "source": [
    "for pair, u1, u2 in results:\n",
    "    print(\"{} {} {}\".format(pair, u1, u2))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "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.6.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}