DataRoutingMultiProcessing.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import findspark\n",
    "findspark.init() # this must be executed before the below import"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "from pyspark import SparkContext, SparkConf\n",
    "from pyspark.sql import SQLContext\n",
    "from pyspark.sql import SparkSession\n",
    "from pyspark import SparkFiles"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pandas as pd\n",
    "import time\n",
    "import rtree\n",
    "from rtree import index\n",
    "import numpy as np\n",
    "from numpy import genfromtxt\n",
    "from multiprocessing import Pool\n",
    "import threading"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "from DRProcess import *\n",
    "from DDProcess import *"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "conf = SparkConf().setAll([(\"spark.executor.memory\", \"8g\"),(\"spark.driver.memory\",\"8g\"),\n",
    "                           (\"spark.memory.offHeap.enabled\",True),(\"spark.memory.offHeap.size\",\"8g\")])\n",
    "\n",
    "sc = SparkContext(conf=conf)\n",
    "sqlContext = SQLContext(sc)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[('spark.memory.offHeap.size', '8g'),\n",
       " ('spark.executor.id', 'driver'),\n",
       " ('spark.app.name', 'pyspark-shell'),\n",
       " ('spark.rdd.compress', 'True'),\n",
       " ('spark.driver.memory', '8g'),\n",
       " ('spark.serializer.objectStreamReset', '100'),\n",
       " ('spark.master', 'local[*]'),\n",
       " ('spark.executor.memory', '8g'),\n",
       " ('spark.submit.pyFiles', ''),\n",
       " ('spark.submit.deployMode', 'client'),\n",
       " ('spark.driver.host', 'namenode.novalocal'),\n",
       " ('spark.memory.offHeap.enabled', 'True'),\n",
       " ('spark.driver.port', '44996'),\n",
       " ('spark.app.id', 'local-1603522830599'),\n",
       " ('spark.ui.showConsoleProgress', 'true')]"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "sc.getConf().getAll()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "class DumpThread(threading.Thread):\n",
    "    def __init__(self, thread_id, name, parameters):\n",
    "        threading.Thread.__init__(self)\n",
    "        self.thread_id = thread_id\n",
    "        self.name = name\n",
    "        self.parameters = parameters\n",
    "        \n",
    "    def run(self):\n",
    "        print('start dumping thread: ', self.thread_id, self.name)\n",
    "        start_index, end_index, pids, pid_data_dict, hdfs_path, column_names = self.parameters\n",
    "        for pid in pids[start_index: end_index]:\n",
    "            path = hdfs_path + 'partition_' + str(pid)+'.parquet'\n",
    "            pdf = pd.DataFrame(pid_data_dict[pid], columns=column_names)\n",
    "            df = sqlContext.createDataFrame(pdf)\n",
    "            df.write.mode('append').parquet(path)\n",
    "            pid_data_dict[pid] = []\n",
    "        print('exit dumping thread: ', self.thread_id, self.name)\n",
    "        \n",
    "def dump_dict_data_2_hdfs(pid_data_dicts, column_names, hdfs_path, num_threads = 8):\n",
    "    \n",
    "    # first merge all the dicts\n",
    "    base_dict = pid_data_dicts[0]\n",
    "    for k in range(1, len(pid_data_dicts)):\n",
    "        for key, val in pid_data_dicts[k].items():\n",
    "            if key in base_dict:\n",
    "                base_dict[key] += val\n",
    "            else:\n",
    "                base_dict.update({key:val})\n",
    "        pid_data_dicts[k].clear()\n",
    "    \n",
    "    if num_threads == 1:\n",
    "        print('start dumping single thread (main)')\n",
    "        pids = list(base_dict.keys())\n",
    "        for pid in pids:\n",
    "            path = hdfs_path + 'partition_' + str(pid)+'.parquet'\n",
    "            pdf = pd.DataFrame(base_dict[pid], columns=column_names)\n",
    "            df = sqlContext.createDataFrame(pdf)\n",
    "            df.write.mode('append').parquet(path)\n",
    "            base_dict[pid] = []\n",
    "        print('finish dumping single thread (main)')\n",
    "    \n",
    "    else:\n",
    "        # apply multi-threading to save\n",
    "        pids = list(base_dict.keys())\n",
    "        step = int(len(pids) / num_threads) + 1\n",
    "        threads = []\n",
    "        for i in range(num_threads):\n",
    "            start_index = i * step\n",
    "            end_index = (i+1) * step\n",
    "            parameters = [start_index, end_index, pids, base_dict, hdfs_path, column_names]\n",
    "            thread = DumpThread(i, 'dump_thread_'+str(i), parameters)\n",
    "            thread.start()\n",
    "            threads.append(thread)\n",
    "            if start_index >= len(pids):\n",
    "                break   \n",
    "        for t in threads:\n",
    "            t.join()\n",
    "\n",
    "# used for multi-process wirting\n",
    "def merge_dicts(pid_data_dicts, num_process):\n",
    "    base_dict = pid_data_dicts[0]\n",
    "    for k in range(1, len(pid_data_dicts)):\n",
    "        for key, val in pid_data_dicts[k].items():\n",
    "            if key in base_dict:\n",
    "                base_dict[key] += val\n",
    "            else:\n",
    "                base_dict.update({key:val})\n",
    "        pid_data_dicts[k].clear()\n",
    "    \n",
    "    # re allocate to non-overlap dicts\n",
    "    pids = list(base_dict.keys())\n",
    "    step = int(len(pids) / num_process) + 1\n",
    "    non_overlap_dicts = [{} for i in range(num_process)]\n",
    "    \n",
    "    for key, val in base_dict.items():\n",
    "        dict_index = key // step\n",
    "        non_overlap_dicts[dict_index][key] = val\n",
    "        \n",
    "    return non_overlap_dicts"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "def batch_data_parallel(table_path, partition_path, chunk_size, used_dims, hdfs_path, \n",
    "                        num_dims, dump_threshold = 1000000, num_process = 8):\n",
    "    \n",
    "    begin_time = time.time()\n",
    "    \n",
    "    col_names = ['_c'+str(i) for i in range(num_dims)]\n",
    "    cols = [i for i in range(num_dims)]\n",
    "    \n",
    "    pid_data_dicts = []\n",
    "    for i in range(num_process):\n",
    "        pid_data_dicts.append({})\n",
    "    \n",
    "    chunks = []\n",
    "    \n",
    "    count = 0\n",
    "    epochs = 0\n",
    "    processed_data = 0\n",
    "    pool = Pool(processes = num_process) # the pool should be reused, or incur memory leak!\n",
    "    pids_each_process = [set() for k in range(num_process)] # used for final merge\n",
    "    \n",
    "    for chunk in pd.read_table(table_path, delimiter='|', usecols=cols, names=col_names, chunksize=chunk_size):\n",
    "    #for chunk in pd.read_csv(table_path, usecols=cols, names=col_names, chunksize=chunk_size):\n",
    "        print('current chunk: ', count)\n",
    "        chunks.append(chunk)\n",
    "        if count % num_process == num_process - 1:\n",
    "            paras = [[chunks[k], used_dims, partition_path, pid_data_dicts[k]] for k in range(num_process)]\n",
    "            pid_data_dicts = pool.map(process_chunk, [para for para in paras])\n",
    "            print('===================')\n",
    "            chunks = []\n",
    "            processed_data += chunk_size * num_process\n",
    "            \n",
    "            # dump data to file\n",
    "            if processed_data >= dump_threshold:\n",
    "                print(\"= = = dumping epoch:\", epochs, \"= = =\")\n",
    "                # parquet write is not thread safe, avoid concurent write\n",
    "                pid_data_dicts = merge_dicts(pid_data_dicts, num_process) # make it non-overlap\n",
    "                for k in range(num_process):\n",
    "                    pids_each_process[k].update(list(pid_data_dicts[k].keys()))\n",
    "                paras = [[epochs, pid_data_dicts[k], col_names, hdfs_path] for k in range(num_process)]\n",
    "                pool.map(dump_data, [para for para in paras])\n",
    "                #dump_dict_data_2_hdfs(pid_data_dicts, col_names, hdfs_path) # multi-thread\n",
    "                processed_data = 0\n",
    "                epochs += 1\n",
    "                for i in range(num_process):\n",
    "                    pid_data_dicts[i].clear()\n",
    "        count += 1\n",
    "        \n",
    "    dict_size = [len(pid_data_dicts[i]) for i in range(num_process)]\n",
    "    print('after exit, chunks size: ', len(chunks))\n",
    "    print('after exit, each dict size: ', dict_size)\n",
    "    # process the last batch\n",
    "    if len(chunks) != 0:\n",
    "        paras = [[chunks[k], used_dims, partition_path, pid_data_dicts[k]] for k in range(len(chunks))]\n",
    "        pid_data_dicts[0:len(chunks)] = pool.map(process_chunk, [para for para in paras])\n",
    "    \n",
    "    \n",
    "    dict_size = [len(pid_data_dicts[i]) for i in range(num_process)]\n",
    "    print('after last chunk, each dict size: ', dict_size)\n",
    "    \n",
    "    if len(pid_data_dicts[0]) != 0:\n",
    "        pid_data_dicts = merge_dicts(pid_data_dicts, num_process) # make it non-overlap\n",
    "        paras = [[epochs, pid_data_dicts[k], col_names, hdfs_path] for k in range(num_process)]\n",
    "        pool.map(dump_data, [para for para in paras])\n",
    "        #dump_dict_data_2_hdfs(pid_data_dicts, col_names, hdfs_path)\n",
    "        for k in range(num_process):\n",
    "            pids_each_process[k].update(list(pid_data_dicts[k].keys()))\n",
    "    \n",
    "    pid_data_dicts.clear() # release memory\n",
    "    \n",
    "    # final merge\n",
    "    epochs += 1\n",
    "    paras = [[epochs, pids_each_process[k], hdfs_path] for k in range(num_process)]\n",
    "    pool.map(merge_parquets, [para for para in paras])\n",
    "        \n",
    "    pool.close()\n",
    "    pool.join()\n",
    "    \n",
    "    finish_time = time.time()\n",
    "    print('total data routing and persisting time: ', finish_time - begin_time)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "# # = = = Configuration (COMP Cloud Ubuntu) = = =\n",
    "# scale_factor = 100\n",
    "\n",
    "# table_base_path = '/home/ubuntu/TPCH/dbgen/'\n",
    "# table_path = table_base_path + 'lineitem_' + str(scale_factor) + '.tbl'\n",
    "\n",
    "# num_process = 12\n",
    "# chunk_size = 200000\n",
    "# dump_threshold = 12000000 # 6M rows = about 1GB raw data\n",
    "\n",
    "# num_dims = 16\n",
    "# used_dims = [1,2]\n",
    "\n",
    "# # base path of HDFS\n",
    "# hdfs_base_path = 'hdfs://10.88.88.103:9000/user/cloudray/'\n",
    "\n",
    "# nora_hdfs = hdfs_base_path + 'NORA/scale' + str(scale_factor) + '/'\n",
    "# qdtree_hdfs = hdfs_base_path + 'QdTree/scale' + str(scale_factor) + '/'\n",
    "# kdtree_hdfs = hdfs_base_path + 'KDTree/scale' + str(scale_factor) + '/'\n",
    "\n",
    "# # base path of Partition\n",
    "# partition_base_path = '/home/ubuntu/PartitionLayout/'\n",
    "\n",
    "# nora_partition = partition_base_path + 'nora_partitions_' + str(scale_factor)\n",
    "# qdtree_partition = partition_base_path + 'qdtree_partitions_' + str(scale_factor)\n",
    "# kdtree_partition = partition_base_path + 'kdt_partitions_' + str(scale_factor)\n",
    "\n",
    "# # Legacy\n",
    "# # table_path = '/home/cloudray/Downloads/TPCH_12M_8Field.csv'\n",
    "# # table_path = '/home/cloudray/TPCH/2.18.0_rc2/dbgen/lineitem.tbl'\n",
    "\n",
    "# # partition_path = '/home/cloudray/NORA_Partitions/nora_partitions'\n",
    "# # partition_path = '/home/cloudray/NORA_Partitions/qd_tree_partitions'\n",
    "\n",
    "# # hdfs_path = 'hdfs://localhost:9000/user/cloudray/NORA/'\n",
    "# # hdfs_path = 'hdfs://localhost:9000/user/cloudray/QdTree/'\n",
    "\n",
    "# # partition_path = '/home/cloudray/NORA_Partitions/nora_test'\n",
    "# # partition_path = '/home/cloudray/NORA_Partitions/qd_tree_test'\n",
    "\n",
    "# # hdfs_path = 'hdfs://localhost:9000/user/cloudray/NORA_Test/'\n",
    "# # hdfs_path = 'hdfs://localhost:9000/user/cloudray/QdTree_Test/'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "# = = = Configuration (UBDA Cloud Centos) = = =\n",
    "scale_factor = 100\n",
    "\n",
    "table_base_path = '/media/datadrive1/TPCH/dbgen/'\n",
    "table_path = table_base_path + 'lineitem_' + str(scale_factor) + '.tbl'\n",
    "\n",
    "num_process = 8\n",
    "chunk_size = 2000000\n",
    "dump_threshold = 16000000 # 6M rows = about 1GB raw data, each dump includes 4G data\n",
    "\n",
    "num_dims = 16\n",
    "used_dims = [1,2]\n",
    "\n",
    "# base path of HDFS\n",
    "hdfs_base_path = 'hdfs://192.168.6.62:9000/user/cloudray/'\n",
    "\n",
    "nora_hdfs = hdfs_base_path + 'NORA/scale' + str(scale_factor) + '/'\n",
    "qdtree_hdfs = hdfs_base_path + 'QdTree/scale' + str(scale_factor) + '/'\n",
    "kdtree_hdfs = hdfs_base_path + 'KDTree/scale' + str(scale_factor) + '/'\n",
    "\n",
    "# base path of Partition\n",
    "partition_base_path = '/home/centos/PartitionLayout/'\n",
    "\n",
    "nora_partition = partition_base_path + 'nora_partitions_' + str(scale_factor)\n",
    "qdtree_partition = partition_base_path + 'qdtree_partitions_' + str(scale_factor)\n",
    "kdtree_partition = partition_base_path + 'kdt_partitions_' + str(scale_factor)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "current chunk:  0\n",
      "current chunk:  1\n",
      "current chunk:  2\n",
      "current chunk:  3\n",
      "current chunk:  4\n",
      "current chunk:  5\n",
      "current chunk:  6\n",
      "current chunk:  7\n",
      "===================\n",
      "= = = dumping epoch: 0 = = =\n",
      "exit dumping process\n",
      "exit dumping process\n",
      "exit dumping process\n",
      "exit dumping process\n",
      "exit dumping process\n",
      "exit dumping process\n",
      "exit dumping process\n",
      "exit dumping process\n",
      "current chunk:  8\n",
      "current chunk:  9\n",
      "current chunk:  10\n",
      "current chunk:  11\n",
      "current chunk:  12\n",
      "current chunk:  13\n",
      "current chunk:  14\n",
      "current chunk:  15\n",
      "===================\n",
      "= = = dumping epoch: 1 = = =\n",
      "exit dumping process\n",
      "exit dumping process\n",
      "exit dumping process\n",
      "exit dumping process\n",
      "exit dumping process\n",
      "exit dumping process\n",
      "exit dumping process\n",
      "exit dumping process\n",
      "current chunk:  16\n",
      "current chunk:  17\n",
      "current chunk:  18\n",
      "current chunk:  19\n",
      "current chunk:  20\n",
      "current chunk:  21\n",
      "current chunk:  22\n",
      "current chunk:  23\n",
      "===================\n",
      "= = = dumping epoch: 2 = = =\n",
      "exit dumping process\n",
      "exit dumping process\n",
      "exit dumping process\n",
      "exit dumping process\n",
      "exit dumping process\n",
      "exit dumping process\n",
      "exit dumping process\n",
      "exit dumping process\n",
      "current chunk:  24\n",
      "current chunk:  25\n",
      "current chunk:  26\n",
      "current chunk:  27\n",
      "current chunk:  28\n",
      "current chunk:  29\n",
      "current chunk:  30\n",
      "current chunk:  31\n",
      "===================\n",
      "= = = dumping epoch: 3 = = =\n",
      "exit dumping process\n",
      "exit dumping process\n",
      "exit dumping process\n",
      "exit dumping process\n",
      "exit dumping process\n",
      "exit dumping process\n",
      "exit dumping process\n",
      "exit dumping process\n",
      "current chunk:  32\n",
      "current chunk:  33\n",
      "current chunk:  34\n",
      "current chunk:  35\n",
      "current chunk:  36\n",
      "current chunk:  37\n",
      "current chunk:  38\n",
      "current chunk:  39\n",
      "===================\n",
      "= = = dumping epoch: 4 = = =\n",
      "exit dumping process\n",
      "exit dumping process\n",
      "exit dumping process\n",
      "exit dumping process\n",
      "exit dumping process\n",
      "exit dumping process\n",
      "exit dumping process\n",
      "exit dumping process\n",
      "current chunk:  40\n",
      "current chunk:  41\n",
      "current chunk:  42\n",
      "current chunk:  43\n",
      "current chunk:  44\n",
      "current chunk:  45\n",
      "current chunk:  46\n",
      "current chunk:  47\n",
      "===================\n",
      "= = = dumping epoch: 5 = = =\n",
      "exit dumping process\n",
      "exit dumping process\n",
      "exit dumping process\n",
      "exit dumping process\n",
      "exit dumping process\n",
      "exit dumping process\n",
      "exit dumping process\n",
      "exit dumping process\n",
      "current chunk:  48\n",
      "current chunk:  49\n",
      "current chunk:  50\n",
      "current chunk:  51\n",
      "current chunk:  52\n",
      "current chunk:  53\n",
      "current chunk:  54\n",
      "current chunk:  55\n",
      "===================\n",
      "= = = dumping epoch: 6 = = =\n",
      "exit dumping process\n",
      "exit dumping process\n",
      "exit dumping process\n",
      "exit dumping process\n",
      "exit dumping process\n",
      "exit dumping process\n",
      "exit dumping process\n",
      "exit dumping process\n",
      "current chunk:  56\n",
      "current chunk:  57\n",
      "current chunk:  58\n",
      "current chunk:  59\n",
      "current chunk:  60\n",
      "current chunk:  61\n",
      "current chunk:  62\n",
      "current chunk:  63\n",
      "===================\n",
      "= = = dumping epoch: 7 = = =\n",
      "exit dumping process\n",
      "exit dumping process\n",
      "exit dumping process\n",
      "exit dumping process\n",
      "exit dumping process\n",
      "exit dumping process\n",
      "exit dumping process\n",
      "exit dumping process\n",
      "current chunk:  64\n",
      "current chunk:  65\n",
      "current chunk:  66\n",
      "current chunk:  67\n",
      "current chunk:  68\n",
      "current chunk:  69\n",
      "current chunk:  70\n",
      "current chunk:  71\n",
      "===================\n",
      "= = = dumping epoch: 8 = = =\n",
      "exit dumping process\n",
      "exit dumping process\n",
      "exit dumping process\n",
      "exit dumping process\n",
      "exit dumping process\n",
      "exit dumping process\n",
      "exit dumping process\n",
      "exit dumping process\n",
      "current chunk:  72\n",
      "current chunk:  73\n",
      "current chunk:  74\n",
      "current chunk:  75\n",
      "current chunk:  76\n",
      "current chunk:  77\n",
      "current chunk:  78\n",
      "current chunk:  79\n",
      "===================\n",
      "= = = dumping epoch: 9 = = =\n",
      "exit dumping process\n",
      "exit dumping process\n",
      "exit dumping process\n",
      "exit dumping process\n",
      "exit dumping process\n",
      "exit dumping process\n",
      "exit dumping process\n",
      "exit dumping process\n",
      "current chunk:  80\n",
      "current chunk:  81\n",
      "current chunk:  82\n",
      "current chunk:  83\n",
      "current chunk:  84\n",
      "current chunk:  85\n",
      "current chunk:  86\n",
      "current chunk:  87\n",
      "===================\n",
      "= = = dumping epoch: 10 = = =\n",
      "exit dumping process\n",
      "exit dumping process\n",
      "exit dumping process\n",
      "exit dumping process\n",
      "exit dumping process\n",
      "exit dumping process\n",
      "exit dumping process\n",
      "exit dumping process\n",
      "current chunk:  88\n",
      "current chunk:  89\n",
      "current chunk:  90\n",
      "current chunk:  91\n",
      "current chunk:  92\n",
      "current chunk:  93\n",
      "current chunk:  94\n",
      "current chunk:  95\n",
      "===================\n",
      "= = = dumping epoch: 11 = = =\n",
      "exit dumping process\n",
      "exit dumping process\n",
      "exit dumping process\n",
      "exit dumping process\n",
      "exit dumping process\n",
      "exit dumping process\n",
      "exit dumping process\n",
      "exit dumping process\n",
      "current chunk:  96\n",
      "current chunk:  97\n",
      "current chunk:  98\n",
      "current chunk:  99\n",
      "current chunk:  100\n",
      "current chunk:  101\n",
      "current chunk:  102\n",
      "current chunk:  103\n",
      "===================\n",
      "= = = dumping epoch: 12 = = =\n",
      "exit dumping process\n",
      "exit dumping process\n",
      "exit dumping process\n",
      "exit dumping process\n",
      "exit dumping process\n",
      "exit dumping process\n",
      "exit dumping process\n",
      "exit dumping process\n",
      "current chunk:  104\n",
      "current chunk:  105\n",
      "current chunk:  106\n",
      "current chunk:  107\n",
      "current chunk:  108\n",
      "current chunk:  109\n",
      "current chunk:  110\n",
      "current chunk:  111\n",
      "===================\n",
      "= = = dumping epoch: 13 = = =\n",
      "exit dumping process\n",
      "exit dumping process\n",
      "exit dumping process\n",
      "exit dumping process\n",
      "exit dumping process\n",
      "exit dumping process\n",
      "exit dumping process\n",
      "exit dumping process\n",
      "current chunk:  112\n",
      "current chunk:  113\n",
      "current chunk:  114\n",
      "current chunk:  115\n",
      "current chunk:  116\n",
      "current chunk:  117\n",
      "current chunk:  118\n",
      "current chunk:  119\n",
      "===================\n",
      "= = = dumping epoch: 14 = = =\n",
      "exit dumping process\n",
      "exit dumping process\n",
      "exit dumping process\n",
      "exit dumping process\n",
      "exit dumping process\n",
      "exit dumping process\n",
      "exit dumping process\n",
      "exit dumping process\n",
      "current chunk:  120\n",
      "current chunk:  121\n",
      "current chunk:  122\n",
      "current chunk:  123\n",
      "current chunk:  124\n",
      "current chunk:  125\n",
      "current chunk:  126\n",
      "current chunk:  127\n",
      "===================\n",
      "= = = dumping epoch: 15 = = =\n",
      "exit dumping process\n",
      "exit dumping process\n",
      "exit dumping process\n",
      "exit dumping process\n",
      "exit dumping process\n",
      "exit dumping process\n",
      "exit dumping process\n",
      "exit dumping process\n",
      "current chunk:  128\n",
      "current chunk:  129\n",
      "current chunk:  130\n",
      "current chunk:  131\n",
      "current chunk:  132\n",
      "current chunk:  133\n",
      "current chunk:  134\n",
      "current chunk:  135\n",
      "===================\n",
      "= = = dumping epoch: 16 = = =\n",
      "exit dumping process\n",
      "exit dumping process\n",
      "exit dumping process\n",
      "exit dumping process\n",
      "exit dumping process\n",
      "exit dumping process\n",
      "exit dumping process\n",
      "exit dumping process\n",
      "current chunk:  136\n",
      "current chunk:  137\n",
      "current chunk:  138\n",
      "current chunk:  139\n",
      "current chunk:  140\n",
      "current chunk:  141\n",
      "current chunk:  142\n",
      "current chunk:  143\n",
      "===================\n",
      "= = = dumping epoch: 17 = = =\n",
      "exit dumping process\n",
      "exit dumping process\n",
      "exit dumping process\n",
      "exit dumping process\n",
      "exit dumping process\n",
      "exit dumping process\n",
      "exit dumping process\n",
      "exit dumping process\n",
      "current chunk:  144\n",
      "current chunk:  145\n",
      "current chunk:  146\n",
      "current chunk:  147\n",
      "current chunk:  148\n",
      "current chunk:  149\n",
      "current chunk:  150\n",
      "current chunk:  151\n",
      "===================\n",
      "= = = dumping epoch: 18 = = =\n",
      "exit dumping process\n",
      "exit dumping process\n",
      "exit dumping process\n",
      "exit dumping process\n",
      "exit dumping process\n",
      "exit dumping process\n",
      "exit dumping process\n",
      "exit dumping process\n",
      "current chunk:  152\n",
      "current chunk:  153\n",
      "current chunk:  154\n",
      "current chunk:  155\n",
      "current chunk:  156\n",
      "current chunk:  157\n",
      "current chunk:  158\n",
      "current chunk:  159\n",
      "===================\n",
      "= = = dumping epoch: 19 = = =\n",
      "exit dumping process\n",
      "exit dumping process\n",
      "exit dumping process\n",
      "exit dumping process\n",
      "exit dumping process\n",
      "exit dumping process\n",
      "exit dumping process\n",
      "exit dumping process\n",
      "current chunk:  160\n",
      "current chunk:  161\n",
      "current chunk:  162\n",
      "current chunk:  163\n",
      "current chunk:  164\n",
      "current chunk:  165\n",
      "current chunk:  166\n",
      "current chunk:  167\n",
      "===================\n",
      "= = = dumping epoch: 20 = = =\n",
      "exit dumping process\n",
      "exit dumping process\n",
      "exit dumping process\n",
      "exit dumping process\n",
      "exit dumping process\n",
      "exit dumping process\n",
      "exit dumping process\n",
      "exit dumping process\n",
      "current chunk:  168\n",
      "current chunk:  169\n",
      "current chunk:  170\n",
      "current chunk:  171\n",
      "current chunk:  172\n",
      "current chunk:  173\n",
      "current chunk:  174\n",
      "current chunk:  175\n",
      "===================\n",
      "= = = dumping epoch: 21 = = =\n",
      "exit dumping process\n",
      "exit dumping process\n",
      "exit dumping process\n",
      "exit dumping process\n",
      "exit dumping process\n",
      "exit dumping process\n",
      "exit dumping process\n",
      "exit dumping process\n",
      "current chunk:  176\n",
      "current chunk:  177\n",
      "current chunk:  178\n",
      "current chunk:  179\n",
      "current chunk:  180\n",
      "current chunk:  181\n",
      "current chunk:  182\n",
      "current chunk:  183\n",
      "===================\n",
      "= = = dumping epoch: 22 = = =\n",
      "exit dumping process\n",
      "exit dumping process\n",
      "exit dumping process\n",
      "exit dumping process\n",
      "exit dumping process\n",
      "exit dumping process\n",
      "exit dumping process\n",
      "exit dumping process\n",
      "current chunk:  184\n",
      "current chunk:  185\n",
      "current chunk:  186\n",
      "current chunk:  187\n",
      "current chunk:  188\n",
      "current chunk:  189\n",
      "current chunk:  190\n",
      "current chunk:  191\n",
      "===================\n",
      "= = = dumping epoch: 23 = = =\n",
      "exit dumping process\n",
      "exit dumping process\n",
      "exit dumping process\n",
      "exit dumping process\n",
      "exit dumping process\n",
      "exit dumping process\n",
      "exit dumping process\n",
      "exit dumping process\n",
      "current chunk:  192\n",
      "current chunk:  193\n",
      "current chunk:  194\n",
      "current chunk:  195\n",
      "current chunk:  196\n",
      "current chunk:  197\n",
      "current chunk:  198\n",
      "current chunk:  199\n",
      "===================\n",
      "= = = dumping epoch: 24 = = =\n",
      "exit dumping process\n",
      "exit dumping process\n",
      "exit dumping process\n",
      "exit dumping process\n",
      "exit dumping process\n",
      "exit dumping process\n",
      "exit dumping process\n",
      "exit dumping process\n",
      "current chunk:  200\n",
      "current chunk:  201\n",
      "current chunk:  202\n",
      "current chunk:  203\n",
      "current chunk:  204\n",
      "current chunk:  205\n",
      "current chunk:  206\n",
      "current chunk:  207\n",
      "===================\n",
      "= = = dumping epoch: 25 = = =\n",
      "exit dumping process\n",
      "exit dumping process\n",
      "exit dumping process\n",
      "exit dumping process\n",
      "exit dumping process\n",
      "exit dumping process\n",
      "exit dumping process\n",
      "exit dumping process\n",
      "current chunk:  208\n",
      "current chunk:  209\n",
      "current chunk:  210\n",
      "current chunk:  211\n",
      "current chunk:  212\n",
      "current chunk:  213\n",
      "current chunk:  214\n",
      "current chunk:  215\n",
      "===================\n",
      "= = = dumping epoch: 26 = = =\n",
      "exit dumping process\n",
      "exit dumping process\n",
      "exit dumping process\n",
      "exit dumping process\n",
      "exit dumping process\n",
      "exit dumping process\n",
      "exit dumping process\n",
      "exit dumping process\n",
      "current chunk:  216\n",
      "current chunk:  217\n",
      "current chunk:  218\n",
      "current chunk:  219\n",
      "current chunk:  220\n",
      "current chunk:  221\n",
      "current chunk:  222\n",
      "current chunk:  223\n",
      "===================\n",
      "= = = dumping epoch: 27 = = =\n",
      "exit dumping process\n",
      "exit dumping process\n",
      "exit dumping process\n",
      "exit dumping process\n",
      "exit dumping process\n",
      "exit dumping process\n",
      "exit dumping process\n",
      "exit dumping process\n",
      "current chunk:  224\n",
      "current chunk:  225\n",
      "current chunk:  226\n",
      "current chunk:  227\n",
      "current chunk:  228\n",
      "current chunk:  229\n",
      "current chunk:  230\n",
      "current chunk:  231\n"
     ]
    }
   ],
   "source": [
    "# = = = Execution = = =\n",
    "# batch_data_parallel(table_path, nora_partition, chunk_size, used_dims, nora_hdfs, num_dims, dump_threshold, num_process)\n",
    "# print('finish nora data routing..')\n",
    "# batch_data_parallel(table_path, qdtree_partition, chunk_size, used_dims, qdtree_hdfs, num_dims, dump_threshold, num_process)\n",
    "# print('finish qdtree data routing..')\n",
    "batch_data_parallel(table_path, kdtree_partition, chunk_size, used_dims, kdtree_hdfs, num_dims, dump_threshold, num_process)\n",
    "print('finish kdtree data routing..')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# it seems the process of merging QdTree partitions are stuck, we re generate the merged data\n",
    "# but the below cannot work, it will also stuck at some point, I can't find out why\n",
    "\n",
    "# pool = Pool(processes = 3)\n",
    "\n",
    "# # totally 68 partitions for QdTree\n",
    "# # pids = [[k * 8 + i for i in range(8)] for k in range(num_process)]\n",
    "# # pids[-1] += [64, 65, 66, 67]\n",
    "\n",
    "# pids = [i for i in range(68)] # 0 - 67\n",
    "\n",
    "# batch = 0\n",
    "# while batch < 3:\n",
    "\n",
    "#     pids_each_process = [set(pids[batch*24+k*8: batch*24+(k+1)*8]) for k in range(3)]\n",
    "#     # totally 94 epochs\n",
    "#     paras = [[94, pids_each_process[k], qdtree_hdfs] for k in range(3)]\n",
    "#     pool.map(merge_parquets, [para for para in paras])\n",
    "#     batch += 1\n",
    "    \n",
    "# pool.close()\n",
    "# pool.join()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# import pandas as pd\n",
    "# import pyarrow as pa\n",
    "# import pyarrow.parquet as pq\n",
    "# import numpy as np\n",
    "\n",
    "# pids = [i for i in range(68)] # 0 - 67\n",
    "# batches = 94\n",
    "# hdfs_path = qdtree_hdfs\n",
    "\n",
    "# start_time = time.time()\n",
    "\n",
    "# # using single process to handle data merge\n",
    "# fs = pa.hdfs.connect()\n",
    "# for pid in pids:\n",
    "#     parquets = []\n",
    "#     print('= = = process pid: ', pid, '= = =')\n",
    "#     for batch in range(batches):\n",
    "#         path = hdfs_path + str(batch) + '/partition_' + str(pid)+'.parquet'\n",
    "#         print(batch)\n",
    "#         try:\n",
    "#             par = pq.read_table(path)\n",
    "#             parquets.append(par)\n",
    "#         except:\n",
    "#             continue\n",
    "#     merged_parquet = pa.concat_tables(parquets)\n",
    "#     merge_path = hdfs_path + 'merged/partition_' + str(pid)+'.parquet'\n",
    "#     fw = fs.open(merge_path, 'wb')\n",
    "#     pq.write_table(merged_parquet, fw)\n",
    "#     fw.close()\n",
    "# print('exit merge process')\n",
    "\n",
    "# end_time = time.time()\n",
    "# print('time usage: ', end_time - start_time) # 2347s"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "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.8.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}