_dataset_table_querying_construction.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pandas as pd\n",
    "from tqdm import tqdm\n",
    "import pickle"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_table_set(df: pd.DataFrame) -> set:\n",
    "    out = []\n",
    "    for r in tqdm(range(df.shape[0])):\n",
    "        out.append(df.iloc[r]['r_id'])\n",
    "        out.append(df.iloc[r]['s_id'])\n",
    "    return set(out)\n",
    "\n",
    "def check_no_repetitions(s1: set, s2: set) -> bool:\n",
    "    for t in tqdm(s1):\n",
    "        if t in s2:\n",
    "            return True\n",
    "    return False"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "df = pd.read_csv('/home/francesco.pugnaloni/GNNTE/Datasets/1_Gittables/baseline_performances/test_set_similarities_gittables_with_armadillo_predictions.csv')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 100000/100000 [00:09<00:00, 10455.33it/s]\n"
     ]
    }
   ],
   "source": [
    "table_count = {}\n",
    "for r in tqdm(range(df.shape[0])):\n",
    "    r_name = df.iloc[r]['r_id']\n",
    "    s_name = df.iloc[r]['s_id']\n",
    "    try:\n",
    "        table_count[r_name].append(s_name)\n",
    "    except:\n",
    "        table_count[r_name] = []\n",
    "        table_count[r_name].append(s_name)\n",
    "    \n",
    "    try:\n",
    "        table_count[s_name].append(r_name)\n",
    "    except:\n",
    "        table_count[s_name] = []\n",
    "        table_count[s_name].append(r_name)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 47607/47607 [00:00<00:00, 1867163.79it/s]\n"
     ]
    }
   ],
   "source": [
    "at_least_10 = []\n",
    "for k in tqdm(table_count.keys()):\n",
    "    count = len(table_count[k])\n",
    "    if count >= 10:\n",
    "        at_least_10.append(k)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "5550"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "len(at_least_10)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 5550/5550 [00:00<00:00, 298846.99it/s]\n"
     ]
    }
   ],
   "source": [
    "at_least_10 = set(at_least_10)\n",
    "new_table_count = {}\n",
    "for k in tqdm(at_least_10):\n",
    "    new_table_count[k] = [i for i in table_count[k] if i not in at_least_10]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "5550"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "len(new_table_count)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "query_tables = list(new_table_count.keys())[:1000]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "1000"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "len(query_tables)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "target_tables = []\n",
    "for k in query_tables:\n",
    "    target_tables+=new_table_count[k]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "3220"
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "len(set(target_tables))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "q_set = set(query_tables)\n",
    "t_set = set(target_tables)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "  0%|          | 0/100000 [00:00<?, ?it/s]"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 100000/100000 [00:09<00:00, 10945.23it/s]\n"
     ]
    }
   ],
   "source": [
    "all_tables = get_table_set(df)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [],
   "source": [
    "target = 10_000 - len(t_set)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "6780"
      ]
     },
     "execution_count": 16,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "target"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " 16%|█▌        | 7543/47607 [00:00<00:00, 1067144.57it/s]\n"
     ]
    }
   ],
   "source": [
    "count = 0\n",
    "new_tabs = []\n",
    "for t in tqdm(all_tables):\n",
    "    if (t not in q_set) and (t not in t_set):\n",
    "        new_tabs.append(t)\n",
    "        count+=1\n",
    "        if count == target:\n",
    "            break"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "6780"
      ]
     },
     "execution_count": 18,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "len(new_tabs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [],
   "source": [
    "t_set = list(t_set)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [],
   "source": [
    "t_set += new_tabs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [],
   "source": [
    "t_set = set(t_set)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "10000"
      ]
     },
     "execution_count": 22,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "len(t_set)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "1000"
      ]
     },
     "execution_count": 23,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "len(q_set)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 10000/10000 [00:00<00:00, 1816659.74it/s]\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "False"
      ]
     },
     "execution_count": 24,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "check_no_repetitions(t_set, q_set)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {},
   "outputs": [],
   "source": [
    "with open('/home/francesco.pugnaloni/GNNTE/Datasets/1_Gittables/table_querying/data_lake_10k.pkl', 'wb') as f:\n",
    "    pickle.dump(t_set, f)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {},
   "outputs": [],
   "source": [
    "with open('/home/francesco.pugnaloni/GNNTE/Datasets/1_Gittables/table_querying/query_set_1k.pkl', 'wb') as f:\n",
    "    pickle.dump(q_set, f)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "10000"
      ]
     },
     "execution_count": 28,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "with open('/home/francesco.pugnaloni/GNNTE/Datasets/1_Gittables/table_querying/data_lake_10k.pkl', 'rb') as f:\n",
    "    a = pickle.load(f)\n",
    "len(a)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "1000"
      ]
     },
     "execution_count": 29,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "with open('/home/francesco.pugnaloni/GNNTE/Datasets/1_Gittables/table_querying/query_set_1k.pkl', 'rb') as f:\n",
    "    b = pickle.load(f)\n",
    "len(b)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "GNNTE",
   "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.11.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}