scripts/process-benchmarks.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Notebook to process benchmark results\n",
    "\n",
    "Please run this notebook after running all the benchmarks and storing them in the `results` dir. This will export them in the desired format for the single node benchmark plots of [qdrant.tech/benchmarks](https://qdrant.tech/benchmarks)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2022-08-05T10:03:50.900734Z",
     "start_time": "2022-08-05T10:03:50.892651Z"
    },
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "from pathlib import Path\n",
    "import re\n",
    "import json\n",
    "import pandas as pd\n",
    "from datetime import datetime, timezone"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2022-08-05T10:03:50.982398Z",
     "start_time": "2022-08-05T10:03:50.903447Z"
    },
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "DATA_DIR = Path().resolve().parent / \"results\"\n",
    "DATA_DIR, list(DATA_DIR.glob(\"*.json\"))[0].name"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2022-08-05T10:03:51.482299Z",
     "start_time": "2022-08-05T10:03:51.477275Z"
    },
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "PATH_REGEX = re.compile(r\"(?P<engine_name>(\"\n",
    "                        r\"?P<engine>[a-z\\-]+)\"\n",
    "                        r\"\\-m\\-(?P<m>[0-9]+)\"\n",
    "                        r\"\\-ef\\-(?P<ef>[0-9]+)\"\n",
    "                        r\")\"\n",
    "                        r\"\\-(?P<dataset>[a-zA-Z0-9\\-]+)\"\n",
    "                        r\"\\-(?P<operation>(search)|(upload))\"\n",
    "                        r\"(\\-(?P<search_index>[0-9]{1,2})\\-)?\"\n",
    "                        r\"\\-?(?P<date>.*)\\.json\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2022-08-05T10:03:54.150582Z",
     "start_time": "2022-08-05T10:03:51.558766Z"
    },
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "upload_results, search_results = [], []\n",
    "\n",
    "for path in DATA_DIR.glob(\"*.json\"):\n",
    "    match = PATH_REGEX.match(path.name)\n",
    "    if match is None:\n",
    "        continue\n",
    "\n",
    "    experiment = match.groupdict()\n",
    "\n",
    "    with open(path, \"r\") as fp:\n",
    "        stats = json.load(fp)\n",
    "\n",
    "    params = stats[\"params\"]\n",
    "    dataset = params.pop(\"dataset\")\n",
    "    engine = params.pop(\"engine\")\n",
    "\n",
    "    entry = {\n",
    "        \"dataset\": dataset,\n",
    "        \"engine\": engine,\n",
    "        \"m\": match[\"m\"],\n",
    "        \"ef\": match[\"ef\"],\n",
    "        \"date\": match[\"date\"],\n",
    "        \"params\": params,\n",
    "        \"results\": stats[\"results\"],\n",
    "    }\n",
    "\n",
    "    if experiment[\"operation\"] == \"search\":\n",
    "        entry.update({\"search_index\": match[\"search_index\"]})\n",
    "        search_results.append(entry)\n",
    "    elif experiment[\"operation\"] == \"upload\":\n",
    "        upload_results.append(entry)\n",
    "    else:\n",
    "        raise Exception(\"Unknown operation\")\n",
    "\n",
    "len(upload_results), len(search_results)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "upload_results, search_results[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2022-08-05T11:31:17.192306Z",
     "start_time": "2022-08-05T11:31:17.125766Z"
    },
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "upload_df = pd.DataFrame(upload_results)\n",
    "upload_df[\"date\"] = pd.to_datetime(upload_df[\"date\"], format=\"%Y-%m-%d-%H-%M-%S\")\n",
    "upload_df = upload_df.sort_values(\"date\", ascending=False) \\\n",
    "    .groupby([\"engine\", \"m\", \"ef\", \"dataset\"]) \\\n",
    "    .first()\n",
    "\n",
    "temp_df = upload_df.copy()\n",
    "temp_df[\"total_time\"] = temp_df[\"results\"].apply(lambda x: x[\"total_time\"])\n",
    "temp_df.sort_values(\"total_time\", ascending=True).head(n=5)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2022-08-05T12:06:58.459530Z",
     "start_time": "2022-08-05T12:06:57.908842Z"
    },
    "pycharm": {
     "name": "#%%\n"
    }
   },
   "outputs": [],
   "source": [
    "search_df = pd.DataFrame(search_results)\n",
    "search_df[\"date\"] = pd.to_datetime(search_df[\"date\"], format=\"%Y-%m-%d-%H-%M-%S\")\n",
    "search_df = search_df.sort_values(\"date\", ascending=False) \\\n",
    "    .groupby([\"engine\", \"m\", \"ef\", \"dataset\", \"search_index\"]) \\\n",
    "    .first()\n",
    "\n",
    "temp_df = search_df.copy()\n",
    "temp_df['rps'] = temp_df['results'].apply(lambda x: x[\"rps\"])\n",
    "temp_df.sort_values(\"rps\", ascending=False).head(n=10)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "_search = search_df.reset_index()\n",
    "_upload = upload_df.reset_index()\n",
    "\n",
    "joined_df = _search.merge(_upload, on=[\"engine\", \"m\", \"ef\", \"dataset\"], how=\"left\", suffixes=(\"_search\", \"_upload\"))\n",
    "print(len(joined_df))\n",
    "joined_df"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "json_results = []\n",
    "\n",
    "for index, row in joined_df.reset_index().iterrows():\n",
    "    engine_params = {}\n",
    "    \n",
    "    if isinstance(row['params_upload'], dict):\n",
    "        engine_params.update(row['params_upload'])\n",
    "    if isinstance(row['params_search'], dict):\n",
    "        search_params = row['params_search']\n",
    "        engine_params.update(search_params.get('config', {}))\n",
    "        engine_params.update(search_params.get('params', {}))\n",
    "        engine_params.update(search_params.get('search_params', {}))\n",
    "        engine_params.update(search_params.get('vectorIndexConfig', {}))\n",
    "\n",
    "    engine_params.pop('experiment')\n",
    "    engine_params.pop('parallel')\n",
    "\n",
    "    engine_name = row['engine']\n",
    "\n",
    "    if engine_name.startswith(\"qdrant-\"):\n",
    "        engine_name = \"qdrant\"\n",
    "\n",
    "    json_object = {\n",
    "        \"engine_name\": engine_name,\n",
    "        \"setup_name\": f\"{row['params_search']['experiment']}\",\n",
    "        \"dataset_name\": row['dataset'],\n",
    "        \"search_idx\": row['search_index'],\n",
    "        \"upload_time\": row['results_upload']['upload_time'],\n",
    "        \"total_upload_time\": row['results_upload']['total_time'],\n",
    "        \"p95_time\": row['results_search']['p95_time'],\n",
    "        \"rps\": row['results_search']['rps'],\n",
    "        \"parallel\": row['params_search']['parallel'],\n",
    "        \"p99_time\": row['results_search']['p99_time'],\n",
    "        \"mean_time\": row['results_search']['mean_time'],\n",
    "        \"mean_precisions\": row['results_search']['mean_precisions'],\n",
    "        \"engine_params\": engine_params,\n",
    "    }\n",
    "    json_results.append(json_object)\n",
    "\n",
    "format = '%Y-%M-%dT%H:%M:%S'\n",
    "now = datetime.now().replace(tzinfo=timezone.utc).strftime(format)\n",
    "\n",
    "Path(\"results.json\").write_text(json.dumps(json_results, indent=2))\n",
    "Path(f\"results-{now}.json\").write_text(json.dumps(json_results, indent=2))\n",
    "\n",
    "print(json_results[-1], len(json_results))\n",
    "\n",
    "results_df = pd.DataFrame(json_results).sort_values(\"p99_time\", ascending=True)\n",
    "# results_df.to_csv('results.csv')\n",
    "results_df"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.9.18"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}