blogs/speed-up-pandas-query-10x-with-dask.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "3cde9f68-64c4-41ca-adf5-10a733cdb89a",
   "metadata": {},
   "source": [
    "# How to speed up a pandas query 10x with Dask\n",
    "\n",
    "This notebook requires you to download a dataset from S3 to your local machine.\n",
    "\n",
    "This notebook works with the `crt-004` environment.  Run `conda env create -f envs/crt-004.yml` to create this environment.\n",
    "\n",
    "Here is the [full blog post](https://coiled.io/blog/speed-up-pandas-query-10x-with-dask/) that describes the computations in this notebook."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d5a782c0-8f63-4109-920a-c2987c5276b3",
   "metadata": {},
   "source": [
    "## Localhost data setup\n",
    "\n",
    "You need to download some data from S3 to run the examples in this notebook.  You need to either setup the AWS CLI on your machine and run the following commands or download the data using another approach."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "30542f7b-e243-4d14-b206-cd3476f7d127",
   "metadata": {},
   "outputs": [],
   "source": [
    "!mkdir data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "7786bb7a-821f-4363-bddc-6c4cb505b76b",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "download: s3://coiled-datasets/h2o-benchmark/N_1e8_K_1e2_single.csv to data/N_1e8_K_1e2_single.csv\n"
     ]
    }
   ],
   "source": [
    "!aws s3 cp s3://coiled-datasets/h2o-benchmark/N_1e8_K_1e2_single.csv data/"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3cb38c7b-2e58-40bd-85b9-09063eb545eb",
   "metadata": {},
   "source": [
    "The N_1e8_K_1e2_single.csv file contains XX rows and is 4.8 GB.  It's not a huge amount of data, but large enough for Dask to provide meaningful performance improvements compared to pandas."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "37235910-24fd-4903-9976-25724fc555a2",
   "metadata": {},
   "source": [
    "## Baseline: Just use Pandas"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "id": "3de5b665-876b-4579-afd7-8b51e3d2c6bc",
   "metadata": {},
   "outputs": [],
   "source": [
    "import pandas as pd"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "id": "7aa4e3a7-856a-42c3-81ed-20339a6ecfc5",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 58.9 s, sys: 32.8 s, total: 1min 31s\n",
      "Wall time: 3min 38s\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "<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>v1</th>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>id1</th>\n",
       "      <th></th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>id001</th>\n",
       "      <td>1999081</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>id002</th>\n",
       "      <td>2002870</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>id003</th>\n",
       "      <td>1998389</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>id004</th>\n",
       "      <td>2001091</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>id005</th>\n",
       "      <td>1999770</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>...</th>\n",
       "      <td>...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>id096</th>\n",
       "      <td>2001015</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>id097</th>\n",
       "      <td>1995014</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>id098</th>\n",
       "      <td>2000933</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>id099</th>\n",
       "      <td>2003447</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>id100</th>\n",
       "      <td>2001151</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "<p>100 rows × 1 columns</p>\n",
       "</div>"
      ],
      "text/plain": [
       "            v1\n",
       "id1           \n",
       "id001  1999081\n",
       "id002  2002870\n",
       "id003  1998389\n",
       "id004  2001091\n",
       "id005  1999770\n",
       "...        ...\n",
       "id096  2001015\n",
       "id097  1995014\n",
       "id098  2000933\n",
       "id099  2003447\n",
       "id100  2001151\n",
       "\n",
       "[100 rows x 1 columns]"
      ]
     },
     "execution_count": 26,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "%%time\n",
    "df = pd.read_csv(\"data/N_1e8_K_1e2_single.csv\")\n",
    "df.groupby(\"id1\", dropna=False, observed=True).agg({\"v1\": \"sum\"})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "id": "d1acd6f0-e327-4cae-b600-54546d753946",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "id1     object\n",
       "id2     object\n",
       "id3     object\n",
       "id4      int64\n",
       "id5      int64\n",
       "id6      int64\n",
       "v1       int64\n",
       "v2       int64\n",
       "v3     float64\n",
       "dtype: object"
      ]
     },
     "execution_count": 27,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "df.dtypes"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "225dbd8e-0102-4ed3-b17e-93e490c5596b",
   "metadata": {},
   "source": [
    "## Optimization #1: Use Dask & inefficient dtypes"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "4596e72c-304b-4bb5-bfbe-65f4c299d364",
   "metadata": {},
   "outputs": [],
   "source": [
    "import dask.dataframe as dd"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "2baf24cb-9e88-42b8-ae9c-9650bcbf8e54",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "id1     object\n",
       "id2     object\n",
       "id3     object\n",
       "id4      int64\n",
       "id5      int64\n",
       "id6      int64\n",
       "v1       int64\n",
       "v2       int64\n",
       "v3     float64\n",
       "dtype: object"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "dd.read_csv(\"data/N_1e8_K_1e2_single.csv\").dtypes"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "7a12e1f7-55c4-46aa-956f-69efe023abcc",
   "metadata": {},
   "outputs": [],
   "source": [
    "dtypes = {\n",
    "    \"id1\": \"object\",\n",
    "    \"id2\": \"object\",\n",
    "    \"id3\": \"object\",\n",
    "    \"id4\": \"object\",\n",
    "    \"id5\": \"object\",\n",
    "    \"id6\": \"object\",\n",
    "    \"v1\": \"object\",\n",
    "    \"v2\": \"object\",\n",
    "    \"v3\": \"object\",\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "d66607bc-b1b5-47a7-8c73-efe634d502a6",
   "metadata": {},
   "outputs": [],
   "source": [
    "ddf = dd.read_csv(\"data/N_1e8_K_1e2_single.csv\", dtype=dtypes)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "e71b9b07-7f71-4603-8b27-5687a596764f",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<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",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>id048</td>\n",
       "      <td>id035</td>\n",
       "      <td>id0000608844</td>\n",
       "      <td>22</td>\n",
       "      <td>63</td>\n",
       "      <td>812586</td>\n",
       "      <td>3</td>\n",
       "      <td>3</td>\n",
       "      <td>95.569069</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>id080</td>\n",
       "      <td>id099</td>\n",
       "      <td>id0000466449</td>\n",
       "      <td>40</td>\n",
       "      <td>79</td>\n",
       "      <td>682109</td>\n",
       "      <td>2</td>\n",
       "      <td>13</td>\n",
       "      <td>97.012438</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>id035</td>\n",
       "      <td>id041</td>\n",
       "      <td>id0000573987</td>\n",
       "      <td>81</td>\n",
       "      <td>10</td>\n",
       "      <td>800483</td>\n",
       "      <td>0</td>\n",
       "      <td>1</td>\n",
       "      <td>74.776325</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>id086</td>\n",
       "      <td>id047</td>\n",
       "      <td>id0000776204</td>\n",
       "      <td>53</td>\n",
       "      <td>72</td>\n",
       "      <td>212890</td>\n",
       "      <td>3</td>\n",
       "      <td>14</td>\n",
       "      <td>45.157498</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>id009</td>\n",
       "      <td>id091</td>\n",
       "      <td>id0000608718</td>\n",
       "      <td>93</td>\n",
       "      <td>65</td>\n",
       "      <td>655745</td>\n",
       "      <td>2</td>\n",
       "      <td>0</td>\n",
       "      <td>51.132162</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "     id1    id2           id3 id4 id5     id6 v1  v2         v3\n",
       "0  id048  id035  id0000608844  22  63  812586  3   3  95.569069\n",
       "1  id080  id099  id0000466449  40  79  682109  2  13  97.012438\n",
       "2  id035  id041  id0000573987  81  10  800483  0   1  74.776325\n",
       "3  id086  id047  id0000776204  53  72  212890  3  14  45.157498\n",
       "4  id009  id091  id0000608718  93  65  655745  2   0  51.132162"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "ddf.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "4ade8120-1b19-4035-8ede-7785e8751583",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "id1    object\n",
       "id2    object\n",
       "id3    object\n",
       "id4    object\n",
       "id5    object\n",
       "id6    object\n",
       "v1     object\n",
       "v2     object\n",
       "v3     object\n",
       "dtype: object"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "ddf.dtypes"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "c89c6f34-ef59-4906-924b-6551153b6c0e",
   "metadata": {},
   "outputs": [],
   "source": [
    "ddf[\"v1\"] = ddf[\"v1\"].astype(\"int64\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "618359fe-68ec-4bb8-8a26-16bcfa9767a8",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 1min 50s, sys: 28 s, total: 2min 18s\n",
      "Wall time: 2min 17s\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "<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>v1</th>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>id1</th>\n",
       "      <th></th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>id001</th>\n",
       "      <td>1999081</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>id002</th>\n",
       "      <td>2002870</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>id003</th>\n",
       "      <td>1998389</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>id004</th>\n",
       "      <td>2001091</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>id005</th>\n",
       "      <td>1999770</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>...</th>\n",
       "      <td>...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>id096</th>\n",
       "      <td>2001015</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>id097</th>\n",
       "      <td>1995014</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>id098</th>\n",
       "      <td>2000933</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>id099</th>\n",
       "      <td>2003447</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>id100</th>\n",
       "      <td>2001151</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "<p>100 rows × 1 columns</p>\n",
       "</div>"
      ],
      "text/plain": [
       "            v1\n",
       "id1           \n",
       "id001  1999081\n",
       "id002  2002870\n",
       "id003  1998389\n",
       "id004  2001091\n",
       "id005  1999770\n",
       "...        ...\n",
       "id096  2001015\n",
       "id097  1995014\n",
       "id098  2000933\n",
       "id099  2003447\n",
       "id100  2001151\n",
       "\n",
       "[100 rows x 1 columns]"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "%%time\n",
    "ddf.groupby(\"id1\", dropna=False, observed=True).agg({\"v1\": \"sum\"}).compute()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6bcf053d-649c-44cf-b341-1d1ffa4da661",
   "metadata": {},
   "source": [
    "### Optimization #2: Use better dtypes"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "9234fff9-f2d4-472e-a2e5-2be5ab435906",
   "metadata": {},
   "outputs": [],
   "source": [
    "better_dtypes = {\n",
    "    \"id1\": \"string[pyarrow]\",\n",
    "    \"id2\": \"string[pyarrow]\",\n",
    "    \"id3\": \"string[pyarrow]\",\n",
    "    \"id4\": \"int64\",\n",
    "    \"id5\": \"int64\",\n",
    "    \"id6\": \"int64\",\n",
    "    \"v1\": \"int64\",\n",
    "    \"v2\": \"int64\",\n",
    "    \"v3\": \"float64\",\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "86af8e08-c3a0-4790-b347-6ab53a119681",
   "metadata": {},
   "outputs": [],
   "source": [
    "ddf = dd.read_csv(\"data/N_1e8_K_1e2_single.csv\", dtype=better_dtypes)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "1d77337a-d37a-4fc2-8052-960099b00b37",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 1min 28s, sys: 22.4 s, total: 1min 50s\n",
      "Wall time: 1min 24s\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "<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>v1</th>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>id1</th>\n",
       "      <th></th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>id001</th>\n",
       "      <td>1999081</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>id002</th>\n",
       "      <td>2002870</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>id003</th>\n",
       "      <td>1998389</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>id004</th>\n",
       "      <td>2001091</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>id005</th>\n",
       "      <td>1999770</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>...</th>\n",
       "      <td>...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>id096</th>\n",
       "      <td>2001015</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>id097</th>\n",
       "      <td>1995014</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>id098</th>\n",
       "      <td>2000933</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>id099</th>\n",
       "      <td>2003447</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>id100</th>\n",
       "      <td>2001151</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "<p>100 rows × 1 columns</p>\n",
       "</div>"
      ],
      "text/plain": [
       "            v1\n",
       "id1           \n",
       "id001  1999081\n",
       "id002  2002870\n",
       "id003  1998389\n",
       "id004  2001091\n",
       "id005  1999770\n",
       "...        ...\n",
       "id096  2001015\n",
       "id097  1995014\n",
       "id098  2000933\n",
       "id099  2003447\n",
       "id100  2001151\n",
       "\n",
       "[100 rows x 1 columns]"
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "%%time\n",
    "ddf.groupby(\"id1\", dropna=False, observed=True).agg({\"v1\": \"sum\"}).compute()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0773ab32-4d21-4fe6-b03b-988433f4ae4b",
   "metadata": {},
   "source": [
    "## Optimization #3: Split dataset into many CSV files"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "c142776c-f08c-442b-9bb7-32d892d5fb14",
   "metadata": {
    "collapsed": true,
    "jupyter": {
     "outputs_hidden": true
    },
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "['/Users/powers/Documents/code/coiled/coiled-resources/blogs/data/csvs/000.part',\n",
       " '/Users/powers/Documents/code/coiled/coiled-resources/blogs/data/csvs/001.part',\n",
       " '/Users/powers/Documents/code/coiled/coiled-resources/blogs/data/csvs/002.part',\n",
       " '/Users/powers/Documents/code/coiled/coiled-resources/blogs/data/csvs/003.part',\n",
       " '/Users/powers/Documents/code/coiled/coiled-resources/blogs/data/csvs/004.part',\n",
       " '/Users/powers/Documents/code/coiled/coiled-resources/blogs/data/csvs/005.part',\n",
       " '/Users/powers/Documents/code/coiled/coiled-resources/blogs/data/csvs/006.part',\n",
       " '/Users/powers/Documents/code/coiled/coiled-resources/blogs/data/csvs/007.part',\n",
       " '/Users/powers/Documents/code/coiled/coiled-resources/blogs/data/csvs/008.part',\n",
       " '/Users/powers/Documents/code/coiled/coiled-resources/blogs/data/csvs/009.part',\n",
       " '/Users/powers/Documents/code/coiled/coiled-resources/blogs/data/csvs/010.part',\n",
       " '/Users/powers/Documents/code/coiled/coiled-resources/blogs/data/csvs/011.part',\n",
       " '/Users/powers/Documents/code/coiled/coiled-resources/blogs/data/csvs/012.part',\n",
       " '/Users/powers/Documents/code/coiled/coiled-resources/blogs/data/csvs/013.part',\n",
       " '/Users/powers/Documents/code/coiled/coiled-resources/blogs/data/csvs/014.part',\n",
       " '/Users/powers/Documents/code/coiled/coiled-resources/blogs/data/csvs/015.part',\n",
       " '/Users/powers/Documents/code/coiled/coiled-resources/blogs/data/csvs/016.part',\n",
       " '/Users/powers/Documents/code/coiled/coiled-resources/blogs/data/csvs/017.part',\n",
       " '/Users/powers/Documents/code/coiled/coiled-resources/blogs/data/csvs/018.part',\n",
       " '/Users/powers/Documents/code/coiled/coiled-resources/blogs/data/csvs/019.part',\n",
       " '/Users/powers/Documents/code/coiled/coiled-resources/blogs/data/csvs/020.part',\n",
       " '/Users/powers/Documents/code/coiled/coiled-resources/blogs/data/csvs/021.part',\n",
       " '/Users/powers/Documents/code/coiled/coiled-resources/blogs/data/csvs/022.part',\n",
       " '/Users/powers/Documents/code/coiled/coiled-resources/blogs/data/csvs/023.part',\n",
       " '/Users/powers/Documents/code/coiled/coiled-resources/blogs/data/csvs/024.part',\n",
       " '/Users/powers/Documents/code/coiled/coiled-resources/blogs/data/csvs/025.part',\n",
       " '/Users/powers/Documents/code/coiled/coiled-resources/blogs/data/csvs/026.part',\n",
       " '/Users/powers/Documents/code/coiled/coiled-resources/blogs/data/csvs/027.part',\n",
       " '/Users/powers/Documents/code/coiled/coiled-resources/blogs/data/csvs/028.part',\n",
       " '/Users/powers/Documents/code/coiled/coiled-resources/blogs/data/csvs/029.part',\n",
       " '/Users/powers/Documents/code/coiled/coiled-resources/blogs/data/csvs/030.part',\n",
       " '/Users/powers/Documents/code/coiled/coiled-resources/blogs/data/csvs/031.part',\n",
       " '/Users/powers/Documents/code/coiled/coiled-resources/blogs/data/csvs/032.part',\n",
       " '/Users/powers/Documents/code/coiled/coiled-resources/blogs/data/csvs/033.part',\n",
       " '/Users/powers/Documents/code/coiled/coiled-resources/blogs/data/csvs/034.part',\n",
       " '/Users/powers/Documents/code/coiled/coiled-resources/blogs/data/csvs/035.part',\n",
       " '/Users/powers/Documents/code/coiled/coiled-resources/blogs/data/csvs/036.part',\n",
       " '/Users/powers/Documents/code/coiled/coiled-resources/blogs/data/csvs/037.part',\n",
       " '/Users/powers/Documents/code/coiled/coiled-resources/blogs/data/csvs/038.part',\n",
       " '/Users/powers/Documents/code/coiled/coiled-resources/blogs/data/csvs/039.part',\n",
       " '/Users/powers/Documents/code/coiled/coiled-resources/blogs/data/csvs/040.part',\n",
       " '/Users/powers/Documents/code/coiled/coiled-resources/blogs/data/csvs/041.part',\n",
       " '/Users/powers/Documents/code/coiled/coiled-resources/blogs/data/csvs/042.part',\n",
       " '/Users/powers/Documents/code/coiled/coiled-resources/blogs/data/csvs/043.part',\n",
       " '/Users/powers/Documents/code/coiled/coiled-resources/blogs/data/csvs/044.part',\n",
       " '/Users/powers/Documents/code/coiled/coiled-resources/blogs/data/csvs/045.part',\n",
       " '/Users/powers/Documents/code/coiled/coiled-resources/blogs/data/csvs/046.part',\n",
       " '/Users/powers/Documents/code/coiled/coiled-resources/blogs/data/csvs/047.part',\n",
       " '/Users/powers/Documents/code/coiled/coiled-resources/blogs/data/csvs/048.part',\n",
       " '/Users/powers/Documents/code/coiled/coiled-resources/blogs/data/csvs/049.part',\n",
       " '/Users/powers/Documents/code/coiled/coiled-resources/blogs/data/csvs/050.part',\n",
       " '/Users/powers/Documents/code/coiled/coiled-resources/blogs/data/csvs/051.part',\n",
       " '/Users/powers/Documents/code/coiled/coiled-resources/blogs/data/csvs/052.part',\n",
       " '/Users/powers/Documents/code/coiled/coiled-resources/blogs/data/csvs/053.part',\n",
       " '/Users/powers/Documents/code/coiled/coiled-resources/blogs/data/csvs/054.part',\n",
       " '/Users/powers/Documents/code/coiled/coiled-resources/blogs/data/csvs/055.part',\n",
       " '/Users/powers/Documents/code/coiled/coiled-resources/blogs/data/csvs/056.part',\n",
       " '/Users/powers/Documents/code/coiled/coiled-resources/blogs/data/csvs/057.part',\n",
       " '/Users/powers/Documents/code/coiled/coiled-resources/blogs/data/csvs/058.part',\n",
       " '/Users/powers/Documents/code/coiled/coiled-resources/blogs/data/csvs/059.part',\n",
       " '/Users/powers/Documents/code/coiled/coiled-resources/blogs/data/csvs/060.part',\n",
       " '/Users/powers/Documents/code/coiled/coiled-resources/blogs/data/csvs/061.part',\n",
       " '/Users/powers/Documents/code/coiled/coiled-resources/blogs/data/csvs/062.part',\n",
       " '/Users/powers/Documents/code/coiled/coiled-resources/blogs/data/csvs/063.part',\n",
       " '/Users/powers/Documents/code/coiled/coiled-resources/blogs/data/csvs/064.part',\n",
       " '/Users/powers/Documents/code/coiled/coiled-resources/blogs/data/csvs/065.part',\n",
       " '/Users/powers/Documents/code/coiled/coiled-resources/blogs/data/csvs/066.part',\n",
       " '/Users/powers/Documents/code/coiled/coiled-resources/blogs/data/csvs/067.part',\n",
       " '/Users/powers/Documents/code/coiled/coiled-resources/blogs/data/csvs/068.part',\n",
       " '/Users/powers/Documents/code/coiled/coiled-resources/blogs/data/csvs/069.part',\n",
       " '/Users/powers/Documents/code/coiled/coiled-resources/blogs/data/csvs/070.part',\n",
       " '/Users/powers/Documents/code/coiled/coiled-resources/blogs/data/csvs/071.part',\n",
       " '/Users/powers/Documents/code/coiled/coiled-resources/blogs/data/csvs/072.part',\n",
       " '/Users/powers/Documents/code/coiled/coiled-resources/blogs/data/csvs/073.part',\n",
       " '/Users/powers/Documents/code/coiled/coiled-resources/blogs/data/csvs/074.part',\n",
       " '/Users/powers/Documents/code/coiled/coiled-resources/blogs/data/csvs/075.part',\n",
       " '/Users/powers/Documents/code/coiled/coiled-resources/blogs/data/csvs/076.part',\n",
       " '/Users/powers/Documents/code/coiled/coiled-resources/blogs/data/csvs/077.part',\n",
       " '/Users/powers/Documents/code/coiled/coiled-resources/blogs/data/csvs/078.part',\n",
       " '/Users/powers/Documents/code/coiled/coiled-resources/blogs/data/csvs/079.part',\n",
       " '/Users/powers/Documents/code/coiled/coiled-resources/blogs/data/csvs/080.part',\n",
       " '/Users/powers/Documents/code/coiled/coiled-resources/blogs/data/csvs/081.part',\n",
       " '/Users/powers/Documents/code/coiled/coiled-resources/blogs/data/csvs/082.part',\n",
       " '/Users/powers/Documents/code/coiled/coiled-resources/blogs/data/csvs/083.part',\n",
       " '/Users/powers/Documents/code/coiled/coiled-resources/blogs/data/csvs/084.part',\n",
       " '/Users/powers/Documents/code/coiled/coiled-resources/blogs/data/csvs/085.part',\n",
       " '/Users/powers/Documents/code/coiled/coiled-resources/blogs/data/csvs/086.part',\n",
       " '/Users/powers/Documents/code/coiled/coiled-resources/blogs/data/csvs/087.part',\n",
       " '/Users/powers/Documents/code/coiled/coiled-resources/blogs/data/csvs/088.part',\n",
       " '/Users/powers/Documents/code/coiled/coiled-resources/blogs/data/csvs/089.part',\n",
       " '/Users/powers/Documents/code/coiled/coiled-resources/blogs/data/csvs/090.part',\n",
       " '/Users/powers/Documents/code/coiled/coiled-resources/blogs/data/csvs/091.part',\n",
       " '/Users/powers/Documents/code/coiled/coiled-resources/blogs/data/csvs/092.part',\n",
       " '/Users/powers/Documents/code/coiled/coiled-resources/blogs/data/csvs/093.part',\n",
       " '/Users/powers/Documents/code/coiled/coiled-resources/blogs/data/csvs/094.part',\n",
       " '/Users/powers/Documents/code/coiled/coiled-resources/blogs/data/csvs/095.part',\n",
       " '/Users/powers/Documents/code/coiled/coiled-resources/blogs/data/csvs/096.part',\n",
       " '/Users/powers/Documents/code/coiled/coiled-resources/blogs/data/csvs/097.part',\n",
       " '/Users/powers/Documents/code/coiled/coiled-resources/blogs/data/csvs/098.part',\n",
       " '/Users/powers/Documents/code/coiled/coiled-resources/blogs/data/csvs/099.part',\n",
       " '/Users/powers/Documents/code/coiled/coiled-resources/blogs/data/csvs/100.part',\n",
       " '/Users/powers/Documents/code/coiled/coiled-resources/blogs/data/csvs/101.part',\n",
       " '/Users/powers/Documents/code/coiled/coiled-resources/blogs/data/csvs/102.part',\n",
       " '/Users/powers/Documents/code/coiled/coiled-resources/blogs/data/csvs/103.part',\n",
       " '/Users/powers/Documents/code/coiled/coiled-resources/blogs/data/csvs/104.part',\n",
       " '/Users/powers/Documents/code/coiled/coiled-resources/blogs/data/csvs/105.part',\n",
       " '/Users/powers/Documents/code/coiled/coiled-resources/blogs/data/csvs/106.part',\n",
       " '/Users/powers/Documents/code/coiled/coiled-resources/blogs/data/csvs/107.part',\n",
       " '/Users/powers/Documents/code/coiled/coiled-resources/blogs/data/csvs/108.part',\n",
       " '/Users/powers/Documents/code/coiled/coiled-resources/blogs/data/csvs/109.part',\n",
       " '/Users/powers/Documents/code/coiled/coiled-resources/blogs/data/csvs/110.part',\n",
       " '/Users/powers/Documents/code/coiled/coiled-resources/blogs/data/csvs/111.part',\n",
       " '/Users/powers/Documents/code/coiled/coiled-resources/blogs/data/csvs/112.part',\n",
       " '/Users/powers/Documents/code/coiled/coiled-resources/blogs/data/csvs/113.part',\n",
       " '/Users/powers/Documents/code/coiled/coiled-resources/blogs/data/csvs/114.part',\n",
       " '/Users/powers/Documents/code/coiled/coiled-resources/blogs/data/csvs/115.part',\n",
       " '/Users/powers/Documents/code/coiled/coiled-resources/blogs/data/csvs/116.part',\n",
       " '/Users/powers/Documents/code/coiled/coiled-resources/blogs/data/csvs/117.part',\n",
       " '/Users/powers/Documents/code/coiled/coiled-resources/blogs/data/csvs/118.part',\n",
       " '/Users/powers/Documents/code/coiled/coiled-resources/blogs/data/csvs/119.part',\n",
       " '/Users/powers/Documents/code/coiled/coiled-resources/blogs/data/csvs/120.part',\n",
       " '/Users/powers/Documents/code/coiled/coiled-resources/blogs/data/csvs/121.part',\n",
       " '/Users/powers/Documents/code/coiled/coiled-resources/blogs/data/csvs/122.part',\n",
       " '/Users/powers/Documents/code/coiled/coiled-resources/blogs/data/csvs/123.part',\n",
       " '/Users/powers/Documents/code/coiled/coiled-resources/blogs/data/csvs/124.part',\n",
       " '/Users/powers/Documents/code/coiled/coiled-resources/blogs/data/csvs/125.part',\n",
       " '/Users/powers/Documents/code/coiled/coiled-resources/blogs/data/csvs/126.part',\n",
       " '/Users/powers/Documents/code/coiled/coiled-resources/blogs/data/csvs/127.part',\n",
       " '/Users/powers/Documents/code/coiled/coiled-resources/blogs/data/csvs/128.part',\n",
       " '/Users/powers/Documents/code/coiled/coiled-resources/blogs/data/csvs/129.part',\n",
       " '/Users/powers/Documents/code/coiled/coiled-resources/blogs/data/csvs/130.part',\n",
       " '/Users/powers/Documents/code/coiled/coiled-resources/blogs/data/csvs/131.part',\n",
       " '/Users/powers/Documents/code/coiled/coiled-resources/blogs/data/csvs/132.part',\n",
       " '/Users/powers/Documents/code/coiled/coiled-resources/blogs/data/csvs/133.part',\n",
       " '/Users/powers/Documents/code/coiled/coiled-resources/blogs/data/csvs/134.part',\n",
       " '/Users/powers/Documents/code/coiled/coiled-resources/blogs/data/csvs/135.part',\n",
       " '/Users/powers/Documents/code/coiled/coiled-resources/blogs/data/csvs/136.part',\n",
       " '/Users/powers/Documents/code/coiled/coiled-resources/blogs/data/csvs/137.part',\n",
       " '/Users/powers/Documents/code/coiled/coiled-resources/blogs/data/csvs/138.part',\n",
       " '/Users/powers/Documents/code/coiled/coiled-resources/blogs/data/csvs/139.part',\n",
       " '/Users/powers/Documents/code/coiled/coiled-resources/blogs/data/csvs/140.part',\n",
       " '/Users/powers/Documents/code/coiled/coiled-resources/blogs/data/csvs/141.part',\n",
       " '/Users/powers/Documents/code/coiled/coiled-resources/blogs/data/csvs/142.part',\n",
       " '/Users/powers/Documents/code/coiled/coiled-resources/blogs/data/csvs/143.part',\n",
       " '/Users/powers/Documents/code/coiled/coiled-resources/blogs/data/csvs/144.part',\n",
       " '/Users/powers/Documents/code/coiled/coiled-resources/blogs/data/csvs/145.part',\n",
       " '/Users/powers/Documents/code/coiled/coiled-resources/blogs/data/csvs/146.part',\n",
       " '/Users/powers/Documents/code/coiled/coiled-resources/blogs/data/csvs/147.part',\n",
       " '/Users/powers/Documents/code/coiled/coiled-resources/blogs/data/csvs/148.part',\n",
       " '/Users/powers/Documents/code/coiled/coiled-resources/blogs/data/csvs/149.part',\n",
       " '/Users/powers/Documents/code/coiled/coiled-resources/blogs/data/csvs/150.part',\n",
       " '/Users/powers/Documents/code/coiled/coiled-resources/blogs/data/csvs/151.part',\n",
       " '/Users/powers/Documents/code/coiled/coiled-resources/blogs/data/csvs/152.part',\n",
       " '/Users/powers/Documents/code/coiled/coiled-resources/blogs/data/csvs/153.part',\n",
       " '/Users/powers/Documents/code/coiled/coiled-resources/blogs/data/csvs/154.part',\n",
       " '/Users/powers/Documents/code/coiled/coiled-resources/blogs/data/csvs/155.part',\n",
       " '/Users/powers/Documents/code/coiled/coiled-resources/blogs/data/csvs/156.part',\n",
       " '/Users/powers/Documents/code/coiled/coiled-resources/blogs/data/csvs/157.part',\n",
       " '/Users/powers/Documents/code/coiled/coiled-resources/blogs/data/csvs/158.part',\n",
       " '/Users/powers/Documents/code/coiled/coiled-resources/blogs/data/csvs/159.part',\n",
       " '/Users/powers/Documents/code/coiled/coiled-resources/blogs/data/csvs/160.part',\n",
       " '/Users/powers/Documents/code/coiled/coiled-resources/blogs/data/csvs/161.part']"
      ]
     },
     "execution_count": 14,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "ddf.repartition(partition_size=\"100MB\").to_csv(\"data/csvs\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "0060aba7-6ed2-4464-8ff1-4ee455bb4291",
   "metadata": {},
   "outputs": [],
   "source": [
    "ddf = dd.read_csv(\"data/csvs/*\", dtype=better_dtypes)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "e175f02e-6b75-4d34-ab8f-77e79930fe1f",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 1min 36s, sys: 10.8 s, total: 1min 46s\n",
      "Wall time: 1min 4s\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "<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>v1</th>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>id1</th>\n",
       "      <th></th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>id001</th>\n",
       "      <td>1999081</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>id002</th>\n",
       "      <td>2002870</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>id003</th>\n",
       "      <td>1998389</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>id004</th>\n",
       "      <td>2001091</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>id005</th>\n",
       "      <td>1999770</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>...</th>\n",
       "      <td>...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>id096</th>\n",
       "      <td>2001015</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>id097</th>\n",
       "      <td>1995014</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>id098</th>\n",
       "      <td>2000933</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>id099</th>\n",
       "      <td>2003447</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>id100</th>\n",
       "      <td>2001151</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "<p>100 rows × 1 columns</p>\n",
       "</div>"
      ],
      "text/plain": [
       "            v1\n",
       "id1           \n",
       "id001  1999081\n",
       "id002  2002870\n",
       "id003  1998389\n",
       "id004  2001091\n",
       "id005  1999770\n",
       "...        ...\n",
       "id096  2001015\n",
       "id097  1995014\n",
       "id098  2000933\n",
       "id099  2003447\n",
       "id100  2001151\n",
       "\n",
       "[100 rows x 1 columns]"
      ]
     },
     "execution_count": 16,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "%%time\n",
    "ddf.groupby(\"id1\", dropna=False, observed=True).agg({\"v1\": \"sum\"}).compute()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "aca771a8-4f87-4a62-b6fc-fef343a3b571",
   "metadata": {},
   "source": [
    "## Optimization #4: Split dataset into multiple uncompressed Parquet files"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "ab970a4d-6026-4aca-affb-65ee61faf146",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(None,)"
      ]
     },
     "execution_count": 17,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "ddf.to_parquet(\"data/parquet\", engine=\"pyarrow\", compression=None)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "c28ac347-df8e-4cb7-8a22-8d41def76a6b",
   "metadata": {},
   "outputs": [],
   "source": [
    "ddf = dd.read_parquet(\"data/parquet\", engine=\"pyarrow\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "50a494d9-ff38-48bb-b1c6-48f0d07adab5",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 40.7 s, sys: 13.3 s, total: 54 s\n",
      "Wall time: 41.3 s\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "<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>v1</th>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>id1</th>\n",
       "      <th></th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>id001</th>\n",
       "      <td>1999081</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>id002</th>\n",
       "      <td>2002870</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>id003</th>\n",
       "      <td>1998389</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>id004</th>\n",
       "      <td>2001091</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>id005</th>\n",
       "      <td>1999770</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>...</th>\n",
       "      <td>...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>id096</th>\n",
       "      <td>2001015</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>id097</th>\n",
       "      <td>1995014</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>id098</th>\n",
       "      <td>2000933</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>id099</th>\n",
       "      <td>2003447</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>id100</th>\n",
       "      <td>2001151</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "<p>100 rows × 1 columns</p>\n",
       "</div>"
      ],
      "text/plain": [
       "            v1\n",
       "id1           \n",
       "id001  1999081\n",
       "id002  2002870\n",
       "id003  1998389\n",
       "id004  2001091\n",
       "id005  1999770\n",
       "...        ...\n",
       "id096  2001015\n",
       "id097  1995014\n",
       "id098  2000933\n",
       "id099  2003447\n",
       "id100  2001151\n",
       "\n",
       "[100 rows x 1 columns]"
      ]
     },
     "execution_count": 19,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "%%time\n",
    "ddf.groupby(\"id1\", dropna=False, observed=True).agg({\"v1\": \"sum\"}).compute()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6905394b-8107-4ea2-83b9-11e826f8ea77",
   "metadata": {},
   "source": [
    "## Optimization #5: Using Snappy compressed Parquet"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "id": "30da4d9f-df97-4873-819f-6ca62bf99ea8",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(None,)"
      ]
     },
     "execution_count": 20,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "ddf.to_parquet(\"data/snappy-parquet\", engine=\"pyarrow\", compression=\"snappy\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "id": "65e22e26-afae-4e2b-a7f0-73afa095cd2a",
   "metadata": {},
   "outputs": [],
   "source": [
    "ddf = dd.read_parquet(\"data/snappy-parquet\", engine=\"pyarrow\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "id": "14895f9c-4e7a-4c2a-8ab7-f6e4b24295a2",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 44.8 s, sys: 14.8 s, total: 59.6 s\n",
      "Wall time: 42.4 s\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "<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>v1</th>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>id1</th>\n",
       "      <th></th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>id001</th>\n",
       "      <td>1999081</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>id002</th>\n",
       "      <td>2002870</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>id003</th>\n",
       "      <td>1998389</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>id004</th>\n",
       "      <td>2001091</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>id005</th>\n",
       "      <td>1999770</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>...</th>\n",
       "      <td>...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>id096</th>\n",
       "      <td>2001015</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>id097</th>\n",
       "      <td>1995014</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>id098</th>\n",
       "      <td>2000933</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>id099</th>\n",
       "      <td>2003447</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>id100</th>\n",
       "      <td>2001151</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "<p>100 rows × 1 columns</p>\n",
       "</div>"
      ],
      "text/plain": [
       "            v1\n",
       "id1           \n",
       "id001  1999081\n",
       "id002  2002870\n",
       "id003  1998389\n",
       "id004  2001091\n",
       "id005  1999770\n",
       "...        ...\n",
       "id096  2001015\n",
       "id097  1995014\n",
       "id098  2000933\n",
       "id099  2003447\n",
       "id100  2001151\n",
       "\n",
       "[100 rows x 1 columns]"
      ]
     },
     "execution_count": 22,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "%%time\n",
    "ddf.groupby(\"id1\", dropna=False, observed=True).agg({\"v1\": \"sum\"}).compute()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "490ad913-1178-4745-a1cf-21bb7a2340ac",
   "metadata": {},
   "source": [
    "## Optimization #6: Leverage Parquet column pruning"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "id": "52a7889b-bfdd-48e7-90f4-f422f062c602",
   "metadata": {},
   "outputs": [],
   "source": [
    "ddf = dd.read_parquet(\"data/snappy-parquet\", engine=\"pyarrow\", columns=[\"id1\", \"v1\"])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "id": "6f356a8d-bbf1-41bd-96de-3dfd5c633737",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 21.4 s, sys: 958 ms, total: 22.3 s\n",
      "Wall time: 16.8 s\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "<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>v1</th>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>id1</th>\n",
       "      <th></th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>id001</th>\n",
       "      <td>1999081</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>id002</th>\n",
       "      <td>2002870</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>id003</th>\n",
       "      <td>1998389</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>id004</th>\n",
       "      <td>2001091</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>id005</th>\n",
       "      <td>1999770</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>...</th>\n",
       "      <td>...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>id096</th>\n",
       "      <td>2001015</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>id097</th>\n",
       "      <td>1995014</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>id098</th>\n",
       "      <td>2000933</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>id099</th>\n",
       "      <td>2003447</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>id100</th>\n",
       "      <td>2001151</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "<p>100 rows × 1 columns</p>\n",
       "</div>"
      ],
      "text/plain": [
       "            v1\n",
       "id1           \n",
       "id001  1999081\n",
       "id002  2002870\n",
       "id003  1998389\n",
       "id004  2001091\n",
       "id005  1999770\n",
       "...        ...\n",
       "id096  2001015\n",
       "id097  1995014\n",
       "id098  2000933\n",
       "id099  2003447\n",
       "id100  2001151\n",
       "\n",
       "[100 rows x 1 columns]"
      ]
     },
     "execution_count": 24,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "%%time\n",
    "ddf.groupby(\"id1\", dropna=False, observed=True).agg({\"v1\": \"sum\"}).compute()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1eae5019-3dbc-4c00-aeb1-b0a0b865b120",
   "metadata": {},
   "source": [
    "## Cleanup (optional)\n",
    "\n",
    "You don't need to keep the CSV data file you downloaded for this example on your machine anymore if you're not going to use it again.  I personally love having this data on my machine for experimentation purposes, but you can delete it if you don't need it anymore."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e406130d-bf3a-4d80-881c-b8d0335c7cb6",
   "metadata": {},
   "outputs": [],
   "source": [
    "!rm -rf data"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1e9fdbff-49b4-4589-bb83-ad740cd9e35f",
   "metadata": {},
   "source": [
    "## Querying a large dataset with the cloud"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "5bbc41b4-1b4e-4e74-8fab-1d57ac2d9d72",
   "metadata": {},
   "outputs": [],
   "source": [
    "import coiled\n",
    "import dask\n",
    "import dask.dataframe as dd"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "c96de150-8a57-4f42-8e48-a98c86f55910",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "197b5296d9334698a13a8fc8d24cb390",
       "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=\"powers-crt-004\",\n",
    "    software=\"crt-004\",\n",
    "    n_workers=16,\n",
    "    worker_vm_types=[\"i3.large\"],\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "a5623f4a-c347-48b2-a339-2e0243976d22",
   "metadata": {},
   "outputs": [],
   "source": [
    "client = dask.distributed.Client(cluster)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "51daf6f5-0d35-419c-8ec8-b572c5ad497b",
   "metadata": {},
   "outputs": [],
   "source": [
    "ddf = dd.read_parquet(\n",
    "    \"s3://coiled-datasets/h2o-benchmark/N_1e9_K_1e2_parquet\",\n",
    "    storage_options={\"anon\": True, \"use_ssl\": True},\n",
    "    columns=[\"id1\", \"id2\", \"v1\"],\n",
    "    engine=\"pyarrow\",\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "f0c59f6d-c62e-4b5d-8265-02eb1ecbbb53",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "1000000000"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "len(ddf)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "52c26327-9bf3-48cb-81c9-8d46b4198ec5",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 231 ms, sys: 26.5 ms, total: 257 ms\n",
      "Wall time: 29.5 s\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "<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>v1</th>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>id1</th>\n",
       "      <th></th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>id086</th>\n",
       "      <td>20002086</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>id090</th>\n",
       "      <td>19984380</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>id089</th>\n",
       "      <td>19994350</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>id038</th>\n",
       "      <td>20002100</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>id077</th>\n",
       "      <td>19990122</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>...</th>\n",
       "      <td>...</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>id027</th>\n",
       "      <td>20002119</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>id009</th>\n",
       "      <td>19989071</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>id013</th>\n",
       "      <td>19997400</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>id097</th>\n",
       "      <td>19997973</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>id015</th>\n",
       "      <td>19988461</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "<p>100 rows × 1 columns</p>\n",
       "</div>"
      ],
      "text/plain": [
       "             v1\n",
       "id1            \n",
       "id086  20002086\n",
       "id090  19984380\n",
       "id089  19994350\n",
       "id038  20002100\n",
       "id077  19990122\n",
       "...         ...\n",
       "id027  20002119\n",
       "id009  19989071\n",
       "id013  19997400\n",
       "id097  19997973\n",
       "id015  19988461\n",
       "\n",
       "[100 rows x 1 columns]"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "%%time\n",
    "ddf.groupby(\"id1\", dropna=False, observed=True).agg({\"v1\": \"sum\"}).compute()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d12e0c01-8201-491c-895d-13000223b446",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python [conda env:crt-004] *",
   "language": "python",
   "name": "conda-env-crt-004-py"
  },
  "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
}