{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Importing and Initializing\n",
    "\n",
    "We need to import the `BTree` class that is located in file `btree.py`. Note that when we import we don't put the `.py` extension.\n",
    "\n",
    "To call the constructor of the `BTree` method we need to use the `super()` function.\n",
    "\n",
    "To simplify testing, we will always build the `BTree` using a `split_threshold` equal to `2`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "from btree import BTree"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "class KVStore(BTree):\n",
    "    \n",
    "    def __init__(self):\n",
    "        super().__init__(split_threshold=2)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Overriding the Add Method\n",
    "\n",
    "To override a method from the superclass we need to declare a method with the same name as the one we are overriding. \n",
    "\n",
    "We override the `add()` method because we want a different behavior in the `KVStore` than the one inherited from the `BTree`. Namely, we want to have no duplicates.\n",
    "\n",
    "To implement the new `add()` method we will need to use the `BTree._find_node()` and `BTree.add()` methods."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "class KVStore(BTree):\n",
    "    \n",
    "    def __init__(self):\n",
    "        super().__init__(split_threshold=2)\n",
    "        \n",
    "    def add(self, key, value):\n",
    "        node = self._find_node(self.root, key)\n",
    "        if node is None:\n",
    "            super().add(key, value)\n",
    "        else:\n",
    "            # Replace the old value by the new\n",
    "            for i, node_key in enumerate(node.keys):\n",
    "                if node_key == key:\n",
    "                    node.values[i] = value"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Testing\n",
    "\n",
    "Let's test the current implementation. We want to make sure that:\n",
    "\n",
    "1. The split threshold is correct.\n",
    "2. We can add entries.\n",
    "3. We can retrieve a value given a key.\n",
    "4. If we add two entries with the same key, the value is updated."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "kv = KVStore()\n",
    "\n",
    "# Check the split threshold\n",
    "assert kv.split_threshold == 2, \"Split threshold is 2\"\n",
    "\n",
    "# Add the entries (i, i) for i from 0 to 9\n",
    "for i in range(10):\n",
    "    kv.add(i, i)\n",
    "\n",
    "# Check the values\n",
    "for i in range(10):\n",
    "    assert kv.get_value(i) == i, \"Value of i is i\"\n",
    "\n",
    "# Add again with different values\n",
    "for i in range(10):\n",
    "    kv.add(i, i + 1)\n",
    "\n",
    "# Check the new values\n",
    "for i in range(10):\n",
    "    assert kv.get_value(i) == i + 1, \"Value of i is i + 1\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Implementing the Item Getter and Setter\n",
    "\n",
    "To allow users to use the bracket notation, we need to implement the `__getitem__()` and `__setitem__()` methods.\n",
    "\n",
    "These methods are already implement but are named `get_value()` and `add()`, respectively. We can thus implement them by calling the corresponding method."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "class KVStore(BTree):\n",
    "    \n",
    "    def __init__(self):\n",
    "        super().__init__(split_threshold=2)\n",
    "        \n",
    "    def add(self, key, value):\n",
    "        node = self._find_node(self.root, key)\n",
    "        if node is None:\n",
    "            super().add(key, value)\n",
    "        else:\n",
    "            # Replace the old value by the new\n",
    "            for i, node_key in enumerate(node.keys):\n",
    "                if node_key == key:\n",
    "                    node.values[i] = value\n",
    "                    \n",
    "    def __setitem__(self, key, value):\n",
    "        self.add(key, value)\n",
    "        \n",
    "    def __getitem__(self, key):\n",
    "        return self.get_value(key)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Testing Getter and Setter\n",
    "\n",
    "Let's redo the same tests but using bracket syntax."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "kv = KVStore()\n",
    "\n",
    "# Check the split threshold\n",
    "assert kv.split_threshold == 2, \"Split threshold is 2\"\n",
    "\n",
    "# Add the entries (i, i) for i from 0 to 9\n",
    "for i in range(10):\n",
    "    kv[i] = i\n",
    "\n",
    "# Check the values\n",
    "for i in range(10):\n",
    "    assert kv[i] == i, \"Value of i is i\"\n",
    "\n",
    "# Add again with different values\n",
    "for i in range(10):\n",
    "    kv[i] = i + 1\n",
    "\n",
    "# Check the new values\n",
    "for i in range(10):\n",
    "    assert kv[i] == i + 1, \"Value of i is i + 1\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Enhancing the Contains Method\n",
    "\n",
    "To enable the `in` operator, we need wrap the `BTree.contains()` method inside a new method named `__contains__()`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "class KVStore(BTree):\n",
    "    \n",
    "    def __init__(self):\n",
    "        super().__init__(split_threshold=2)\n",
    "\n",
    "    def add(self, key, value):\n",
    "        node = self._find_node(self.root, key)\n",
    "        if node is None:\n",
    "            super().add(key, value)\n",
    "        else:\n",
    "            # Replace the old value by the new\n",
    "            for i, node_key in enumerate(node.keys):\n",
    "                if node_key == key:\n",
    "                    node.values[i] = value\n",
    "                    \n",
    "    def __setitem__(self, key, value):\n",
    "        self.add(key, value)\n",
    "        \n",
    "    def __getitem__(self, key):\n",
    "        return self.get_value(key)\n",
    "    \n",
    "    def __contains__(self, key):\n",
    "        return self.contains(key)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Testing the In Operator\n",
    "\n",
    "Let's do a test to see if we can use the `in` operator. We'll use alphabet letters as keys to see if other kind of keys are supported."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "kv = KVStore()\n",
    "for c in 'abcdefghijklmnopqrstuvwxyz':\n",
    "    kv[c] = c\n",
    "\n",
    "for c in 'abcdefghijklmnopqrstuvwxyz':\n",
    "    assert c in kv, \"Character is in the key-value store\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Range Queries\n",
    "\n",
    "Our solution consisted in replacing both `float('-inf')` and `float('inf')` by `None`. Then we created a method named `_range_intersects` that checks whether the query range intersects with the node range.\n",
    "\n",
    "We make the condition work in a way such that if `min_key` is `None` then it is always considered smaller than any other key and if `max_key` is `None` then it is always considered larger than any other key."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "class KVStore(BTree):\n",
    "    \n",
    "    def __init__(self):\n",
    "        super().__init__(split_threshold=2)\n",
    "        \n",
    "    def add(self, key, value):\n",
    "        node = self._find_node(self.root, key)\n",
    "        if node is None:\n",
    "            super().add(key, value)\n",
    "        else:\n",
    "            # Replace the old value by the new\n",
    "            for i, node_key in enumerate(node.keys):\n",
    "                if node_key == key:\n",
    "                    node.values[i] = value\n",
    "                    \n",
    "    def __setitem__(self, key, value):\n",
    "        self.add(key, value)\n",
    "        \n",
    "    def __getitem__(self, key):\n",
    "        return self.get_value(key)\n",
    "\n",
    "    def __contains__(self, key):\n",
    "        return self.contains(key)\n",
    "    \n",
    "    def _range_query(self, range_start, range_end, current_node, min_key, max_key):\n",
    "        if not self._range_intersects(range_start, range_end, min_key, max_key):\n",
    "            return []\n",
    "        results = []\n",
    "        for i, key in enumerate(current_node.keys):\n",
    "            if range_start <= key and key <= range_end:\n",
    "                results.append(current_node.values[i])\n",
    "        if not current_node.is_leaf():\n",
    "            for i, child in enumerate(current_node.children):\n",
    "                new_min_key = current_node.keys[i - 1] if i > 0 else min_key\n",
    "                new_max_key = current_node.keys[i] if i < len(current_node) else max_key\n",
    "                results += self._range_query(range_start, range_end, child, new_min_key, new_max_key)\n",
    "        return results \n",
    "\n",
    "    def range_query(self, range_start, range_end):\n",
    "        return self._range_query(range_start, range_end, self.root, float('-inf'), float('inf'))\n",
    "    \n",
    "    def _range_intersects(self, range_start, range_end, node_min, node_max):\n",
    "        if not node_min is None and node_min > range_end:\n",
    "            return False\n",
    "        if not node_max is None and node_max < range_start:\n",
    "            return False\n",
    "        return True"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "class DictKVStore(dict):\n",
    "    \n",
    "    def range_query(self, range_start, range_end):\n",
    "        result = []\n",
    "        for key in self.keys():\n",
    "            if range_start <= key and key <= range_end:\n",
    "                result.append(self[key])\n",
    "        return result"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Test range query\n",
    "\n",
    "Test the range against the dictionary range query."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "dict_kv = DictKVStore()\n",
    "our_kv = KVStore()\n",
    "for i in range(10):\n",
    "    dict_kv[i] = i\n",
    "    our_kv[i] = i\n",
    "\n",
    "for range_start, range_end in [(1, 3), (4, 6), (1, 10), (5, 5)]:\n",
    "    dict_res = sorted(dict_kv.range_query(range_start, range_end))\n",
    "    our_res = sorted(our_kv.range_query(range_start, range_end))\n",
    "    assert dict_res == our_res, \"Both data structures return the same range query result.\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Random Tests"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Testing Insertion\n",
      "Testing Length\n",
      "Testing Values\n",
      "Testing in Operator\n",
      "Testing Range Queries\n"
     ]
    }
   ],
   "source": [
    "import random\n",
    "random.seed(0)\n",
    "\n",
    "NUM_INSERTS = 10\n",
    "NUM_CONTAINS = 10\n",
    "NUM_RANGE_QUERIES = 10\n",
    "\n",
    "dict_kv = DictKVStore()\n",
    "\n",
    "kv = KVStore()\n",
    "\n",
    "print(\"Testing Insertion\")\n",
    "for _ in range(NUM_INSERTS):\n",
    "    key = random.randint(0, 100)\n",
    "    value = random.randint(0, 1000000)\n",
    "    dict_kv[key] = value\n",
    "    kv[key] = value\n",
    "    \n",
    "print(\"Testing Length\")\n",
    "assert len(dict_kv) == len(kv), \"Wrong length. Length should be {} but is {}.\".format(len(dict_kv), len(kv))\n",
    "    \n",
    "\n",
    "print(\"Testing Values\")\n",
    "for key in dict_kv:\n",
    "    assert dict_kv[key] == kv[key], \"Wrong value for key {}. Expected value {} but found value {}.\".format(key, dict_kv[key], kv[key])\n",
    "\n",
    "    \n",
    "print(\"Testing in Operator\")\n",
    "for i in range(NUM_CONTAINS):\n",
    "    key = random.randint(0, 1000)\n",
    "    assert (key in dict_kv) == (key in kv), \"Contains method did not return the correct value for key {}.\".format(key)\n",
    "\n",
    "\n",
    "print(\"Testing Range Queries\")\n",
    "for _ in range(NUM_RANGE_QUERIES):\n",
    "    range_start = random.randint(0, 100)\n",
    "    range_end = random.randint(range_start, 100)\n",
    "    dict_results = dict_kv.range_query(range_start, range_end)\n",
    "    kv_results = kv.range_query(range_start, range_end)\n",
    "    assert len(dict_results) == len(kv_results), \"Wrong number of reuslt in range query [{}, {}]. Should be {} but was {}.\".format(range_start, range_end, len(dict_result), len(kv_result))\n",
    "    dict_results.sort()\n",
    "    kv_results.sort()\n",
    "    assert dict_results == kv_results, \"Wrong number of reuslt in range query [{}, {}]. Should be {} but was {}.\".format(range_start, range_end, len(dict_result), len(kv_result))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Speed Tests\n",
    "\n",
    "To perform the speed tests we start by creating an empty data structure of each type.\n",
    "\n",
    "Then we load all entries from the `entries.csv` file.\n",
    "\n",
    "After that, loop over each query in the `queries.csv` file. For each query, we measure its execution time on both data structure. Then we compute the execution time ratio between the dictionary solution and our solution.\n",
    "\n",
    "In the end we plot the result."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 15480 characters]\n",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "import time\n",
    "import csv\n",
    "\n",
    "dict_kv = DictKVStore()\n",
    "our_kv = KVStore()\n",
    "\n",
    "# Load the entries\n",
    "with open('entries.csv', 'r') as f:\n",
    "    rows = list(csv.reader(f))[1:]\n",
    "    for row in rows:\n",
    "        key = int(row[0])\n",
    "        value = int(row[1])\n",
    "        dict_kv[key] = value\n",
    "        our_kv[key] = value\n",
    "\n",
    "# Measure query times\n",
    "time_ratios = []\n",
    "with open('queries.csv', 'r') as f:\n",
    "    rows = list(csv.reader(f))[1:]\n",
    "    for row in rows:\n",
    "        range_start = int(row[0])\n",
    "        range_end = int(row[1])\n",
    "        \n",
    "        start = time.time()\n",
    "        dict_kv.range_query(range_start, range_end)\n",
    "        end = time.time()\n",
    "        time_dict = end - start\n",
    "\n",
    "        start = time.time()\n",
    "        our_kv.range_query(range_start, range_end)\n",
    "        end = time.time()\n",
    "        time_kv = end - start\n",
    "\n",
    "        time_ratios.append(time_dict / time_kv)\n",
    "\n",
    "# Plot results\n",
    "import matplotlib.pyplot as plt\n",
    "plt.plot(time_ratios)\n",
    "plt.xlabel('Query range result size')\n",
    "plt.ylabel('Runtime ratio')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Conclusion\n",
    "\n",
    "For 50,000 entries, we get a performance boost of at most 50 times. \n",
    "\n",
    "We see that the performance boost decreases as the size of the of query increases. This is expected since the more result we return the closer we get to having to iterate of all entries in the tree.\n",
    "\n",
    "With 100,000 entries the performance boost can go up to about 120 times."
   ]
  }
 ],
 "metadata": {
  "anaconda-cloud": {},
  "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.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
