run-qdrant.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "46ef55c3-ba47-4250-bf90-4092db6bddf2",
   "metadata": {},
   "outputs": [],
   "source": [
    "import gzip\n",
    "import json\n",
    "import os\n",
    "import pprint\n",
    "import re\n",
    "import subprocess\n",
    "import time\n",
    "from datetime import timedelta, datetime\n",
    "from pathlib import Path\n",
    "from dataclasses import dataclass, asdict\n",
    "import multiprocessing\n",
    "\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "import requests"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cbaa8bfa-c307-44dc-82c4-415ad36984d8",
   "metadata": {},
   "outputs": [],
   "source": [
    "@dataclass\n",
    "class DataSetConfig:\n",
    "    content_path: Path\n",
    "    embedding_path: Path\n",
    "    num_of_docs: int\n",
    "    index_size: int\n",
    "    bulk_size: int\n",
    "    index_name: str\n",
    "    distance: str\n",
    "    dimension: int\n",
    "    exact: bool\n",
    "    hnsw_m: int\n",
    "    hnsw_ef_construction: int\n",
    "    hnsw_ef: int\n",
    "    update_docs_per_sec: int\n",
    "    quantization: str\n",
    "\n",
    "    qdrant_name: str = \"benchmark_qdrant\"\n",
    "    qdrant_host: str = \"localhost\"\n",
    "    qdrant_port: int = 6344\n",
    "    qdrant_version: str = \"1.9.1\"\n",
    "\n",
    "\n",
    "def get_dataset_config(target_name):\n",
    "    setting = {\n",
    "        \"100k-768-m32-efc200-ef100-ip\": {\n",
    "            \"content_path\": Path(\"dataset/passages-c400-jawiki-20230403\"),\n",
    "            \"embedding_path\": Path(\"dataset/passages-c400-jawiki-20230403/multilingual-e5-base-passage\"),\n",
    "            \"num_of_docs\": 5555583,\n",
    "            \"index_size\": 100000,\n",
    "            \"bulk_size\": 1000,\n",
    "            \"index_name\": \"contents\",\n",
    "            \"distance\": \"Dot\", # \"Cosine\"\n",
    "            \"dimension\": 768,\n",
    "            \"exact\": False,\n",
    "            \"hnsw_m\": 32,\n",
    "            \"hnsw_ef_construction\": 200,\n",
    "            \"hnsw_ef\": 100,\n",
    "            \"update_docs_per_sec\": 0,\n",
    "            \"quantization\": \"int8\",\n",
    "        },\n",
    "        \"1m-768-m48-efc200-ef100-ip\": {\n",
    "            \"content_path\": Path(\"dataset/passages-c400-jawiki-20230403\"),\n",
    "            \"embedding_path\": Path(\"dataset/passages-c400-jawiki-20230403/multilingual-e5-base-passage\"),\n",
    "            \"num_of_docs\": 5555583,\n",
    "            \"index_size\": 1000000,\n",
    "            \"bulk_size\": 1000,\n",
    "            \"index_name\": \"contents\",\n",
    "            \"distance\": \"Dot\", # \"Cosine\"\n",
    "            \"dimension\": 768,\n",
    "            \"exact\": False,\n",
    "            \"hnsw_m\": 48,\n",
    "            \"hnsw_ef_construction\": 200,\n",
    "            \"hnsw_ef\": 100,\n",
    "            \"update_docs_per_sec\": 0,\n",
    "            \"quantization\": \"int8\",\n",
    "        },\n",
    "        \"5m-768-m48-efc200-ef100-ip\": {\n",
    "            \"content_path\": Path(\"dataset/passages-c400-jawiki-20230403\"),\n",
    "            \"embedding_path\": Path(\"dataset/passages-c400-jawiki-20230403/multilingual-e5-base-passage\"),\n",
    "            \"num_of_docs\": 5555583,\n",
    "            \"index_size\": 5000000,\n",
    "            \"bulk_size\": 1000,\n",
    "            \"index_name\": \"contents\",\n",
    "            \"distance\": \"Dot\", # \"Cosine\"\n",
    "            \"dimension\": 768,\n",
    "            \"exact\": False,\n",
    "            \"hnsw_m\": 48,\n",
    "            \"hnsw_ef_construction\": 200,\n",
    "            \"hnsw_ef\": 100,\n",
    "            \"update_docs_per_sec\": 0,\n",
    "            \"quantization\": \"int8\",\n",
    "        },\n",
    "    }\n",
    "    return DataSetConfig(**setting.get(target_name))\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e3e93d3d-7b74-4a5c-8bf9-b44a0d33c9ac",
   "metadata": {},
   "outputs": [],
   "source": [
    "def run_qdrant(config):\n",
    "    volume_dir = os.getenv(\"VOLUME_DIR\", \"./data\")\n",
    "    print(f\"Starting {config.qdrant_name}... \", end=\"\")\n",
    "    docker_cmd = [\n",
    "        # \"sudo\",\n",
    "        \"docker\", \"run\", \"-d\",\n",
    "        \"--name\", config.qdrant_name,\n",
    "        \"-p\", f\"{config.qdrant_port}:6333\",\n",
    "        # \"-v\", f\"{volume_dir}:/qdrant/storage\",\n",
    "        f\"qdrant/qdrant:v{config.qdrant_version}\"\n",
    "    ]\n",
    "    result = subprocess.run(docker_cmd, capture_output=True, text=True)\n",
    "    if result.returncode == 0:\n",
    "        print(\"[OK]\")\n",
    "    else:\n",
    "        print(f\"[FAIL]\")\n",
    "        print(\"STDOUT:\")\n",
    "        print(result.stdout)\n",
    "        print(\"STDERR:\")\n",
    "        print(result.stderr)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "277c817c-fddc-4516-872c-6bedb0746a96",
   "metadata": {},
   "outputs": [],
   "source": [
    "def stop_qdrant(config):\n",
    "    print(f\"Stopping {config.qdrant_name}... \", end=\"\")\n",
    "    docker_cmd = [\n",
    "        # \"sudo\",\n",
    "        \"docker\", \"stop\", config.qdrant_name\n",
    "    ]\n",
    "    result = subprocess.run(docker_cmd, capture_output=True, text=True)\n",
    "    if result.returncode == 0:\n",
    "        print(\"[OK]\")\n",
    "    else:\n",
    "        print(f\"[FAIL]\")\n",
    "        print(\"STDOUT:\")\n",
    "        print(result.stdout)\n",
    "        print(\"STDERR:\")\n",
    "        print(result.stderr)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b53c5315-4dbe-4c6d-92f8-14c8f50dd49d",
   "metadata": {},
   "outputs": [],
   "source": [
    "def prune_docker(config):\n",
    "    print(f\"Cleaning up... \", end=\"\")\n",
    "    docker_cmd = [\n",
    "        # \"sudo\",\n",
    "        \"docker\", \"system\", \"prune\", \"-f\"\n",
    "    ]\n",
    "    result = subprocess.run(docker_cmd, capture_output=True, text=True)\n",
    "    if result.returncode == 0:\n",
    "        print(\"[OK]\")\n",
    "    else:\n",
    "        print(f\"[FAIL]\")\n",
    "        print(\"STDOUT:\")\n",
    "        print(result.stdout)\n",
    "        print(\"STDERR:\")\n",
    "        print(result.stderr)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7b1104ef-f760-423f-a066-6eb22f743cd8",
   "metadata": {},
   "outputs": [],
   "source": [
    "def print_docker_system_df(config):\n",
    "    docker_cmd = [\n",
    "        # \"sudo\",\n",
    "        \"docker\", \"system\", \"df\"\n",
    "    ]\n",
    "    result = subprocess.run(docker_cmd, capture_output=True, text=True)\n",
    "    if result.returncode == 0:\n",
    "        print(result.stdout)\n",
    "    else:\n",
    "        print(result.stderr)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "58bda7d9-308a-41f8-89de-79566cd7e641",
   "metadata": {},
   "outputs": [],
   "source": [
    "def print_docker_container_stats(config):\n",
    "    docker_cmd = [\n",
    "        # \"sudo\",\n",
    "        \"docker\", \"container\", \"stats\", \"--no-stream\"\n",
    "    ]\n",
    "    result = subprocess.run(docker_cmd, capture_output=True, text=True)\n",
    "    containers = {}\n",
    "    if result.returncode == 0:\n",
    "        print(result.stdout)\n",
    "        for line in result.stdout.split(\"\\n\"):\n",
    "            if line.startswith(\"CONTAINER\") or len(line) == 0:\n",
    "                continue\n",
    "            values = line.split()\n",
    "            containers[values[1]] = {\n",
    "                \"container_id\": values[0],\n",
    "                \"cpu\": values[2],\n",
    "                \"mem\": values[6],\n",
    "                \"mem_usage\": values[3],\n",
    "                \"mem_limit\": values[5],\n",
    "                \"net_in\": values[7],\n",
    "                \"net_out\": values[9],\n",
    "                \"block_in\": values[10],\n",
    "                \"block_out\": values[12],\n",
    "                \"pids\": values[13],\n",
    "            }\n",
    "    else:\n",
    "        print(result.stderr)\n",
    "    return containers\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a2842a4c-16ed-4b78-ae50-a066254000cd",
   "metadata": {},
   "outputs": [],
   "source": [
    "def create_index(config):\n",
    "    print(F\"Creating Collection {config.index_name} with {config.quantization}... \", end=\"config.quantization\")\n",
    "    schema = {\n",
    "        \"vectors\": {\n",
    "            \"size\": config.dimension,\n",
    "            \"distance\": config.distance,\n",
    "            \"hnsw_config\": {\n",
    "                \"m\": config.hnsw_m,\n",
    "                \"ef_construct\": config.hnsw_ef_construction,\n",
    "            }\n",
    "        },\n",
    "        \"hnsw_config\": {\n",
    "            \"m\": config.hnsw_m,\n",
    "            \"ef_construct\": config.hnsw_ef_construction,\n",
    "        }\n",
    "    }\n",
    "    if config.quantization == \"int8\":\n",
    "        schema[\"quantization_config\"] = {\n",
    "            \"scalar\": {\n",
    "                \"type\": \"int8\",\n",
    "                \"quantile\": 0.99,\n",
    "                \"always_ram\": True\n",
    "            }\n",
    "        }\n",
    "    response = requests.put(f\"http://{config.qdrant_host}:{config.qdrant_port}/collections/{config.index_name}\",\n",
    "                            headers={\"Content-Type\": \"application/json\"},\n",
    "                            json=schema)\n",
    "    if response.status_code == 200:\n",
    "        print(\"[OK]\")\n",
    "    else:\n",
    "        print(f\"[FAIL]\\n{response.text}\")\n",
    "\n",
    "    for field_name in [\"page_id\", \"rev_id\"]:\n",
    "        print(F\"Creating Payload integer:{config.index_name}... \", end=\"\")\n",
    "        response = requests.put(f\"http://{config.qdrant_host}:{config.qdrant_port}/collections/{config.index_name}/index\",\n",
    "                                headers={\"Content-Type\": \"application/json\"},\n",
    "                                json={\n",
    "                                    \"field_name\": field_name,\n",
    "                                    \"field_schema\": \"integer\"\n",
    "                                })\n",
    "        if response.status_code == 200:\n",
    "            print(\"[OK]\")\n",
    "        else:\n",
    "            print(f\"[FAIL]\\n{response.text}\")\n",
    "\n",
    "    for field_name in [\"section\"]:\n",
    "        print(F\"Creating Payload keyword:{config.index_name}... \", end=\"\")\n",
    "        response = requests.put(f\"http://{config.qdrant_host}:{config.qdrant_port}/collections/{config.index_name}/index\",\n",
    "                                headers={\"Content-Type\": \"application/json\"},\n",
    "                                json={\n",
    "                                    \"field_name\": field_name,\n",
    "                                    \"field_schema\": \"keyword\"\n",
    "                                })\n",
    "        if response.status_code == 200:\n",
    "            print(\"[OK]\")\n",
    "        else:\n",
    "            print(f\"[FAIL]\\n{response.text}\")\n",
    "\n",
    "    for field_name in [\"title\", \"text\"]:\n",
    "        print(F\"Creating Payload text:{config.index_name}... \", end=\"\")\n",
    "        response = requests.put(f\"http://{config.qdrant_host}:{config.qdrant_port}/collections/{config.index_name}/index\",\n",
    "                                headers={\"Content-Type\": \"application/json\"},\n",
    "                                json={\n",
    "                                    \"field_name\": field_name,\n",
    "                                    \"field_schema\": {\n",
    "                                        \"type\": \"text\",\n",
    "                                        \"tokenizer\": \"word\",\n",
    "                                        \"min_token_len\": 2,\n",
    "                                        \"max_token_len\": 2,\n",
    "                                        \"lowercase\": True\n",
    "                                    }\n",
    "                                })\n",
    "        if response.status_code == 200:\n",
    "            print(\"[OK]\")\n",
    "        else:\n",
    "            print(f\"[FAIL]\\n{response.text}\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "433e9452-3e3d-4a95-a26d-136f746e4c31",
   "metadata": {},
   "outputs": [],
   "source": [
    "def delete_index(config):\n",
    "    print(F\"Deleting Collection {config.index_name}... \", end=\"\")\n",
    "    response = requests.delete(f\"http://{config.qdrant_host}:{config.qdrant_port}/collections/{config.index_name}\")\n",
    "    if response.status_code == 200:\n",
    "        print(\"[OK]\")\n",
    "    else:\n",
    "        print(f\"[FAIL]\\n{response.text}\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "174c3528-858a-4d35-ad72-602c2c33dd84",
   "metadata": {},
   "outputs": [],
   "source": [
    "def print_indices(config):\n",
    "    response = requests.get(f\"http://{config.qdrant_host}:{config.qdrant_port}/collections/{config.index_name}\")\n",
    "    obj = json.loads(response.text)\n",
    "    pprint.pprint(obj)\n",
    "    return {\n",
    "        \"num_of_docs\": obj.get(\"result\").get(\"points_count\"),\n",
    "    }\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "05ab5ca4-9f55-42d4-8e04-82baba08af6c",
   "metadata": {},
   "outputs": [],
   "source": [
    "def wait_for_qdrant(config, retry_count=60):\n",
    "    print(f\"Waiting for {config.qdrant_name}\", end=\"\")\n",
    "    for i in range(retry_count):\n",
    "        try:\n",
    "            response = requests.get(f\"http://{config.qdrant_host}:{config.qdrant_port}/cluster\")\n",
    "            if response.status_code == 200:\n",
    "                print(\"[OK]\")        \n",
    "                return\n",
    "        except:\n",
    "            pass\n",
    "        print(\".\", end=\"\")\n",
    "        time.sleep(1)\n",
    "    print(\"[FAIL]\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ae8e027a-f07b-480f-b5b6-799e0b83d97d",
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_embedding(config, embedding_index, embedding_data, id):\n",
    "    emb_index = int(id / 100000) * 100000\n",
    "    if embedding_data is None or embedding_index != emb_index:\n",
    "        with np.load(config.embedding_path / f\"{emb_index}.npz\") as data:\n",
    "            embedding_data = data[\"embs\"]\n",
    "    embedding = embedding_data[id - emb_index]\n",
    "    if config.distance == \"Dot\":\n",
    "        embedding = embedding / np.linalg.norm(embedding)\n",
    "    return emb_index, embedding_data, embedding\n",
    "\n",
    "\n",
    "def get_section_values(config, df, min_count=10000):\n",
    "    tmp_df = df[[\"id\",\"section\"]].groupby(\"section\").count().reset_index()\n",
    "    tmp_df = tmp_df[tmp_df[\"id\"] >= min_count]\n",
    "    return tmp_df[\"section\"].values.tolist()\n",
    "\n",
    "\n",
    "def insert_data(config, max_size, bulk_size, controller=None, query_data=None):\n",
    "    start_time = time.time()\n",
    "\n",
    "    ids = []\n",
    "    vectors = []\n",
    "    payloads = []\n",
    "    def send_data(pos):\n",
    "        print(F\"Sending {int(len(ids))} docs ({pos}/{max_size})... \", end=\"\")\n",
    "        response = requests.put(f\"http://{config.qdrant_host}:{config.qdrant_port}/collections/{config.index_name}/points\",\n",
    "                                 headers={\"Content-Type\": \"application/json\"},\n",
    "                                 params={\n",
    "                                     \"wait\": \"true\",\n",
    "                                 },\n",
    "                                 data=json.dumps({\n",
    "                                     \"batch\": {\n",
    "                                         \"ids\": ids,\n",
    "                                         \"vectors\": vectors,\n",
    "                                         \"payloads\": payloads,\n",
    "                                     }\n",
    "                                 }))\n",
    "        t = json.loads(response.text).get(\"time\")\n",
    "        if response.status_code == 200:\n",
    "            print(f\"[OK] {t}\")\n",
    "        else:\n",
    "            print(f\"[FAIL] {t} {response.status_code} {response.text}\")\n",
    "        return t\n",
    "\n",
    "    total_time = 0\n",
    "    count = 0\n",
    "    embedding_index = -1\n",
    "    embedding_data = None\n",
    "    for content_file in sorted(config.content_path.glob(\"*.parquet\")):\n",
    "        if count >= max_size:\n",
    "            break\n",
    "        df = pd.read_parquet(content_file)\n",
    "        if query_data is not None:\n",
    "            query_data[\"section_values\"].extend(get_section_values(config, df))\n",
    "        for i,row in df.iterrows():\n",
    "            if count >= max_size:\n",
    "                break\n",
    "            embedding_index, embedding_data, embedding = get_embedding(config, embedding_index, embedding_data, row.id)\n",
    "            count += 1\n",
    "            ids.append(count)\n",
    "            vectors.append(embedding.tolist())\n",
    "            payloads.append({\n",
    "                \"page_id\": row.pageid,\n",
    "                \"rev_id\": row.revid,\n",
    "                # \"title\": row.title,\n",
    "                \"section\": row.section,\n",
    "                # \"text\": row.text,\n",
    "            })\n",
    "            if len(ids) >= bulk_size:\n",
    "                total_time += send_data(count)\n",
    "                ids = []\n",
    "                vectors = []\n",
    "                payloads = []\n",
    "                if controller is not None and not controller.run():\n",
    "                    return\n",
    "\n",
    "    if len(ids) > 0:\n",
    "        total_time += send_data(count)\n",
    "\n",
    "    green_count = 0\n",
    "    while green_count < 30:\n",
    "        response = requests.get(f\"http://{config.qdrant_host}:{config.qdrant_port}/collections/{config.index_name}\")\n",
    "        obj = json.loads(response.text)\n",
    "        if obj.get(\"result\").get(\"status\") == \"green\":\n",
    "            green_count += 1\n",
    "        else:\n",
    "            green_count = 0 # reset\n",
    "        print(\".\", end=\"\")\n",
    "        time.sleep(1)\n",
    "    print(\".\")\n",
    "\n",
    "    execution_time = time.time() - start_time\n",
    "    hours, remainder = divmod(execution_time, 3600)\n",
    "    minutes, seconds = divmod(remainder, 60)\n",
    "    print(f\"Execution Time: {int(hours):02d}:{int(minutes):02d}:{seconds:02.2f} ({timedelta(seconds=total_time)})\")\n",
    "    return {\n",
    "        \"execution_time\": execution_time,\n",
    "        \"process_time\": total_time,\n",
    "    }\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8b19c777-6cea-4899-ab8c-51ef814f093a",
   "metadata": {},
   "outputs": [],
   "source": [
    "class IndexingController:\n",
    "\n",
    "    def __init__(self, stop_event):\n",
    "        self._previous_time = time.time()\n",
    "        self._stop_event = stop_event\n",
    "\n",
    "    def run(self):\n",
    "        if self._stop_event.is_set():\n",
    "            return False\n",
    "\n",
    "        wait_time = 1 - ( time.time() - self._previous_time)\n",
    "        if wait_time > 0:\n",
    "            # print(f\"Waiting for {wait_time} sec\")\n",
    "            time.sleep(wait_time)\n",
    "        self._previous_time =  time.time()\n",
    "        return True\n",
    "\n",
    "\n",
    "def update_data(target_config, stop_event):\n",
    "    dataset_config = get_dataset_config(target_config)\n",
    "    if dataset_config.update_docs_per_sec > 0:\n",
    "        print(f\"Starting update for {target_config}\")\n",
    "        insert_data(dataset_config, \n",
    "                    max_size=dataset_config.index_size,\n",
    "                    bulk_size=dataset_config.update_docs_per_sec,\n",
    "                    controller=IndexingController(stop_event))\n",
    "        print(f\"Stopping update for {target_config}\")\n",
    "    else:\n",
    "        print(f\"No background updates\")\n",
    "\n",
    "\n",
    "def start_update(target_config):\n",
    "    stop_event = multiprocessing.Event()\n",
    "\n",
    "    p = multiprocessing.Process(target=update_data, args=(target_config, stop_event))\n",
    "    p.start()\n",
    "\n",
    "    def stop_update():\n",
    "        stop_event.set()\n",
    "        p.join()\n",
    "\n",
    "    return stop_update\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4156e80f-ef11-4015-a39e-c772e994ac65",
   "metadata": {},
   "outputs": [],
   "source": [
    "def search(config, query):\n",
    "    response = requests.post(f\"http://{config.qdrant_host}:{config.qdrant_port}/collections/{config.index_name}/points/search\",\n",
    "                             headers={\"Content-Type\": \"application/json\"},\n",
    "                             json=query)\n",
    "    # print(response.text)\n",
    "\n",
    "    if response.status_code == 200:\n",
    "        obj = json.loads(response.text)\n",
    "        if obj.get(\"status\") != \"ok\":\n",
    "            print(f\"[FAIL] {response.text}\")\n",
    "            return -1, -1, [], [], []\n",
    "        product_ids = [x.get(\"id\") for x in obj.get(\"result\")]\n",
    "        scores = [x.get(\"score\") for x in obj.get(\"result\")]\n",
    "        return obj.get(\"time\") * 1000, len(obj.get(\"result\")), product_ids, scores\n",
    "    print(f\"[FAIL][{response.status_code}] {response.text}\")\n",
    "    return -1, -1, [], []\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4bd7acf0-c0f7-4eb4-87c7-8fff9e334ea1",
   "metadata": {},
   "outputs": [],
   "source": [
    "def search_with_knn_queries(config, output_path, pre_filter=None, max_size=10000, page_size=100, offset=0, max_error_count=100, exact=False):\n",
    "    print(\"Sending knn queries...\")\n",
    "    start_time = time.time()\n",
    "    pos = offset\n",
    "    doc_id = 0\n",
    "    count = 0\n",
    "    running = True\n",
    "    error_count = 0\n",
    "    with gzip.open(output_path, \"wt\", encoding=\"utf-8\") as f:\n",
    "        while running:\n",
    "            with np.load(config.embedding_path / f\"{pos}.npz\") as data:\n",
    "                embedding_data = data[\"embs\"]\n",
    "            for embedding in embedding_data:\n",
    "                doc_id += 1\n",
    "                if count >= max_size:\n",
    "                    running = False\n",
    "                    break\n",
    "                if config.distance == \"Dot\":\n",
    "                    embedding = embedding / np.linalg.norm(embedding)\n",
    "                query = {\n",
    "                    \"vector\": embedding.tolist(),\n",
    "                    \"limit\": page_size,\n",
    "                    # \"with_payload\": True,\n",
    "                    \"params\": {\n",
    "                        \"hnsw_ef\": config.hnsw_ef,\n",
    "                        \"exact\": exact,\n",
    "                    },\n",
    "                }\n",
    "                if pre_filter is not None:\n",
    "                    query[\"filter\"] = next(pre_filter)\n",
    "                took, hits, ids, scores = search(config, query=query)\n",
    "                # print(f\"{took}, {ids}, {scores}\")\n",
    "                if took == -1:\n",
    "                    error_count += 1\n",
    "                    if error_count >= max_error_count:\n",
    "                        running = False\n",
    "                        break\n",
    "                    continue\n",
    "                result = {\n",
    "                    \"id\": doc_id,\n",
    "                    \"took\": took,\n",
    "                    \"hits\": hits,\n",
    "                    \"ids\": ids,\n",
    "                    \"scores\": scores,\n",
    "                }\n",
    "                f.write(json.dumps(result, ensure_ascii=False))\n",
    "                f.write(\"\\n\")\n",
    "                count += 1\n",
    "                if count % 10000 == 0:\n",
    "                    print(f\"Sent {count}/{max_size} queries.\")\n",
    "\n",
    "            pos += 100000\n",
    "            if pos > config.num_of_docs:\n",
    "                pos = 0\n",
    "\n",
    "    execution_time = time.time() - start_time\n",
    "    hours, remainder = divmod(execution_time, 3600)\n",
    "    minutes, seconds = divmod(remainder, 60)\n",
    "    print(f\"Execution Time: {int(hours):02d}:{int(minutes):02d}:{seconds:02.2f}\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cee339fd-e622-45cb-a17a-e3b19492f92e",
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_output_filename(version, name, exact):\n",
    "    filename = f\"output/qdrant{version.replace('.', '_')}_{name}\"\n",
    "    if exact:\n",
    "        filename += \"_exact\"\n",
    "    filename += \".jsonl.gz\"\n",
    "    return filename\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a9292d64-c116-4548-aaf9-6865613e2562",
   "metadata": {},
   "outputs": [],
   "source": [
    "def print_took_and_total_hits(k, filename, truth_filename):\n",
    "    query_ids = []\n",
    "    tooks = []\n",
    "    hits = []\n",
    "    df = pd.read_json(filename, lines=True)\n",
    "    result = {\n",
    "        \"num_of_queries\": len(df),\n",
    "        \"took\": {\n",
    "            \"mean\": df.took.mean(),\n",
    "            \"std\": df.took.std(),\n",
    "            \"min\": df.took.min(),\n",
    "            \"25%\": df.took.quantile(0.25),\n",
    "            \"50%\": df.took.quantile(0.5),\n",
    "            \"75%\": df.took.quantile(0.75),\n",
    "            \"90%\": df.took.quantile(0.9),\n",
    "            \"99%\": df.took.quantile(0.99),\n",
    "            \"max\": df.took.max(),\n",
    "        },\n",
    "        \"hits\": {\n",
    "            \"mean\": df.hits.mean(),\n",
    "            \"std\": df.hits.std(),\n",
    "            \"min\": df.hits.min(),\n",
    "            \"25%\": df.hits.quantile(0.25),\n",
    "            \"50%\": df.hits.quantile(0.5),\n",
    "            \"75%\": df.hits.quantile(0.75),\n",
    "            \"max\": df.hits.max(),\n",
    "        },\n",
    "    }\n",
    "    if os.path.exists(truth_filename):\n",
    "        df = pd.merge(df,\n",
    "                      pd.read_json(truth_filename, lines=True)[[\"id\", \"ids\"]].rename(columns={\"ids\":\"truth_ids\"}),\n",
    "                      on=\"id\", how=\"inner\")\n",
    "        def get_precision(row):\n",
    "            size = len(row[\"truth_ids\"])\n",
    "            if size > k:\n",
    "                size = k\n",
    "            return len(set(row[\"ids\"]).intersection(set(row[\"truth_ids\"]))) / size\n",
    "        df[\"precision\"] = df.apply(get_precision, axis=1)\n",
    "        result[\"precision\"] = {\n",
    "            \"mean\": df.precision.mean(),\n",
    "            \"std\": df.precision.std(),\n",
    "            \"min\": df.precision.min(),\n",
    "            \"25%\": df.precision.quantile(0.25),\n",
    "            \"50%\": df.precision.quantile(0.5),\n",
    "            \"75%\": df.precision.quantile(0.75),\n",
    "            \"90%\": df.precision.quantile(0.9),\n",
    "            \"99%\": df.precision.quantile(0.99),\n",
    "            \"max\": df.precision.max(),\n",
    "        }\n",
    "    print(df.describe().to_markdown())\n",
    "    return result\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "55aa4a6b-7eae-4fba-aabd-59dbaa746981",
   "metadata": {},
   "outputs": [],
   "source": [
    "def save_results(target_config, config, results):\n",
    "    with open(\"results.json\", \"wt\", encoding=\"utf-8\") as f:\n",
    "        json.dump({\n",
    "            \"target\": target_config,\n",
    "            \"version\": config.qdrant_version,\n",
    "            \"settings\": asdict(config),\n",
    "            \"results\": results,\n",
    "            \"timestamp\": datetime.now().isoformat(),\n",
    "        }, f, ensure_ascii=False, default=lambda x: int(x) if isinstance(x, np.int64) else None)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ef1e70da-5f07-4659-bde2-af40786f76c8",
   "metadata": {},
   "outputs": [],
   "source": [
    "query_data = {\"section_values\": []}\n",
    "results = {}\n",
    "target_config = os.getenv(\"TARGET_CONFIG\", \"100k-768-m32-efc200-ef100-ip\")\n",
    "dataset_config = get_dataset_config(target_config)\n",
    "pprint.pprint(dataset_config)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "65a5f319-12dd-443e-9dec-aefbacccf75f",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "prune_docker(dataset_config)\n",
    "print(f\"<<<Qdrant {dataset_config.qdrant_version}>>>\")\n",
    "run_qdrant(dataset_config)\n",
    "wait_for_qdrant(dataset_config)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8a9720ed-c07c-48cb-bad3-bb4b49d482ae",
   "metadata": {},
   "outputs": [],
   "source": [
    "print_docker_container_stats(dataset_config)\n",
    "print_docker_system_df(dataset_config)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5127ae0c-ac36-4141-a643-3492163b8f84",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "create_index(dataset_config)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bcbdfa51-5bdf-494f-a9aa-3f265b8ec20d",
   "metadata": {},
   "outputs": [],
   "source": [
    "print_docker_container_stats(dataset_config)\n",
    "print_indices(dataset_config)\n",
    "print_docker_system_df(dataset_config)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3dc02726-3045-4f24-9a66-f018ab3f2d93",
   "metadata": {},
   "outputs": [],
   "source": [
    "results[\"indexing\"] = insert_data(dataset_config,\n",
    "                                  max_size=dataset_config.index_size,\n",
    "                                  bulk_size=dataset_config.bulk_size,\n",
    "                                  query_data=query_data)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c30e3946-3158-4e86-9aa6-5bb5e06a2368",
   "metadata": {},
   "outputs": [],
   "source": [
    "results[\"indexing\"][\"container\"] = print_docker_container_stats(dataset_config)\n",
    "print_indices(dataset_config)\n",
    "print_docker_system_df(dataset_config)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cfbef882-944e-4e5d-8297-e85d9981ebce",
   "metadata": {},
   "outputs": [],
   "source": [
    "for page_size in [10, 100]:\n",
    "    print(f\"page size: {page_size}\")\n",
    "    filename = get_output_filename(dataset_config.qdrant_version, f\"knn_{page_size}\", exact=dataset_config.exact)\n",
    "    stop_update = start_update(target_config)\n",
    "    search_with_knn_queries(dataset_config, filename, page_size=page_size, max_size=1000) # warmup\n",
    "    search_with_knn_queries(dataset_config, filename, page_size=page_size, offset=dataset_config.index_size, exact=dataset_config.exact)\n",
    "    stop_update()\n",
    "    results[f\"top_{page_size}\"] = print_took_and_total_hits(page_size, filename,  f\"dataset/ground_truth/{re.sub(r'-m.*', '', target_config)}/knn_{page_size}.jsonl.gz\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5c312bff-d8a8-4e4c-a5d6-2caac4916c04",
   "metadata": {},
   "outputs": [],
   "source": [
    "def pre_filter_generator():\n",
    "    section_values = query_data[\"section_values\"]\n",
    "    if len(section_values) > 0:\n",
    "        while True:\n",
    "            for s in section_values:\n",
    "                yield {\n",
    "                    \"must\": [\n",
    "                        {\n",
    "                            \"key\": \"section\",\n",
    "                            \"match\": {\n",
    "                                \"value\": s\n",
    "                            }\n",
    "                        }\n",
    "                    ]\n",
    "                }\n",
    "\n",
    "results[\"num_of_filtered_words\"] = len(query_data[\"section_values\"])\n",
    "for page_size in [10, 100]:\n",
    "    print(f\"page size: {page_size}\")\n",
    "    filename = get_output_filename(dataset_config.qdrant_version, f\"knn_{page_size}_filtered\", exact=dataset_config.exact)\n",
    "    stop_update = start_update(target_config)\n",
    "    search_with_knn_queries(dataset_config, filename, page_size=page_size, max_size=1000, pre_filter=pre_filter_generator()) # warmup\n",
    "    search_with_knn_queries(dataset_config, filename, page_size=page_size, offset=dataset_config.index_size, pre_filter=pre_filter_generator(), exact=dataset_config.exact)\n",
    "    stop_update()\n",
    "    results[f\"top_{page_size}_filtered\"] = print_took_and_total_hits(page_size, filename, f\"dataset/ground_truth/{re.sub(r'-m.*', '', target_config)}/knn_{page_size}_filtered.jsonl.gz\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "caf22d76-2ff2-4b6d-b567-6a1117403154",
   "metadata": {},
   "outputs": [],
   "source": [
    "results[\"indexing\"][\"container\"] = print_docker_container_stats(dataset_config)\n",
    "print_indices(dataset_config)\n",
    "print_docker_system_df(dataset_config)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9662cbf9-79ae-4159-b012-39327d02198d",
   "metadata": {},
   "outputs": [],
   "source": [
    "save_results(target_config, dataset_config, results)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "086b270b-2dba-4f0e-9793-e315b3279aac",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "delete_index(dataset_config)\n",
    "stop_qdrant(dataset_config)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fdc6c0c1-a508-4c04-a30c-6f14958d7eb7",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "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.10.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}