notebooks/03_ann_es_docker.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Approximate Nearest Neighbors on Elastic Search with Docker\n",
    "[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/stephenleo/adventures-with-ann/blob/main/ann_es_docker.ipynb)\n",
    "\n",
    "Scaling ANNs to \"Big\" Data Volumes\n",
    "\n",
    "![Header Image](images/ann_es_docker.png)\n",
    "\n",
    "Docker containers are crucial for Data Science at Scale [Link](https://blogs.nvidia.com/blog/2020/09/03/what-is-mlops/). That's very well the case for Approximate Nearest Neighbors (ANNs) on \"big\" data too!\n",
    "\n",
    "> Everything must run in a container\n",
    "\n",
    "Speed and Accuracy (or Recall) are the top two considerations while choosing a Nearest Neighbors or Similarity Search algorithm. In my previous post, [KNN is Dead](https://medium.com/towards-artificial-intelligence/knn-k-nearest-neighbors-is-dead-fc16507eb3e), I have proven the tremendous (>300X) speed advantage ANNs have over KNN at comparable accuracy. I've also discussed how you can choose the fastest, most accurate ANN algorithm on your own dataset [Link](https://medium.com/towards-artificial-intelligence/how-to-choose-the-best-nearest-neighbors-algorithm-8d75d42b16ab). \n",
    "\n",
    "However, sometimes, in addition to speed and accuracy, you also need the ability to scale to large data volumes. The ease of distributing the data across multiple machines is a third consideration in these cases. Let's solve all these three considerations concurrently with the fantastic OpenDistro for Elastic Search [Link](https://opendistro.github.io/for-elasticsearch/) in this post.\n",
    "\n",
    "## Why Elastic Search (ES)?\n",
    "ES might be a foreign concept to many Data Scientists reading this post, so let me introduce why it is important beyond its typical usage by your IT team.\n",
    "\n",
    "ES is a search engine database that famously powers search capabilities across the massive volume of Wikipedia [Link](https://opendistro.github.io/for-elasticsearch-docs/docs/elasticsearch/). It allows for full-text search [Link](https://en.wikipedia.org/wiki/Full-text_search) similar to Apache Solr for those from the Hadoop world. ES's distributed nature means it can be scaled to handle huge data volumes by adding more servers/nodes similar to Hadoop. \n",
    "\n",
    "The following are the three reasons why Data Scientists would be interested in Elastic Seach.\n",
    "1. ***Scalable ANN***: OpenDistro ES distribution has implemented the HNSW ANN algorithm as a plugin [Link](https://opendistro.github.io/for-elasticsearch/features/knn.html). ES's distributed nature enables it to scale this high-speed, high-accuracy HNSW ANN search to many millions of records.\n",
    "2. ***Support***: Most modern IT and Infrastructure teams are already familiar with (and heavily using) ES. As a Data Scientist, you probably want the IT Infra team to set up and maintain the server hardware, so using a technology that they are already familiar with increases your chance of securing their support. This support can literally make or break the productionizing of your Data Science project, so don't underestimate it! \n",
    "3. ***Simpler Stack***: ES can be used as a database to store additional fields, which can then be queried together with the nearest neighbors. For example, if you search for nearest neighbors using product name embeddings, you can also get the product price, category, date added, etc. As long as you stored it in ES, you can retrieve it during the ANN search. This simplifies your application stack because you can get all the information you need from one place instead of querying the nearest neighbors from one location and then hitting another database to get the other fields for all these neighbors."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Installing OpenDistro ES (which includes HNSW)\n",
    "This is super-duper simple using Docker. We'll follow the instructions from the official documentation [Link](https://opendistro.github.io/for-elasticsearch-docs/docs/install/docker/). \n",
    "\n",
    "***Important!*** This setup is only for your experimentation of HNSW ANN on ES! For Production setup, please ask your IT Infra team to set up the necessary security protocols.\n",
    "\n",
    "## Single Machine\n",
    "1. Install `docker` on your machine following instructions from the docker website [Link](https://docs.docker.com/engine/install/)\n",
    "2. Pull the docker image for OpenDistro: `sudo docker pull amazon/opendistro-for-elasticsearch:1.12.0`\n",
    "3. Run the docker image: `sudo docker run --rm -p 9200:9200 -p 9600:9600 -e \"discovery.type=single-node\" amazon/opendistro-for-elasticsearch:1.12.0`\n",
    "\n",
    "That's it! You can now interact with the ES service."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%bash\n",
    "sudo docker pull amazon/opendistro-for-elasticsearch:1.12.0\n",
    "sudo docker run --rm -p 9200:9200 -p 9600:9600 -e \"discovery.type=single-node\" amazon/opendistro-for-elasticsearch:1.12.0"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Cluster of Multiple Machines\n",
    "Run the following steps in **each of the machines** you want to set up in your cluster. ***Again Important!*** This is just for experimentation, DO NOT use it in production without the proper security protocols.\n",
    "\n",
    "1. Install `docker` on your machine following instructions from the docker website [Link](https://docs.docker.com/engine/install/)\n",
    "2. Pull the docker image for OpenDistro: `sudo docker pull amazon/opendistro-for-elasticsearch:1.12.0`\n",
    "3. Install `docker-compose` following instructions from the docker website [Link](https://docs.docker.com/compose/install/)\n",
    "4. Create a `docker-compose.yml` file\n",
    "5. Create a `elasticsearch.yml` file\n",
    "6. Start the ES service: `sudo docker-compose up`\n",
    "\n",
    "If you encounter any errors, try updating the VM map count as per [Link](https://github.com/opendistro-for-elasticsearch/opendistro-build/issues/329): `sudo sysctl -w vm.max_map_count=262144`\n",
    "\n",
    "The whole process is shown below"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%bash\n",
    "# Run this script in each of the machine on your cluster\n",
    "# Provide a unique node_name for each machine on your cluster\n",
    "node_name=odfe-node01\n",
    "\n",
    "# JVM memory, update as per your machine's hardware\n",
    "Xms=10g\n",
    "Xmx=10g\n",
    "\n",
    "# Fixed\n",
    "cluster_name=odfe-cluster01\n",
    "initial_master=odfe-node01\n",
    "internal_ip=hostname -I | awk '{print $1}'\n",
    "\n",
    "# Pull the docker container\n",
    "sudo docker pull amazon/opendistro-for-elasticsearch:1.12.0\n",
    "\n",
    "# Create docker-compose.yml\n",
    "echo '# Instructions from\n",
    "# https://opendistro.github.io/for-elasticsearch-docs/docs/install/docker/\n",
    "# https://www.elastic.co/guide/en/elastic-stack-get-started/current/get-started-docker.html\n",
    "\n",
    "version: \"3\"\n",
    "services:\n",
    "    '$node_name':\n",
    "    image: amazon/opendistro-for-elasticsearch:1.12.0\n",
    "    container_name: odfe-node\n",
    "    environment:\n",
    "        - cluster.name='$cluster_name'\n",
    "        - bootstrap.memory_lock=true # along with the memlock settings below, disables swapping\n",
    "        - \"ES_JAVA_OPTS=-Xms'$Xms' -Xmx'$Xmx'\" # minimum and maximum Java heap size, recommend setting both to 50% of system RAM\n",
    "    ulimits:\n",
    "        memlock:\n",
    "            soft: -1\n",
    "            hard: -1\n",
    "        nofile:\n",
    "            soft: 65536 # maximum number of open files for the Elasticsearch user, set to at least 65536 on modern systems\n",
    "            hard: 65536\n",
    "    volumes:\n",
    "        - odfe-data:/usr/share/elasticsearch/data\n",
    "        - ./elasticsearch.yml:/usr/share/elasticsearch/config/elasticsearch.yml\n",
    "    ports:\n",
    "        - 9200:9200\n",
    "        - 9300:9300 # required for adding more nodes\n",
    "        - 9600:9600 # required for Performance Analyzer\n",
    "    networks:\n",
    "        - odfe-net\n",
    "\n",
    "volumes:\n",
    "    odfe-data:\n",
    "        driver: local\n",
    "\n",
    "networks:\n",
    "    odfe-net:\n",
    "        driver: bridge' > docker-compose.yml\n",
    "\n",
    "# Create elasticsearch.yml\n",
    "echo 'cluster.name: '$cluster_name'\n",
    "node.name: '$node_name'\n",
    "node.roles: [master, data]\n",
    "opendistro_security.disabled: true\n",
    "http.host: 0.0.0.0\n",
    "transport.host: 0.0.0.0\n",
    "transport.publish_host: '$internal_ip'\n",
    "http.publish_host: '$internal_ip'\n",
    "http.port: 9200\n",
    "transport.tcp.port: 9300\n",
    "network.host: [127.0.0.1, '$internal_ip']\n",
    "cluster.initial_master_nodes:\n",
    "    - '$initial_master'\n",
    "discovery.seed_hosts:\n",
    "    - '$internal_ip'\n",
    "    - <ip address of 2nd machine on this cluster>\n",
    "    - <ip address of 3rd machine on this cluster>\n",
    "path:\n",
    "    data: /usr/share/elasticsearch/data\n",
    "' > elasticsearch.yml\n",
    "\n",
    "# Start the ES service\n",
    "sudo docker-compose up"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Check ES Status\n",
    "After you complete running the steps above on each of your machines, you can check the status of your cluster by using: `curl -XGET <ip_address_of_any_node>:9200/_cat/nodes?v -u <username>:<password> --insecure`. You should see all the nodes in your cluster."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "ip         heap.percent ram.percent cpu load_1m load_5m load_15m node.role master name\n",
      "172.17.0.2           28          18   1    0.39    0.20     0.13 dimr      *      b8d08220be95\n"
     ]
    }
   ],
   "source": [
    "!curl -XGET https://localhost:9200/_cat/nodes?v -u 'admin:admin' --insecure"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Using ES for ANN"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Open a DataFrame to load into ES\n",
    "We shall use the same Amazon 500K product dataset used in my previous post [KNN is Dead!](https://medium.com/towards-artificial-intelligence/knn-k-nearest-neighbors-is-dead-fc16507eb3e)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(527543, 2)\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>title</th>\n",
       "      <th>emb</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>Puppies Faceplate Hard Case Protector for Net1...</td>\n",
       "      <td>[0.01443001, 0.0070197457, 0.011907284, 0.0197...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>White Wolf Faceplate Protector Hard Case for S...</td>\n",
       "      <td>[0.007030556, 0.03295864, 0.028047869, 0.01396...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>Camo Duck Grass Rubberized Hard Case Phone Fac...</td>\n",
       "      <td>[0.0013388951, 0.03315278, 0.03617841, -0.0105...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>Camoflague Camo Usa Deer Combo Hybrid Hard Cas...</td>\n",
       "      <td>[0.01922137, -0.0041188085, 0.022043534, -0.00...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>Motorola H700 Black - Non-Retail Packaging</td>\n",
       "      <td>[0.004890033, 0.019337852, 0.026887875, -0.037...</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "                                               title  \\\n",
       "0  Puppies Faceplate Hard Case Protector for Net1...   \n",
       "1  White Wolf Faceplate Protector Hard Case for S...   \n",
       "2  Camo Duck Grass Rubberized Hard Case Phone Fac...   \n",
       "3  Camoflague Camo Usa Deer Combo Hybrid Hard Cas...   \n",
       "4         Motorola H700 Black - Non-Retail Packaging   \n",
       "\n",
       "                                                 emb  \n",
       "0  [0.01443001, 0.0070197457, 0.011907284, 0.0197...  \n",
       "1  [0.007030556, 0.03295864, 0.028047869, 0.01396...  \n",
       "2  [0.0013388951, 0.03315278, 0.03617841, -0.0105...  \n",
       "3  [0.01922137, -0.0041188085, 0.022043534, -0.00...  \n",
       "4  [0.004890033, 0.019337852, 0.026887875, -0.037...  "
      ]
     },
     "execution_count": 1,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import pandas as pd\n",
    "df = pd.read_pickle(\"df.pkl\")\n",
    "embedding_col = \"emb\"\n",
    "\n",
    "print(df.shape)\n",
    "df.head()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Creating an ES Index with HNSW\n",
    "Once you have setup OpenDistro ES, we need to create an ES Index that will hold all our data. This can be done either by using the Python requests library or using the Python [ElasticSearch library](https://elasticsearch-py.readthedocs.io/en/7.10.0/): `pip install elasticsearch`. I'll use the ElasticSearch library in this post. \n",
    "\n",
    "In this step, we need to specify the HNSW parameters `ef_construction` and `M` as part of the ES index settings. We also need to indicate whether to use `l2` (euclidean) or `cosinesimil` (angular) distance to find neighbors. Some of the settings are best practices from [Link](https://medium.com/@kumon/how-to-realize-similarity-search-with-elasticsearch-3dd5641b9adb)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Collecting elasticsearch\n",
      "  Downloading elasticsearch-7.10.1-py2.py3-none-any.whl (322 kB)\n",
      "\u001b[K     |████████████████████████████████| 322 kB 9.5 MB/s eta 0:00:01\n",
      "\u001b[?25hRequirement already satisfied: certifi in /opt/conda/lib/python3.7/site-packages (from elasticsearch) (2020.11.8)\n",
      "Requirement already satisfied: urllib3<2,>=1.21.1 in /opt/conda/lib/python3.7/site-packages (from elasticsearch) (1.25.11)\n",
      "Installing collected packages: elasticsearch\n",
      "Successfully installed elasticsearch-7.10.1\n"
     ]
    }
   ],
   "source": [
    "!pip install elasticsearch"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Imports\n",
    "from elasticsearch import Elasticsearch\n",
    "\n",
    "# Suppress the insecure warning while using verify = False\n",
    "import requests\n",
    "from requests.packages.urllib3.exceptions import InsecureRequestWarning\n",
    "requests.packages.urllib3.disable_warnings(InsecureRequestWarning)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/opt/conda/lib/python3.7/site-packages/elasticsearch/connection/http_urllib3.py:206: UserWarning: Connecting to https://localhost:9200 using SSL with verify_certs=False is insecure.\n",
      "  % self.host\n"
     ]
    }
   ],
   "source": [
    "# ES constants\n",
    "index_name = \"amazon-500k\"\n",
    "\n",
    "es = Elasticsearch([\"https://localhost:9200\"], \n",
    "                   http_auth=(\"admin\", \"admin\"), \n",
    "                   verify_certs=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "# ES settings\n",
    "body = {\n",
    "  \"settings\": {\n",
    "    \"index\": {\n",
    "      \"knn\": True,\n",
    "      \"knn.space_type\": \"l2\", \n",
    "      \"knn.algo_param.ef_construction\": 100, \n",
    "      \"knn.algo_param.m\": 16\n",
    "    },\n",
    "    \"number_of_shards\": 10, \n",
    "    \"number_of_replicas\": 0,\n",
    "    \"refresh_interval\": -1,\n",
    "    \"translog.flush_threshold_size\": \"10gb\"\n",
    "  }\n",
    "}\n",
    "\n",
    "mapping = {\n",
    "  \"properties\": {\n",
    "    embedding_col: {\n",
    "      \"type\": \"knn_vector\", \n",
    "      \"dimension\": len(df.loc[0,embedding_col])\n",
    "    }\n",
    "  }\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'acknowledged': True}"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Create the Index\n",
    "es.indices.create(index_name, body=body)\n",
    "es.indices.put_mapping(mapping, index_name)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "After you complete running the code above, we can check the available indices on our ES cluster using : `curl -XGET <ip_address_of_any_node>:9200/_cat/indices?v -u <username>:<password> --insecure`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "health status index                        uuid                   pri rep docs.count docs.deleted store.size pri.store.size\n",
      "yellow open   security-auditlog-2020.12.20 c9IG9Uc5Sn2cF68UAiQJTw   1   1          2            0     30.2kb         30.2kb\n",
      "green  open   amazon-500k                  Mb3C3OFvSvKIseAINNTAIQ  10   0          0            0       624b           624b\n",
      "green  open   .opendistro_security         BEb84-71RcSEZ3oAJUaPJg   1   0          9            0     55.4kb         55.4kb\n"
     ]
    }
   ],
   "source": [
    "!curl -XGET https://localhost:9200/_cat/indices?v -u 'admin:admin' --insecure"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Uploading Data\n",
    "\n",
    "Now that the index has been created, we can upload data to it either by using the Python requests library or using the Python [ElasticSearch library](https://elasticsearch-py.readthedocs.io/en/7.10.0/), which is shown below."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Imports\n",
    "from elasticsearch.helpers import bulk\n",
    "from tqdm import tqdm"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 527543/527543 [04:29<00:00, 1955.64it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{\n",
      "  \"_shards\" : {\n",
      "    \"total\" : 10,\n",
      "    \"successful\" : 10,\n",
      "    \"failed\" : 0\n",
      "  }\n",
      "}\n",
      "\n"
     ]
    }
   ],
   "source": [
    "# Data Generator\n",
    "def gen():\n",
    "    for row in tqdm(df.itertuples(), total=len(df)):\n",
    "        output_yield = { \n",
    "            \"_op_type\": \"index\", \n",
    "            \"_index\": index_name\n",
    "        }\n",
    "        output_yield.update(row._asdict())\n",
    "        output_yield.update({\n",
    "            embedding_col: output_yield[embedding_col].tolist()\n",
    "        })\n",
    "\n",
    "        yield output_yield\n",
    "\n",
    "# Upload data to ES in bulk\n",
    "_, errors = bulk(es, gen(), chunk_size=500, max_retries=2)\n",
    "assert len(errors) == 0, errors\n",
    "\n",
    "# Refresh the data\n",
    "es.indices.refresh(index_name, request_timeout=1000)   \n",
    "\n",
    "# Warmup API    \n",
    "res = requests.get(\"https://localhost:9200/_opendistro/_knn/warmup/\"+index_name+\"?pretty\", auth=(\"admin\", \"admin\"), verify=False)\n",
    "print(res.text)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Querying Nearest Neighbors\n",
    "Once we have uploaded all our data into the ES index, we can then query the \"K\" nearest neighbors for any new data point as follows."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Query parameters\n",
    "query_df = df.sample(1000).copy()\n",
    "K = 5 # Number of neighbors\n",
    "step = 200 # Number of items to query at a time\n",
    "\n",
    "cols_to_query = [\"Index\", \"title\"]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'acknowledged': True}"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Update the search settings\n",
    "body = {\n",
    "  \"settings\": {\n",
    "    \"index\": {\"knn.algo_param.ef_search\": 100}\n",
    "  }        \n",
    "}   \n",
    "\n",
    "es.indices.put_settings(body=body, index=index_name)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 5/5 [00:34<00:00,  6.85s/it]\n",
      "100%|██████████| 1000/1000 [00:06<00:00, 163.84it/s]\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>title</th>\n",
       "      <th>emb</th>\n",
       "      <th>es_neighbors_Index</th>\n",
       "      <th>es_neighbors_title</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>137539</th>\n",
       "      <td>MOTOROLA DROID RAZR M XT907 IMPACT CASE COVER ...</td>\n",
       "      <td>[0.031467807, -0.048276246, 0.014576058, -0.02...</td>\n",
       "      <td>[137539, 135533, 135529, 135540, 192993]</td>\n",
       "      <td>[MOTOROLA DROID RAZR M XT907 IMPACT CASE COVER...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>186522</th>\n",
       "      <td>Forever Collectibles NHL Dual Hybrid iPhone 4/...</td>\n",
       "      <td>[0.023334723, 0.017239397, 0.03823308, -0.0193...</td>\n",
       "      <td>[158876, 158881, 158875, 186522, 186523]</td>\n",
       "      <td>[Forever Collectibles NHL Dual Hybrid iPhone 4...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>344437</th>\n",
       "      <td>Samsung Galaxy S6 Edge, White Pearl 64GB (Sprint)</td>\n",
       "      <td>[-0.0013037474, -0.0038629419, 0.05036281, 0.0...</td>\n",
       "      <td>[344437, 344439, 344445, 478780, 344444]</td>\n",
       "      <td>[Samsung Galaxy S6 Edge, White Pearl 64GB (Spr...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>27490</th>\n",
       "      <td>Amzer Super Clear Screen Protector with Cleani...</td>\n",
       "      <td>[-0.021655105, -0.007051119, 0.024543112, -0.0...</td>\n",
       "      <td>[27490, 27484, 27483, 23949, 33349]</td>\n",
       "      <td>[Amzer Super Clear Screen Protector with Clean...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>387250</th>\n",
       "      <td>[Benly-5] 1/8&amp;quot; 3.5mm Stereo Mini Jack / F...</td>\n",
       "      <td>[0.025798462, 0.028967546, 0.013909815, 0.0072...</td>\n",
       "      <td>[387250, 401900, 440491, 440501, 440465]</td>\n",
       "      <td>[[Benly-5] 1/8&amp;quot; 3.5mm Stereo Mini Jack / ...</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "                                                    title  \\\n",
       "137539  MOTOROLA DROID RAZR M XT907 IMPACT CASE COVER ...   \n",
       "186522  Forever Collectibles NHL Dual Hybrid iPhone 4/...   \n",
       "344437  Samsung Galaxy S6 Edge, White Pearl 64GB (Sprint)   \n",
       "27490   Amzer Super Clear Screen Protector with Cleani...   \n",
       "387250  [Benly-5] 1/8&quot; 3.5mm Stereo Mini Jack / F...   \n",
       "\n",
       "                                                      emb  \\\n",
       "137539  [0.031467807, -0.048276246, 0.014576058, -0.02...   \n",
       "186522  [0.023334723, 0.017239397, 0.03823308, -0.0193...   \n",
       "344437  [-0.0013037474, -0.0038629419, 0.05036281, 0.0...   \n",
       "27490   [-0.021655105, -0.007051119, 0.024543112, -0.0...   \n",
       "387250  [0.025798462, 0.028967546, 0.013909815, 0.0072...   \n",
       "\n",
       "                              es_neighbors_Index  \\\n",
       "137539  [137539, 135533, 135529, 135540, 192993]   \n",
       "186522  [158876, 158881, 158875, 186522, 186523]   \n",
       "344437  [344437, 344439, 344445, 478780, 344444]   \n",
       "27490        [27490, 27484, 27483, 23949, 33349]   \n",
       "387250  [387250, 401900, 440491, 440501, 440465]   \n",
       "\n",
       "                                       es_neighbors_title  \n",
       "137539  [MOTOROLA DROID RAZR M XT907 IMPACT CASE COVER...  \n",
       "186522  [Forever Collectibles NHL Dual Hybrid iPhone 4...  \n",
       "344437  [Samsung Galaxy S6 Edge, White Pearl 64GB (Spr...  \n",
       "27490   [Amzer Super Clear Screen Protector with Clean...  \n",
       "387250  [[Benly-5] 1/8&quot; 3.5mm Stereo Mini Jack / ...  "
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Run the Query\n",
    "responses = []\n",
    "for n in tqdm(range(0, len(query_df), step)):\n",
    "    subset_df = query_df.iloc[n:n+step,:]\n",
    "    request = []\n",
    "\n",
    "    for row in subset_df.itertuples():\n",
    "        req_head = {\"index\": index_name}\n",
    "        req_body = {\n",
    "            \"query\": {\n",
    "                \"knn\": {\n",
    "                    embedding_col: {\n",
    "                        \"vector\": getattr(row, embedding_col).tolist(), \n",
    "                        \"k\": K}\n",
    "                }\n",
    "            },\n",
    "            \"size\": K,\n",
    "            \"_source\": cols_to_query\n",
    "        }\n",
    "\n",
    "        request.extend([req_head, req_body])\n",
    "\n",
    "    r = es.msearch(body=request)\n",
    "    responses.extend(r['responses'])\n",
    "\n",
    "# Convert the responses to dataframe columns\n",
    "nn_data = {f'es_neighbors_{key}': [] for key in cols_to_query}\n",
    "\n",
    "for item in tqdm(responses):\n",
    "    nn = pd.concat(map(pd.DataFrame.from_dict, \n",
    "                       item['hits']['hits']), \n",
    "                   axis=1)['_source'].T.reset_index(drop=True)\n",
    "    for key in cols_to_query:\n",
    "        nn_data[f'es_neighbors_{key}'].append(nn[key]\n",
    "                                              .values\n",
    "                                              .tolist())\n",
    "\n",
    "query_df = query_df.assign(**nn_data)\n",
    "query_df.head()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let's take a closer look at the neighbors for one of the rows in query_df. They're all phone cases for the same phone indicating just how good the nearest neighbor search is!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "['MOTOROLA DROID RAZR M XT907 IMPACT CASE COVER + BLACK RUBBER SKIN PROTECTOR ACCESSORY CAMO DEER',\n",
       " 'MOTOROLA DROID RAZR M XT907 CAMO DEER HEAVY DUTY CASE + BLACK GEL SKIN SNAP-ON PROTECTOR ACCESSORY',\n",
       " 'MOTOROLA DROID RAZR M XT907 CAMO DEER HEAVY DUTY CASE + LIME GREEN GEL SKIN SNAP-ON PROTECTOR ACCESSORY',\n",
       " 'MOTOROLA DROID RAZR M XT907 CAMO DRIED LEAVES HEAVY DUTY CASE + BLACK GEL SKIN SNAP-ON PROTECTOR ACCESSORY',\n",
       " 'MOTOROLA DROID RAZR M XT907 NON SLIP GRAY HEAVY DUTY CASE + RED GEL SKIN SNAP-ON PROTECTOR ACCESSORY']"
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "query_df[\"es_neighbors_title\"].iloc[0]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Delete Index\n",
    "Finally, after our experimentation is done (or if we made a mistake and need to delete everything to start from scratch), we can delete the whole index with a single line of code as below."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'acknowledged': True}"
      ]
     },
     "execution_count": 15,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "es.indices.delete(index=index_name)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# OpenDistro's HNSW vs. Other ES ANNs"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "OpenDistro ES is now available in the [ann-benchmarks](https://github.com/erikbern/ann-benchmarks) repository that I introduced in [How to Choose the Best Nearest Neighbors Algorithm](https://medium.com/towards-artificial-intelligence/how-to-choose-the-best-nearest-neighbors-algorithm-8d75d42b16ab). OpenDistro's HNSW is ~3X faster at comparable Recall than other ES ANN algorithms, as shown below. Though ANN on ES cannot compete on single-core performance with standalone ANN algorithms due to ES overheads, the scalability to huge data makes OpenDistro ES's HNSW ANN plugin worthwhile to consider for large datasets!\n",
    "\n",
    "![ANN Benchmarks comparison between various ES ANNs](images/custom-euclidean-onlyES.png)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Thank you for reading! All the code in this post is available on my [GitHub repository](https://github.com/stephenleo/adventures-with-ann/blob/main/ann_es_docker.ipynb)."
   ]
  }
 ],
 "metadata": {
  "environment": {
   "name": "tf2-2-3-gpu.2-3.m59",
   "type": "gcloud",
   "uri": "gcr.io/deeplearning-platform-release/tf2-2-3-gpu.2-3:m59"
  },
  "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.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}