h2o_benchmarks_2022-10-0_pyarrow_str_50GB.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "e3468eac",
   "metadata": {},
   "source": [
    "## Grouby queries from h2o-benchmarks (parquet-pyarrow str)\n",
    "\n",
    "The dataset \"50GB_pyarr\" was generated in the same way than the other ones, except that when reading the multiple csv files, id1, id2 and id3 where casted as pyarroew_strings, then performed a repartition using 100MB as target size, and then we wrote to parquet. \n",
    "\n",
    "Once reading the data back and casting id1, id2 and id3 as pyarrow strings each partition size in memory is around 72MB\n",
    "\n",
    "\n",
    "- Dask = 2022.10.0\n",
    "- pyarrow = 10.0.0.dev497\n",
    "- pandas = 1.5.1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "ae69b689",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "\n",
    "import coiled\n",
    "import dask.dataframe as dd\n",
    "from dask.distributed import Client, performance_report\n",
    "from distributed.diagnostics import MemorySampler\n",
    "import pandas as pd"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "e3d5a996",
   "metadata": {},
   "outputs": [],
   "source": [
    "import dask\n",
    "import distributed"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "2560ca17",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'2022.10.0'"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "dask.__version__"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "fc158264",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'2022.10.0'"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "distributed.__version__"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "7427273b",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'0.2.38'"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import coiled\n",
    "coiled.__version__"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "f7870aeb",
   "metadata": {},
   "outputs": [],
   "source": [
    "#cluster = coiled.Cluster(name=\"h2o-benchmarks\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "d63f04c8",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "65bfd21b855f4a8baabf0596cecc87ac",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Output()"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/html": [
       "<pre style=\"white-space:pre;overflow-x:auto;line-height:normal;font-family:Menlo,'DejaVu Sans Mono',consolas,'Courier New',monospace\"></pre>\n"
      ],
      "text/plain": []
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "cluster = coiled.Cluster(\n",
    "    name=\"h2o-benchmarks\",\n",
    "    n_workers=10,\n",
    "    worker_vm_types=[\"t3.large\"],  # 2CPU, 8GiB\n",
    "    scheduler_vm_types=[\"t3.large\"],\n",
    "    software=\"dask-engineering/h2o_pyarrow_nightly\",\n",
    "#     package_sync=True,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "8b80b3ca",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div>\n",
       "    <div style=\"width: 24px; height: 24px; background-color: #e1e1e1; border: 3px solid #9D9D9D; border-radius: 5px; position: absolute;\"> </div>\n",
       "    <div style=\"margin-left: 48px;\">\n",
       "        <h3 style=\"margin-bottom: 0px;\">Client</h3>\n",
       "        <p style=\"color: #9D9D9D; margin-bottom: 0px;\">Client-69b0f404-56d4-11ed-a13b-9eae289ce62a</p>\n",
       "        <table style=\"width: 100%; text-align: left;\">\n",
       "\n",
       "        <tr>\n",
       "        \n",
       "            <td style=\"text-align: left;\"><strong>Connection method:</strong> Cluster object</td>\n",
       "            <td style=\"text-align: left;\"><strong>Cluster type:</strong> coiled.ClusterBeta</td>\n",
       "        \n",
       "        </tr>\n",
       "\n",
       "        \n",
       "            <tr>\n",
       "                <td style=\"text-align: left;\">\n",
       "                    <strong>Dashboard: </strong> <a href=\"http://18.116.80.128:8787\" target=\"_blank\">http://18.116.80.128:8787</a>\n",
       "                </td>\n",
       "                <td style=\"text-align: left;\"></td>\n",
       "            </tr>\n",
       "        \n",
       "\n",
       "        </table>\n",
       "\n",
       "        \n",
       "            <details>\n",
       "            <summary style=\"margin-bottom: 20px;\"><h3 style=\"display: inline;\">Cluster Info</h3></summary>\n",
       "            <div class=\"jp-RenderedHTMLCommon jp-RenderedHTML jp-mod-trusted jp-OutputArea-output\">\n",
       "    <div style=\"width: 24px; height: 24px; background-color: #e1e1e1; border: 3px solid #9D9D9D; border-radius: 5px; position: absolute;\">\n",
       "    </div>\n",
       "    <div style=\"margin-left: 48px;\">\n",
       "        <h3 style=\"margin-bottom: 0px; margin-top: 0px;\">ClusterBeta</h3>\n",
       "        <p style=\"color: #9D9D9D; margin-bottom: 0px;\">h2o-benchmarks</p>\n",
       "        <table style=\"width: 100%; text-align: left;\">\n",
       "            <tr>\n",
       "                <td style=\"text-align: left;\">\n",
       "                    <strong>Dashboard:</strong> <a href=\"http://18.116.80.128:8787\" target=\"_blank\">http://18.116.80.128:8787</a>\n",
       "                </td>\n",
       "                <td style=\"text-align: left;\">\n",
       "                    <strong>Workers:</strong> 10\n",
       "                </td>\n",
       "            </tr>\n",
       "            <tr>\n",
       "                <td style=\"text-align: left;\">\n",
       "                    <strong>Total threads:</strong> 20\n",
       "                </td>\n",
       "                <td style=\"text-align: left;\">\n",
       "                    <strong>Total memory:</strong> 72.02 GiB\n",
       "                </td>\n",
       "            </tr>\n",
       "            \n",
       "        </table>\n",
       "\n",
       "        <details>\n",
       "            <summary style=\"margin-bottom: 20px;\">\n",
       "                <h3 style=\"display: inline;\">Scheduler Info</h3>\n",
       "            </summary>\n",
       "\n",
       "            <div style=\"\">\n",
       "    <div>\n",
       "        <div style=\"width: 24px; height: 24px; background-color: #FFF7E5; border: 3px solid #FF6132; border-radius: 5px; position: absolute;\"> </div>\n",
       "        <div style=\"margin-left: 48px;\">\n",
       "            <h3 style=\"margin-bottom: 0px;\">Scheduler</h3>\n",
       "            <p style=\"color: #9D9D9D; margin-bottom: 0px;\">Scheduler-c22ebc78-8b8f-4d51-97ff-c9406ab43892</p>\n",
       "            <table style=\"width: 100%; text-align: left;\">\n",
       "                <tr>\n",
       "                    <td style=\"text-align: left;\">\n",
       "                        <strong>Comm:</strong> tls://10.0.15.238:8786\n",
       "                    </td>\n",
       "                    <td style=\"text-align: left;\">\n",
       "                        <strong>Workers:</strong> 10\n",
       "                    </td>\n",
       "                </tr>\n",
       "                <tr>\n",
       "                    <td style=\"text-align: left;\">\n",
       "                        <strong>Dashboard:</strong> <a href=\"http://10.0.15.238:8787/status\" target=\"_blank\">http://10.0.15.238:8787/status</a>\n",
       "                    </td>\n",
       "                    <td style=\"text-align: left;\">\n",
       "                        <strong>Total threads:</strong> 20\n",
       "                    </td>\n",
       "                </tr>\n",
       "                <tr>\n",
       "                    <td style=\"text-align: left;\">\n",
       "                        <strong>Started:</strong> Just now\n",
       "                    </td>\n",
       "                    <td style=\"text-align: left;\">\n",
       "                        <strong>Total memory:</strong> 72.02 GiB\n",
       "                    </td>\n",
       "                </tr>\n",
       "            </table>\n",
       "        </div>\n",
       "    </div>\n",
       "\n",
       "    <details style=\"margin-left: 48px;\">\n",
       "        <summary style=\"margin-bottom: 20px;\">\n",
       "            <h3 style=\"display: inline;\">Workers</h3>\n",
       "        </summary>\n",
       "\n",
       "        \n",
       "        <div style=\"margin-bottom: 20px;\">\n",
       "            <div style=\"width: 24px; height: 24px; background-color: #DBF5FF; border: 3px solid #4CC9FF; border-radius: 5px; position: absolute;\"> </div>\n",
       "            <div style=\"margin-left: 48px;\">\n",
       "            <details>\n",
       "                <summary>\n",
       "                    <h4 style=\"margin-bottom: 0px; display: inline;\">Worker: h2o-benchmarks-worker-1cbae68403</h4>\n",
       "                </summary>\n",
       "                <table style=\"width: 100%; text-align: left;\">\n",
       "                    <tr>\n",
       "                        <td style=\"text-align: left;\">\n",
       "                            <strong>Comm: </strong> tls://10.0.9.210:33039\n",
       "                        </td>\n",
       "                        <td style=\"text-align: left;\">\n",
       "                            <strong>Total threads: </strong> 2\n",
       "                        </td>\n",
       "                    </tr>\n",
       "                    <tr>\n",
       "                        <td style=\"text-align: left;\">\n",
       "                            <strong>Dashboard: </strong> <a href=\"http://10.0.9.210:8787/status\" target=\"_blank\">http://10.0.9.210:8787/status</a>\n",
       "                        </td>\n",
       "                        <td style=\"text-align: left;\">\n",
       "                            <strong>Memory: </strong> 7.21 GiB\n",
       "                        </td>\n",
       "                    </tr>\n",
       "                    <tr>\n",
       "                        <td style=\"text-align: left;\">\n",
       "                            <strong>Nanny: </strong> tls://10.0.9.210:35915\n",
       "                        </td>\n",
       "                        <td style=\"text-align: left;\"></td>\n",
       "                    </tr>\n",
       "                    <tr>\n",
       "                        <td colspan=\"2\" style=\"text-align: left;\">\n",
       "                            <strong>Local directory: </strong> /scratch/dask-worker-space/worker-oyf45xvt\n",
       "                        </td>\n",
       "                    </tr>\n",
       "\n",
       "                    \n",
       "\n",
       "                    \n",
       "\n",
       "                </table>\n",
       "            </details>\n",
       "            </div>\n",
       "        </div>\n",
       "        \n",
       "        <div style=\"margin-bottom: 20px;\">\n",
       "            <div style=\"width: 24px; height: 24px; background-color: #DBF5FF; border: 3px solid #4CC9FF; border-radius: 5px; position: absolute;\"> </div>\n",
       "            <div style=\"margin-left: 48px;\">\n",
       "            <details>\n",
       "                <summary>\n",
       "                    <h4 style=\"margin-bottom: 0px; display: inline;\">Worker: h2o-benchmarks-worker-1ee943eba9</h4>\n",
       "                </summary>\n",
       "                <table style=\"width: 100%; text-align: left;\">\n",
       "                    <tr>\n",
       "                        <td style=\"text-align: left;\">\n",
       "                            <strong>Comm: </strong> tls://10.0.6.47:34001\n",
       "                        </td>\n",
       "                        <td style=\"text-align: left;\">\n",
       "                            <strong>Total threads: </strong> 2\n",
       "                        </td>\n",
       "                    </tr>\n",
       "                    <tr>\n",
       "                        <td style=\"text-align: left;\">\n",
       "                            <strong>Dashboard: </strong> <a href=\"http://10.0.6.47:8787/status\" target=\"_blank\">http://10.0.6.47:8787/status</a>\n",
       "                        </td>\n",
       "                        <td style=\"text-align: left;\">\n",
       "                            <strong>Memory: </strong> 7.20 GiB\n",
       "                        </td>\n",
       "                    </tr>\n",
       "                    <tr>\n",
       "                        <td style=\"text-align: left;\">\n",
       "                            <strong>Nanny: </strong> tls://10.0.6.47:37639\n",
       "                        </td>\n",
       "                        <td style=\"text-align: left;\"></td>\n",
       "                    </tr>\n",
       "                    <tr>\n",
       "                        <td colspan=\"2\" style=\"text-align: left;\">\n",
       "                            <strong>Local directory: </strong> /scratch/dask-worker-space/worker-370d2l3n\n",
       "                        </td>\n",
       "                    </tr>\n",
       "\n",
       "                    \n",
       "\n",
       "                    \n",
       "\n",
       "                </table>\n",
       "            </details>\n",
       "            </div>\n",
       "        </div>\n",
       "        \n",
       "        <div style=\"margin-bottom: 20px;\">\n",
       "            <div style=\"width: 24px; height: 24px; background-color: #DBF5FF; border: 3px solid #4CC9FF; border-radius: 5px; position: absolute;\"> </div>\n",
       "            <div style=\"margin-left: 48px;\">\n",
       "            <details>\n",
       "                <summary>\n",
       "                    <h4 style=\"margin-bottom: 0px; display: inline;\">Worker: h2o-benchmarks-worker-4d18072d13</h4>\n",
       "                </summary>\n",
       "                <table style=\"width: 100%; text-align: left;\">\n",
       "                    <tr>\n",
       "                        <td style=\"text-align: left;\">\n",
       "                            <strong>Comm: </strong> tls://10.0.4.158:32797\n",
       "                        </td>\n",
       "                        <td style=\"text-align: left;\">\n",
       "                            <strong>Total threads: </strong> 2\n",
       "                        </td>\n",
       "                    </tr>\n",
       "                    <tr>\n",
       "                        <td style=\"text-align: left;\">\n",
       "                            <strong>Dashboard: </strong> <a href=\"http://10.0.4.158:8787/status\" target=\"_blank\">http://10.0.4.158:8787/status</a>\n",
       "                        </td>\n",
       "                        <td style=\"text-align: left;\">\n",
       "                            <strong>Memory: </strong> 7.21 GiB\n",
       "                        </td>\n",
       "                    </tr>\n",
       "                    <tr>\n",
       "                        <td style=\"text-align: left;\">\n",
       "                            <strong>Nanny: </strong> tls://10.0.4.158:33309\n",
       "                        </td>\n",
       "                        <td style=\"text-align: left;\"></td>\n",
       "                    </tr>\n",
       "                    <tr>\n",
       "                        <td colspan=\"2\" style=\"text-align: left;\">\n",
       "                            <strong>Local directory: </strong> /scratch/dask-worker-space/worker-u27kby59\n",
       "                        </td>\n",
       "                    </tr>\n",
       "\n",
       "                    \n",
       "\n",
       "                    \n",
       "\n",
       "                </table>\n",
       "            </details>\n",
       "            </div>\n",
       "        </div>\n",
       "        \n",
       "        <div style=\"margin-bottom: 20px;\">\n",
       "            <div style=\"width: 24px; height: 24px; background-color: #DBF5FF; border: 3px solid #4CC9FF; border-radius: 5px; position: absolute;\"> </div>\n",
       "            <div style=\"margin-left: 48px;\">\n",
       "            <details>\n",
       "                <summary>\n",
       "                    <h4 style=\"margin-bottom: 0px; display: inline;\">Worker: h2o-benchmarks-worker-5adad28fa2</h4>\n",
       "                </summary>\n",
       "                <table style=\"width: 100%; text-align: left;\">\n",
       "                    <tr>\n",
       "                        <td style=\"text-align: left;\">\n",
       "                            <strong>Comm: </strong> tls://10.0.6.35:41143\n",
       "                        </td>\n",
       "                        <td style=\"text-align: left;\">\n",
       "                            <strong>Total threads: </strong> 2\n",
       "                        </td>\n",
       "                    </tr>\n",
       "                    <tr>\n",
       "                        <td style=\"text-align: left;\">\n",
       "                            <strong>Dashboard: </strong> <a href=\"http://10.0.6.35:8787/status\" target=\"_blank\">http://10.0.6.35:8787/status</a>\n",
       "                        </td>\n",
       "                        <td style=\"text-align: left;\">\n",
       "                            <strong>Memory: </strong> 7.20 GiB\n",
       "                        </td>\n",
       "                    </tr>\n",
       "                    <tr>\n",
       "                        <td style=\"text-align: left;\">\n",
       "                            <strong>Nanny: </strong> tls://10.0.6.35:45171\n",
       "                        </td>\n",
       "                        <td style=\"text-align: left;\"></td>\n",
       "                    </tr>\n",
       "                    <tr>\n",
       "                        <td colspan=\"2\" style=\"text-align: left;\">\n",
       "                            <strong>Local directory: </strong> /scratch/dask-worker-space/worker-3wp7ita5\n",
       "                        </td>\n",
       "                    </tr>\n",
       "\n",
       "                    \n",
       "\n",
       "                    \n",
       "\n",
       "                </table>\n",
       "            </details>\n",
       "            </div>\n",
       "        </div>\n",
       "        \n",
       "        <div style=\"margin-bottom: 20px;\">\n",
       "            <div style=\"width: 24px; height: 24px; background-color: #DBF5FF; border: 3px solid #4CC9FF; border-radius: 5px; position: absolute;\"> </div>\n",
       "            <div style=\"margin-left: 48px;\">\n",
       "            <details>\n",
       "                <summary>\n",
       "                    <h4 style=\"margin-bottom: 0px; display: inline;\">Worker: h2o-benchmarks-worker-8b0ad999e7</h4>\n",
       "                </summary>\n",
       "                <table style=\"width: 100%; text-align: left;\">\n",
       "                    <tr>\n",
       "                        <td style=\"text-align: left;\">\n",
       "                            <strong>Comm: </strong> tls://10.0.14.68:44901\n",
       "                        </td>\n",
       "                        <td style=\"text-align: left;\">\n",
       "                            <strong>Total threads: </strong> 2\n",
       "                        </td>\n",
       "                    </tr>\n",
       "                    <tr>\n",
       "                        <td style=\"text-align: left;\">\n",
       "                            <strong>Dashboard: </strong> <a href=\"http://10.0.14.68:8787/status\" target=\"_blank\">http://10.0.14.68:8787/status</a>\n",
       "                        </td>\n",
       "                        <td style=\"text-align: left;\">\n",
       "                            <strong>Memory: </strong> 7.19 GiB\n",
       "                        </td>\n",
       "                    </tr>\n",
       "                    <tr>\n",
       "                        <td style=\"text-align: left;\">\n",
       "                            <strong>Nanny: </strong> tls://10.0.14.68:33073\n",
       "                        </td>\n",
       "                        <td style=\"text-align: left;\"></td>\n",
       "                    </tr>\n",
       "                    <tr>\n",
       "                        <td colspan=\"2\" style=\"text-align: left;\">\n",
       "                            <strong>Local directory: </strong> /scratch/dask-worker-space/worker-0_zv6ysv\n",
       "                        </td>\n",
       "                    </tr>\n",
       "\n",
       "                    \n",
       "\n",
       "                    \n",
       "\n",
       "                </table>\n",
       "            </details>\n",
       "            </div>\n",
       "        </div>\n",
       "        \n",
       "        <div style=\"margin-bottom: 20px;\">\n",
       "            <div style=\"width: 24px; height: 24px; background-color: #DBF5FF; border: 3px solid #4CC9FF; border-radius: 5px; position: absolute;\"> </div>\n",
       "            <div style=\"margin-left: 48px;\">\n",
       "            <details>\n",
       "                <summary>\n",
       "                    <h4 style=\"margin-bottom: 0px; display: inline;\">Worker: h2o-benchmarks-worker-9440ef1918</h4>\n",
       "                </summary>\n",
       "                <table style=\"width: 100%; text-align: left;\">\n",
       "                    <tr>\n",
       "                        <td style=\"text-align: left;\">\n",
       "                            <strong>Comm: </strong> tls://10.0.9.227:43897\n",
       "                        </td>\n",
       "                        <td style=\"text-align: left;\">\n",
       "                            <strong>Total threads: </strong> 2\n",
       "                        </td>\n",
       "                    </tr>\n",
       "                    <tr>\n",
       "                        <td style=\"text-align: left;\">\n",
       "                            <strong>Dashboard: </strong> <a href=\"http://10.0.9.227:8787/status\" target=\"_blank\">http://10.0.9.227:8787/status</a>\n",
       "                        </td>\n",
       "                        <td style=\"text-align: left;\">\n",
       "                            <strong>Memory: </strong> 7.21 GiB\n",
       "                        </td>\n",
       "                    </tr>\n",
       "                    <tr>\n",
       "                        <td style=\"text-align: left;\">\n",
       "                            <strong>Nanny: </strong> tls://10.0.9.227:40387\n",
       "                        </td>\n",
       "                        <td style=\"text-align: left;\"></td>\n",
       "                    </tr>\n",
       "                    <tr>\n",
       "                        <td colspan=\"2\" style=\"text-align: left;\">\n",
       "                            <strong>Local directory: </strong> /scratch/dask-worker-space/worker-cvi30uka\n",
       "                        </td>\n",
       "                    </tr>\n",
       "\n",
       "                    \n",
       "\n",
       "                    \n",
       "\n",
       "                </table>\n",
       "            </details>\n",
       "            </div>\n",
       "        </div>\n",
       "        \n",
       "        <div style=\"margin-bottom: 20px;\">\n",
       "            <div style=\"width: 24px; height: 24px; background-color: #DBF5FF; border: 3px solid #4CC9FF; border-radius: 5px; position: absolute;\"> </div>\n",
       "            <div style=\"margin-left: 48px;\">\n",
       "            <details>\n",
       "                <summary>\n",
       "                    <h4 style=\"margin-bottom: 0px; display: inline;\">Worker: h2o-benchmarks-worker-bc8610f279</h4>\n",
       "                </summary>\n",
       "                <table style=\"width: 100%; text-align: left;\">\n",
       "                    <tr>\n",
       "                        <td style=\"text-align: left;\">\n",
       "                            <strong>Comm: </strong> tls://10.0.9.108:39389\n",
       "                        </td>\n",
       "                        <td style=\"text-align: left;\">\n",
       "                            <strong>Total threads: </strong> 2\n",
       "                        </td>\n",
       "                    </tr>\n",
       "                    <tr>\n",
       "                        <td style=\"text-align: left;\">\n",
       "                            <strong>Dashboard: </strong> <a href=\"http://10.0.9.108:8787/status\" target=\"_blank\">http://10.0.9.108:8787/status</a>\n",
       "                        </td>\n",
       "                        <td style=\"text-align: left;\">\n",
       "                            <strong>Memory: </strong> 7.20 GiB\n",
       "                        </td>\n",
       "                    </tr>\n",
       "                    <tr>\n",
       "                        <td style=\"text-align: left;\">\n",
       "                            <strong>Nanny: </strong> tls://10.0.9.108:45301\n",
       "                        </td>\n",
       "                        <td style=\"text-align: left;\"></td>\n",
       "                    </tr>\n",
       "                    <tr>\n",
       "                        <td colspan=\"2\" style=\"text-align: left;\">\n",
       "                            <strong>Local directory: </strong> /scratch/dask-worker-space/worker-zwh2vl7c\n",
       "                        </td>\n",
       "                    </tr>\n",
       "\n",
       "                    \n",
       "\n",
       "                    \n",
       "\n",
       "                </table>\n",
       "            </details>\n",
       "            </div>\n",
       "        </div>\n",
       "        \n",
       "        <div style=\"margin-bottom: 20px;\">\n",
       "            <div style=\"width: 24px; height: 24px; background-color: #DBF5FF; border: 3px solid #4CC9FF; border-radius: 5px; position: absolute;\"> </div>\n",
       "            <div style=\"margin-left: 48px;\">\n",
       "            <details>\n",
       "                <summary>\n",
       "                    <h4 style=\"margin-bottom: 0px; display: inline;\">Worker: h2o-benchmarks-worker-c757868e0f</h4>\n",
       "                </summary>\n",
       "                <table style=\"width: 100%; text-align: left;\">\n",
       "                    <tr>\n",
       "                        <td style=\"text-align: left;\">\n",
       "                            <strong>Comm: </strong> tls://10.0.8.101:44417\n",
       "                        </td>\n",
       "                        <td style=\"text-align: left;\">\n",
       "                            <strong>Total threads: </strong> 2\n",
       "                        </td>\n",
       "                    </tr>\n",
       "                    <tr>\n",
       "                        <td style=\"text-align: left;\">\n",
       "                            <strong>Dashboard: </strong> <a href=\"http://10.0.8.101:8787/status\" target=\"_blank\">http://10.0.8.101:8787/status</a>\n",
       "                        </td>\n",
       "                        <td style=\"text-align: left;\">\n",
       "                            <strong>Memory: </strong> 7.20 GiB\n",
       "                        </td>\n",
       "                    </tr>\n",
       "                    <tr>\n",
       "                        <td style=\"text-align: left;\">\n",
       "                            <strong>Nanny: </strong> tls://10.0.8.101:43155\n",
       "                        </td>\n",
       "                        <td style=\"text-align: left;\"></td>\n",
       "                    </tr>\n",
       "                    <tr>\n",
       "                        <td colspan=\"2\" style=\"text-align: left;\">\n",
       "                            <strong>Local directory: </strong> /scratch/dask-worker-space/worker-q56m9eeg\n",
       "                        </td>\n",
       "                    </tr>\n",
       "\n",
       "                    \n",
       "\n",
       "                    \n",
       "\n",
       "                </table>\n",
       "            </details>\n",
       "            </div>\n",
       "        </div>\n",
       "        \n",
       "        <div style=\"margin-bottom: 20px;\">\n",
       "            <div style=\"width: 24px; height: 24px; background-color: #DBF5FF; border: 3px solid #4CC9FF; border-radius: 5px; position: absolute;\"> </div>\n",
       "            <div style=\"margin-left: 48px;\">\n",
       "            <details>\n",
       "                <summary>\n",
       "                    <h4 style=\"margin-bottom: 0px; display: inline;\">Worker: h2o-benchmarks-worker-c7fa48f7f1</h4>\n",
       "                </summary>\n",
       "                <table style=\"width: 100%; text-align: left;\">\n",
       "                    <tr>\n",
       "                        <td style=\"text-align: left;\">\n",
       "                            <strong>Comm: </strong> tls://10.0.2.161:36669\n",
       "                        </td>\n",
       "                        <td style=\"text-align: left;\">\n",
       "                            <strong>Total threads: </strong> 2\n",
       "                        </td>\n",
       "                    </tr>\n",
       "                    <tr>\n",
       "                        <td style=\"text-align: left;\">\n",
       "                            <strong>Dashboard: </strong> <a href=\"http://10.0.2.161:8787/status\" target=\"_blank\">http://10.0.2.161:8787/status</a>\n",
       "                        </td>\n",
       "                        <td style=\"text-align: left;\">\n",
       "                            <strong>Memory: </strong> 7.20 GiB\n",
       "                        </td>\n",
       "                    </tr>\n",
       "                    <tr>\n",
       "                        <td style=\"text-align: left;\">\n",
       "                            <strong>Nanny: </strong> tls://10.0.2.161:38415\n",
       "                        </td>\n",
       "                        <td style=\"text-align: left;\"></td>\n",
       "                    </tr>\n",
       "                    <tr>\n",
       "                        <td colspan=\"2\" style=\"text-align: left;\">\n",
       "                            <strong>Local directory: </strong> /scratch/dask-worker-space/worker-ppbbw_3_\n",
       "                        </td>\n",
       "                    </tr>\n",
       "\n",
       "                    \n",
       "\n",
       "                    \n",
       "\n",
       "                </table>\n",
       "            </details>\n",
       "            </div>\n",
       "        </div>\n",
       "        \n",
       "        <div style=\"margin-bottom: 20px;\">\n",
       "            <div style=\"width: 24px; height: 24px; background-color: #DBF5FF; border: 3px solid #4CC9FF; border-radius: 5px; position: absolute;\"> </div>\n",
       "            <div style=\"margin-left: 48px;\">\n",
       "            <details>\n",
       "                <summary>\n",
       "                    <h4 style=\"margin-bottom: 0px; display: inline;\">Worker: h2o-benchmarks-worker-dfb601bc32</h4>\n",
       "                </summary>\n",
       "                <table style=\"width: 100%; text-align: left;\">\n",
       "                    <tr>\n",
       "                        <td style=\"text-align: left;\">\n",
       "                            <strong>Comm: </strong> tls://10.0.15.89:40063\n",
       "                        </td>\n",
       "                        <td style=\"text-align: left;\">\n",
       "                            <strong>Total threads: </strong> 2\n",
       "                        </td>\n",
       "                    </tr>\n",
       "                    <tr>\n",
       "                        <td style=\"text-align: left;\">\n",
       "                            <strong>Dashboard: </strong> <a href=\"http://10.0.15.89:8787/status\" target=\"_blank\">http://10.0.15.89:8787/status</a>\n",
       "                        </td>\n",
       "                        <td style=\"text-align: left;\">\n",
       "                            <strong>Memory: </strong> 7.20 GiB\n",
       "                        </td>\n",
       "                    </tr>\n",
       "                    <tr>\n",
       "                        <td style=\"text-align: left;\">\n",
       "                            <strong>Nanny: </strong> tls://10.0.15.89:34827\n",
       "                        </td>\n",
       "                        <td style=\"text-align: left;\"></td>\n",
       "                    </tr>\n",
       "                    <tr>\n",
       "                        <td colspan=\"2\" style=\"text-align: left;\">\n",
       "                            <strong>Local directory: </strong> /scratch/dask-worker-space/worker-cbw2oy5_\n",
       "                        </td>\n",
       "                    </tr>\n",
       "\n",
       "                    \n",
       "\n",
       "                    \n",
       "\n",
       "                </table>\n",
       "            </details>\n",
       "            </div>\n",
       "        </div>\n",
       "        \n",
       "\n",
       "    </details>\n",
       "</div>\n",
       "\n",
       "        </details>\n",
       "    </div>\n",
       "</div>\n",
       "            </details>\n",
       "        \n",
       "\n",
       "    </div>\n",
       "</div>"
      ],
      "text/plain": [
       "<Client: 'tls://10.0.15.238:8786' processes=10 threads=20, memory=72.02 GiB>"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "client = Client(cluster)\n",
    "client"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "599d3096",
   "metadata": {},
   "outputs": [],
   "source": [
    "ms = MemorySampler()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "f55fd544",
   "metadata": {},
   "outputs": [],
   "source": [
    "#\"50GB_pyarr\" I save the data as pyarrow strings, did repartitioning using pyarrow\n",
    "\n",
    "data_size = {\n",
    "    \"50GB_pyarr\": \"s3://coiled-datasets/h2o-benchmark/pyarrow_strings/N_1e9_K_1e2/*.parquet\",\n",
    "\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "96004376",
   "metadata": {},
   "outputs": [],
   "source": [
    "ds = \"50GB_pyarr\" #run next with \"50GB_orig\"\n",
    "report_dir = \"performance-reports-pyarr_str-50GB\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "d3ee8cfa",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<div><strong>Dask DataFrame Structure:</strong></div>\n",
       "<div>\n",
       "<style scoped>\n",
       "    .dataframe tbody tr th:only-of-type {\n",
       "        vertical-align: middle;\n",
       "    }\n",
       "\n",
       "    .dataframe tbody tr th {\n",
       "        vertical-align: top;\n",
       "    }\n",
       "\n",
       "    .dataframe thead th {\n",
       "        text-align: right;\n",
       "    }\n",
       "</style>\n",
       "<table border=\"1\" class=\"dataframe\">\n",
       "  <thead>\n",
       "    <tr style=\"text-align: right;\">\n",
       "      <th></th>\n",
       "      <th>id1</th>\n",
       "      <th>id2</th>\n",
       "      <th>id3</th>\n",
       "      <th>id4</th>\n",
       "      <th>id5</th>\n",
       "      <th>id6</th>\n",
       "      <th>v1</th>\n",
       "      <th>v2</th>\n",
       "      <th>v3</th>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>npartitions=1000</th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "      <th></th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th></th>\n",
       "      <td>string</td>\n",
       "      <td>string</td>\n",
       "      <td>string</td>\n",
       "      <td>Int32</td>\n",
       "      <td>Int32</td>\n",
       "      <td>Int32</td>\n",
       "      <td>Int32</td>\n",
       "      <td>Int32</td>\n",
       "      <td>float64</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th></th>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>...</th>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th></th>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th></th>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "      <td>...</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>\n",
       "<div>Dask Name: read-parquet, 1 graph layer</div>"
      ],
      "text/plain": [
       "Dask DataFrame Structure:\n",
       "                     id1     id2     id3    id4    id5    id6     v1     v2       v3\n",
       "npartitions=1000                                                                    \n",
       "                  string  string  string  Int32  Int32  Int32  Int32  Int32  float64\n",
       "                     ...     ...     ...    ...    ...    ...    ...    ...      ...\n",
       "...                  ...     ...     ...    ...    ...    ...    ...    ...      ...\n",
       "                     ...     ...     ...    ...    ...    ...    ...    ...      ...\n",
       "                     ...     ...     ...    ...    ...    ...    ...    ...      ...\n",
       "Dask Name: read-parquet, 1 graph layer"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "ddf = dd.read_parquet(\n",
    "    data_size[ds],\n",
    "    engine=\"pyarrow\",\n",
    "    storage_options={\"anon\": True},\n",
    ")\n",
    "ddf"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3226a723",
   "metadata": {},
   "source": [
    "### Q1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "36df277e",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 289 ms, sys: 75.7 ms, total: 364 ms\n",
      "Wall time: 1min 13s\n"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "with performance_report(filename=os.path.join(report_dir, f\"q1_{ds}.html\")):\n",
    "    with ms.sample(f\"q1_22-10_{ds}\"):\n",
    "        ddf_q1 = ddf[[\"id1\", \"v1\"]].astype({\"id1\": \"string[pyarrow]\"})\n",
    "        ddf_q1.groupby(\"id1\", dropna=False, observed=True).agg({\"v1\": \"sum\"}).compute()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "c84be5fc",
   "metadata": {},
   "outputs": [],
   "source": [
    "client.restart();"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8c36c4cb",
   "metadata": {},
   "source": [
    "### Q2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "46ed6e0f",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 292 ms, sys: 54.3 ms, total: 347 ms\n",
      "Wall time: 1min 39s\n"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "with performance_report(filename=os.path.join(report_dir, f\"q2_{ds}.html\")):\n",
    "    with ms.sample(f\"q2_22-10_{ds}\"):\n",
    "        ddf_q2 = ddf[[\"id1\", \"id2\", \"v1\"]].astype({\"id1\": \"string[pyarrow]\",\n",
    "                                                   \"id2\": \"string[pyarrow]\"})\n",
    "        (\n",
    "            ddf_q2.groupby([\"id1\", \"id2\"], dropna=False, observed=True)\n",
    "            .agg({\"v1\": \"sum\"})\n",
    "            .compute()\n",
    "        )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "db7e7643",
   "metadata": {},
   "outputs": [],
   "source": [
    "client.restart();"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "52376148",
   "metadata": {},
   "source": [
    "### Q3"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "ed1da782",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 7.78 s, sys: 2.57 s, total: 10.3 s\n",
      "Wall time: 15min 17s\n"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "with performance_report(filename=os.path.join(report_dir, f\"q3_{ds}-p2p.html\")):\n",
    "    \n",
    "    with ms.sample(f\"q3_22-10_{ds}_p2p\"):\n",
    "        ddf_q3 = ddf[[\"id3\", \"v1\", \"v3\"]].astype({\"id3\": \"string[pyarrow]\"})\n",
    "        (\n",
    "            ddf_q3.groupby(\"id3\", dropna=False, observed=True)\n",
    "            .agg({\"v1\": \"sum\", \"v3\": \"mean\"}, shuffle=\"p2p\")\n",
    "            .compute()\n",
    "        )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "3ae0138a",
   "metadata": {},
   "outputs": [],
   "source": [
    "client.restart();"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "242e7b80",
   "metadata": {},
   "source": [
    "### Q7"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "82212572",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 3.62 s, sys: 1.25 s, total: 4.87 s\n",
      "Wall time: 11min 6s\n"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "with performance_report(filename=os.path.join(report_dir, f\"q7_{ds}-p2p.html\")):\n",
    "    \n",
    "    with ms.sample(f\"q7_22-10_{ds}_p2p\"):\n",
    "\n",
    "        ddf_q7 = ddf[[\"id3\", \"v1\", \"v2\"]].astype({\"id3\": \"string[pyarrow]\"})\n",
    "        (\n",
    "            ddf_q7.groupby(\"id3\", dropna=False, observed=True)\n",
    "            .agg({\"v1\": \"max\", \"v2\": \"min\"}, shuffle=\"p2p\")\n",
    "            .assign(range_v1_v2=lambda x: x[\"v1\"] - x[\"v2\"])[[\"range_v1_v2\"]]\n",
    "            .compute()\n",
    "        )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "5333c82c",
   "metadata": {},
   "outputs": [],
   "source": [
    "client.restart();"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b62084e3",
   "metadata": {},
   "source": [
    "### Q9"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "id": "24f29fae",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 2.72 s, sys: 467 ms, total: 3.19 s\n",
      "Wall time: 5min 54s\n"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "with performance_report(filename=os.path.join(report_dir, f\"q9_{ds}.html\")):\n",
    "    \n",
    "    with ms.sample(f\"q9_22-10_{ds}\"):\n",
    "        \n",
    "        ddf_q9 = ddf[[\"id2\", \"id4\", \"v1\", \"v2\"]].astype({\"id2\": \"string[pyarrow]\"})\n",
    "        (\n",
    "            ddf_q9[[\"id2\", \"id4\", \"v1\", \"v2\"]]\n",
    "            .groupby([\"id2\", \"id4\"], dropna=False, observed=True)\n",
    "            .apply(\n",
    "                lambda x: pd.Series({\"r2\": x.corr()[\"v1\"][\"v2\"] ** 2}),\n",
    "                meta={\"r2\": \"float64\"},\n",
    "            )\n",
    "            .compute()\n",
    "        )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "id": "0ff5fc27",
   "metadata": {},
   "outputs": [],
   "source": [
    "client.restart();"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "id": "a4c5b9a7",
   "metadata": {},
   "outputs": [],
   "source": [
    "ms.to_pandas(align=True).to_csv(\"id1id2id3_pyarrow_str_2022-10.csv\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ceb899cf",
   "metadata": {},
   "source": [
    "### Q4, Q5, Q6 and Q8\n",
    "Note: Columns involved in these queries are not affected by pyarrow strings. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "08dbe5f3",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.9.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}