End2EndCostEstimator_FeatureExtraction.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import json\n",
    "import re\n",
    "import pandas as pd\n",
    "import re\n",
    "import pypred"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Load Dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def load_dataset(dir_path):\n",
    "    data={}\n",
    "    data[\"aka_name\"] = pd.read_csv(dir_path+'/aka_name.csv',header=None)\n",
    "    data[\"aka_title\"] = pd.read_csv(dir_path+'/aka_title.csv',header=None)\n",
    "    data[\"cast_info\"] = pd.read_csv(dir_path+'/cast_info.csv',header=None)\n",
    "    data[\"char_name\"] = pd.read_csv(dir_path+'/char_name.csv',header=None)\n",
    "    data[\"company_name\"] = pd.read_csv(dir_path+'/company_name.csv',header=None)\n",
    "    data[\"company_type\"] = pd.read_csv(dir_path+'/company_type.csv',header=None)\n",
    "    data[\"comp_cast_type\"] = pd.read_csv(dir_path+'/comp_cast_type.csv',header=None)\n",
    "    data[\"complete_cast\"] = pd.read_csv(dir_path+'/complete_cast.csv',header=None)\n",
    "    data[\"info_type\"] = pd.read_csv(dir_path+'/info_type.csv',header=None)\n",
    "    data[\"keyword\"] = pd.read_csv(dir_path+'/keyword.csv',header=None)\n",
    "    data[\"kind_type\"] = pd.read_csv(dir_path+'/kind_type.csv',header=None)\n",
    "    data[\"link_type\"] = pd.read_csv(dir_path+'/link_type.csv',header=None)\n",
    "    data[\"movie_companies\"] = pd.read_csv(dir_path+'/movie_companies.csv',header=None)\n",
    "    data[\"movie_info\"] = pd.read_csv(dir_path+'/movie_info.csv',header=None)\n",
    "    data[\"movie_info_idx\"] = pd.read_csv(dir_path+'/movie_info_idx.csv',header=None)\n",
    "    data[\"movie_keyword\"] = pd.read_csv(dir_path+'/movie_keyword.csv',header=None)\n",
    "    data[\"movie_link\"] = pd.read_csv(dir_path+'/movie_link.csv',header=None)\n",
    "    data[\"name\"] = pd.read_csv(dir_path+'/name.csv',header=None)\n",
    "    data[\"person_info\"] = pd.read_csv(dir_path+'/person_info.csv',header=None)\n",
    "    data[\"role_type\"] = pd.read_csv(dir_path+'/role_type.csv',header=None)\n",
    "    data[\"title\"] = pd.read_csv(dir_path+'/title.csv',header=None)\n",
    "\n",
    "    aka_name_column = {\n",
    "        'id':0,\n",
    "        'person_id':1,\n",
    "        'name':2,\n",
    "        'imdb_index':3,\n",
    "        'name_pcode_cf':4,\n",
    "        'name_pcode_nf':5,\n",
    "        'surname_pcode':6,\n",
    "        'md5sum':7\n",
    "    }\n",
    "\n",
    "    aka_title_column = {\n",
    "        'id':0,\n",
    "        'movie_id':1,\n",
    "        'title':2,\n",
    "        'imdb_index':3,\n",
    "        'kind_id':4,\n",
    "        'production_year':5,\n",
    "        'phonetic_code':6,\n",
    "        'episode_of_id':7,\n",
    "        'season_nr':8,\n",
    "        'episode_nr':9,\n",
    "        'note':10,\n",
    "        'md5sum':11\n",
    "    }\n",
    "\n",
    "    cast_info_column = {\n",
    "        'id':0,\n",
    "        'person_id':1,\n",
    "        'movie_id':2,\n",
    "        'person_role_id':3,\n",
    "        'note':4,\n",
    "        'nr_order':5,\n",
    "        'role_id':6\n",
    "    }\n",
    "\n",
    "    char_name_column = {\n",
    "        'id':0,\n",
    "        'name':1,\n",
    "        'imdb_index':2,\n",
    "        'imdb_id':3,\n",
    "        'name_pcode_nf':4,\n",
    "        'surname_pcode':5,\n",
    "        'md5sum':6\n",
    "    }\n",
    "\n",
    "    comp_cast_type_column = {\n",
    "        'id':0,\n",
    "        'kind':1\n",
    "    }\n",
    "\n",
    "    company_name_column = {\n",
    "        'id':0,\n",
    "        'name':1,\n",
    "        'country_code':2,\n",
    "        'imdb_id':3,\n",
    "        'name_pcode_nf':4,\n",
    "        'name_pcode_sf':5,\n",
    "        'md5sum':6\n",
    "    }\n",
    "\n",
    "    company_type_column = {\n",
    "        'id':0,\n",
    "        'kind':1\n",
    "    }\n",
    "\n",
    "    complete_cast_column = {\n",
    "        'id':0,\n",
    "        'movie_id':1,\n",
    "        'subject_id':2,\n",
    "        'status_id':3\n",
    "    }\n",
    "\n",
    "    info_type_column = {\n",
    "        'id':0,\n",
    "        'info':1\n",
    "    }\n",
    "\n",
    "    keyword_column = {\n",
    "        'id':0,\n",
    "        'keyword':1,\n",
    "        'phonetic_code':2\n",
    "    }\n",
    "\n",
    "    kind_type_column = {\n",
    "        'id':0,\n",
    "        'kind':1\n",
    "    }\n",
    "\n",
    "    link_type_column = {\n",
    "        'id':0,\n",
    "        'link':1\n",
    "    }\n",
    "\n",
    "    movie_companies_column = {\n",
    "        'id':0,\n",
    "        'movie_id':1,\n",
    "        'company_id':2,\n",
    "        'company_type_id':3,\n",
    "        'note':4\n",
    "    }\n",
    "\n",
    "    movie_info_idx_column = {\n",
    "        'id':0,\n",
    "        'movie_id':1,\n",
    "        'info_type_id':2,\n",
    "        'info':3,\n",
    "        'note':4\n",
    "    }\n",
    "\n",
    "    movie_keyword_column = {\n",
    "        'id':0,\n",
    "        'movie_id':1,\n",
    "        'keyword_id':2\n",
    "    }\n",
    "\n",
    "    movie_link_column = {\n",
    "        'id':0,\n",
    "        'movie_id':1,\n",
    "        'linked_movie_id':2,\n",
    "        'link_type_id':3\n",
    "    }\n",
    "\n",
    "    name_column = {\n",
    "        'id':0,\n",
    "        'name':1,\n",
    "        'imdb_index':2,\n",
    "        'imdb_id':3,\n",
    "        'gender':4,\n",
    "        'name_pcode_cf':5,\n",
    "        'name_pcode_nf':6,\n",
    "        'surname_pcode':7,\n",
    "        'md5sum':8\n",
    "    }\n",
    "\n",
    "    role_type_column = {\n",
    "        'id':0,\n",
    "        'role':1\n",
    "    }\n",
    "\n",
    "    title_column = {\n",
    "        'id':0,\n",
    "        'title':1,\n",
    "        'imdb_index':2,\n",
    "        'kind_id':3,\n",
    "        'production_year':4,\n",
    "        'imdb_id':5,\n",
    "        'phonetic_code':6,\n",
    "        'episode_of_id':7,\n",
    "        'season_nr':8,\n",
    "        'episode_nr':9,\n",
    "        'series_years':10,\n",
    "        'md5sum':11\n",
    "    }\n",
    "\n",
    "    movie_info_column = {\n",
    "        'id':0,\n",
    "        'movie_id':1,\n",
    "        'info_type_id':2,\n",
    "        'info':3,\n",
    "        'note':4\n",
    "    }\n",
    "\n",
    "    person_info_column = {\n",
    "        'id':0,\n",
    "        'person_id':1,\n",
    "        'info_type_id':2,\n",
    "        'info':3,\n",
    "        'note':4\n",
    "    }\n",
    "    data[\"aka_name\"].columns = aka_name_column\n",
    "    data[\"aka_title\"].columns = aka_title_column\n",
    "    data[\"cast_info\"].columns = cast_info_column\n",
    "    data[\"char_name\"].columns = char_name_column\n",
    "    data[\"company_name\"].columns = company_name_column\n",
    "    data[\"company_type\"].columns = company_type_column\n",
    "    data[\"comp_cast_type\"].columns = comp_cast_type_column\n",
    "    data[\"complete_cast\"].columns = complete_cast_column\n",
    "    data[\"info_type\"].columns = info_type_column\n",
    "    data[\"keyword\"].columns = keyword_column\n",
    "    data[\"kind_type\"].columns = kind_type_column\n",
    "    data[\"link_type\"].columns = link_type_column\n",
    "    data[\"movie_companies\"].columns = movie_companies_column\n",
    "    data[\"movie_info\"].columns = movie_info_column\n",
    "    data[\"movie_info_idx\"].columns = movie_info_idx_column\n",
    "    data[\"movie_keyword\"].columns = movie_keyword_column\n",
    "    data[\"movie_link\"].columns = movie_link_column\n",
    "    data[\"name\"].columns = name_column\n",
    "    data[\"person_info\"].columns = person_info_column\n",
    "    data[\"role_type\"].columns = role_type_column\n",
    "    data[\"title\"].columns = title_column\n",
    "    return data\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Prepare Samples"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def prepare_samples(data, sample_num):\n",
    "    sample = {}\n",
    "    sample['aka_name'] = data['aka_name'].sample(n=min(sample_num,len(data['aka_name'])))\n",
    "    sample['aka_title'] = data['aka_title'].sample(n=min(sample_num,len(data['aka_title'])))\n",
    "    sample['cast_info'] = data['cast_info'].sample(n=min(sample_num,len(data['cast_info'])))\n",
    "    sample['char_name'] = data['char_name'].sample(n=min(sample_num,len(data['char_name'])))\n",
    "    sample['company_name'] = data['company_name'].sample(n=min(sample_num,len(data['company_name'])))\n",
    "    sample['company_type'] = data['company_type'].sample(n=min(sample_num,len(data['company_type'])))\n",
    "    sample['comp_cast_type'] = data['comp_cast_type'].sample(n=min(sample_num,len(data['comp_cast_type'])))\n",
    "    sample['complete_cast'] = data['complete_cast'].sample(n=min(sample_num,len(data['complete_cast'])))\n",
    "    sample['info_type'] = data['info_type'].sample(n=min(sample_num,len(data['info_type'])))\n",
    "    sample['keyword'] = data['keyword'].sample(n=min(sample_num,len(data['keyword'])))\n",
    "    sample['kind_type'] = data['kind_type'].sample(n=min(sample_num,len(data['kind_type'])))\n",
    "    sample['link_type'] = data['link_type'].sample(n=min(sample_num,len(data['link_type'])))\n",
    "    sample['movie_companies'] = data['movie_companies'].sample(n=min(sample_num,len(data['movie_companies'])))\n",
    "    sample['movie_info'] = data['movie_info'].sample(n=min(sample_num,len(data['movie_info'])))\n",
    "    sample['movie_info_idx'] = data['movie_info_idx'].sample(n=min(sample_num,len(data['movie_info_idx'])))\n",
    "    sample['movie_keyword'] = data['movie_keyword'].sample(n=min(sample_num,len(data['movie_keyword'])))\n",
    "    sample['movie_link'] = data['movie_link'].sample(n=min(sample_num,len(data['movie_link'])))\n",
    "    sample['name'] = data['name'].sample(n=min(sample_num,len(data['name'])))\n",
    "    sample['person_info'] = data['person_info'].sample(n=min(sample_num,len(data['person_info'])))\n",
    "    sample['role_type'] = data['role_type'].sample(n=min(sample_num,len(data['role_type'])))\n",
    "    sample['title'] = data['title'].sample(n=min(sample_num,len(data['title'])))\n",
    "    return sample\n",
    "    \n",
    "    "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Extract Features From Predicates"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Operator(object):\n",
    "    def __init__(self, opt):\n",
    "        self.op_type = 'Bool'\n",
    "        self.operator = opt\n",
    "        \n",
    "    def __str__(self):\n",
    "        return 'Operator: ' + self.operator\n",
    "\n",
    "class Comparison(object):\n",
    "    def __init__(self, opt, left_value, right_value):\n",
    "        self.op_type = 'Compare'\n",
    "        self.operator = opt\n",
    "        self.left_value = left_value\n",
    "        self.right_value = right_value\n",
    "    \n",
    "    def __str__(self):\n",
    "        return 'Comparison: ' + self.left_value + ' ' + self.operator + ' ' + self.right_value\n",
    "\n",
    "def remove_invalid_tokens(predicate):\n",
    "    x = re.sub(r'\\(\\(([a-zA-Z_]+)\\)::text ~~ \\'(((?!::text).)*)\\'::text\\)', r\"(\\1 = '__LIKE__\\2')\", predicate)\n",
    "    x = re.sub(r'\\(\\(([a-zA-Z_]+)\\)::text !~~ \\'(((?!::text).)*)\\'::text\\)', r\"(\\1 = '__NOTLIKE__\\2')\", x)\n",
    "    x = re.sub(r'\\(\\(([a-zA-Z_]+)\\)::text <> \\'(((?!::text).)*)\\'::text\\)', r\"(\\1 = '__NOTEQUAL__\\2')\", x)\n",
    "    x = re.sub(r'\\(([a-zA-Z_]+) ~~ \\'(((?!::text).)*)\\'::text\\)', r\"(\\1 = '__LIKE__\\2')\", x)\n",
    "    x = re.sub(r'\\(([a-zA-Z_]+) !~~ \\'(((?!::text).)*)\\'::text\\)', r\"(\\1 = '__NOTLIKE__\\2')\", x)\n",
    "    x = re.sub(r'\\(([a-zA-Z_]+) <> \\'(((?!::text).)*)\\'::text\\)', r\"(\\1 = '__NOTEQUAL__\\2')\", x)\n",
    "    x = re.sub(r'(\\'[^\\']*\\')::[a-z_]+', r'\\1', x)\n",
    "    x = re.sub(r'\\(([^\\(]+)\\)::[a-z_]+', r'\\1', x)\n",
    "    x = re.sub(r'\\(([a-z_0-9A-Z\\-]+) = ANY \\(\\'(\\{.+\\})\\'\\[\\]\\)\\)', r\"(\\1 = '__ANY__\\2')\", x)\n",
    "    return x\n",
    "    \n",
    "def predicates2seq(pre_tree, alias2table, relation_name, index_name, is_join_condition):\n",
    "    current_level = -1\n",
    "    current_line = 0\n",
    "    sequence = []\n",
    "    while current_line < len(pre_tree):\n",
    "        operator_str = pre_tree[current_line]\n",
    "        level = len(re.findall(r'\\t', operator_str))\n",
    "        operator_seq = operator_str.strip('\\t').split(' ')\n",
    "        operator_type = operator_seq[1]\n",
    "        operator = operator_seq[0]\n",
    "        if level <= current_level:\n",
    "            for i in range(current_level - level + 1):\n",
    "                sequence.append(None)\n",
    "        current_level = level\n",
    "        if operator_type == 'operator':\n",
    "            sequence.append(Operator(operator))\n",
    "            current_line += 1\n",
    "        elif operator_type == 'comparison':\n",
    "            operator = operator_seq[0]\n",
    "            current_line += 1\n",
    "            operator_str = pre_tree[current_line]\n",
    "            operator_seq = operator_str.strip('\\t').split(' ')\n",
    "            left_type = operator_seq[0]\n",
    "            left_value = operator_seq[1]\n",
    "            current_line += 1\n",
    "            operator_str = pre_tree[current_line]\n",
    "            operator_seq = operator_str.strip('\\t').split(' ')\n",
    "            right_type = operator_seq[0]\n",
    "            if right_type == 'Number':\n",
    "                right_value = operator_seq[1]\n",
    "            elif right_type == 'Literal':\n",
    "                p = re.compile(\"Literal (.*) at line:\")\n",
    "                result = p.search(operator_str)\n",
    "                right_value = result.group(1)\n",
    "            elif right_type == 'Constant':\n",
    "                p = re.compile(\"Constant (.*) at line:\")\n",
    "                result = p.search(operator_str)\n",
    "                right_value = result.group(1)\n",
    "            else:\n",
    "                raise \"Unsupport Value Type: \"+right_type\n",
    "            if re.match(r'^[a-z][a-z0-9_]*\\.[a-z][a-z0-9_]*$', left_value) != None:\n",
    "                left_relation = left_value.split('.')[0]\n",
    "                left_column = left_value.split('.')[1]\n",
    "                if left_relation in alias2table:\n",
    "                    left_relation = alias2table[left_relation]\n",
    "                left_value = left_relation + '.' + left_column\n",
    "            else:\n",
    "                if relation_name == None:\n",
    "                    relation = index_name.replace(left_value+'_', '')\n",
    "                else:\n",
    "                    relation = relation_name\n",
    "                left_value = relation + '.' + left_value\n",
    "            if re.match(r'^[a-z][a-z0-9_]*\\.[a-z][a-z0-9_]*$', right_value) != None:\n",
    "                right_relation = right_value.split('.')[0]\n",
    "                right_column = right_value.split('.')[1]\n",
    "                if right_relation in alias2table:\n",
    "                    right_relation = alias2table[right_relation]\n",
    "                right_value = right_relation + '.' + right_column\n",
    "            sequence.append(Comparison(operator, left_value, right_value.strip('\\'')))\n",
    "            current_line += 1\n",
    "    return sequence\n",
    "\n",
    "def pre2seq(predicates, alias2table, relation_name, index_name, is_join_condition = False):\n",
    "    pr = remove_invalid_tokens(predicates)\n",
    "    pr = pr.replace(\"''\",\" \")\n",
    "    p = pypred.Predicate(pr)\n",
    "    try:\n",
    "        predicates = predicates2seq(p.description().strip('\\n').split('\\n'), alias2table, relation_name, index_name, is_join_condition)\n",
    "    except:\n",
    "        raise\n",
    "    return predicates"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Extract Features from Plans"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Materialize(object):\n",
    "    def __init__(self):\n",
    "        self.node_type = 'Materialize'\n",
    "    def __str__(self):\n",
    "        return 'Materialize'\n",
    "    \n",
    "class Aggregate(object):\n",
    "    def __init__(self, strategy, keys):\n",
    "        self.node_type = 'Aggregate'\n",
    "        self.strategy = strategy\n",
    "        self.group_keys = keys\n",
    "    def __str__(self):\n",
    "        return 'Aggregate ON: ' + ','.join(self.group_keys)\n",
    "\n",
    "class Sort(object):\n",
    "    def __init__(self, sort_keys):\n",
    "        self.sort_keys = sort_keys\n",
    "        self.node_type = 'Sort'\n",
    "    def __str__(self):\n",
    "        return 'Sort by: ' + ','.join(self.sort_keys)\n",
    "\n",
    "class Hash(object):\n",
    "    def __init__(self):\n",
    "        self.node_type = 'Hash'\n",
    "    def __str__(self):\n",
    "        return 'Hash'\n",
    "\n",
    "class Join(object):\n",
    "    def __init__(self, node_type, condition_seq):\n",
    "        self.node_type = node_type\n",
    "        self.condition = condition_seq\n",
    "    def __str__(self):\n",
    "        return self.node_type+' ON ' + ','.join([str(i) for i in self.condition])\n",
    "    \n",
    "class Scan(object):\n",
    "    def __init__(self, node_type, condition_seq_filter, condition_seq_index, relation_name, index_name):\n",
    "        self.node_type = node_type\n",
    "        self.condition_filter = condition_seq_filter\n",
    "        self.condition_index = condition_seq_index\n",
    "        self.relation_name = relation_name\n",
    "        self.index_name = index_name\n",
    "    def __str__(self):\n",
    "        return self.node_type+' ON ' + ','.join([str(i) for i in self.condition_filter]) + '; ' + ','.join([str(i) for i in self.condition_index])\n",
    "\n",
    "class BitmapCombine(object):\n",
    "    def __init__(self, operator):\n",
    "        self.node_type = operator\n",
    "    def __str__(self):\n",
    "        return self.node_type\n",
    "    \n",
    "class Result(object):\n",
    "    def __init__(self):\n",
    "        self.node_type = 'Result'\n",
    "    def __str__(self):\n",
    "        return 'Result'\n",
    "\n",
    "def class2json(instance):\n",
    "    if instance == None:\n",
    "        return json.dumps({})\n",
    "    else:\n",
    "        return json.dumps(todict(instance))\n",
    "    \n",
    "def todict(obj, classkey=None):\n",
    "    if isinstance(obj, dict):\n",
    "        data = {}\n",
    "        for (k, v) in obj.items():\n",
    "            data[k] = todict(v, classkey)\n",
    "        return data\n",
    "    elif hasattr(obj, \"_ast\"):\n",
    "        return todict(obj._ast())\n",
    "    elif hasattr(obj, \"__iter__\") and not isinstance(obj, str):\n",
    "        return [todict(v, classkey) for v in obj]\n",
    "    elif hasattr(obj, \"__dict__\"):\n",
    "        data = dict([(key, todict(value, classkey)) \n",
    "            for key, value in obj.__dict__.items() \n",
    "            if not callable(value) and not key.startswith('_')])\n",
    "        if classkey is not None and hasattr(obj, \"__class__\"):\n",
    "            data[classkey] = obj.__class__.__name__\n",
    "        return data\n",
    "    else:\n",
    "        return obj\n",
    "    \n",
    "def change_alias2table(column, alias2table):\n",
    "    relation_name = column.split('.')[0]\n",
    "    column_name = column.split('.')[1]\n",
    "    if relation_name in alias2table:\n",
    "        return alias2table[relation_name]+'.'+column_name\n",
    "    else:\n",
    "        return column\n",
    "    \n",
    "def extract_info_from_node(node, alias2table):\n",
    "    relation_name, index_name = None, None\n",
    "    if 'Relation Name' in node:\n",
    "        relation_name = node['Relation Name']\n",
    "    if 'Index Name' in node:\n",
    "        index_name = node['Index Name']\n",
    "    if node['Node Type'] == 'Materialize':\n",
    "        return Materialize(), None\n",
    "    elif node['Node Type'] == 'Hash':\n",
    "        return Hash(), None\n",
    "    elif node['Node Type'] == 'Sort':\n",
    "        keys = [change_alias2table(key, alias2table) for key in node['Sort Key']]\n",
    "        return Sort(keys), None\n",
    "    elif node['Node Type'] == 'BitmapAnd':\n",
    "        return BitmapCombine('BitmapAnd'), None\n",
    "    elif node['Node Type'] == 'BitmapOr':\n",
    "        return BitmapCombine('BitmapOr'), None\n",
    "    elif node['Node Type'] == 'Result':\n",
    "        return Result(), None\n",
    "    elif node['Node Type'] == 'Hash Join':\n",
    "        return Join('Hash Join', pre2seq(node[\"Hash Cond\"], alias2table, relation_name, index_name, True)), None\n",
    "    elif node['Node Type'] == 'Merge Join':\n",
    "        return Join('Merge Join', pre2seq(node[\"Merge Cond\"], alias2table, relation_name, index_name, True)), None\n",
    "    elif node['Node Type'] == 'Nested Loop':\n",
    "        if 'Join Filter' in node:\n",
    "            condition = pre2seq(node['Join Filter'], alias2table, relation_name, index_name, True)\n",
    "        else:\n",
    "            condition = []\n",
    "        return Join('Nested Loop', condition), None\n",
    "    elif node['Node Type'] == 'Aggregate':\n",
    "        if 'Group Key' in node:\n",
    "            keys = [change_alias2table(key, alias2table) for key in node['Group Key']]\n",
    "        else:\n",
    "            keys = []\n",
    "        return Aggregate(node['Strategy'], keys), None\n",
    "    elif node['Node Type'] == 'Seq Scan':\n",
    "        if 'Filter' in node:\n",
    "            condition_seq_filter = pre2seq(node['Filter'], alias2table, relation_name, index_name)\n",
    "        else:\n",
    "            condition_seq_filter = []\n",
    "        condition_seq_index, relation_name, index_name = [], node[\"Relation Name\"], None\n",
    "        return Scan('Seq Scan', condition_seq_filter, condition_seq_index, relation_name, index_name), None\n",
    "    elif node['Node Type'] == 'Bitmap Heap Scan':\n",
    "        if 'Filter' in node:\n",
    "            condition_seq_filter = pre2seq(node['Filter'], alias2table, relation_name, index_name)\n",
    "        else:\n",
    "            condition_seq_filter = []\n",
    "        condition_seq_index, relation_name, index_name = [], node[\"Relation Name\"], None\n",
    "        return Scan('Bitmap Heap Scan', condition_seq_filter, condition_seq_index, relation_name, index_name), None\n",
    "    elif node['Node Type'] == 'Index Scan':\n",
    "        if 'Filter' in node:\n",
    "            condition_seq_filter = pre2seq(node['Filter'], alias2table, relation_name, index_name)\n",
    "        else:\n",
    "            condition_seq_filter = []\n",
    "        if 'Index Cond' in node:\n",
    "            condition_seq_index = pre2seq(node['Index Cond'], alias2table, relation_name, index_name)\n",
    "        else:\n",
    "            condition_seq_index = []\n",
    "        relation_name, index_name = node[\"Relation Name\"], node['Index Name']\n",
    "        if len(condition_seq_index) == 1 and re.match(r'[a-zA-Z]+', condition_seq_index[0].right_value) != None:\n",
    "            return Scan('Index Scan', condition_seq_filter, condition_seq_index, relation_name, index_name), condition_seq_index\n",
    "        else:\n",
    "            return Scan('Index Scan', condition_seq_filter, condition_seq_index, relation_name, index_name), None\n",
    "    elif node['Node Type'] == 'Bitmap Index Scan':\n",
    "        if 'Index Cond' in node:\n",
    "            condition_seq_index = pre2seq(node['Index Cond'], alias2table, relation_name, index_name)\n",
    "        else:\n",
    "            condition_seq_index = []\n",
    "        condition_seq_filter, relation_name, index_name = [], None, node['Index Name']\n",
    "        if len(condition_seq_index) == 1 and re.match(r'[a-zA-Z]+', condition_seq_index[0].right_value) != None:\n",
    "            return Scan('Bitmap Index Scan', condition_seq_filter, condition_seq_index, relation_name, index_name), condition_seq_index\n",
    "        else:\n",
    "            return Scan('Bitmap Index Scan', condition_seq_filter, condition_seq_index, relation_name, index_name), None\n",
    "    elif node['Node Type'] == 'Index Only Scan':\n",
    "        if 'Index Cond' in node:\n",
    "            condition_seq_index = pre2seq(node['Index Cond'], alias2table, relation_name, index_name)\n",
    "        else:\n",
    "            condition_seq_index = []\n",
    "        condition_seq_filter, relation_name, index_name = [], None, node['Index Name']\n",
    "        if len(condition_seq_index) == 1 and re.match(r'[a-zA-Z]+', condition_seq_index[0].right_value) != None:\n",
    "            return Scan('Index Only Scan', condition_seq_filter, condition_seq_index, relation_name, index_name), condition_seq_index\n",
    "        else:\n",
    "            return Scan('Index Only Scan', condition_seq_filter, condition_seq_index, relation_name, index_name), None\n",
    "    else:\n",
    "        raise Exception('Unsupported Node Type: '+node['Node Type'])\n",
    "        return None, None\n",
    "\n",
    "def plan2seq(root, alias2table):\n",
    "    sequence = []\n",
    "    join_conditions = []\n",
    "    node, join_condition = extract_info_from_node(root, alias2table)\n",
    "    if join_condition != None:\n",
    "        join_conditions += join_condition\n",
    "    sequence.append(node)\n",
    "    if 'Plans' in root:\n",
    "        for plan in root['Plans']:\n",
    "            next_sequence, next_join_conditions = plan2seq(plan, alias2table)\n",
    "            sequence += next_sequence\n",
    "            join_conditions += next_join_conditions\n",
    "    sequence.append(None)\n",
    "    return sequence, join_conditions"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Get Feature Except for Sample Bitmap"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_subplan(root):\n",
    "    results = []\n",
    "    if root.has_key('Actual Rows') and root.has_key('Actual Total Time') and root['Actual Rows'] > 0:\n",
    "        results.append((root, root['Actual Total Time'], root['Actual Rows']))\n",
    "    if root.has_key('Plans'):\n",
    "        for plan in root['Plans']:\n",
    "            results += get_subplan(plan)\n",
    "    return results\n",
    "\n",
    "def get_plan(root):\n",
    "    return (root, root['Actual Total Time'], root['Actual Rows'])\n",
    "\n",
    "class PlanInSeq(object):\n",
    "    def __init__(self, seq, cost, cardinality):\n",
    "        self.seq = seq\n",
    "        self.cost = cost\n",
    "        self.cardinality = cardinality\n",
    "        \n",
    "def get_alias2table(root, alias2table):\n",
    "    if root.has_key('Relation Name') and root.has_key('Alias'):\n",
    "        alias2table[root['Alias']] = root['Relation Name']\n",
    "    if root.has_key('Plans'):\n",
    "        for child in root['Plans']:\n",
    "            get_alias2table(child, alias2table)\n",
    "def feature_extractor(input_path, out_path):\n",
    "    with open(out_path, 'w') as out:\n",
    "        with open(input_path, 'r') as f:\n",
    "            for index, plan in enumerate(f.readlines()):\n",
    "                print (index)\n",
    "                if plan != 'null\\n':\n",
    "                    plan = json.loads(plan)['Plan']\n",
    "                    alias2table = {}\n",
    "                    get_alias2table(plan, alias2table)\n",
    "                    subplan, cost, cardinality = get_plan(plan)\n",
    "                    seq, _ = plan2seq(subplan, alias2table)\n",
    "                    seqs = PlanInSeq(seq, cost, cardinality)\n",
    "                    out.write(class2json(seqs)+'\\n')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Add Sample Bitmap"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class TreeNode(object):\n",
    "    def __init__(self, current_vec, parent):\n",
    "        self.item = current_vec\n",
    "        self.parent = parent\n",
    "        self.children = []\n",
    "    def get_parent(self):\n",
    "        return self.parent\n",
    "    def get_item(self):\n",
    "        return self.item\n",
    "    def get_children(self):\n",
    "        return self.children\n",
    "    def add_child(self, child):\n",
    "        self.children.append(child)\n",
    "        \n",
    "def recover_tree(vecs, parent):\n",
    "    if len(vecs) == 0:\n",
    "        return vecs\n",
    "    if vecs[0] == None:\n",
    "        return vecs[1:]\n",
    "    node = TreeNode(vecs[0], parent)\n",
    "    while True:\n",
    "        vecs = recover_tree(vecs[1:], node)\n",
    "        parent.add_child(node)\n",
    "        if len(vecs) == 0:\n",
    "            return vecs\n",
    "        if vecs[0] == None:\n",
    "            return vecs[1:]\n",
    "        node = TreeNode(vecs[0], parent)\n",
    "        \n",
    "def bitand(bit1, bit2):\n",
    "    if len(bit1) > 0 and len(bit2) > 0:\n",
    "        result = []\n",
    "        for i in range(len(bit1)):\n",
    "            result.append(min(bit1[i], bit2[i]))\n",
    "        return result\n",
    "    elif len(bit1) > 0:\n",
    "        return bit1\n",
    "    elif len(bit2) > 0:\n",
    "        return bit2\n",
    "    else:\n",
    "        return []\n",
    "    \n",
    "def bitor(bit1, bit2):\n",
    "    if len(bit1) > 0 and len(bit2) > 0:\n",
    "        result = []\n",
    "        for i in range(len(bit1)):\n",
    "            result.append(max(bit1[i], bit2[i]))\n",
    "        return result\n",
    "    elif len(bit1) > 0:\n",
    "        return bit1\n",
    "    elif len(bit2) > 0:\n",
    "        return bit2\n",
    "    else:\n",
    "        return []\n",
    "\n",
    "def get_bitmap(root):\n",
    "    predicate = root.get_item()\n",
    "    if predicate != None and predicate['op_type'] == 'Compare':\n",
    "        table_name = predicate['left_value'].split('.')[0]\n",
    "        column = predicate['left_value'].split('.')[1]\n",
    "        vec = []\n",
    "        pattern = re.compile(r'^[a-z_]+\\.[a-z][a-z0-9_]*$')\n",
    "        result = pattern.match(predicate['right_value'])\n",
    "        if result == None:\n",
    "            dtype = data[table_name].dtypes[column]\n",
    "            for value in sample[table_name][column].tolist():\n",
    "                if isSelected(value, predicate, dtype):\n",
    "                    vec.append(1)\n",
    "                else:\n",
    "                    vec.append(0)\n",
    "            for i in range(len(vec), sample_num):\n",
    "                vec.append(0)\n",
    "        elif not predicate['right_value'].split('.')[0] in data:\n",
    "            dtype = data[table_name].dtypes[column]\n",
    "            for value in sample[table_name][column].tolist():\n",
    "                if isSelected(value, predicate, dtype):\n",
    "                    vec.append(1)\n",
    "                else:\n",
    "                    vec.append(0)\n",
    "            for i in range(len(vec), sample_num):\n",
    "                vec.append(0)\n",
    "        return vec\n",
    "    elif predicate != None and predicate['op_type'] == 'Bool':\n",
    "        bitmap = []\n",
    "        if predicate['operator'] == 'AND':\n",
    "            for child in root.get_children():\n",
    "                vec = get_bitmap(child)\n",
    "                bitmap = bitand(bitmap, vec)\n",
    "        elif predicate['operator'] == 'OR':\n",
    "            for child in root.get_children():\n",
    "                vec = get_bitmap(child)\n",
    "                bitmap = bitor(bitmap, vec)\n",
    "        else:\n",
    "            print (predicate['operator'])\n",
    "            raise\n",
    "        return bitmap\n",
    "    else:\n",
    "        return []\n",
    "            \n",
    "\n",
    "# Second, we should encode each training instance one by one except for the label\n",
    "def isSelected(row_value, predicate, dtype):\n",
    "    if dtype == 'int64':\n",
    "        row_value = int(row_value)\n",
    "        value = int(predicate['right_value'])\n",
    "        op = predicate['operator']\n",
    "        if op == '=':\n",
    "            if row_value != value:\n",
    "                return False\n",
    "        elif op == '!=':\n",
    "            if row_value == value:\n",
    "                return False\n",
    "        elif op == '<':\n",
    "            if row_value >= value:\n",
    "                return False\n",
    "        elif op == '>':\n",
    "            if row_value <= value:\n",
    "                return False\n",
    "        elif op == '<=':\n",
    "            if row_value > value:\n",
    "                return False\n",
    "        elif op == '>=':\n",
    "            if row_value < value:\n",
    "                return False\n",
    "        else:\n",
    "            print (op)\n",
    "            raise\n",
    "    elif dtype == 'float64':\n",
    "        row_value = float(row_value)\n",
    "        value = float(predicate['right_value'])\n",
    "        op = predicate['operator']\n",
    "        if op == '=':\n",
    "            if row_value != value:\n",
    "                return False\n",
    "        elif op == '!=':\n",
    "            if row_value == value:\n",
    "                return False\n",
    "        elif op == '<':\n",
    "            if row_value >= value:\n",
    "                return False\n",
    "        elif op == '>':\n",
    "            if row_value <= value:\n",
    "                return False\n",
    "        elif op == '<=':\n",
    "            if row_value > value:\n",
    "                return False\n",
    "        elif op == '>=':\n",
    "            if row_value < value:\n",
    "                return False\n",
    "        else:\n",
    "            print (op)\n",
    "            raise\n",
    "    elif dtype == 'object':\n",
    "        value = predicate['right_value']\n",
    "        op = predicate['operator']\n",
    "        if pd.isnull(row_value):\n",
    "            row_value = ''\n",
    "        else:\n",
    "            row_value = str(row_value)\n",
    "        if op == '=':\n",
    "            if value.startswith('__LIKE__'):\n",
    "                v = value[8:]\n",
    "                pattern = r'^'\n",
    "                for idx, token in enumerate(v.split('%')):\n",
    "                    if len(token) == 0:\n",
    "                        pattern += r'.*'\n",
    "                    else:\n",
    "                        pattern += re.escape(token)\n",
    "                        if idx < len(v.split('%')) - 1:\n",
    "                            pattern += r'.*'\n",
    "                pattern += r'$'\n",
    "                if re.match(pattern, row_value) == None:\n",
    "                    return False\n",
    "            elif value.startswith('__NOTLIKE__'):\n",
    "                v = value[11:]\n",
    "                pattern = r'^'\n",
    "                for idx, token in enumerate(v.split('%')):\n",
    "                    if len(token) == 0:\n",
    "                        pattern += r'.*'\n",
    "                    else:\n",
    "                        pattern += re.escape(token)\n",
    "                        if idx < len(v.split('%')) - 1:\n",
    "                            pattern += r'.*'\n",
    "                pattern += r'$'\n",
    "                if re.match(pattern, row_value) != None:\n",
    "                    return False\n",
    "            elif value.startswith('__NOTEQUAL__'):\n",
    "                pattern = value[12:]\n",
    "                if row_value == pattern:\n",
    "                    return False\n",
    "            elif value.startswith('__ANY__'):\n",
    "                pattern = value.strip('__ANY__')\n",
    "                pattern = pattern.strip('{}')\n",
    "                for token in pattern.split(','):\n",
    "                    token = token.strip('\"').strip('\\'')\n",
    "                    if row_value == token:\n",
    "                        return True\n",
    "                return False\n",
    "            elif value == 'None':\n",
    "                if len(row_value) > 0:\n",
    "                    return False\n",
    "            else:\n",
    "                if row_value != value:\n",
    "                    return False\n",
    "        elif op == 'IS':\n",
    "            if value == 'None':\n",
    "                if len(row_value) > 0:\n",
    "                    return False\n",
    "            else:\n",
    "                print (value)\n",
    "                raise\n",
    "        elif op == '!=':\n",
    "            if value == 'None':\n",
    "                if len(row_value) == 0:\n",
    "                    return False\n",
    "            else:\n",
    "                if row_value == value:\n",
    "                    return False\n",
    "        elif op == '<':\n",
    "            if row_value >= value:\n",
    "                return False\n",
    "        elif op == '>':\n",
    "            if row_value <= value:\n",
    "                return False\n",
    "        elif op == '<=':\n",
    "            if row_value > value:\n",
    "                return False\n",
    "        elif op == '>=':\n",
    "            if row_value < value:\n",
    "                return False\n",
    "        else:\n",
    "            print (op)\n",
    "            raise\n",
    "    else:\n",
    "        print (dtype)\n",
    "        raise\n",
    "    return True\n",
    "\n",
    "def chunks(l, n):\n",
    "    \"\"\"Yield successive n-sized chunks from l.\"\"\"\n",
    "    for i in range(0, len(l), n):\n",
    "        yield l[i:i + n]\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def add_sample_bitmap(input_path, output_path):\n",
    "    with open(input_path, 'r') as ff:\n",
    "        with open(output_path, 'w') as f:\n",
    "            for count, plan in enumerate(ff.readlines()):\n",
    "                print (count)\n",
    "                parsed_plan = json.loads(plan)\n",
    "                nodes_with_sample = []\n",
    "                for node in parsed_plan['seq']:\n",
    "                    bitmap_filter = []\n",
    "                    bitmap_index = []\n",
    "                    bitmap_other = []\n",
    "                    if node != None and node.has_key('condition'):\n",
    "                        predicates = node['condition']\n",
    "                        if len(predicates) > 0:\n",
    "                            root = TreeNode(predicates[0], None)\n",
    "                            if len(predicates) > 1:\n",
    "                                recover_tree(predicates[1:], root)\n",
    "                            bitmap_other = get_bitmap(root)\n",
    "                    if node != None and node.has_key('condition_filter'):\n",
    "                        predicates = node['condition_filter']\n",
    "                        if len(predicates) > 0:\n",
    "                            root = TreeNode(predicates[0], None)\n",
    "                            if len(predicates) > 1:\n",
    "                                recover_tree(predicates[1:], root)\n",
    "                            bitmap_filter = get_bitmap(root)\n",
    "                    if node != None and node.has_key('condition_index'):\n",
    "                        predicates = node['condition_index']\n",
    "                        if len(predicates) > 0:\n",
    "                            root = TreeNode(predicates[0], None)\n",
    "                            if len(predicates) > 1:\n",
    "                                recover_tree(predicates[1:], root)\n",
    "                            bitmap_index = get_bitmap(root)\n",
    "                    if len(bitmap_filter) > 0 or len(bitmap_index) > 0 or len(bitmap_other) > 0:\n",
    "                        bitmap = [1 for _ in range(sample_num)]\n",
    "                        bitmap = bitand(bitmap, bitmap_filter)\n",
    "                        bitmap = bitand(bitmap, bitmap_index)\n",
    "                        bitmap = bitand(bitmap, bitmap_other)\n",
    "                        node['bitmap'] = ''.join([str(x) for x in bitmap])\n",
    "                    nodes_with_sample.append(node)\n",
    "                parsed_plan['seq'] = nodes_with_sample\n",
    "                f.write(json.dumps(parsed_plan))\n",
    "                f.write('\\n')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Main"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sample_num = 1000\n",
    "data = load_dataset('dir_path')\n",
    "sample = prepare_samples(data, sample_num)\n",
    "feature_extractor('input_path', 'out_path')\n",
    "add_sample_bitmap('input_path', 'output_path')"
   ]
  }
 ],
 "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.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}