notebooks/elasticsearch-spark-recommender.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "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, Apache Spark and the Elasticsearch Spark connector detailed in the [setup steps](https://github.com/IBM/elasticsearch-spark-recommender/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's script score query together with vector functions"
   ]
  },
  {
   "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",
    "* Follow the [Code Pattern instructions](https://github.com/IBM/elasticsearch-spark-recommender/tree/master#5-download-the-data) to 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": {},
   "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": {},
   "outputs": [],
   "source": [
    "# if you unzipped the data to a different location than that specified in the Code Pattern 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: {}\".format(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": {},
   "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": {},
   "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": [
    "For indexing into Elasticsearch, we would prefer to represent the genres as a list. Create a `DataFrame` user-defined function (UDF) to extract this delimited string into a list of genres."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "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": {},
   "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": {},
   "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": {},
   "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` in your console (see the [Code Pattern Steps](https://github.com/IBM/elasticsearch-spark-recommender/tree/master#6-launch-the-notebook))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "hide_input": true
   },
   "outputs": [],
   "source": [
    "try:\n",
    "    import tmdbsimple as tmdb\n",
    "    import json\n",
    "    from requests.exceptions import HTTPError\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 ImportError:\n",
    "    print(\"Cannot import tmdbsimple as it is not installed, no movie posters will be displayed!\")\n",
    "except HTTPError as e:\n",
    "    if e.response.status_code == 401:\n",
    "        j = json.loads(e.response.text)\n",
    "        print(\"TMdb API call failed: {}\".format(j['status_message']))"
   ]
  },
  {
   "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 Elasticsearch indices, with mappings for users, movies and rating events\n",
    "\n",
    "In Elasticsearch, an \"index\" is roughly similar to a \"database\" or \"database table\". The schema for an index 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 the vector field that will hold the recommendation \"model\" (that is, the factor vectors). When creating a vector field, you need to provide the dimension of the vector explicitly, so it cannot be a dynamic mapping.\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 [Code Pattern Links section](https://github.com/IBM/elasticsearch-spark-recommender/blob/master/README.md#links) for more detail.\n",
    "\n",
    "__References:__\n",
    "* [Create index API](https://www.elastic.co/guide/en/elasticsearch/reference/current/indices-create-index.html)\n",
    "* [Index mapping](https://www.elastic.co/guide/en/elasticsearch/reference/current/mapping.html)\n",
    "* [Dense vector types](https://www.elastic.co/guide/en/elasticsearch/reference/current/dense-vector.html)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "> **_Optional_**\n",
    "\n",
    "> If you are re-running the notebook and have previously created the `movies`, `users` and `ratings`indices in Elasticsearch, you should first delete them by un-commenting and running the next cell, before running the index creation cell that follows."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# es.indices.delete(index=\"ratings,users,movies\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now you're ready to create your indices."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# set the factor vector dimension for the recommendation model\n",
    "VECTOR_DIM = 20\n",
    "\n",
    "create_ratings = {\n",
    "    # this mapping definition sets up the fields for the rating events\n",
    "    \"mappings\": {\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",
    "}\n",
    "\n",
    "create_users = {\n",
    "    # this mapping definition sets up the metadata fields for the users\n",
    "    \"mappings\": {\n",
    "        \"properties\": {\n",
    "            \"userId\": {\n",
    "                \"type\": \"integer\"\n",
    "            },\n",
    "            # the following fields define our model factor vectors and metadata\n",
    "            \"model_factor\": {\n",
    "                \"type\": \"dense_vector\",\n",
    "                \"dims\" : VECTOR_DIM\n",
    "            },\n",
    "            \"model_version\": {\n",
    "                \"type\": \"keyword\"\n",
    "            },\n",
    "            \"model_timestamp\": {\n",
    "                \"type\": \"date\"\n",
    "            }\n",
    "        }\n",
    "    }\n",
    "}\n",
    "\n",
    "create_movies = {\n",
    "    # this mapping definition sets up the metadata fields for the movies\n",
    "    \"mappings\": {\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",
    "            # the following fields define our model factor vectors and metadata\n",
    "            \"model_factor\": {\n",
    "                \"type\": \"dense_vector\",\n",
    "                \"dims\" : VECTOR_DIM\n",
    "            },\n",
    "            \"model_version\": {\n",
    "                \"type\": \"keyword\"\n",
    "            },\n",
    "            \"model_timestamp\": {\n",
    "                \"type\": \"date\"\n",
    "            }          \n",
    "        }\n",
    "    }\n",
    "}\n",
    "\n",
    "# create indices with the settings and mappings above\n",
    "res_ratings = es.indices.create(index=\"ratings\", body=create_ratings)\n",
    "res_users = es.indices.create(index=\"users\", body=create_users)\n",
    "res_movies = es.indices.create(index=\"movies\", body=create_movies)\n",
    "\n",
    "print(\"Created indices:\")\n",
    "print(res_ratings)\n",
    "print(res_users)\n",
    "print(res_movies)"
   ]
  },
  {
   "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": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# write ratings data\n",
    "ratings.write.format(\"es\").save(\"ratings\")\n",
    "num_ratings_es = es.count(index=\"ratings\")['count']\n",
    "num_ratings_df = ratings.count()\n",
    "# check write went ok\n",
    "print(\"Dataframe count: {}\".format(num_ratings_df))\n",
    "print(\"ES index count:  {}\".format(num_ratings_es))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# test things out by retrieving a few rating event documents from Elasticsearch\n",
    "es.search(index=\"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": {},
   "outputs": [],
   "source": [
    "es.count(index=\"ratings\", q=\"timestamp:[2018-01-01 TO 2018-02-01]\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Next write the movie metadata"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "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(\"movies\")\n",
    "num_movies_df = movie_data.count()\n",
    "num_movies_es = es.count(index=\"movies\")['count']\n",
    "# check load went ok\n",
    "print(\"Movie DF count: {}\".format(num_movies_df))\n",
    "print(\"ES index count: {}\".format(num_movies_es))"
   ]
  },
  {
   "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": {},
   "outputs": [],
   "source": [
    "# test things out by searching for movies containing \"matrix\" in the title\n",
    "es.search(index=\"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": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "ratings_from_es = spark.read.format(\"es\").load(\"ratings\")\n",
    "ratings_from_es.show(5)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "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.02, rank=VECTOR_DIM, seed=54)\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",
    "We can export the model factor vector directly to Elasticsearch, since it is an array and the `dense_vector` field expects an array as input.\n",
    "\n",
    "For illustrative purposes, we will also export model metadata (such as the Spark model id and a timestamp)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Write the model factor vectors, model version and model timestamp to Elasticsearch"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "from pyspark.sql.functions import lit, current_timestamp, unix_timestamp\n",
    "ver = model.uid\n",
    "ts = unix_timestamp(current_timestamp())\n",
    "movie_vectors = model.itemFactors.select(\"id\",\\\n",
    "                                         col(\"features\").alias(\"model_factor\"),\\\n",
    "                                         lit(ver).alias(\"model_version\"),\\\n",
    "                                         ts.alias(\"model_timestamp\"))\n",
    "movie_vectors.show(5)\n",
    "user_vectors = model.userFactors.select(\"id\",\\\n",
    "                                        col(\"features\").alias(\"model_factor\"),\\\n",
    "                                        lit(ver).alias(\"model_version\"),\\\n",
    "                                        ts.alias(\"model_timestamp\"))\n",
    "user_vectors.show(5)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "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(\"movies\", mode=\"append\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "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(\"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..188..., ]` field in the returned movie document, as well as a `model_version` and `model_timestamp` field)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "# search for a particular sci-fi movie\n",
    "es.search(index=\"movies\", q=\"force awakens\")['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 [script score query for vector functions](https://www.elastic.co/guide/en/elasticsearch/reference/current/query-dsl-script-score-query.html#vector-functions) 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": {
    "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",
    "    \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 vector_query(query_vec, vector_field, q=\"*\", cosine=False):\n",
    "    \"\"\"\n",
    "    Construct an Elasticsearch script score query using `dense_vector` fields\n",
    "    \n",
    "    The script score query takes as parameters the query vector (as a Python list)\n",
    "    \n",
    "    Parameters\n",
    "    ----------\n",
    "    query_vec : list\n",
    "        The query vector\n",
    "    vector_field : str\n",
    "        The field name in the document against which to score `query_vec`\n",
    "    q : str, optional\n",
    "        Query string for the search query (default: '*' to search across all documents)\n",
    "    cosine : bool, optional\n",
    "        Whether to compute cosine similarity. If `False` then the dot product is computed (default: False)\n",
    "     \n",
    "    Note: Elasticsearch cannot rank negative scores. Therefore, in the case of the dot product, a sigmoid transform\n",
    "    is applied. In the case of cosine similarity, 1.0 is added to the score. In both cases, documents with no \n",
    "    factor vectors are ignored by applying a 0.0 score.\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",
    "    \n",
    "    if cosine:\n",
    "        score_fn = \"doc['{v}'].size() == 0 ? 0 : cosineSimilarity(params.vector, '{v}') + 1.0\"\n",
    "    else:\n",
    "        score_fn = \"doc['{v}'].size() == 0 ? 0 : sigmoid(1, Math.E, -dotProduct(params.vector, '{v}'))\"\n",
    "       \n",
    "    score_fn = score_fn.format(v=vector_field, fn=score_fn)\n",
    "    \n",
    "    return {\n",
    "    \"query\": {\n",
    "        \"script_score\": {\n",
    "            \"query\" : { \n",
    "                \"query_string\": {\n",
    "                    \"query\": q\n",
    "                }\n",
    "            },\n",
    "            \"script\": {\n",
    "                \"source\": score_fn,\n",
    "                \"params\": {\n",
    "                    \"vector\": query_vec\n",
    "                }\n",
    "            }\n",
    "        }\n",
    "    }\n",
    "}\n",
    "\n",
    "\n",
    "def get_similar(the_id, q=\"*\", num=10, index=\"movies\", vector_field='model_factor'):\n",
    "    \"\"\"\n",
    "    Given a movie id, execute the recommendation script score query to find similar movies,\n",
    "    ranked by cosine similarity. We return the `num` most similar, excluding the movie itself.\n",
    "    \"\"\"\n",
    "    response = es.get(index=index, id=the_id)\n",
    "    src = response['_source']\n",
    "    if vector_field in src:\n",
    "        query_vec = src[vector_field]\n",
    "        q = vector_query(query_vec, vector_field, q=q, cosine=True)\n",
    "        results = es.search(index=index, 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, users=\"users\", movies=\"movies\", vector_field='model_factor'):\n",
    "    \"\"\"\n",
    "    Given a user id, execute the recommendation script score query to find top movies, ranked by predicted rating\n",
    "    \"\"\"\n",
    "    response = es.get(index=users, id=the_id)\n",
    "    src = response['_source']\n",
    "    if vector_field in src:\n",
    "        query_vec = src[vector_field]\n",
    "        q = vector_query(query_vec, vector_field, 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, ratings=\"ratings\", movies=\"movies\"):\n",
    "    \"\"\"\n",
    "    Given a user id, get the movies rated by that user, from highest- to lowest-rated.\n",
    "    \"\"\"\n",
    "    response = es.search(index=ratings, q=\"userId:{}\".format(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=movies, _source_includes=['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, users=\"users\", movies=\"movies\", ratings=\"ratings\"):\n",
    "    user, recs = get_user_recs(the_id, q, num, users, movies)\n",
    "    user_movies = get_movies_for_user(the_id, num_last, ratings, movies)\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, movies=\"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, movies)\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(858, num=5)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "So we see that people who like the original Godfather movie tend like other crime drama movies (including another Godfather film), as well as some general drama movies.\n",
    "\n",
    "> _Note_ since we are using a relatively small and sparse 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",
    "#### Filter recommendations based on title\n",
    "\n",
    "For example, perhaps you want to remove any movies with \"godfather\" 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(858, num=5, q=\"title:(NOT godfather)\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We see that we now only have similar movies that are _not also part of the Godfather trilogy_.\n",
    "\n",
    "#### Filter recommendations based on genre\n",
    "\n",
    "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 the same manner as for 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(72, num=5, num_last=5)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Again, note that since we are using a relatively small and sparse dataset, the results may not be too good. However, we can see that this user seems to like some sci-fi and some drama films. The recommended movies fall broadly into these categories and seem to be somewhat reasonable.\n",
    "\n",
    "#### Filter recommendations based on release date\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 3 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(72, num=5, num_last=5, q=\"release_date:[2017 TO *]\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "You can see that the recommendation include only recent movies, again including mostly drama and sci-fi / fantasy titles.\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 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}