Example Use.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "d9c5213b",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import os\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import time\n",
    "import pandas as pd\n",
    "from scipy.stats import pearsonr"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "037bec4a",
   "metadata": {},
   "outputs": [],
   "source": [
    "from model.util import Normalizer\n",
    "from model.database_util import get_hist_file, get_job_table_sample, collator\n",
    "from model.model import QueryFormer\n",
    "from model.database_util import Encoding\n",
    "from model.dataset import PlanTreeDataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "822fdcaf",
   "metadata": {},
   "outputs": [],
   "source": [
    "data_path = './data/imdb/'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "fbcd4773",
   "metadata": {},
   "outputs": [],
   "source": [
    "class Args:\n",
    "    pass"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "aace65f6",
   "metadata": {},
   "outputs": [],
   "source": [
    "hist_file = get_hist_file(data_path + 'histogram_string.csv')\n",
    "cost_norm = Normalizer(-3.61192, 12.290855)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "4e5f421a",
   "metadata": {},
   "outputs": [],
   "source": [
    "encoding_ckpt = torch.load('checkpoints/encoding.pt')\n",
    "encoding = encoding_ckpt['encoding']\n",
    "checkpoint = torch.load('checkpoints/cost_model.pt', map_location='cpu')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "71759b41",
   "metadata": {},
   "outputs": [],
   "source": [
    "from model.util import seed_everything\n",
    "seed_everything()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fc547102",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "8de39e12",
   "metadata": {},
   "outputs": [],
   "source": [
    "args = checkpoint['args']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "f9592f60",
   "metadata": {},
   "outputs": [],
   "source": [
    "model = QueryFormer(emb_size = args.embed_size ,ffn_dim = args.ffn_dim, head_size = args.head_size, \\\n",
    "                 dropout = args.dropout, n_layers = args.n_layers, \\\n",
    "                 use_sample = True, use_hist = True, \\\n",
    "                 pred_hid = args.pred_hid\n",
    "                )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "b2ae2829",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<All keys matched successfully>"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "model.load_state_dict(checkpoint['model'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "85e27584",
   "metadata": {},
   "outputs": [],
   "source": [
    "device = 'cuda:0'\n",
    "_ = model.to(device).eval()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "6c8a7d25",
   "metadata": {},
   "outputs": [],
   "source": [
    "to_predict = 'cost'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "bfa517f2",
   "metadata": {},
   "outputs": [],
   "source": [
    "methods = {\n",
    "    'get_sample' : get_job_table_sample,\n",
    "    'encoding': encoding,\n",
    "    'cost_norm': cost_norm,\n",
    "    'hist_file': hist_file,\n",
    "    'model': model,\n",
    "    'device': device,\n",
    "    'bs': 512,\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bf2b14a9",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "bd9e7796",
   "metadata": {},
   "outputs": [],
   "source": [
    "def print_qerror(preds_unnorm, labels_unnorm):\n",
    "    qerror = []\n",
    "    for i in range(len(preds_unnorm)):\n",
    "        if preds_unnorm[i] > float(labels_unnorm[i]):\n",
    "            qerror.append(preds_unnorm[i] / float(labels_unnorm[i]))\n",
    "        else:\n",
    "            qerror.append(float(labels_unnorm[i]) / float(preds_unnorm[i]))\n",
    "\n",
    "    e_50, e_90 = np.median(qerror), np.percentile(qerror,90)    \n",
    "    e_mean = np.mean(qerror)\n",
    "    print(\"Median: {}\".format(e_50))\n",
    "    print(\"90th percentile: {}\".format(e_90))\n",
    "    print(\"Mean: {}\".format(e_mean))\n",
    "    return \n",
    "\n",
    "def get_corr(ps, ls): # unnormalised\n",
    "    ps = np.array(ps)\n",
    "    ls = np.array(ls)\n",
    "    corr, _ = pearsonr(np.log(ps), np.log(ls))\n",
    "    \n",
    "    return corr"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "0b1cee2c",
   "metadata": {},
   "outputs": [],
   "source": [
    "def evaluate(model, ds, bs, norm, device):\n",
    "    model.eval()\n",
    "    cost_predss = np.empty(0)\n",
    "\n",
    "    with torch.no_grad():\n",
    "        for i in range(0, len(ds), bs):\n",
    "            batch, batch_labels = collator(list(zip(*[ds[j] for j in range(i,min(i+bs, len(ds)) ) ])))\n",
    "\n",
    "            batch = batch.to(device)\n",
    "\n",
    "            cost_preds, _ = model(batch)\n",
    "            cost_preds = cost_preds.squeeze()\n",
    "\n",
    "            cost_predss = np.append(cost_predss, cost_preds.cpu().detach().numpy())\n",
    "\n",
    "            \n",
    "    print_qerror(norm.unnormalize_labels(cost_predss), ds.costs)\n",
    "    corr = get_corr(norm.unnormalize_labels(cost_predss), ds.costs)\n",
    "    print('Corr: ', corr)\n",
    "\n",
    "    return "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "49f2e71c",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "13c13677",
   "metadata": {},
   "outputs": [],
   "source": [
    "def eval_workload(workload, methods):\n",
    "\n",
    "    get_table_sample = methods['get_sample']\n",
    "\n",
    "    workload_file_name = './data/imdb/workloads/' + workload\n",
    "    table_sample = get_table_sample(workload_file_name)\n",
    "    plan_df = pd.read_csv('./data/imdb/{}_plan.csv'.format(workload))\n",
    "    workload_csv = pd.read_csv('./data/imdb/workloads/{}.csv'.format(workload),sep='#',header=None)\n",
    "    workload_csv.columns = ['table','join','predicate','cardinality']\n",
    "    \n",
    "    ds = PlanTreeDataset(plan_df, workload_csv, \\\n",
    "        methods['encoding'], methods['hist_file'], methods['cost_norm'], \\\n",
    "        methods['cost_norm'], 'cost', table_sample)\n",
    "\n",
    "    evaluate(methods['model'], ds, methods['bs'], methods['cost_norm'], methods['device'])\n",
    "    return "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "fd89df96",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Loaded queries with len  70\n",
      "Loaded bitmaps\n",
      "Median: 1.4075482846905147\n",
      "90th percentile: 20.276220153301075\n",
      "Mean: 12.494164553182886\n",
      "Corr:  0.9028569487410075\n"
     ]
    }
   ],
   "source": [
    "eval_workload('job-light', methods)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "9e40c30c",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Loaded queries with len  5000\n",
      "Loaded bitmaps\n",
      "Median: 1.0650682558884779\n",
      "90th percentile: 1.717860179422548\n",
      "Mean: 1.5365826871030694\n",
      "Corr:  0.9866923820645216\n"
     ]
    }
   ],
   "source": [
    "eval_workload('synthetic', methods)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b30aceed",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5e47dfb9",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f0622ba9",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "21b92aa3",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a3ceb39d",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}