2_dask_laptop.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Getting up to speed with Dask\n",
    "\n",
    "## Part 2: Dask to the rescue!\n",
    "\n",
    "We will do the same analysis as Part 1 with the same laptop-sized machine, except we will be using a rougly 10x larger dataset\n",
    "\n",
    "AWS EC2 instance type: r5.xlarge (2 CPU, 16GB RAM)\n",
    "\n",
    "Dask will work without intializing a distributed backend, but it's [recommended to do so](https://docs.dask.org/en/latest/setup/single-distributed.html) even on a laptop. To initialize, start a `Client` like so:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<table style=\"border: 2px solid white;\">\n",
       "<tr>\n",
       "<td style=\"vertical-align: top; border: 0px solid white\">\n",
       "<h3 style=\"text-align: left;\">Client</h3>\n",
       "<ul style=\"text-align: left; list-style: none; margin: 0; padding: 0;\">\n",
       "  <li><b>Scheduler: </b>tcp://127.0.0.1:46179</li>\n",
       "  <li><b>Dashboard: </b><a href='http://127.0.0.1:8787/status' target='_blank'>http://127.0.0.1:8787/status</a></li>\n",
       "</ul>\n",
       "</td>\n",
       "<td style=\"vertical-align: top; border: 0px solid white\">\n",
       "<h3 style=\"text-align: left;\">Cluster</h3>\n",
       "<ul style=\"text-align: left; list-style:none; margin: 0; padding: 0;\">\n",
       "  <li><b>Workers: </b>2</li>\n",
       "  <li><b>Cores: </b>2</li>\n",
       "  <li><b>Memory: </b>16.25 GB</li>\n",
       "</ul>\n",
       "</td>\n",
       "</tr>\n",
       "</table>"
      ],
      "text/plain": [
       "<Client: 'tcp://127.0.0.1:46179' processes=2 threads=2, memory=16.25 GB>"
      ]
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from dask.distributed import Client\n",
    "\n",
    "client = Client()\n",
    "client"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pandas as pd\n",
    "import numpy as np\n",
    "import datetime\n",
    "import s3fs\n",
    "from pathlib import Path\n",
    "\n",
    "# change this path if you changed in Part 0\n",
    "data_path = Path('data')\n",
    "seed = 42"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Load and explore data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "taxi_dtypes = {\n",
    "    'store_and_fwd_flag': str,\n",
    "    'RatecodeID': 'float64',\n",
    "    'VendorID': 'float64',\n",
    "    'passenger_count': 'float64',\n",
    "    'payment_type': 'float64',\n",
    "}\n",
    "\n",
    "def load_csv(f):\n",
    "    \"\"\" Same as Part 1\"\"\"\n",
    "    return pd.read_csv(\n",
    "        f,\n",
    "        dtype=taxi_dtypes, \n",
    "        parse_dates=['tpep_pickup_datetime', 'tpep_dropoff_datetime'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "# new dask imports\n",
    "from dask import delayed\n",
    "import dask.dataframe as dd"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## The hard way\n",
    "\n",
    "We could parallelize our for-loop that blew up in Part 1 using [dask.delayed](https://docs.dask.org/en/latest/delayed.html).\n",
    "\n",
    "NOTE: There is a better way below! This is just for educational purposes."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 436 ms, sys: 64 ms, total: 500 ms\n",
      "Wall time: 15.2 s\n"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "\n",
    "dfs = []\n",
    "for f in data_path.iterdir():\n",
    "    df = delayed(load_csv)(f)  # see the delayed() call\n",
    "    dfs.append(df)\n",
    "taxi = dd.from_delayed(dfs)  # from_delayed rather than pd.concat"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[Delayed('load_csv-863adce6-516b-46c6-b156-fdea6e5a113f'),\n",
       " Delayed('load_csv-0b46a36e-f5ff-4755-a3df-463e3dafba00'),\n",
       " Delayed('load_csv-556f630c-7049-4274-8328-10c83792566e'),\n",
       " Delayed('load_csv-eafa309f-5067-422f-9441-e77b9c6f9034'),\n",
       " Delayed('load_csv-ef165669-ca8f-4db4-9108-76ad14adf30f'),\n",
       " Delayed('load_csv-26533d43-621b-47d0-b99d-9581a00f689a'),\n",
       " Delayed('load_csv-5cca96e7-58b3-4fa5-ac9e-c0a104a1621b'),\n",
       " Delayed('load_csv-458953cf-ddc5-4f2b-a03c-882b3c5b12db'),\n",
       " Delayed('load_csv-7c92b38c-6f20-4ea9-a945-66936763b159'),\n",
       " Delayed('load_csv-4f8194eb-a35e-47d0-9bad-bfb35d915164'),\n",
       " Delayed('load_csv-9f0ba641-eb6d-4691-b9dd-be8a9c35d325'),\n",
       " Delayed('load_csv-5cf9bd5f-65e2-4505-a0b5-5c39d0fc1da0')]"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "dfs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 2.04 s, sys: 2.62 s, total: 4.65 s\n",
      "Wall time: 18.7 s\n"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "df0 = dfs[0].compute()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "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>VendorID</th>\n",
       "      <th>tpep_pickup_datetime</th>\n",
       "      <th>tpep_dropoff_datetime</th>\n",
       "      <th>passenger_count</th>\n",
       "      <th>trip_distance</th>\n",
       "      <th>RatecodeID</th>\n",
       "      <th>store_and_fwd_flag</th>\n",
       "      <th>PULocationID</th>\n",
       "      <th>DOLocationID</th>\n",
       "      <th>payment_type</th>\n",
       "      <th>fare_amount</th>\n",
       "      <th>extra</th>\n",
       "      <th>mta_tax</th>\n",
       "      <th>tip_amount</th>\n",
       "      <th>tolls_amount</th>\n",
       "      <th>improvement_surcharge</th>\n",
       "      <th>total_amount</th>\n",
       "      <th>congestion_surcharge</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>1.0</td>\n",
       "      <td>2019-12-01 00:26:58</td>\n",
       "      <td>2019-12-01 00:41:45</td>\n",
       "      <td>1.0</td>\n",
       "      <td>4.2</td>\n",
       "      <td>1.0</td>\n",
       "      <td>N</td>\n",
       "      <td>142</td>\n",
       "      <td>116</td>\n",
       "      <td>2.0</td>\n",
       "      <td>14.5</td>\n",
       "      <td>3.0</td>\n",
       "      <td>0.5</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.3</td>\n",
       "      <td>18.3</td>\n",
       "      <td>2.5</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>1.0</td>\n",
       "      <td>2019-12-01 00:12:08</td>\n",
       "      <td>2019-12-01 00:12:14</td>\n",
       "      <td>1.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>1.0</td>\n",
       "      <td>N</td>\n",
       "      <td>145</td>\n",
       "      <td>145</td>\n",
       "      <td>2.0</td>\n",
       "      <td>2.5</td>\n",
       "      <td>0.5</td>\n",
       "      <td>0.5</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.3</td>\n",
       "      <td>3.8</td>\n",
       "      <td>0.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>1.0</td>\n",
       "      <td>2019-12-01 00:25:53</td>\n",
       "      <td>2019-12-01 00:26:04</td>\n",
       "      <td>1.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>1.0</td>\n",
       "      <td>N</td>\n",
       "      <td>145</td>\n",
       "      <td>145</td>\n",
       "      <td>2.0</td>\n",
       "      <td>2.5</td>\n",
       "      <td>0.5</td>\n",
       "      <td>0.5</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.3</td>\n",
       "      <td>3.8</td>\n",
       "      <td>0.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>1.0</td>\n",
       "      <td>2019-12-01 00:12:03</td>\n",
       "      <td>2019-12-01 00:33:19</td>\n",
       "      <td>2.0</td>\n",
       "      <td>9.4</td>\n",
       "      <td>1.0</td>\n",
       "      <td>N</td>\n",
       "      <td>138</td>\n",
       "      <td>25</td>\n",
       "      <td>1.0</td>\n",
       "      <td>28.5</td>\n",
       "      <td>0.5</td>\n",
       "      <td>0.5</td>\n",
       "      <td>10.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.3</td>\n",
       "      <td>39.8</td>\n",
       "      <td>0.0</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>1.0</td>\n",
       "      <td>2019-12-01 00:05:27</td>\n",
       "      <td>2019-12-01 00:16:32</td>\n",
       "      <td>2.0</td>\n",
       "      <td>1.6</td>\n",
       "      <td>1.0</td>\n",
       "      <td>N</td>\n",
       "      <td>161</td>\n",
       "      <td>237</td>\n",
       "      <td>2.0</td>\n",
       "      <td>9.0</td>\n",
       "      <td>3.0</td>\n",
       "      <td>0.5</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.3</td>\n",
       "      <td>12.8</td>\n",
       "      <td>2.5</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "   VendorID tpep_pickup_datetime tpep_dropoff_datetime  passenger_count  \\\n",
       "0       1.0  2019-12-01 00:26:58   2019-12-01 00:41:45              1.0   \n",
       "1       1.0  2019-12-01 00:12:08   2019-12-01 00:12:14              1.0   \n",
       "2       1.0  2019-12-01 00:25:53   2019-12-01 00:26:04              1.0   \n",
       "3       1.0  2019-12-01 00:12:03   2019-12-01 00:33:19              2.0   \n",
       "4       1.0  2019-12-01 00:05:27   2019-12-01 00:16:32              2.0   \n",
       "\n",
       "   trip_distance  RatecodeID store_and_fwd_flag  PULocationID  DOLocationID  \\\n",
       "0            4.2         1.0                  N           142           116   \n",
       "1            0.0         1.0                  N           145           145   \n",
       "2            0.0         1.0                  N           145           145   \n",
       "3            9.4         1.0                  N           138            25   \n",
       "4            1.6         1.0                  N           161           237   \n",
       "\n",
       "   payment_type  fare_amount  extra  mta_tax  tip_amount  tolls_amount  \\\n",
       "0           2.0         14.5    3.0      0.5         0.0           0.0   \n",
       "1           2.0          2.5    0.5      0.5         0.0           0.0   \n",
       "2           2.0          2.5    0.5      0.5         0.0           0.0   \n",
       "3           1.0         28.5    0.5      0.5        10.0           0.0   \n",
       "4           2.0          9.0    3.0      0.5         0.0           0.0   \n",
       "\n",
       "   improvement_surcharge  total_amount  congestion_surcharge  \n",
       "0                    0.3          18.3                   2.5  \n",
       "1                    0.3           3.8                   0.0  \n",
       "2                    0.3           3.8                   0.0  \n",
       "3                    0.3          39.8                   0.0  \n",
       "4                    0.3          12.8                   2.5  "
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "df0.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "pandas.core.frame.DataFrame"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "type(df0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "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>VendorID</th>\n",
       "      <th>tpep_pickup_datetime</th>\n",
       "      <th>tpep_dropoff_datetime</th>\n",
       "      <th>passenger_count</th>\n",
       "      <th>trip_distance</th>\n",
       "      <th>RatecodeID</th>\n",
       "      <th>store_and_fwd_flag</th>\n",
       "      <th>PULocationID</th>\n",
       "      <th>DOLocationID</th>\n",
       "      <th>payment_type</th>\n",
       "      <th>fare_amount</th>\n",
       "      <th>extra</th>\n",
       "      <th>mta_tax</th>\n",
       "      <th>tip_amount</th>\n",
       "      <th>tolls_amount</th>\n",
       "      <th>improvement_surcharge</th>\n",
       "      <th>total_amount</th>\n",
       "      <th>congestion_surcharge</th>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>npartitions=12</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",
       "      <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>float64</td>\n",
       "      <td>datetime64[ns]</td>\n",
       "      <td>datetime64[ns]</td>\n",
       "      <td>float64</td>\n",
       "      <td>float64</td>\n",
       "      <td>float64</td>\n",
       "      <td>object</td>\n",
       "      <td>int64</td>\n",
       "      <td>int64</td>\n",
       "      <td>float64</td>\n",
       "      <td>float64</td>\n",
       "      <td>float64</td>\n",
       "      <td>float64</td>\n",
       "      <td>float64</td>\n",
       "      <td>float64</td>\n",
       "      <td>float64</td>\n",
       "      <td>float64</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",
       "      <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",
       "      <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",
       "      <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",
       "      <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: from-delayed, 24 tasks</div>"
      ],
      "text/plain": [
       "Dask DataFrame Structure:\n",
       "               VendorID tpep_pickup_datetime tpep_dropoff_datetime passenger_count trip_distance RatecodeID store_and_fwd_flag PULocationID DOLocationID payment_type fare_amount    extra  mta_tax tip_amount tolls_amount improvement_surcharge total_amount congestion_surcharge\n",
       "npartitions=12                                                                                                                                                                                                                                                                     \n",
       "                float64       datetime64[ns]        datetime64[ns]         float64       float64    float64             object        int64        int64      float64     float64  float64  float64    float64      float64               float64      float64              float64\n",
       "                    ...                  ...                   ...             ...           ...        ...                ...          ...          ...          ...         ...      ...      ...        ...          ...                   ...          ...                  ...\n",
       "...                 ...                  ...                   ...             ...           ...        ...                ...          ...          ...          ...         ...      ...      ...        ...          ...                   ...          ...                  ...\n",
       "                    ...                  ...                   ...             ...           ...        ...                ...          ...          ...          ...         ...      ...      ...        ...          ...                   ...          ...                  ...\n",
       "                    ...                  ...                   ...             ...           ...        ...                ...          ...          ...          ...         ...      ...      ...        ...          ...                   ...          ...                  ...\n",
       "Dask Name: from-delayed, 24 tasks"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "taxi"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 4.8 s, sys: 139 ms, total: 4.94 s\n",
      "Wall time: 2min 22s\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "84399019"
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "%%time\n",
    "len(taxi)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Okay, let's do it the better way\n",
    "\n",
    "`dask.dataframe.read_csv` has the same API as `pandas.dataframe.read_csv`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 20.5 ms, sys: 950 µs, total: 21.4 ms\n",
      "Wall time: 21.1 ms\n"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "\n",
    "taxi = dd.read_csv(\n",
    "    'data/yellow_tripdata_2019-*.csv',\n",
    "    dtype=taxi_dtypes, \n",
    "    parse_dates=['tpep_pickup_datetime', 'tpep_dropoff_datetime'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 4.89 s, sys: 400 ms, total: 5.29 s\n",
      "Wall time: 2min 17s\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "84399019"
      ]
     },
     "execution_count": 14,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "%%time\n",
    "len(taxi)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<br>\n",
    "\n",
    "Note the calls to `.compute()`. Dask executes _lazily_ meaning it won't do something unless you explicitly tell it to. Peek at the dashboard while these are executing to see the progress!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 4.95 s, sys: 624 ms, total: 5.58 s\n",
      "Wall time: 2min 26s\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "16.367014316"
      ]
     },
     "execution_count": 15,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "%%time\n",
    "taxi.memory_usage(deep=True).sum().compute() / 1e9"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<br>\n",
    "This takes a while"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 3min 9s, sys: 10.4 s, total: 3min 19s\n",
      "Wall time: 17min 55s\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>count</th>\n",
       "      <th>mean</th>\n",
       "      <th>std</th>\n",
       "      <th>min</th>\n",
       "      <th>25%</th>\n",
       "      <th>50%</th>\n",
       "      <th>75%</th>\n",
       "      <th>max</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>VendorID</th>\n",
       "      <td>84152418.0</td>\n",
       "      <td>1.645</td>\n",
       "      <td>0.498</td>\n",
       "      <td>1.00</td>\n",
       "      <td>1.00</td>\n",
       "      <td>2.00</td>\n",
       "      <td>2.00</td>\n",
       "      <td>4.00</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>passenger_count</th>\n",
       "      <td>84152418.0</td>\n",
       "      <td>1.563</td>\n",
       "      <td>1.208</td>\n",
       "      <td>0.00</td>\n",
       "      <td>1.00</td>\n",
       "      <td>1.00</td>\n",
       "      <td>2.00</td>\n",
       "      <td>9.00</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>trip_distance</th>\n",
       "      <td>84399019.0</td>\n",
       "      <td>3.001</td>\n",
       "      <td>8.091</td>\n",
       "      <td>-37264.53</td>\n",
       "      <td>1.07</td>\n",
       "      <td>1.93</td>\n",
       "      <td>8.82</td>\n",
       "      <td>45977.22</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>RatecodeID</th>\n",
       "      <td>84152418.0</td>\n",
       "      <td>1.061</td>\n",
       "      <td>0.760</td>\n",
       "      <td>1.00</td>\n",
       "      <td>1.00</td>\n",
       "      <td>1.00</td>\n",
       "      <td>1.00</td>\n",
       "      <td>99.00</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>PULocationID</th>\n",
       "      <td>84399019.0</td>\n",
       "      <td>163.158</td>\n",
       "      <td>66.016</td>\n",
       "      <td>1.00</td>\n",
       "      <td>132.00</td>\n",
       "      <td>162.00</td>\n",
       "      <td>234.00</td>\n",
       "      <td>265.00</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>DOLocationID</th>\n",
       "      <td>84399019.0</td>\n",
       "      <td>161.353</td>\n",
       "      <td>70.251</td>\n",
       "      <td>1.00</td>\n",
       "      <td>116.00</td>\n",
       "      <td>163.00</td>\n",
       "      <td>236.00</td>\n",
       "      <td>265.00</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>payment_type</th>\n",
       "      <td>84152418.0</td>\n",
       "      <td>1.289</td>\n",
       "      <td>0.479</td>\n",
       "      <td>1.00</td>\n",
       "      <td>1.00</td>\n",
       "      <td>1.00</td>\n",
       "      <td>2.00</td>\n",
       "      <td>5.00</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>fare_amount</th>\n",
       "      <td>84399019.0</td>\n",
       "      <td>13.344</td>\n",
       "      <td>174.375</td>\n",
       "      <td>-1856.00</td>\n",
       "      <td>7.00</td>\n",
       "      <td>11.00</td>\n",
       "      <td>32.04</td>\n",
       "      <td>943274.80</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>extra</th>\n",
       "      <td>84399019.0</td>\n",
       "      <td>1.087</td>\n",
       "      <td>1.249</td>\n",
       "      <td>-60.00</td>\n",
       "      <td>0.00</td>\n",
       "      <td>1.00</td>\n",
       "      <td>3.00</td>\n",
       "      <td>535.38</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>mta_tax</th>\n",
       "      <td>84399019.0</td>\n",
       "      <td>0.495</td>\n",
       "      <td>0.067</td>\n",
       "      <td>-0.50</td>\n",
       "      <td>0.50</td>\n",
       "      <td>0.50</td>\n",
       "      <td>0.50</td>\n",
       "      <td>212.42</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>tip_amount</th>\n",
       "      <td>84399019.0</td>\n",
       "      <td>2.195</td>\n",
       "      <td>15.657</td>\n",
       "      <td>-221.00</td>\n",
       "      <td>0.00</td>\n",
       "      <td>2.06</td>\n",
       "      <td>3.26</td>\n",
       "      <td>141492.02</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>tolls_amount</th>\n",
       "      <td>84399019.0</td>\n",
       "      <td>0.383</td>\n",
       "      <td>1.817</td>\n",
       "      <td>-70.00</td>\n",
       "      <td>0.00</td>\n",
       "      <td>0.00</td>\n",
       "      <td>0.00</td>\n",
       "      <td>3288.00</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>improvement_surcharge</th>\n",
       "      <td>84399019.0</td>\n",
       "      <td>0.299</td>\n",
       "      <td>0.028</td>\n",
       "      <td>-0.30</td>\n",
       "      <td>0.30</td>\n",
       "      <td>0.30</td>\n",
       "      <td>0.30</td>\n",
       "      <td>1.00</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>total_amount</th>\n",
       "      <td>84399019.0</td>\n",
       "      <td>19.124</td>\n",
       "      <td>184.087</td>\n",
       "      <td>-1871.80</td>\n",
       "      <td>11.76</td>\n",
       "      <td>15.80</td>\n",
       "      <td>38.16</td>\n",
       "      <td>1084772.17</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>congestion_surcharge</th>\n",
       "      <td>79543038.0</td>\n",
       "      <td>2.188</td>\n",
       "      <td>0.837</td>\n",
       "      <td>-2.50</td>\n",
       "      <td>2.50</td>\n",
       "      <td>2.50</td>\n",
       "      <td>2.50</td>\n",
       "      <td>4.50</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "                            count     mean      std       min     25%     50%  \\\n",
       "VendorID               84152418.0    1.645    0.498      1.00    1.00    2.00   \n",
       "passenger_count        84152418.0    1.563    1.208      0.00    1.00    1.00   \n",
       "trip_distance          84399019.0    3.001    8.091 -37264.53    1.07    1.93   \n",
       "RatecodeID             84152418.0    1.061    0.760      1.00    1.00    1.00   \n",
       "PULocationID           84399019.0  163.158   66.016      1.00  132.00  162.00   \n",
       "DOLocationID           84399019.0  161.353   70.251      1.00  116.00  163.00   \n",
       "payment_type           84152418.0    1.289    0.479      1.00    1.00    1.00   \n",
       "fare_amount            84399019.0   13.344  174.375  -1856.00    7.00   11.00   \n",
       "extra                  84399019.0    1.087    1.249    -60.00    0.00    1.00   \n",
       "mta_tax                84399019.0    0.495    0.067     -0.50    0.50    0.50   \n",
       "tip_amount             84399019.0    2.195   15.657   -221.00    0.00    2.06   \n",
       "tolls_amount           84399019.0    0.383    1.817    -70.00    0.00    0.00   \n",
       "improvement_surcharge  84399019.0    0.299    0.028     -0.30    0.30    0.30   \n",
       "total_amount           84399019.0   19.124  184.087  -1871.80   11.76   15.80   \n",
       "congestion_surcharge   79543038.0    2.188    0.837     -2.50    2.50    2.50   \n",
       "\n",
       "                          75%         max  \n",
       "VendorID                 2.00        4.00  \n",
       "passenger_count          2.00        9.00  \n",
       "trip_distance            8.82    45977.22  \n",
       "RatecodeID               1.00       99.00  \n",
       "PULocationID           234.00      265.00  \n",
       "DOLocationID           236.00      265.00  \n",
       "payment_type             2.00        5.00  \n",
       "fare_amount             32.04   943274.80  \n",
       "extra                    3.00      535.38  \n",
       "mta_tax                  0.50      212.42  \n",
       "tip_amount               3.26   141492.02  \n",
       "tolls_amount             0.00     3288.00  \n",
       "improvement_surcharge    0.30        1.00  \n",
       "total_amount            38.16  1084772.17  \n",
       "congestion_surcharge     2.50        4.50  "
      ]
     },
     "execution_count": 16,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "%%time\n",
    "np.round(taxi.describe().compute(), 3).T"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Feature engineering"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [],
   "source": [
    "def make_features(df):\n",
    "    \"\"\" Same code from Part 1 \"\"\"\n",
    "    df['pickup_weekday'] = df.tpep_pickup_datetime.dt.weekday\n",
    "    df['pickup_weekofyear'] = df.tpep_pickup_datetime.dt.isocalendar().week.astype(int)\n",
    "    df['pickup_hour'] = df.tpep_pickup_datetime.dt.hour\n",
    "    df['pickup_minute'] = df.tpep_pickup_datetime.dt.minute\n",
    "    df['pickup_year_seconds'] = (df.tpep_pickup_datetime - datetime.datetime(2019, 1, 1, 0, 0, 0)).dt.seconds\n",
    "    df['pickup_week_hour'] = (df.pickup_weekday * 24) + df.pickup_hour\n",
    "    df['store_and_fwd_flag'] = (df.store_and_fwd_flag == 'Y').astype(int)\n",
    "    df['VendorID'] = df.VendorID.fillna(-1)\n",
    "    df['RatecodeID'] = df.RatecodeID.fillna(-1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 54.2 ms, sys: 4.33 ms, total: 58.5 ms\n",
      "Wall time: 75.1 ms\n"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "\n",
    "make_features(taxi)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 101 ms, sys: 18.5 ms, total: 120 ms\n",
      "Wall time: 2.06 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>VendorID</th>\n",
       "      <th>tpep_pickup_datetime</th>\n",
       "      <th>tpep_dropoff_datetime</th>\n",
       "      <th>passenger_count</th>\n",
       "      <th>trip_distance</th>\n",
       "      <th>RatecodeID</th>\n",
       "      <th>store_and_fwd_flag</th>\n",
       "      <th>PULocationID</th>\n",
       "      <th>DOLocationID</th>\n",
       "      <th>payment_type</th>\n",
       "      <th>...</th>\n",
       "      <th>tolls_amount</th>\n",
       "      <th>improvement_surcharge</th>\n",
       "      <th>total_amount</th>\n",
       "      <th>congestion_surcharge</th>\n",
       "      <th>pickup_weekday</th>\n",
       "      <th>pickup_weekofyear</th>\n",
       "      <th>pickup_hour</th>\n",
       "      <th>pickup_minute</th>\n",
       "      <th>pickup_year_seconds</th>\n",
       "      <th>pickup_week_hour</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>1.0</td>\n",
       "      <td>2019-01-01 00:46:40</td>\n",
       "      <td>2019-01-01 00:53:20</td>\n",
       "      <td>1.0</td>\n",
       "      <td>1.5</td>\n",
       "      <td>1.0</td>\n",
       "      <td>0</td>\n",
       "      <td>151</td>\n",
       "      <td>239</td>\n",
       "      <td>1.0</td>\n",
       "      <td>...</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.3</td>\n",
       "      <td>9.95</td>\n",
       "      <td>NaN</td>\n",
       "      <td>1</td>\n",
       "      <td>1</td>\n",
       "      <td>0</td>\n",
       "      <td>46</td>\n",
       "      <td>2800</td>\n",
       "      <td>24</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>1.0</td>\n",
       "      <td>2019-01-01 00:59:47</td>\n",
       "      <td>2019-01-01 01:18:59</td>\n",
       "      <td>1.0</td>\n",
       "      <td>2.6</td>\n",
       "      <td>1.0</td>\n",
       "      <td>0</td>\n",
       "      <td>239</td>\n",
       "      <td>246</td>\n",
       "      <td>1.0</td>\n",
       "      <td>...</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.3</td>\n",
       "      <td>16.30</td>\n",
       "      <td>NaN</td>\n",
       "      <td>1</td>\n",
       "      <td>1</td>\n",
       "      <td>0</td>\n",
       "      <td>59</td>\n",
       "      <td>3587</td>\n",
       "      <td>24</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>2.0</td>\n",
       "      <td>2018-12-21 13:48:30</td>\n",
       "      <td>2018-12-21 13:52:40</td>\n",
       "      <td>3.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>1.0</td>\n",
       "      <td>0</td>\n",
       "      <td>236</td>\n",
       "      <td>236</td>\n",
       "      <td>1.0</td>\n",
       "      <td>...</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.3</td>\n",
       "      <td>5.80</td>\n",
       "      <td>NaN</td>\n",
       "      <td>4</td>\n",
       "      <td>51</td>\n",
       "      <td>13</td>\n",
       "      <td>48</td>\n",
       "      <td>49710</td>\n",
       "      <td>109</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>3</th>\n",
       "      <td>2.0</td>\n",
       "      <td>2018-11-28 15:52:25</td>\n",
       "      <td>2018-11-28 15:55:45</td>\n",
       "      <td>5.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>1.0</td>\n",
       "      <td>0</td>\n",
       "      <td>193</td>\n",
       "      <td>193</td>\n",
       "      <td>2.0</td>\n",
       "      <td>...</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.3</td>\n",
       "      <td>7.55</td>\n",
       "      <td>NaN</td>\n",
       "      <td>2</td>\n",
       "      <td>48</td>\n",
       "      <td>15</td>\n",
       "      <td>52</td>\n",
       "      <td>57145</td>\n",
       "      <td>63</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>4</th>\n",
       "      <td>2.0</td>\n",
       "      <td>2018-11-28 15:56:57</td>\n",
       "      <td>2018-11-28 15:58:33</td>\n",
       "      <td>5.0</td>\n",
       "      <td>0.0</td>\n",
       "      <td>2.0</td>\n",
       "      <td>0</td>\n",
       "      <td>193</td>\n",
       "      <td>193</td>\n",
       "      <td>2.0</td>\n",
       "      <td>...</td>\n",
       "      <td>0.0</td>\n",
       "      <td>0.3</td>\n",
       "      <td>55.55</td>\n",
       "      <td>NaN</td>\n",
       "      <td>2</td>\n",
       "      <td>48</td>\n",
       "      <td>15</td>\n",
       "      <td>56</td>\n",
       "      <td>57417</td>\n",
       "      <td>63</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "<p>5 rows × 24 columns</p>\n",
       "</div>"
      ],
      "text/plain": [
       "   VendorID tpep_pickup_datetime tpep_dropoff_datetime  passenger_count  \\\n",
       "0       1.0  2019-01-01 00:46:40   2019-01-01 00:53:20              1.0   \n",
       "1       1.0  2019-01-01 00:59:47   2019-01-01 01:18:59              1.0   \n",
       "2       2.0  2018-12-21 13:48:30   2018-12-21 13:52:40              3.0   \n",
       "3       2.0  2018-11-28 15:52:25   2018-11-28 15:55:45              5.0   \n",
       "4       2.0  2018-11-28 15:56:57   2018-11-28 15:58:33              5.0   \n",
       "\n",
       "   trip_distance  RatecodeID  store_and_fwd_flag  PULocationID  DOLocationID  \\\n",
       "0            1.5         1.0                   0           151           239   \n",
       "1            2.6         1.0                   0           239           246   \n",
       "2            0.0         1.0                   0           236           236   \n",
       "3            0.0         1.0                   0           193           193   \n",
       "4            0.0         2.0                   0           193           193   \n",
       "\n",
       "   payment_type  ...  tolls_amount  improvement_surcharge  total_amount  \\\n",
       "0           1.0  ...           0.0                    0.3          9.95   \n",
       "1           1.0  ...           0.0                    0.3         16.30   \n",
       "2           1.0  ...           0.0                    0.3          5.80   \n",
       "3           2.0  ...           0.0                    0.3          7.55   \n",
       "4           2.0  ...           0.0                    0.3         55.55   \n",
       "\n",
       "   congestion_surcharge  pickup_weekday  pickup_weekofyear  pickup_hour  \\\n",
       "0                   NaN               1                  1            0   \n",
       "1                   NaN               1                  1            0   \n",
       "2                   NaN               4                 51           13   \n",
       "3                   NaN               2                 48           15   \n",
       "4                   NaN               2                 48           15   \n",
       "\n",
       "   pickup_minute  pickup_year_seconds  pickup_week_hour  \n",
       "0             46                 2800                24  \n",
       "1             59                 3587                24  \n",
       "2             48                49710               109  \n",
       "3             52                57145                63  \n",
       "4             56                57417                63  \n",
       "\n",
       "[5 rows x 24 columns]"
      ]
     },
     "execution_count": 19,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "%%time\n",
    "\n",
    "taxi.head()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Machine learning"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [],
   "source": [
    "# same as Part 1\n",
    "numeric_feat = [\n",
    "    'pickup_weekday', \n",
    "    'pickup_weekofyear', \n",
    "    'pickup_hour', \n",
    "    'pickup_minute', \n",
    "    'pickup_year_seconds',\n",
    "    'pickup_week_hour', \n",
    "    'passenger_count',\n",
    "]\n",
    "categorical_feat = [\n",
    "    'VendorID', \n",
    "    'RatecodeID', \n",
    "    'store_and_fwd_flag',\n",
    "    'PULocationID',\n",
    "    'DOLocationID',\n",
    "]\n",
    "features = numeric_feat + categorical_feat\n",
    "y_col = 'total_amount'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [],
   "source": [
    "# note the dask_ml imports rather than sklearn\n",
    "from dask_ml.model_selection import train_test_split\n",
    "from dask_ml.metrics import mean_squared_error\n",
    "from xgboost.dask import DaskXGBRegressor"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "CPU times: user 4.68 ms, sys: 3.7 ms, total: 8.38 ms\n",
      "Wall time: 10 ms\n"
     ]
    }
   ],
   "source": [
    "%%time\n",
    "\n",
    "X_train, X_test, y_train, y_test = train_test_split(\n",
    "    taxi[features], taxi[y_col], test_size=0.33, random_state=seed, shuffle=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "<br>\n",
    "\n",
    "Remember, nothing happens yet because of lazy evaluation. If you have the RAM, you can call `df.persist()` to avoid repeated CSV loading\n",
    "\n",
    "(but if you had the RAM, you could do without Dask)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [],
   "source": [
    "# X_train.persist()\n",
    "# y_train.persist()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [],
   "source": [
    "xgb = DaskXGBRegressor(\n",
    "    n_estimators=10, \n",
    "    max_depth=3, \n",
    "    learning_rate=0.1, \n",
    "    random_state=seed, \n",
    "    n_jobs=-1,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%time\n",
    "\n",
    "_ = xgb.fit(X_train, y_train)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%time\n",
    "\n",
    "# get test RMSE\n",
    "preds = xgb.predict(X_test)\n",
    "np.sqrt(mean_squared_error(preds, y_test.to_dask_array()))"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "saturn (Python 3)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.10"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}