notebooks/elasticsearch-spark-recommender.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "# Licensed under the Apache License, Version 2.0 (the \"License\"); you may\n",
    "# not use this file except in compliance with the License. You may obtain\n",
    "# a copy of the License at\n",
    "#\n",
    "#      http://www.apache.org/licenses/LICENSE-2.0\n",
    "#\n",
    "# Unless required by applicable law or agreed to in writing, software\n",
    "# distributed under the License is distributed on an \"AS IS\" BASIS, WITHOUT\n",
    "# WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the\n",
    "# License for the specific language governing permissions and limitations\n",
    "# under the License."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Creating a Scalable Recommender with Apache Spark & Elasticsearch\n",
    "\n",
    "In this notebook, you will create a recommendation engine using Spark and Elasticsearch. Using some movie rating data,\n",
    "you will train a collaborative filtering model in Spark and export the trained model to Elasticsearch. Once exported, \n",
    "you can test your recommendations by querying Elasticsearch and displaying the results.\n",
    "\n",
    "### _Prerequisites_\n",
    "\n",
    "The notebook assumes you have installed Elasticsearch, the Elasticsearch vector-scoring plugin, Apache Spark and the Elasticsearch Spark connector detailed in the [setup steps](https://github.com/MLnick/elasticsearch-spark-recommender-demo/tree/master#steps).\n",
    "\n",
    "> _Optional:_\n",
    "\n",
    "> In order to display the images in the recommendation demo, you will need to access [The Movie Database (TMdb) API](https://www.themoviedb.org/documentation/api). Please follow the [instructions](https://developers.themoviedb.org/3/getting-started) to get an API key.\n",
    "\n",
    "## Overview\n",
    "\n",
    "You will work through the following steps\n",
    "\n",
    "1. Prepare the data\n",
    "2. Use the Elasticsearch Spark connector to save it to Elasticsearch\n",
    "3. Load ratings data and train a collaborative filtering recommendation model using Spark MLlib\n",
    "3. Save the model to Elasticsearch\n",
    "4. Show recommendations using Elasticsearch vector scoring plugin"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Step 1: Prepare the data\n",
    "\n",
    "* This notebook uses the \"small\" version of the latest MovieLens movie rating dataset, containing about 100,000 ratings, 9,000 movies and 700 users\n",
    "* The latest version of the data can be downloaded at https://grouplens.org/datasets/movielens/latest/\n",
    "* Download the `ml-latest-small.zip` file and unzip it to a suitable location on your system.\n",
    "\n",
    "The folder should contain a number of CSV files. We will be using the following files:\n",
    "* `ratings.csv` - movie rating data\n",
    "* `links.csv` - external database ids for each movie\n",
    "* `movies.csv` - movie title and genres"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "# first import a few utility methods that we'll use later on\n",
    "from IPython.display import Image, HTML, display\n",
    "# check PySpark is running\n",
    "spark"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Load rating and movie data"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Ratings**\n",
    "\n",
    "The ratings data consists of around 100,000 ratings given by users to movies. Each row of the `DataFrame` consists of a `userId`, `movieId` and `timestamp` for the event, together with the `rating` given by the user to the movie"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "# if you unzipped the data to a different location than that specified in the Journey setup steps\n",
    "# you can change the path below to point to the correct location\n",
    "PATH_TO_DATA = \"../data/ml-latest-small\"\n",
    "# load ratings data\n",
    "ratings = spark.read.csv(PATH_TO_DATA + \"/ratings.csv\", header=True, inferSchema=True)\n",
    "ratings.cache()\n",
    "print(\"Number of ratings: %i\" % ratings.count())\n",
    "print(\"Sample of ratings:\")\n",
    "ratings.show(5)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "You will see that the `timestamp` field is a UNIX timestamp in seconds. Elasticsearch takes timestamps in milliseconds, so you will use some `DataFrame` operations to convert the timestamps into milliseconds."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "ratings = ratings.select(\n",
    "    ratings.userId, ratings.movieId, ratings.rating, (ratings.timestamp.cast(\"long\") * 1000).alias(\"timestamp\"))\n",
    "ratings.show(5)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Movies**\n",
    "\n",
    "The file `movies.csv` contains the `movieId`, `title` and `genres` for each movie. As you can see, the `genres` field is a bit tricky to use, as the genres are in the form of one string delimited by the `|` character: `Adventure|Animation|Children|Comedy|Fantasy`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "# load raw data from CSV\n",
    "raw_movies = spark.read.csv(PATH_TO_DATA + \"/movies.csv\", header=True, inferSchema=True)\n",
    "print(\"Raw movie data:\")\n",
    "raw_movies.show(5, truncate=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Create a `DataFrame` user-defined function (UDF) to extract this delimited string into a list of genres."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "from pyspark.sql.functions import udf\n",
    "from pyspark.sql.types import *\n",
    "# define a UDF to convert the raw genres string to an array of genres and lowercase\n",
    "extract_genres = udf(lambda x: x.lower().split(\"|\"), ArrayType(StringType()))\n",
    "# test it out\n",
    "raw_movies.select(\"movieId\", \"title\", extract_genres(\"genres\").alias(\"genres\")).show(5, False)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Ok, that looks better!\n",
    "\n",
    "You may also notice that the movie titles contain the year of release. It would be useful to have that as a field in your search index for filtering results (say you want to filter our recommendations to include only more recent movies).\n",
    "\n",
    "Create a UDF to extract the release year from the title using a Python regular expression."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "import re\n",
    "# define a UDF to extract the release year from the title, and return the new title and year in a struct type\n",
    "def extract_year_fn(title):\n",
    "    result = re.search(\"\\(\\d{4}\\)\", title)\n",
    "    try:\n",
    "        if result:\n",
    "            group = result.group()\n",
    "            year = group[1:-1]\n",
    "            start_pos = result.start()\n",
    "            title = title[:start_pos-1]\n",
    "            return (title, year)\n",
    "        else:\n",
    "            return (title, 1970)\n",
    "    except:\n",
    "        print(title)\n",
    "\n",
    "extract_year = udf(extract_year_fn,\\\n",
    "                   StructType([StructField(\"title\", StringType(), True),\\\n",
    "                               StructField(\"release_date\", StringType(), True)]))\n",
    "    \n",
    "# test out our function\n",
    "s = \"Jumanji (1995)\"\n",
    "extract_year_fn(s)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Ok the function works! Now create a new `DataFrame` with the cleaned-up titles, release dates and genres of the movies."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "movies = raw_movies.select(\n",
    "    \"movieId\", extract_year(\"title\").title.alias(\"title\"),\\\n",
    "    extract_year(\"title\").release_date.alias(\"release_date\"),\\\n",
    "    extract_genres(\"genres\").alias(\"genres\"))\n",
    "print(\"Cleaned movie data:\")\n",
    "movies.show(5, truncate=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Next, join the `links.csv` data to `movies` so that there is an id for _The Movie Database_ corresponding to each movie. You can use this id to retrieve movie poster images when displaying your recommendations later."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "link_data = spark.read.csv(PATH_TO_DATA + \"/links.csv\", header=True, inferSchema=True)\n",
    "# join movies with links to get TMDB id\n",
    "movie_data = movies.join(link_data, movies.movieId == link_data.movieId)\\\n",
    "    .select(movies.movieId, movies.title, movies.release_date, movies.genres, link_data.tmdbId)\n",
    "num_movies = movie_data.count()\n",
    "print(\"Cleaned movie data with tmdbId links:\")\n",
    "movie_data.show(5, truncate=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "> **_Optional_**\n",
    "\n",
    "> Run the below cell to test your access to TMDb API. You should see the _Toy Story_ movie poster displayed inline.\n",
    "\n",
    "> To install the Python package run `pip install tmdbsimple`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true,
    "hide_input": true
   },
   "outputs": [],
   "source": [
    "try:\n",
    "    import tmdbsimple as tmdb\n",
    "    # replace this variable with your actual TMdb API key\n",
    "    tmdb.API_KEY = 'YOUR_API_KEY'\n",
    "    print(\"Successfully imported tmdbsimple!\")\n",
    "    # base URL for TMDB poster images\n",
    "    IMAGE_URL = 'https://image.tmdb.org/t/p/w500'\n",
    "    movie_id = movie_data.first().tmdbId\n",
    "    movie_info = tmdb.Movies(movie_id).info()\n",
    "    movie_poster_url = IMAGE_URL + movie_info['poster_path']\n",
    "    display(Image(movie_poster_url, width=200))\n",
    "except Exception:\n",
    "    print(\"Cannot import tmdbsimple, no movie posters will be displayed!\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Step 2: Load data into Elasticsearch\n",
    "\n",
    "Now that you have your dataset processed and prepared, you will load it into Elasticsearch.\n",
    "\n",
    "_Note:_ for the purposes of this demo notebook you have started with an existing example dataset and will load that into Elasticsearch. In practice you may write your event data as well as user and item metadata from your application directly into Elasticsearch.\n",
    "\n",
    "First test that your Elasticsearch instance is running and you can connect to it using the Python Elasticsearch client."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from elasticsearch import Elasticsearch\n",
    "\n",
    "# test your ES instance is running\n",
    "es = Elasticsearch()\n",
    "es.info(pretty=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Create an Elasticsearch index with mappings for users, movies and rating events\n",
    "\n",
    "In Elasticsearch, an \"index\" is roughly similar to a \"database\", while a \"document type\" is roughly similar to a \"table\" in that database. The schema for a document type is called an index mapping.\n",
    "\n",
    "While Elasticsearch supports dynamic mapping, it's advisable to specify the mapping explicitly when creating an index if you know what your data looks like.\n",
    "\n",
    "For the purposes of your recommendation engine, this is also necessary so that you can specify a custom analyzer for the field that will hold the recommendation \"model\" (that is, the factor vectors). This will ensure the vector-scoring plugin will work correctly.\n",
    "\n",
    "> _Note_ This notebook does not go into detail about the underlying scoring mechanism or the relevant Elasticsearch internals. See the talks and slides in the [Journey Links section](https://github.com/MLnick/elasticsearch-spark-recommender-demo/blob/master/README.md#links) for more detail.\n",
    "\n",
    "__References:__\n",
    "* [Create index request](https://www.elastic.co/guide/en/elasticsearch/reference/current/indices-create-index.html)\n",
    "* [Delimited payload filter](https://www.elastic.co/guide/en/elasticsearch/reference/2.4/analysis-delimited-payload-tokenfilter.html)\n",
    "* [Term vectors](https://www.elastic.co/guide/en/elasticsearch/reference/2.4/docs-termvectors.html#_term_information)\n",
    "* [Mapping](https://www.elastic.co/guide/en/elasticsearch/reference/2.4/mapping.html)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "> **_Optional_**\n",
    "\n",
    "> If you are re-running the notebook and have previously created the `demo` index in Elasticsearch, you should first delete it by un-commenting and running the next cell, before running the index creation cell that follows."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "# es.indices.delete(index=\"demo\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now you're ready to create your index."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "create_index = {\n",
    "    \"settings\": {\n",
    "        \"analysis\": {\n",
    "            \"analyzer\": {\n",
    "                # this configures the custom analyzer we need to parse vectors such that the scoring\n",
    "                # plugin will work correctly\n",
    "                \"payload_analyzer\": {\n",
    "                    \"type\": \"custom\",\n",
    "                    \"tokenizer\":\"whitespace\",\n",
    "                    \"filter\":\"delimited_payload_filter\"\n",
    "                }\n",
    "            }\n",
    "        }\n",
    "    },\n",
    "    \"mappings\": {\n",
    "        \"ratings\": {\n",
    "          # this mapping definition sets up the fields for the rating events\n",
    "          \"properties\": {\n",
    "                \"timestamp\": {\n",
    "                    \"type\": \"date\"\n",
    "                },\n",
    "                \"userId\": {\n",
    "                    \"type\": \"integer\"\n",
    "                },\n",
    "                \"movieId\": {\n",
    "                    \"type\": \"integer\"\n",
    "                },\n",
    "                \"rating\": {\n",
    "                    \"type\": \"double\"\n",
    "                }\n",
    "            }  \n",
    "        },\n",
    "        \"users\": {\n",
    "            # this mapping definition sets up the metadata fields for the users\n",
    "            \"properties\": {\n",
    "                \"userId\": {\n",
    "                    \"type\": \"integer\"\n",
    "                },\n",
    "                \"@model\": {\n",
    "                    # this mapping definition sets up the fields for user factor vectors of our model\n",
    "                    \"properties\": {\n",
    "                        \"factor\": {\n",
    "                            \"type\": \"text\",\n",
    "                            \"term_vector\": \"with_positions_offsets_payloads\",\n",
    "                            \"analyzer\" : \"payload_analyzer\"\n",
    "                        },\n",
    "                        \"version\": {\n",
    "                            \"type\": \"keyword\"\n",
    "                        },\n",
    "                        \"timestamp\": {\n",
    "                            \"type\": \"date\"\n",
    "                        }\n",
    "                    }\n",
    "                }\n",
    "            }\n",
    "        },\n",
    "        \"movies\": {\n",
    "            # this mapping definition sets up the metadata fields for the movies\n",
    "            \"properties\": {\n",
    "                \"movieId\": {\n",
    "                    \"type\": \"integer\"\n",
    "                },\n",
    "                \"tmdbId\": {\n",
    "                    \"type\": \"keyword\"\n",
    "                },\n",
    "                \"genres\": {\n",
    "                    \"type\": \"keyword\"\n",
    "                },\n",
    "                \"release_date\": {\n",
    "                    \"type\": \"date\",\n",
    "                    \"format\": \"year\"\n",
    "                },\n",
    "                \"@model\": {\n",
    "                    # this mapping definition sets up the fields for movie factor vectors of our model\n",
    "                    \"properties\": {\n",
    "                        \"factor\": {\n",
    "                            \"type\": \"text\",\n",
    "                            \"term_vector\": \"with_positions_offsets_payloads\",\n",
    "                            \"analyzer\" : \"payload_analyzer\"\n",
    "                        },\n",
    "                        \"version\": {\n",
    "                            \"type\": \"keyword\"\n",
    "                        },\n",
    "                        \"timestamp\": {\n",
    "                            \"type\": \"date\"\n",
    "                        }\n",
    "                    }\n",
    "                }\n",
    "            }\n",
    "        }\n",
    "    }\n",
    "}\n",
    "# create index with the settings and mappings above\n",
    "es.indices.create(index=\"demo\", body=create_index)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Load Ratings and Movies DataFrames into Elasticsearch\n",
    "\n",
    "First you will write the ratings data to Elasticsearch. Notice that you can simply use the Spark Elasticsearch connector to write a `DataFrame` with the native Spark datasource API by specifying `format(\"es\")`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true,
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# write ratings data\n",
    "ratings.write.format(\"es\").save(\"demo/ratings\")\n",
    "# check write went ok\n",
    "print(\"Dataframe count: %d\" % ratings.count())\n",
    "print(\"ES index count:  %d\" % es.count(index=\"demo\", doc_type=\"ratings\")['count'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "# test things out by retrieving a few rating event documents from Elasticsearch\n",
    "es.search(index=\"demo\", doc_type=\"ratings\", q=\"*\", size=3)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Since you've indexed the rating event data into Elasticsearch, you can use all the capabilities of a search engine to query the data. For example, you could count the number of ratings events in a given date range using Elasticsearch's date math in a query string:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "es.count(index=\"demo\", doc_type=\"ratings\", q=\"timestamp:[2016-01-01 TO 2016-02-01]\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Next write the movie metadata"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "# write movie data, specifying the DataFrame column to use as the id mapping\n",
    "movie_data.write.format(\"es\").option(\"es.mapping.id\", \"movieId\").save(\"demo/movies\")\n",
    "# check load went ok\n",
    "print(\"Movie DF count: %d\" % movie_data.count())\n",
    "print(\"ES index count: %d\" % es.count(index=\"demo\", doc_type=\"movies\")['count'])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Again you can harness the power of search to query the movie metadata:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "# test things out by searching for movies containing \"matrix\" in the title\n",
    "es.search(index=\"demo\", doc_type=\"movies\", q=\"title:matrix\", size=3)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Step 3: Train a recommmender model on the ratings data\n",
    "\n",
    "Your data is now stored in Elasticsearch and you will use the ratings data to build a collaborative filtering recommendation model.\n",
    "\n",
    "[Collaborative filtering](https://en.wikipedia.org/wiki/Collaborative_filtering) is a recommendation approach that is effectively based on the \"wisdom of the crowd\". It makes the assumption that, if two people share similar preferences, then the things that one of them prefers could be good recommendations to make to the other. In other words, if user A tends to like certain movies, and user B shares some of these preferences with user A, then the movies that user A likes, that user B _has not yet seen_, may well be movies that user B will also like.\n",
    "\n",
    "In a similar manner, we can think about _items_ as being similar if they tend to be rated highly by the same people, on average. \n",
    "\n",
    "Hence these models are based on the combined, collaborative preferences and behavior of all users in aggregate. They tend to be very effective in practice (provided you have enough preference data to train the model). The ratings data you have is a form of _explicit preference data_, perfect for training collaborative filtering models.\n",
    "\n",
    "### Alternating Least Squares\n",
    "\n",
    "Alternating Least Squares (ALS) is a specific algorithm for solving a type of collaborative filtering model known as [matrix factorization (MF)](https://en.wikipedia.org/wiki/Matrix_decomposition). The core idea of MF is to represent the ratings as a _user-item ratings matrix_. In the diagram below you will see this matrix on the left (with users as _rows_ and movies as _columns_). The entries in this matrix are the ratings given by users to movies.\n",
    "\n",
    "You may also notice that the matrix has _missing entries_ because not all users have rated all movies. In this situation we refer to the data as _sparse_.\n",
    "\n",
    "![als-diagram.png](../doc/source/images/als-diagram.png)\n",
    "\n",
    "MF methods aim to find two much smaller matrices (one representing the _users_ and the other the _items_) that, when multiplied together, re-construct the original ratings matrix as closely as possible. This is know as _factorizing_ the original matrix, hence the name of the technique.\n",
    "\n",
    "The two smaller matrices are called _factor matrices_ (or _latent features_). The user and movie factor matrices are illustrated on the right in the diagram above. The idea is that each user factor vector is a compressed representation of the user's preferences and behavior. Likewise, each item factor vector is a compressed representation of the item. Once the model is trained, the factor vectors can be used to make recommendations, which is what you will do in the following sections.\n",
    "\n",
    "__Further reading:__\n",
    "\n",
    "* [Spark MLlib Collaborative Filtering](http://spark.apache.org/docs/latest/ml-collaborative-filtering.html)\n",
    "* [Alternating Least Squares and collaborative filtering](https://datasciencemadesimpler.wordpress.com/tag/alternating-least-squares/)\n",
    "* [Quora question on Alternating Least Squares](https://www.quora.com/What-is-the-Alternating-Least-Squares-method-in-recommendation-systems-And-why-does-this-algorithm-work-intuition-behind-this)\n",
    "\n",
    "Fortunately, Spark's MLlib machine learning library has a scalable, efficient implementation of matrix factorization built in, which we can use to train our recommendation model. Next, you will use Spark's ALS to train a model on your ratings data from Elasticsearch."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true,
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "ratings_from_es = spark.read.format(\"es\").load(\"demo/ratings\")\n",
    "ratings_from_es.show(5)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true,
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "from pyspark.ml.recommendation import ALS\n",
    "from pyspark.sql.functions import col\n",
    "als = ALS(userCol=\"userId\", itemCol=\"movieId\", ratingCol=\"rating\", regParam=0.01, rank=20, seed=12)\n",
    "model = als.fit(ratings_from_es)\n",
    "model.userFactors.show(5)\n",
    "model.itemFactors.show(5)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Step 4: Export ALS user and item factor vectors to Elasticsearch\n",
    "\n",
    "Congratulations, you've trained a recommendation model! The next step is to export the model factors (shown in the `DataFrames` above) to Elasticsearch.\n",
    "\n",
    "In order to store the model in the correct format for the index mappings set up earlier, you will need to create some utility functions. These functions will allow you to convert the raw vectors (which are equivalent to a Python list in the factor `DataFrames` above) to the correct _delimited string format_. This ensures Elasticsearch will parse the vector field in the model correctly using the delimited token filter custom analyzer you configured earlier.\n",
    "\n",
    "You will also create a function to convert a vector and related metadata (such as the Spark model id and a timestamp) into a `DataFrame` field that matches the `model` field in the Elasticsearch index mapping."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Utility functions for converting factor vectors"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true,
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "from pyspark.sql.types import *\n",
    "from pyspark.sql.functions import udf, lit, current_timestamp, unix_timestamp\n",
    "\n",
    "def convert_vector(x):\n",
    "    '''Convert a list or numpy array to delimited token filter format'''\n",
    "    return \" \".join([\"%s|%s\" % (i, v) for i, v in enumerate(x)])\n",
    "\n",
    "def reverse_convert(s):\n",
    "    '''Convert a delimited token filter format string back to list format'''\n",
    "    return  [float(f.split(\"|\")[1]) for f in s.split(\" \")]\n",
    "\n",
    "def vector_to_struct(x, version, ts):\n",
    "    '''Convert a vector to a SparkSQL Struct with string-format vector and version fields'''\n",
    "    return (convert_vector(x), version, ts)\n",
    "\n",
    "vector_struct = udf(vector_to_struct, \\\n",
    "                    StructType([StructField(\"factor\", StringType(), True), \\\n",
    "                                StructField(\"version\", StringType(), True),\\\n",
    "                                StructField(\"timestamp\", LongType(), True)]))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true,
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "# test out the vector conversion function\n",
    "test_vec = model.userFactors.select(\"features\").first().features\n",
    "print(test_vec)\n",
    "print()\n",
    "print(convert_vector(test_vec))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Convert factor vectors to [factor, version, timestamp] form and write to Elasticsearch"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true,
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "ver = model.uid\n",
    "ts = unix_timestamp(current_timestamp())\n",
    "movie_vectors = model.itemFactors.select(\"id\", vector_struct(\"features\", lit(ver), ts).alias(\"@model\"))\n",
    "movie_vectors.select(\"id\", \"@model.factor\", \"@model.version\", \"@model.timestamp\").show(5)\n",
    "user_vectors = model.userFactors.select(\"id\", vector_struct(\"features\", lit(ver), ts).alias(\"@model\"))\n",
    "user_vectors.select(\"id\", \"@model.factor\", \"@model.version\", \"@model.timestamp\").show(5)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true,
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "# write data to ES, use:\n",
    "# - \"id\" as the column to map to ES movie id\n",
    "# - \"update\" write mode for ES, since you want to update new fields only\n",
    "# - \"append\" write mode for Spark\n",
    "movie_vectors.write.format(\"es\") \\\n",
    "    .option(\"es.mapping.id\", \"id\") \\\n",
    "    .option(\"es.write.operation\", \"update\") \\\n",
    "    .save(\"demo/movies\", mode=\"append\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true,
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "# write data to ES, use:\n",
    "# - \"id\" as the column to map to ES movie id\n",
    "# - \"index\" write mode for ES, since you have not written to the user index previously\n",
    "# - \"append\" write mode for Spark\n",
    "user_vectors.write.format(\"es\") \\\n",
    "    .option(\"es.mapping.id\", \"id\") \\\n",
    "    .option(\"es.write.operation\", \"index\") \\\n",
    "    .save(\"demo/users\", mode=\"append\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Check the data  was written correctly\n",
    "\n",
    "You can search for a movie to see if the model factor vector was written correctly. You should see a `'@model': {'factor': '0|...` field in the returned movie document."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true,
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "# search for a particular sci-fi movie\n",
    "es.search(index=\"demo\", doc_type=\"movies\", q=\"star wars phantom menace\", size=1)['hits']['hits'][0]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Step 5: Recommend using Elasticsearch!\n",
    "\n",
    "Now that you have loaded your recommendation model into Elasticsearch, you will generate some recommendations.\n",
    "First, you will need to create a few utility functions for:\n",
    "\n",
    "* Fetching movie posters from TMdb API (optional)\n",
    "* Constructing the Elasticsearch [function score query](https://www.elastic.co/guide/en/elasticsearch/reference/5.3/query-dsl-function-score-query.html) to generate recommendations from your factor model\n",
    "* Given a movie, use this query to find the movies most similar to it\n",
    "* Given a user, use this query to find the movies with the highest predicted rating, to recommend to the user\n",
    "* Display the results as an HTML table in Jupyter"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true,
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "from IPython.display import Image, HTML, display\n",
    "\n",
    "def get_poster_url(id):\n",
    "    \"\"\"Fetch movie poster image URL from TMDb API given a tmdbId\"\"\"\n",
    "    IMAGE_URL = 'https://image.tmdb.org/t/p/w500'\n",
    "    try:\n",
    "        import tmdbsimple as tmdb\n",
    "        from tmdbsimple import APIKeyError\n",
    "        try:\n",
    "            movie = tmdb.Movies(id).info()\n",
    "            poster_url = IMAGE_URL + movie['poster_path'] if 'poster_path' in movie and movie['poster_path'] is not None else \"\"\n",
    "            return poster_url\n",
    "        except APIKeyError as ae:\n",
    "            return \"KEY_ERR\"\n",
    "    except Exception as me:\n",
    "        return \"NA\"\n",
    "    \n",
    "    \n",
    "def fn_query(query_vec, q=\"*\", cosine=False):\n",
    "    \"\"\"\n",
    "    Construct an Elasticsearch function score query.\n",
    "    \n",
    "    The query takes as parameters:\n",
    "        - the field in the candidate document that contains the factor vector\n",
    "        - the query vector\n",
    "        - a flag indicating whether to use dot product or cosine similarity (normalized dot product) for scores\n",
    "        \n",
    "    The query vector passed in will be the user factor vector (if generating recommended movies for a user)\n",
    "    or movie factor vector (if generating similar movies for a given movie)\n",
    "    \"\"\"\n",
    "    return {\n",
    "    \"query\": {\n",
    "        \"function_score\": {\n",
    "            \"query\" : { \n",
    "                \"query_string\": {\n",
    "                    \"query\": q\n",
    "                }\n",
    "            },\n",
    "            \"script_score\": {\n",
    "                \"script\": {\n",
    "                        \"inline\": \"payload_vector_score\",\n",
    "                        \"lang\": \"native\",\n",
    "                        \"params\": {\n",
    "                            \"field\": \"@model.factor\",\n",
    "                            \"vector\": query_vec,\n",
    "                            \"cosine\" : cosine\n",
    "                        }\n",
    "                    }\n",
    "            },\n",
    "            \"boost_mode\": \"replace\"\n",
    "        }\n",
    "    }\n",
    "}\n",
    "\n",
    "\n",
    "def get_similar(the_id, q=\"*\", num=10, index=\"demo\", dt=\"movies\"):\n",
    "    \"\"\"\n",
    "    Given a movie id, execute the recommendation function score query to find similar movies, ranked by cosine similarity\n",
    "    \"\"\"\n",
    "    response = es.get(index=index, doc_type=dt, id=the_id)\n",
    "    src = response['_source']\n",
    "    if '@model' in src and 'factor' in src['@model']:\n",
    "        raw_vec = src['@model']['factor']\n",
    "        # our script actually uses the list form for the query vector and handles conversion internally\n",
    "        query_vec = reverse_convert(raw_vec)\n",
    "        q = fn_query(query_vec, q=q, cosine=True)\n",
    "        results = es.search(index, dt, body=q)\n",
    "        hits = results['hits']['hits']\n",
    "        return src, hits[1:num+1]\n",
    "    \n",
    "    \n",
    "def get_user_recs(the_id, q=\"*\", num=10, index=\"demo\"):\n",
    "    \"\"\"\n",
    "    Given a user id, execute the recommendation function score query to find top movies, ranked by predicted rating\n",
    "    \"\"\"\n",
    "    response = es.get(index=index, doc_type=\"users\", id=the_id)\n",
    "    src = response['_source']\n",
    "    if '@model' in src and 'factor' in src['@model']:\n",
    "        raw_vec = src['@model']['factor']\n",
    "        # our script actually uses the list form for the query vector and handles conversion internally\n",
    "        query_vec = reverse_convert(raw_vec)\n",
    "        q = fn_query(query_vec, q=q, cosine=False)\n",
    "        results = es.search(index, \"movies\", body=q)\n",
    "        hits = results['hits']['hits']\n",
    "        return src, hits[:num]\n",
    "\n",
    "def get_movies_for_user(the_id, num=10, index=\"demo\"):\n",
    "    \"\"\"\n",
    "    Given a user id, get the movies rated by that user, from highest- to lowest-rated.\n",
    "    \"\"\"\n",
    "    response = es.search(index=index, doc_type=\"ratings\", q=\"userId:%s\" % the_id, size=num, sort=[\"rating:desc\"])\n",
    "    hits = response['hits']['hits']\n",
    "    ids = [h['_source']['movieId'] for h in hits]\n",
    "    movies = es.mget(body={\"ids\": ids}, index=index, doc_type=\"movies\", _source_include=['tmdbId', 'title'])\n",
    "    movies_hits = movies['docs']\n",
    "    tmdbids = [h['_source'] for h in movies_hits]\n",
    "    return tmdbids\n",
    "\n",
    "            \n",
    "def display_user_recs(the_id, q=\"*\", num=10, num_last=10, index=\"demo\"):\n",
    "    user, recs = get_user_recs(the_id, q, num, index)\n",
    "    user_movies = get_movies_for_user(the_id, num_last, index)\n",
    "    # check that posters can be displayed\n",
    "    first_movie = user_movies[0]\n",
    "    first_im_url = get_poster_url(first_movie['tmdbId'])\n",
    "    if first_im_url == \"NA\":\n",
    "        display(HTML(\"<i>Cannot import tmdbsimple. No movie posters will be displayed!</i>\"))\n",
    "    if first_im_url == \"KEY_ERR\":\n",
    "        display(HTML(\"<i>Key error accessing TMDb API. Check your API key. No movie posters will be displayed!</i>\"))\n",
    "        \n",
    "    # display the movies that this user has rated highly\n",
    "    display(HTML(\"<h2>Get recommended movies for user id %s</h2>\" % the_id))\n",
    "    display(HTML(\"<h4>The user has rated the following movies highly:</h4>\"))\n",
    "    user_html = \"<table border=0>\"\n",
    "    i = 0\n",
    "    for movie in user_movies:\n",
    "        movie_im_url = get_poster_url(movie['tmdbId'])\n",
    "        movie_title = movie['title']\n",
    "        user_html += \"<td><h5>%s</h5><img src=%s width=150></img></td>\" % (movie_title, movie_im_url)\n",
    "        i += 1\n",
    "        if i % 5 == 0:\n",
    "            user_html += \"</tr><tr>\"\n",
    "    user_html += \"</tr></table>\"\n",
    "    display(HTML(user_html))\n",
    "    # now display the recommended movies for the user\n",
    "    display(HTML(\"<br>\"))\n",
    "    display(HTML(\"<h2>Recommended movies:</h2>\"))\n",
    "    rec_html = \"<table border=0>\"\n",
    "    i = 0\n",
    "    for rec in recs:\n",
    "        r_im_url = get_poster_url(rec['_source']['tmdbId'])\n",
    "        r_score = rec['_score']\n",
    "        r_title = rec['_source']['title']\n",
    "        rec_html += \"<td><h5>%s</h5><img src=%s width=150></img></td><td><h5>%2.3f</h5></td>\" % (r_title, r_im_url, r_score)\n",
    "        i += 1\n",
    "        if i % 5 == 0:\n",
    "            rec_html += \"</tr><tr>\"\n",
    "    rec_html += \"</tr></table>\"\n",
    "    display(HTML(rec_html))\n",
    "\n",
    "    \n",
    "def display_similar(the_id, q=\"*\", num=10, index=\"demo\", dt=\"movies\"):\n",
    "    \"\"\"\n",
    "    Display query movie, together with similar movies and similarity scores, in a table\n",
    "    \"\"\"\n",
    "    movie, recs = get_similar(the_id, q, num, index, dt)\n",
    "    q_im_url = get_poster_url(movie['tmdbId'])\n",
    "    if q_im_url == \"NA\":\n",
    "        display(HTML(\"<i>Cannot import tmdbsimple. No movie posters will be displayed!</i>\"))\n",
    "    if q_im_url == \"KEY_ERR\":\n",
    "        display(HTML(\"<i>Key error accessing TMDb API. Check your API key. No movie posters will be displayed!</i>\"))\n",
    "        \n",
    "    display(HTML(\"<h2>Get similar movies for:</h2>\"))\n",
    "    display(HTML(\"<h4>%s</h4>\" % movie['title']))\n",
    "    if q_im_url != \"NA\":\n",
    "        display(Image(q_im_url, width=200))\n",
    "    display(HTML(\"<br>\"))\n",
    "    display(HTML(\"<h2>People who liked this movie also liked these:</h2>\"))\n",
    "    sim_html = \"<table border=0>\"\n",
    "    i = 0\n",
    "    for rec in recs:\n",
    "        r_im_url = get_poster_url(rec['_source']['tmdbId'])\n",
    "        r_score = rec['_score']\n",
    "        r_title = rec['_source']['title']\n",
    "        sim_html += \"<td><h5>%s</h5><img src=%s width=150></img></td><td><h5>%2.3f</h5></td>\" % (r_title, r_im_url, r_score)\n",
    "        i += 1\n",
    "        if i % 5 == 0:\n",
    "            sim_html += \"</tr><tr>\"\n",
    "    sim_html += \"</tr></table>\"\n",
    "    display(HTML(sim_html))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now, you're ready to generate some recommendations.\n",
    "\n",
    "### 5(a) Find similar movies for a given movie\n",
    "\n",
    "To start, you can find movies that are _similar_ to a given movie. This similarity score is computed from the model factor vectors for each movie. Recall that the ALS model you trained earlier is a collaborative filtering model, so the similarity between movie vectors will be based on the _rating co-occurrence_ of the movies. In other words, two movies that tend to be rated highly by a user will tend to be more similar. It is common to use the [cosine similarity](https://en.wikipedia.org/wiki/Cosine_similarity) of the movie factor vectors as a measure of the similarity between two movies.\n",
    "\n",
    "Using this similarity you can show recommendations along the lines of _people who liked this movie also liked these_."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "display_similar(2628, num=5)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "So we see that people who like Star Wars tend like other sci-fi movies (including other Star Wars films), as well as some action and drama.\n",
    "\n",
    "> _Note_ since we are using a very small dataset, results may not be as good as those for the same model trained on a larger dataset.\n",
    "\n",
    "Now you will see the power and flexibility that comes from using a search engine to generate recommendations. Elasticsearch allows you to tweak the results returned by the recommendation query using any standard search query or filter - from free text search through to filters based on time and geo-location (or any other piece of metadata you can think of).\n",
    "\n",
    "For example, perhaps you want to remove any movies with \"matrix\" in the title from the recommendations. You can do this by simply passing a valid Elasticsearch query string to the recommendation function."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "display_similar(2628, num=5, q=\"title:(NOT matrix)\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Or you may want to ensure that only valid children's movies are shown to young viewers."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "display_similar(1, num=5, q=\"genres:children\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Feel free to check out the documentation for the Elasticsearch [query string query](https://www.elastic.co/guide/en/elasticsearch/reference/current/query-dsl-query-string-query.html) and play around with the various queries you can construct by passing in a query string as `q` in the recommendation function above!"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "collapsed": true
   },
   "source": [
    "### 5(b) Find movies to recommend to a user\n",
    "\n",
    "Now, you're ready to generate some movie recommendations, personalized for a specific user.\n",
    "\n",
    "Given a user, you can recommend movies to that user based on the predicted ratings from your model. In a similar manner to the similar movie recommendations, this predicted rating score is computed from the model factor vector for the user and the factor vectors for each movie. Recall that the collaborative filtering model means that, at a high level, we will recommend movies _liked by other users who liked the same movies as the given user_."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "display_user_recs(12, num=5, num_last=5)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Again, note that since we are using a very small dataset, the results may not be too good. However, we can see that this user seems to like some sci-fi, some horror and some comedy films. The recommended movies fall broadly into these categories and seem to be somewhat reasonable.\n",
    "\n",
    "Next, you can again apply the power of Elasticsearch's filtering capabilities to your recommendation engine. Let's say you only want to recommend more recent movies (say, from the past 5 years). This can be done by adding a date math query to the recommendation function score query."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "display_user_recs(12, num=5, num_last=5, q=\"release_date:[2012 TO *]\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "You can see that the recommendation include only recent movies, and this time they seem to be heavily tilted to sci-fi and fantasy genres.\n",
    "\n",
    "As you did with the similar movies recommendations, feel free to play around with the various queries you could pass into the user recommendation query."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 2",
   "language": "python",
   "name": "python2"
  },
  "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.1"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}