notebooks/gutenberg/part2-ngram-viewer-EMPTY.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "from __future__ import print_function\n",
    "%matplotlib inline\n",
    "import matplotlib.pylab as plt\n",
    "import sys, os, glob\n",
    "import numpy as np\n",
    "import subprocess\n",
    "\n",
    "from IPython.display import HTML\n",
    "import xml.etree.ElementTree as ET\n",
    "try:\n",
    "    tree = ET.parse(os.environ['HADOOP_CONF_DIR'] + '/yarn-site.xml')\n",
    "except IOError:\n",
    "    raise IOError(\"Can't find the yarn configuration -- is HADOOP_CONF_DIR set?\")\n",
    "root = tree.getroot()\n",
    "yarn_web_app = root.findall(\"./property[name='yarn.resourcemanager.webapp.address']\")[0].find('value').text\n",
    "yarn_web_app_string = \"If this works successfully, you can check the <a target='_blank' href='http://{yarn_web_app}'>YARN application scheduler</a> and you should see your app listed there. Clicking on the 'Application Master' link will bring up the familiar Spark Web UI. \"\n",
    "\n",
    "plt.rcParams['figure.figsize'] = (10,6)\n",
    "plt.rcParams['font.size'] = 18\n",
    "plt.style.use('fivethirtyeight')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Analyzing the Gutenberg Books Corpus"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In this notebook, we will use the cleaned, pre-processed data that we created in the [pre-processing part](gutenberg-preprocessing-SOLUTIONS.ipynb). As a reminder, we ended up with an RDD of `(gid, text)` tuples that has been cleaned and we stored it on HDFS at `/user/<YOUR_USERNAME>/gutenberg/cleaned_rdd`. \n",
    "\n",
    "We have two goals in this notebook:\n",
    "\n",
    "1. generate data that will enable us to create something like the [Google Ngram Viewer](https://books.google.com/ngrams) but for the Gutenberg book corpus. \n",
    "2. Use Spark's machine learning library to perform language classification on the English and German book corpus"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Set up and launch the Spark runtime\n",
    "\n",
    "Remember from the previous notebook that we have a saved configuration in `./spark_config/` -- so all we need to do is set the `SPARK_CONF_DIR` environment variable and our default configuration will be used: "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "os.environ['SPARK_CONF_DIR'] = os.path.realpath('./spark_config')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "As before, we'll also set the amount of memory we want to use for the driver and set up the paths using `findspark`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "# how many cores do we have for the driver\n",
    "ncores = int(os.environ.get('LSB_DJOB_NUMPROC', 1)) \n",
    "\n",
    "# here we set the memory we want spark to use for the driver JVM\n",
    "os.environ['SPARK_DRIVER_MEMORY'] = '%dG'%(ncores*2*0.7)\n",
    "\n",
    "# we have to tell spark which python executable we are using\n",
    "os.environ['PYSPARK_PYTHON'] = subprocess.check_output('which python', shell=True).rstrip()\n",
    "\n",
    "import findspark\n",
    "findspark.init()\n",
    "\n",
    "import pyspark\n",
    "from pyspark import SparkConf, SparkContext"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Note that there is another file in `./spark_config`: "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "!cat ./spark_config/log4j.properties"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This logging set-up quiets down the amount of stuff we see in the console..."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Initializing the `SparkContext`\n",
    "\n",
    "We'll create an empty `SparkConf` here for completeness -- if we needed to, we could alter any of the settings here. Then we launch the `SparkContext` as before:  "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "conf = SparkConf()\n",
    "\n",
    "sc = SparkContext(master='yarn-client', conf=conf)\n",
    "\n",
    "HTML(yarn_web_app_string.format(yarn_web_app=yarn_web_app))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Load the data from HDFS"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "# TODO: load cleaned_rdd from your directory on HDFS\n",
    "cleaned_rdd = sc.pickleFile('/user/roskarr/gutenberg/cleaned_rdd').setName('cleaned_rdd').cache()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "A quick inspection of the data to remind ourselves what it looks like: "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "gid, text = cleaned_rdd.first()\n",
    "print(gid, text[10000:10500]) # somewhere in the middle of the book"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Load in the metadata dictionary and broadcast it\n",
    "Remember that we saved our metadata as a dictionary on disk. We didn't use the broadcasted dictionary in the last notebook, but we will make heavy use of it soon."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "from cPickle import load\n",
    "\n",
    "with open('{home}/gutenberg_metadata.dump'.format(home=os.environ['HOME']), 'r') as f :\n",
    "    meta_dict = load(f)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "# create meta_b by broadcasting meta_dict\n",
    "meta_b = sc.broadcast(meta_dict)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Remember that our `cleaned_rdd` contains `gid`'s as keys and text as values and if we want some other piece of metadata, we can just access it via the lookup table, for example to extract the first names of the authors, we can do something like this: "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "cleaned_rdd.map(lambda (gid, text): meta_b.value[gid]['first_name']).first()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Histogram of book publication years\n",
    "Now we're ready to start asking some questions of the data. To begin with, lets do a simple histogram of the year distribution of the books. Since we don't have original publication dates, we just use the simple formula: \n",
    "\n",
    "$year = max\\left((year_{birth} + year_{death})/2, year_{birth} + offset\\right)$, \n",
    "\n",
    "where $offset$ is a number drawn from a gaussian centered on 40 with $\\sigma = 5$ years. This means that we assume most people write their books around age 40. ;)\n",
    "\n",
    "The function `publication_year` is provided for you and you should use it to *transform* the `year_rdd` into an RDD of publication years. \n",
    "\n",
    "Remember that we created a broadcast variable at the beginning of the notebook, which can be used to efficiently look up metadata."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "from random import gauss\n",
    "\n",
    "def publication_year(gid) : \n",
    "    \"\"\"Returns the publication year for the given gutenberg id (gid)\"\"\"\n",
    "    \n",
    "    # extract the metadata dictionary for this gid\n",
    "    meta = meta_b.value[gid] \n",
    "    \n",
    "    birth_year = int(meta['birth_year'])\n",
    "    if meta['death_year'] is None : \n",
    "        year = birth_year + gauss(40,5)\n",
    "    else :\n",
    "        death_year = int(meta['death_year'])\n",
    "        year = max((birth_year + death_year) / 2.0, birth_year + gauss(40,5))\n",
    "\n",
    "    return min(int(year),2015)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "# TODO: map cleaned_rdd to contain just the publication years by using the publication_year function above\n",
    "year_rdd = cleaned_rdd.<FILL IN>\n",
    "                       \n",
    "year_rdd.count()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "year_rdd.first()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The histogram function actually already exists in the Spark API (but it didn't use to!). However, for fun we will write our own. Calculating the histogram can be split up into two parts:\n",
    "\n",
    "2. calculate the bin that the value maps to and create an RDD of (`bin`, 1) pair\n",
    "3. do a `reduceByKey` where we just add up all the values belonging to each bin. \n",
    "\n",
    "The last step is the simplest for us to execute, but the most complicated in terms of what happens behind the scenes. The reduction happens in two parts - first, it creates a list of `(key, count)` pairs locally by iterating through all the elements of the partition and creating a hash table. Then, these pairs are communicated to other partitions and added up. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "# define a helper function to determine which bin a value falls into\n",
    "from bisect import bisect_right\n",
    "def get_bin(bin_edges, value) : \n",
    "    \"\"\"Returns which bin, specified by `bin_edges`, the `value` falls into.\"\"\"\n",
    "    return bisect_right(bin_edges, value) - 1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "# TODO: fill in the bits and pieces of the histogram function\n",
    "\n",
    "def histogram(rdd, nbins = 100, min_val=None, max_val=None) :\n",
    "    \"\"\"\n",
    "    Calculate a histogram of the data, given the number of bins and the data range. \n",
    "    \n",
    "    Arguments: \n",
    "        \n",
    "        rdd: the data RDD\n",
    "    \n",
    "    Optional Keywords:\n",
    "        \n",
    "        nbins: number of bins (default 100)\n",
    "        \n",
    "        min_val: minimum value to consider (default, min value of the input data)\n",
    "        \n",
    "        max_val: maximum value to consider (default, max value of the input data)\n",
    "    \n",
    "    Returns: \n",
    "        \n",
    "        tuple of bins and array of counts\n",
    "        \n",
    "    \"\"\"\n",
    "    # if either min_val or max_val are missing, get them from the data\n",
    "    if min_val is None : \n",
    "        min_val = rdd.<FILL IN> # FILL IN --> look at the Spark API!\n",
    "    if max_val is None : \n",
    "        max_val = rdd.<FILL IN> # FILL IN --> look at the Spark API!\n",
    "        \n",
    "    # create the array of bin edges\n",
    "    bin_edges = np.linspace(min_val,max_val,nbins+1)\n",
    "    \n",
    "    # create an RDD where each data value is mapped into a tuple (bin, 1) that will \n",
    "    # be used for counting up the bin values (use the get_bin function defined above to generate the bin)\n",
    "    binned_rdd = rdd.<FILL IN> # FILL IN \n",
    "    \n",
    "    # reduce binned_rdd and collect the values into summed_bins \n",
    "    summed_bins = binned_rdd.<FILL IN>.collect() # FILL IN \n",
    "    \n",
    "    # This is a sparse result -- turn into a dense vector for plotting: \n",
    "    res_full = np.zeros(nbins)\n",
    "    overflow = 0\n",
    "    for item in summed_bins : \n",
    "        if item[0] > len(res_full)-1 or item[0] < 0: \n",
    "            continue # ignore\n",
    "        else: res_full[item[0]] = item[1]\n",
    "            \n",
    "    return .5*(bin_edges[:-1]+bin_edges[1:]), res_full"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "%time bins, vals = histogram(year_rdd, min_val=1500, max_val=2015)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Finally, here is the histogram of publication years:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "plt.plot(bins, vals)\n",
    "plt.xlim(1500,2015)\n",
    "plt.xlabel('year')\n",
    "plt.ylabel('Number of books')\n",
    "plt.semilogy();"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Reducing network traffic by using `mapPartition`\n",
    "\n",
    "This kind of operation is a good candidate for first computing a partial result on each partition and then following with a reduce step; that is, we can make a histogram of just the data in each partition first, and then all we need to do is add those histograms together. \n",
    "\n",
    "In our `histogram` function above, the `reduceByKey` step requires a reshuffling of the data which means spending potentially a lot of time in network communication and other overhead associated with creating distributed hash tables for all the keys. \n",
    "\n",
    "However, since our binning function effectively already hashes the values anyway, we can avoid the overhead of `reduceByKey` by instructing first each partition to calculate its local histogram and then simply adding up the histograms in the end. In this way, we are basically combining the `map` operation with a local `reduce` operation, and since the number of bins is always relatively small, we end up sending a trivial amount of data across the network. \n",
    "\n",
    "Here is where our knowledge of generators comes in handy, because the method [`mapPartitions`](http://spark.apache.org/docs/latest/api/python/pyspark.html#pyspark.RDD.mapPartitions) *requires* a generator function. "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Defining the generator function\n",
    "The key here is to define a histogram function that calculates the histogram on each partition locally. `mapPartitions` gives us an *iterator* over the data; we need to extract the data out of this iterator and calculate the local histograms *for the data in each partition*. Below, this is accomplished using the `bin_partition` function.\n",
    "\n",
    "Finally, we define a new `histogram_partition` function that uses `bin_partition` to compute the local histograms and a simple addition in a `reduce` operation to sum up the histograms in the end. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "def bin_partition(iterator, nbins,  min_val, max_val) : \n",
    "    \"\"\"\n",
    "    Perform the binning of data contained in an iterator\n",
    "    \n",
    "    Arguments: \n",
    "        \n",
    "        iterator: the data iterable \n",
    "        \n",
    "        nbins: number of bins\n",
    "        \n",
    "        min_val, max_val: min and max values to consider\n",
    "    \n",
    "    yields the local histogram\n",
    "    \"\"\"\n",
    "    from bisect import bisect_right\n",
    "    \n",
    "    bin_edges = np.linspace(min_val,max_val,nbins+1)\n",
    "    \n",
    "    histogram = np.zeros(len(bin_edges)-1)\n",
    "    \n",
    "    for item in iterator : # iterating over all the items in the partition\n",
    "        try : \n",
    "            ind = get_bin(bin_edges,item)\n",
    "            histogram[ind] += 1\n",
    "        except IndexError : \n",
    "            pass\n",
    "        \n",
    "    yield histogram\n",
    "\n",
    "def histogram_partition(rdd, nbins = 100, min_val=None, max_val=None) :\n",
    "    \"\"\"\n",
    "    Calculate a histogram of the data by using partition methods\n",
    "    \n",
    "    Arguments: \n",
    "    \n",
    "        rdd: the data\n",
    "        \n",
    "    Optional Keywords:\n",
    "    \n",
    "        nbins: number of bins (default 100)\n",
    "        \n",
    "        min_val: minimum value to consider (default, min value of the input data)\n",
    "        \n",
    "        max_val: maximum value to consider (default, max value of the input data)\n",
    "    \n",
    "    Returns: \n",
    "        \n",
    "        tuple of bins and array of counts\n",
    "    \"\"\"\n",
    "    # if either min_val or max_val are missing, get them from the data\n",
    "    if min_val is None : \n",
    "        min_val = rdd.<FILL IN> # FILL IN\n",
    "    if max_val is None : \n",
    "        max_val = rdd.<FILL IN> # FILL IN\n",
    "        \n",
    "    bin_edges = np.linspace(min_val,max_val,nbins+1)\n",
    "    \n",
    "    result = (rdd.mapPartitions(lambda iterator: bin_partition(iterator, nbins, min_val, max_val))\n",
    "                 .reduce(lambda a,b: a+b))\n",
    "    \n",
    "    return .5*(bin_edges[:-1]+bin_edges[1:]), result"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "%time bins, vals = histogram_partition(year_rdd, min_val=1500, max_val=2015)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The difference here doesn't look dramatic because the total amount of data is rather small, but have a look at the Spark Web UI and you will see that the second implementation didn't do any shuffle writing. If you have a large number of keys, the intermediate shuffles that need to take place can have a substantial impact on performance. With a bigger data set, this difference could potentially matter quite a lot!"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Inspecting the metadata some more\n",
    "\n",
    "Lets do a couple more checks and practice using the Spark API. \n",
    "\n",
    "#### How many unique authors are there in the dataset? \n",
    "\n",
    "1. make `author_rdd` that is composed of a string `\"last_name, first_name\"` (use the broadcast variable `meta_b` to get the data for each `gid`)\n",
    "2. keep only the unique author strings (*hint*: look at the Spark API to find an appropriate method)\n",
    "3. count the number of elements remaining\n",
    "\n",
    "**note**: use `meta_b.value` to access the actual metadata dictionary, i.e. to get the metadata for `gid=101`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "meta_b.value[101]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "# TODO: map cleaned_rdd to contain the string \"last_name, first_name\" \n",
    "author_rdd = cleaned_rdd.<FILL IN>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "# TODO: use RDD methods to obtain the distinct author strings and count them\n",
    "n_authors = (author_rdd.<FILL IN>)\n",
    "print(\"Number of distinct authors: %s \" % n_authors)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "assert(n_authors == 10192)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Most-represented authors in the corpus: \n",
    "\n",
    "1. use the `author_rdd` from above\n",
    "2. use the pattern `(key, 1)` to set up an RDD that can be passed to `reduceByKey` \n",
    "3. run `reduceByKey`, yielding an RDD composed of `(author, count)` tuples\n",
    "4. sort by descending order of number of books per author and print out the top 10 (try using `takeOrdered`)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "# TODO: generate a list of authors, reverse-sorted by the number of books they have in the corpus\n",
    "(author_rdd.<FILL IN>\n",
    "           .<FILL IN>\n",
    "           .<FILL IN>)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Finally, lets do the same thing per language, just to get an idea of how much data there is: "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "# FILL IN: generate a lang_rdd that contains just the language of each book\n",
    "lang_rdd = cleaned_rdd.<FILL IN>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "# FILL IN: reduce the `lang_rdd` to yield number of books in each language\n",
    "(lang_rdd.<FILL IN>\n",
    "         .<FILL IN>\n",
    "         .<FILL IN>)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## How many unique words were used in English in these 500+ years? \n",
    "\n",
    "We could have done the above metadata gymnastics without ever invoking a distributed processing framework by simply extracting the years from the metadata -- nevertheless we used the metadata to have a closer look at some of the RDD methods. However, the text body of each data element is where the bulk of the data volume lies. \n",
    "\n",
    "To construct a corpus wide vocabulary, we have to deconstruct each document into a list of words and then extract the unique words from the entire data set. If our dataset fits into memory of a single machine, this is a simple `set` operation. But what if it doesn't? \n",
    "\n",
    "We'll assume this is the case and instead of converting each `gid,text` pair into a `gid,list_of_words` pair, we will simply construct one global RDD of words. Here we aren't necessarily interested in preserving the provenance of words, but just finding the unique words in the whole corpus, so we drop the metadata altogether. \n",
    "\n",
    "The steps are as follows:\n",
    "\n",
    "1. map the entire RDD of text into an RDD of single words (use flatMap -- this returns a different number of elements than it takes in)\n",
    "2. use the `distinct` method of the resulting RDD to transform it into an RDD with only unique words\n",
    "\n",
    "As a reminder, here's an illustration of how `flatMap` differs from `map`:\n",
    "\n",
    "![flatMap](../../slides/figs/flatMap_example.svg)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "*Hint:* In python, splitting a string into a set of words separated by spaces is easy: "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "line = 'splitting a string is super simple'\n",
    "line.split()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Make an RDD `distinct_rdd` which holds the *unique English* words. Consider the steps this will require:\n",
    "\n",
    "* use `cleaned_rdd` but keep only books in the english language (make sure you use the broadcast metadata variable!)\n",
    "* convert each document into individual words\n",
    "* retain only the unique words\n",
    "\n",
    "Which RDD methods can you use to achieve these three steps? (note that this will be a pretty expensive operation so it might take some time...)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "# TODO: create distinct_rdd by filtering for english books and using RDD methods to generate an RDD of distinct words\n",
    "distinct_rdd = (cleaned_rdd.filter(lambda (gid, text): meta_b.value[gid]['lang'] == 'en')\n",
    "                           .flatMap(lambda (gid, text): text.split())\n",
    "                           .distinct())\n",
    "nwords = distinct_rdd.count()\n",
    "print(\"Number of unique English words: \", nwords)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "assert(abs(nwords - 3688479) < 10)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Note that not all of these are actual words, this is just how many character sequences separated by spaces we found. The pre-processing steps are less than perfect so there is some garbage in there. We will trim this down to just the most commonly-used ones later and that will get rid of most of the nonsense. "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## What are the most common words? \n",
    "\n",
    "A \"MapReduce\" tutorial has to include a word counting example -- it's basically the equivalent of a \"Hello World!\"  in a programming tutorial!\n",
    "\n",
    "So, lets count the occurences of all the words across the entire corpus. This is a fairly straightforward operation, but it exposes some very common patterns that can be useful for many tasks. To simplify this a bit, we'll use only the English-language corpus for the moment.\n",
    "\n",
    "Here are the steps we need to take:\n",
    "\n",
    "0. keep only the english language books (use a filter)\n",
    "1. `flatMap` each document into (`word, count`) pairs, but only for words that are not in the `stop_words` set (try with a list comprehension!)\n",
    "2. call `reduceByKey` to sum up all the `count`s for each word\n",
    "3. finally sort it in descending order to see the most common words first\n",
    "\n",
    "The first part here (filtering and `flatMap`) is much like what we did before, but with a twist: for each word, check that it is *not* a member of the `stop_words` set. \"Stop words\" include common words like \"a, the, he\" etc.  "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false,
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "from pickle import load\n",
    "stop_words = load(open('./stop_words.dump')).union(['gutenbergtm', 'gutenberg', 'electronic', 'foundation', 'license', 'copyright', 'donation', 'donations'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "# TODO: create english_rdd which contains only the english books by filtering on the metadata\n",
    "english_rdd = cleaned_rdd.<FILL IN>.cache()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "# TODO: use flatMap to extract the words from each document's text using the english_rdd we made above\n",
    "words_rdd = (english_rdd.<FILL IN>\n",
    "                        .setName('words_rdd')\n",
    "                        .cache())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now that we have our \"flattened\" data set, do the counting by first mapping each word in `words_rdd` into a `(word, 1)` tuple, and then using `reduceByKey` to calculate the word frequencies. At the end of this step, use the `sortBy` method to sort the word counts in descending order. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "# TODO: do the word count!\n",
    "word_count = (words_rdd.<FILL IN>\n",
    "                       .setName('word_count')\n",
    "                       .cache())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "assert(word_count.take(5) == [('said', 4339532), ('man', 2775185), ('time', 2688315), ('little', 2458328), ('like', 2342552)])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "# fifty most common words (excluding stop words)\n",
    "word_count.take(50)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Reduces, Shuffles, and Partitioning\n",
    "During a `reduceByKey`, or any other reduction for that matter, data must be shuffled around the cluster and combined. Other common RDD methods like `join`, `sortByKey` etc. also typically require lots of data shuffling. By default, this is done in an intelligent way by first reducing values locally on each partition, and then combining the results of the partitions. Still, as is the case here, for common keys, every partition will have to send its results to others. This can result in a lot of temporary file IO if the data that needs to be communicated can't all be held in memory on all of the executors. \n",
    "\n",
    "One way around this is to pre-partition the data ahead of time so that the same keys land on the same partition by design. This results in much less data needing to be shipped around the network and can improve the performance. Of course, at the cost of an expensive initial shuffle that takes place during the partitioning step! But if many \"by key\" have to be done on the same data, it might be worth it. "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Lets have a look at these concepts by performing the word count in a few different ways. "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we will create a re-partitioned `words_rdd` using the default partitioning of `partitionBy` (just a hash function). You may also specify your own partitioning function. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "num_partitions = words_rdd.getNumPartitions()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "par = (words_rdd.map(lambda word: (word,1))\n",
    "                .partitionBy(num_partitions)\n",
    "                .cache())\n",
    "par.count() # call count to compute and cache the data"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Since both datasets are cached in memory, we can compare the time it takes the reduce step to complete:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "%%time\n",
    "(words_rdd.map(lambda word: (word,1))\n",
    "          .reduceByKey(lambda a,b: a+b)\n",
    "          .count())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "%%time\n",
    "par.reduceByKey(lambda a,b: a+b).count()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "collapsed": true
   },
   "source": [
    "Quite a speedup, but at the cost of an expensive initial shuffle. If the \"by key\" operation is only done once, then this is not worth it -- but if it's done repeatedly  (i.e. frequent joins) then it may be beneficial. However, we can also expect the difference to depend on the nature of the dataset. If you inspect the Spark UI, you can see that the first `reduceByKey` (i.e. one done on `word_rdd`) shuffled ~390 Mb of data, while the second `reduceByKey` (i.e. done on `par`) only shuffled ~50 Mb of data. This dataset is still pretty small, but when the shuffles are in the Gb range, the differences can be substantial. \n",
    "\n",
    "If you need to do lookups of individual keys, this becomes even more dramatic: "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "%time x = words_rdd.lookup('environment')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "%time x = par.lookup('environment')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Computing word frequency vs. time\n",
    "\n",
    "Now we have all the components to build a simple tool that visualizes the relative word frequency as a function of time in the Gutenberg corpus. For inspiration, see the [Google Ngram viewer](https://books.google.com/ngrams).\n",
    "\n",
    "### Converting documents into vectors\n",
    "\n",
    "To make quantitative analysis of the corpus possible, we will convert each document into a vector that represents the word counts for each word appearing in the document. \n",
    "\n",
    "This will look something like this. Imagine we have have a corpus consisting of two \"documents\"\n",
    "\n",
    "    document 1: \"a dog bit me\"\n",
    "    document 2: \"i bit the dog back\"\n",
    "    \n",
    "Then our corpus vocabulary (of 1-grams) is\n",
    "\n",
    "    [\"a\", \"dog\", \"bit\", \"me\", \"i\", \"the\", \"back\"]\n",
    "    \n",
    "Since this is an array and each word in the array has a unique index, we can \"encode\" the two documents using this index mapping. Our corpus now looks like this: \n",
    "\n",
    "    document 1: [1, 1, 1, 1, 0, 0, 0]\n",
    "    document 2: [0, 1, 1, 0, 1, 1, 1]\n",
    "    \n",
    "In order for the vector indices to remain consistent across the whole corpus, the first step is to build a corpus-wide lookup table, a `word --> index` mapping. \n",
    "\n",
    "\n",
    "### Generating word counts \n",
    "\n",
    "Once each document is converted to a vector, doing the word counts for sub-groups of documents is a simple vector addition operation. In our case, we will reduce the vectors by year, yielding an RDD that will have the total number of occurrences of each word in every year. From there, it is trivial to look up the desired word and plot the relative frequency vs. year. \n",
    "\n",
    "## Create the vocabulary lookup table\n",
    "Create a look-up table of words by attaching a unique index to each word. The Spark API provides a [zipWithIndex](http://spark.apache.org/docs/latest/api/python/pyspark.html#pyspark.RDD.zipWithIndex) method that makes this easy. \n",
    "\n",
    "Above, we created the `word_count` RDD that already contains the unique words and their associated counts. To reduce the size of the lookup table (and the size of the vectors), we will restrict ourselves to using only the first 100k words. \n",
    "\n",
    "You can either create the `(word, index)` pairs in the RDD and collect the top 100,000 as a dictionary using `collectAsMap`, or you can collect the top 100,000 words from `word_count` RDD and turn them into a dictionary locally. \n",
    "\n",
    "These are the steps we need to take to make the vocabulary lookup from an RDD:\n",
    "\n",
    "1. use `zipWithIndex` on the keys of `word_count` to generate a unique index for each word -- we don't care about the counts anymore, so we can get rid of the values and just work with the keys using the `keys()` method\n",
    "2. use filter to retain only the first 100000 words \n",
    "3. finally, use [collectAsMap](http://spark.apache.org/docs/latest/api/python/pyspark.html#pyspark.RDD.collectAsMap) to return the resulting RDD to the driver as a dictionary. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "# TODO: create word_lookup, a dictionary that maps each of the top 100,000 words to a unique integer\n",
    "word_lookup = <FILL IN>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Make a `word_lookup` into a broadcast variable so we can use it on all the workers:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "word_lookup_b = sc.broadcast(word_lookup)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This dictionary is approximately 6 Mb in size - without a broadcast, it would get sent over the network to each task, resulting in a lot of network traffic! As a broadcast variable, it gets sent only *once* to each *executor*, i.e. it's transferred only 20 times (if you are using the defaults, otherwise however many executors you have). "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Vectorize the documents\n",
    "\n",
    "Now that we have a vocabulary lookup table, we can use this to turn each document into a vector. \n",
    "\n",
    "This is done by counting up the occurrences of all words in the document that are also in the global vocabulary. \n",
    "\n",
    "The function `vectorize_doc` below accomplishes this by using a dictionary to keep track of the local word count. Once the counting is done we use the counts to create a sparse vector that represents the document. A sparse vector consists of two arrays, one representing the *locations* of the non-zero values, and the other the values themselves. \n",
    "\n",
    "To return to our contrived example from above, we had \n",
    "\n",
    "    document 1: \"a dog bit me\"\n",
    "    document 2: \"i bit the dog back\"\n",
    "    \n",
    "which turned into \n",
    "\n",
    "    document 1: [1, 1, 1, 1, 0, 0, 0]\n",
    "    document 2: [0, 1, 1, 0, 1, 1, 1]\n",
    "\n",
    "with a vocabulary of \n",
    "\n",
    "    [\"a\", \"dog\", \"bit\", \"me\", \"i\", \"the\", \"back\"]\n",
    "    \n",
    "As sparse vectors, these two documents could be represented with two arrays: \n",
    "\n",
    "    document 1: indices = [0,1,2,3]; values = [1, 1, 1, 1]\n",
    "    document 2: indices = [1,2,4,5,6]; values = [1, 1, 1, 1, 1]\n",
    "\n",
    "We use Spark's own `SparseVector` data type, for which we must specify a size (total number of features), a (sorted) array of indices, and an array of values. This means that we start to save space if sparsity is > 50%. Note that the `SparseVector` provides some nice higher-level methods, but it does not allow simple operations like addition. If a lot of vector arithmetic is needed, you should use the scipy sparse types instead. "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In the next cell, we define two functions: \n",
    "\n",
    "* **`extract_ngrams`** converts a sequence of words or characters into a sequence of n-grams (here we are just using single worde, i.e. 1-grams so we'll postpone talking about ngrams until later)\n",
    "\n",
    "* **`vectorize_doc`** converts a document into a sparse vector by using `extract_ngrams` to tokenize it and a vocabulary mapping to turn each word into a vector component"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "from pyspark.mllib.linalg import SparseVector\n",
    "\n",
    "def extract_ngrams(tokens, ngram_range=[1,1], select_ngrams = None, ngram_type='word'):\n",
    "    \"\"\"\n",
    "    Turn tokens into a sequence of n-grams \n",
    "\n",
    "    Arguments:\n",
    "\n",
    "        tokens: a list of tokens\n",
    "\n",
    "    Optional Keywords:\n",
    "\n",
    "        ngram_range: a tuple with min, max ngram ngram_range\n",
    "    \n",
    "        select_ngrams: the vocabulary to use\n",
    "    \n",
    "        ngram_type: whether to produce word or character ngrams\n",
    "\n",
    "    Output:\n",
    "\n",
    "    Generator yielding a list of ngrams in the desired range\n",
    "    generated from the input list of tokens\n",
    "\n",
    "    \"\"\"\n",
    "    join_str = \"\" if ngram_type=='character' else \" \"\n",
    "    \n",
    "    # handle token n-grams\n",
    "    min_n, max_n = ngram_range\n",
    "    n_tokens = len(tokens)\n",
    "    \n",
    "    for n in xrange(min_n, min(max_n + 1, n_tokens + 1)):\n",
    "        for i in xrange(n_tokens - n + 1):\n",
    "            if n == 1: \n",
    "                res = tokens[i]\n",
    "            else : \n",
    "                res = join_str.join(tokens[i: i+n])\n",
    "           \n",
    "            # if we are using a lookup vocabulary, check for membership here\n",
    "            if select_ngrams is not None : \n",
    "                if res in select_ngrams: \n",
    "                    yield res\n",
    "            else : \n",
    "                yield res\n",
    "            \n",
    "\n",
    "def vectorize_doc(doc, vocab, ngram_range = [1,1], ngram_type='word') : \n",
    "    \"\"\"\n",
    "    Returns a vector representation of `doc` given the reference \n",
    "    vocabulary `vocab` after tokenizing it with `tokenizer`\n",
    "    \n",
    "    Arguments: \n",
    "        \n",
    "        doc: a sequence of tokens (words or characters)\n",
    "        \n",
    "        vocab: the vocabulary mapping\n",
    "        \n",
    "    Keywords:\n",
    "    \n",
    "        ngram_range: the range of ngrams to process\n",
    "        \n",
    "        ngram_type: whether to produce character or word ngrams; default is \n",
    "        \n",
    "    Returns:\n",
    "    \n",
    "        a sparse vector representation of the document given the vocabulary mapping\n",
    "    \"\"\"\n",
    "    from collections import defaultdict\n",
    "    from scipy.sparse import csr_matrix \n",
    "        \n",
    "    # count all the word occurences \n",
    "    data = np.zeros(len(vocab))\n",
    "    \n",
    "    for ngram in extract_ngrams(doc, ngram_range, vocab, ngram_type) : \n",
    "         data[vocab[ngram]] += 1\n",
    "            \n",
    "    # only keep the nonzero indices for the sparse representation\n",
    "    indices = data.nonzero()[0]\n",
    "    values = data[indices]\n",
    "    \n",
    "    return SparseVector(len(vocab), indices, values)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Using these functions to vectorize our two-sentence test corpus: "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "import string\n",
    "s1 = \"a dog bit me\"\n",
    "s2 = \"i bit the dog back\"\n",
    "vocab = [\"a\", \"dog\", \"bit\", \"me\", \"i\", \"the\", \"back\"]\n",
    "vocab_dict = {word:ind for ind, word in enumerate(vocab)}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "print(s1)\n",
    "print(vectorize_doc(s1.split(), vocab_dict))\n",
    "print(s2)\n",
    "print(vectorize_doc(s2.split(), vocab_dict))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we have all the components we need to create an RDD of english-language books vectorized using the most common 100k words. All we need to do is to use `mapValues` to map the text of each document in `english_rdd` into a vector using `vectorize_doc` and our broadcast vocabulary lookup table `word_lookup_b`. Recall that the broadcast variable, `word_lookup_b` is just a wrapper for the lookup table; to pass the actual lookup table to the `vectorize_doc` function, use `word_lookup_b.value`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "# FILL IN \n",
    "vector_rdd = <FILL IN>.cache()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false,
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "vector_rdd.first()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "# make sure the transformation can be carried out for all elements by using count\n",
    "# this is a good way of catching anomalies in the data\n",
    "vector_rdd.count()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Perform the aggregation\n",
    "\n",
    "We now have the entire Gutenberg English book corpus in the form of sparse vectors encoding the most used 100k words. \n",
    "\n",
    "To get the yearly sums, we will turn the metadata of each document into its publication year (i.e. the key will be the year, the value is the vector) and then do an aggregation by year. \n",
    "\n",
    "We will use the powerful [`treeAggregate`](http://spark.apache.org/docs/latest/api/python/pyspark.html?highlight=values#pyspark.RDD.treeAggregate) method, which requires that we specify three different components:\n",
    "\n",
    "1. the starting aggregate\n",
    "2. a function that adds a new value to the aggregate \n",
    "3. a function that adds together two aggregates\n",
    "\n",
    "The way `treeAggregate` works is that it performs the reduction in a tree pattern in order to minimize the strain on the driver. In a \"normal\" reduction, the workers send their results to the driver, which is tasked with putting it all together -- however, if these partial results are large (as is potentially the case here) then the driver can run into memory issues. Furthermore, most of the cluster is sitting idle while the driver performs the aggregation. `treeAggregate` fixes this by performing partial aggregations on the workers and only sending the final stages to the driver. See this [blog post](https://databricks.com/blog/2014/09/22/spark-1-1-mllib-performance-improvements.html) for a bit more description of this method. \n",
    "\n",
    "The aggregation methods are powerful because the \"aggregate\" can be any object -- we can write a class that gets passed around to do the aggregation, for example. Aggregation methods are more general reduction methods because they allow you to change the type of the variables -- in our case here, we are converting the data `(year, vector)` tuples into a dictionary of arrays. \n",
    "\n",
    "Below, we will use an instance of a dictionary as the aggregation object and define two functions that will do the actual aggregation. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "def add_vector(d, data_tuple) : \n",
    "    \"\"\"Add a new vector to the aggregation dictionary\n",
    "    \n",
    "    The vectors in the aggregation dictionary are dense since for most years we can expect that \n",
    "    this will be the case anyway. Note that we use 32-bit floats to save a bit on memory. \n",
    "    \n",
    "    Arguments: \n",
    "        d: the aggregation dictionary\n",
    "        \n",
    "        data_tuple: the (year, vec) tuple\n",
    "        \n",
    "    Returns: \n",
    "        the updated aggregation dictionary \n",
    "    \"\"\"\n",
    "    # expand the data tuple\n",
    "    year, vec = data_tuple\n",
    "    \n",
    "    if year in d : \n",
    "        \n",
    "        d[year][vec.indices] += vec.values\n",
    "    else :\n",
    "        # this is the first time we've encountered this year --> make an empty vector \n",
    "        new_vec = np.zeros(vec.size, dtype=np.float32)\n",
    "        \n",
    "        # now put in the contents of the current vector\n",
    "        new_vec[vec.indices] = vec.values\n",
    "        \n",
    "        # create the year in the dictionary\n",
    "        d[year] = new_vec\n",
    "        \n",
    "    return d\n",
    "\n",
    "def add_dicts(d1, d2) : \n",
    "    \"\"\"Add two dictionaries together\n",
    "    \n",
    "    Arguments: \n",
    "        d1: first dictionary\n",
    "        \n",
    "        d2: second dictionary\n",
    "        \n",
    "    Returns: \n",
    "        merged dictionaries\n",
    "    \"\"\"\n",
    "    \n",
    "    # iterate through all the items in the second dictionary\n",
    "    for year, vec in d2.iteritems() : \n",
    "        # if this year is also in d1, add the vectors together\n",
    "        if year in d1 : \n",
    "            d1[year] += vec\n",
    "        # if not, create a new year entry in d1\n",
    "        else : \n",
    "            d1[year] = vec\n",
    "    return d1"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we create an RDD of `(year, vec)` pairs using the `publication_year` function we defined at the top:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "year_vec = vector_rdd.map(lambda (gid, vec): (publication_year(gid), vec))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Before we perform the aggregation, we can do one final bit of optimization. Passing around dictionaries full of large arrays can get expensive very quickly. The memory footprint of our partial results will depend on how heterogeneous the years on each partition or group of partitions are: if most of the data on a partition is for the same key (year in this case) then the dictionary we create on that partition will only contain a handful of vectors.  We can control this by first partitioning the RDD in a way that groups data with the same keys onto the same partitions. \n",
    "\n",
    "Spark provides a `partitionBy` method that does exactly this -- by default, it uses a hash function to map the keys to partitions, but you can also pass a custom partitioner if you want. If you look at the Spark UI after executing the next cell, you'll see that the partition step caused some shuffling of data, but that the aggregation itself ran very quickly and with minimal data movement. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "n_partitions = year_vec.getNumPartitions()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "# TODO: use an empty dictionary and the two functions defined above as arguments to the treeAggregate method\n",
    "year_sums = (year_vec.partitionBy(n_partitions)\n",
    "                     .<FILL IN>)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Note that `year_sums` is a single \"value\", in this case our aggregate dictionary containing years as keys and vectors representing cummulative word counts as values. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "list(year_sums.iteritems())[0:10]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Gutenberg Project N-gram viewer\n",
    "\n",
    "Lets plot some results!\n",
    "\n",
    "Below we define a plotting function and then make a plot of some interesting examples -- feel free to change the word list..."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "def plot_usage_frequency(words, year_data, word_lookup, plot_range = [1500,2015]) : \n",
    "    years = sorted(year_data.keys())\n",
    "    tot_count = np.array([year_data[year].sum() for year in years])\n",
    "    \n",
    "    if ',' in words:\n",
    "        words = [word.strip() for word in words.split(',')]\n",
    "    elif type(words) is not list: \n",
    "        words = [words]\n",
    "        \n",
    "    n_words = len(words)\n",
    "    \n",
    "    for i, word in enumerate(words) : \n",
    "        word_ind = word_lookup[word]\n",
    "        w_count = np.array([year_data[year][word_ind] for year in years])\n",
    "        \n",
    "        plt.plot(years, smooth(w_count/(tot_count-w_count)),label=word, color = plt.cm.Set1(1.*i/n_words))\n",
    "    \n",
    "    plt.xlim(plot_range)\n",
    "    plt.xlabel('year')\n",
    "    plt.ylabel('relative frequency')\n",
    "    plt.legend(loc='upper left', fontsize = 'small')\n",
    "    \n",
    "    \n",
    "def smooth(x,window_len=11,window='hanning'):\n",
    "        if x.ndim != 1:\n",
    "                raise ValueError, \"smooth only accepts 1 dimension arrays.\"\n",
    "        if x.size < window_len:\n",
    "                raise ValueError, \"Input vector needs to be bigger than window size.\"\n",
    "        if window_len<3:\n",
    "                return x\n",
    "        if not window in ['flat', 'hanning', 'hamming', 'bartlett', 'blackman']:\n",
    "                raise ValueError, \"Window is on of 'flat', 'hanning', 'hamming', 'bartlett', 'blackman'\"\n",
    "        s=np.r_[2*x[0]-x[window_len-1::-1],x,2*x[-1]-x[-1:-window_len:-1]]\n",
    "        if window == 'flat': #moving average\n",
    "                w=np.ones(window_len,'d')\n",
    "        else:  \n",
    "                w=eval('np.'+window+'(window_len)')\n",
    "        y=np.convolve(w/w.sum(),s,mode='same')\n",
    "        return y[window_len:-window_len+1]\n",
    "    \n",
    "from ipywidgets import interact, interactive, fixed\n",
    "import ipywidgets as widgets\n",
    "\n",
    "plot_partial = lambda words: plot_usage_frequency(words, year_sums, word_lookup)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Here are just some illustrative examples -- feel free to try your own..."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false
   },
   "outputs": [],
   "source": [
    "interact(plot_partial, words = 'giveth, environment, machine')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "If you are very motivated, you can adapt the workflow above to work with higher-order n-grams and allow for the lookup of phrases (i.e. \"world war\") instead of just single words. To do this, you have to create a new `word_lookup` table and regenerate the vectors. Since single words (i.e. one-grams) will dominate, it might make sense to build separate list of top N-grams (top two-grams, top three-grams) and then merge them together into a vocabulary map. Beware that the size of the data will increase quickly for N > 1!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "sc.stop()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To continue the exploration of the Gutenberg books corpus, you can move on to the [language classification notebook](part3-lang-classification-EMPTY.ipynb)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 2",
   "language": "python",
   "name": "python2"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 2
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython2",
   "version": "2.7.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 0
}