4-performance-optimization.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "f06a2f61-89b5-4a1f-9928-da4fe14d62a5",
   "metadata": {},
   "source": [
    "<img src=\"images/dask_horizontal.svg\" width=\"45%\" alt=\"Dask logo\\\">\n",
    "\n",
    "# Performance Optimization\n",
    "\n",
    "This notebook walks through a Dask DataFrame ETL workload. We'll demonstrate how to diagnose performance issues, utilize the Dask dashboard, and cover several common DataFrame best practices. \n",
    "\n",
    "## Dataset: Uber/Lyft TLC Trip Records\n",
    "\n",
    "The New York City Taxi and Limousine Commission (TLC) collects trip information for each taxi and for-hire vehicle trip completed by licensed drivers and vehicles; Here we'll analyze a subset of the [High-Volume For-Hire Services](https://www.nyc.gov/site/tlc/businesses/high-volume-for-hire-services.page) datset stored which provides a good example of an out-of-core dataset that's too large for a standard laptop due to memory limitations.\n",
    "\n",
    "Some characteristics of the dataset:\n",
    "\n",
    "- CSV dataset that's ~115 GB in memory\n",
    "- Stored in `s3://coiled-datasets/uber-lyft-tlc-sample/csv-10/`\n",
    "- In region `us-east-2`\n",
    "\n",
    "## Cluster setup\n",
    "\n",
    "Because the dataset is too large for a laptop, we'll create a larger Dask cluster on AWS using [Coiled](https://www.coiled.io).\n",
    "(Disclaimer: Some of the instructors for this tutorial are employed by Coiled.):"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "29f3298c-6626-4eb4-8389-1e3ca7f83f67",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "import coiled\n",
    "\n",
    "cluster = coiled.Cluster(\n",
    "    n_workers=20,\n",
    "    region=\"us-east-2\",  # start workers close to data to minimize costs\n",
    ")\n",
    "client = cluster.get_client()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "547788e0",
   "metadata": {},
   "source": [
    "Once we have initialized a cluster and client, we can easily view the Dask dashboard either through widgets provided by [dask-labextension](https://github.com/dask/dask-labextension), or by visiting the dashboard URL directly:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b3396b6a-6432-4e86-b826-17556efee8ec",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "client"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2a4c4623-e950-4685-849c-9793b5f18e59",
   "metadata": {},
   "source": [
    "Using `dask.dataframe.read_csv()`, we can lazily read this data in and do some low-level exploration before performing more complex computations:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3b486301-0df1-4700-b618-343bbb397013",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "%%time\n",
    "\n",
    "import dask.dataframe as dd\n",
    "\n",
    "ddf = dd.read_csv(\n",
    "    \"s3://coiled-datasets/uber-lyft-tlc-sample/csv-0.2-10/*\", \n",
    "    dtype={\"wav_match_flag\": \"category\"},\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "72a8ba8c-7ebe-4c26-81c4-76ac345e135b",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "ddf.dtypes"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "da2e59c0",
   "metadata": {},
   "source": [
    "After some initial exploration, we see that the columns representing on-scene and pickup times are stored as `object`s. We decide to do some feature engineering by converting these to datetimes and moving relevant date components into separate columns."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "63f2390e-5983-4844-8c37-2651afd8d940",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "%%time\n",
    "\n",
    "# Convert to datetime\n",
    "ddf[\"on_scene_datetime\"] = dd.to_datetime(ddf[\"on_scene_datetime\"], format=\"mixed\")\n",
    "ddf[\"pickup_datetime\"] = dd.to_datetime(ddf[\"pickup_datetime\"], format=\"mixed\")\n",
    "\n",
    "# Unpack columns\n",
    "ddf = ddf.assign(\n",
    "    accessible_vehicle=ddf.on_scene_datetime.isnull(),\n",
    "    pickup_month=ddf.pickup_datetime.dt.month,\n",
    "    pickup_dow=ddf.pickup_datetime.dt.dayofweek,\n",
    "    pickup_hour=ddf.pickup_datetime.dt.hour,\n",
    ")\n",
    "ddf = ddf.drop(columns=[\"on_scene_datetime\", \"pickup_datetime\"])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "166a0eb0",
   "metadata": {},
   "source": [
    "From here, some data sanitization and improvements to readability:\n",
    "\n",
    "- Normalize airport fees to non-null floats\n",
    "- Remove trip time outliers\n",
    "- Rename service codes to their corresponding rideshare companies"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3b502560",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "%%time\n",
    "\n",
    "# Format airport_fee\n",
    "ddf[\"airport_fee\"] = ddf[\"airport_fee\"].fillna(0)\n",
    "\n",
    "# Remove outliers\n",
    "lower_bound = 0\n",
    "Q3 = ddf[\"trip_time\"].quantile(0.75)\n",
    "upper_bound = Q3 + (1.5 * (Q3 - lower_bound))\n",
    "ddf = ddf.loc[(ddf[\"trip_time\"] >= lower_bound) & (ddf[\"trip_time\"] <= upper_bound)]\n",
    "\n",
    "service_names = {\n",
    "    \"HV0002\": \"juno\",\n",
    "    \"HV0005\": \"lyft\",\n",
    "    \"HV0003\": \"uber\",\n",
    "    \"HV0004\": \"via\",\n",
    "}\n",
    "\n",
    "ddf[\"service_names\"] = ddf[\"hvfhs_license_num\"].map(service_names)\n",
    "ddf = ddf.drop(columns=[\"hvfhs_license_num\"])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "24e4e66b-24f7-4f6a-907a-96a34d48b7d0",
   "metadata": {
    "tags": []
   },
   "source": [
    "Now that the data is cleaned up, we can do some computations on our data.\n",
    "\n",
    "First, let's compute the average tip amount across all riders:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "abfae138-c7cc-4bf4-9b83-446e3cfd7492",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "%%time\n",
    "\n",
    "(ddf.tips > 0).mean().compute()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c932453a",
   "metadata": {},
   "source": [
    "Or some metrics of tipping grouped by rideshare company:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7ceab8e1-6696-4ca1-b870-931c93ca684a",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "%%time\n",
    "\n",
    "ddf.loc[lambda x: x.tips > 0].groupby(\"service_names\").tips.sum().compute()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "24e6b445-fa87-40aa-aa53-45e410a34e37",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "%%time\n",
    "\n",
    "ddf.loc[lambda x: x.tips > 0].groupby(\"service_names\").tips.mean().compute()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f3bb9d04-dd26-45bb-87be-3b57ef8358b6",
   "metadata": {},
   "source": [
    "# Persist when possible\n",
    "\n",
    "Looking at the dashboard while performing the above analysis, it should become clear that whenever we compute operations on `ddf`, we must also run through all the dependent operations that read in and sanitize `ddf`, which forces several repeated computation steps.\n",
    "\n",
    "When doing mutliple computations on the same dataset, it can save both time and money to `.persist()` it first - this incurs the time and cost of computing the dataset once, in exchange for future computations on the dataset working with an in-memory copy of the computed data:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4e7489e9-0408-4550-a756-baf7fdeadccc",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "%%time\n",
    "\n",
    "ddf = ddf.persist()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ef63458e-25fa-4be1-bd60-fad2e7b0987d",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "%%time\n",
    "\n",
    "from distributed import wait\n",
    "wait(ddf);"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d5b14a0e",
   "metadata": {},
   "source": [
    "Now that `ddf` has been persisted, we can see that the same analysis as above can be computed much faster, with the initial creation of `ddf` no longer being included:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "830696a8-f007-4be7-a482-375d7abbeb8e",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "%%time\n",
    "\n",
    "(ddf.tips > 0).mean().compute()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b4c6e467-ebb7-4e6f-a56c-b97be7f19f96",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "%%time\n",
    "\n",
    "ddf.loc[lambda x: x.tips > 0].groupby(\"service_names\").tips.sum().compute()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f1f97070-2081-4803-9ae9-bc1ce6a17f70",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "%%time\n",
    "\n",
    "ddf.loc[lambda x: x.tips > 0].groupby(\"service_names\").tips.mean().compute()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8eea0988",
   "metadata": {},
   "source": [
    "Note that the choice to persist data depends on several factors, including:\n",
    "\n",
    "- Whether or not it fits into your clusters memory\n",
    "- If it's being reused in enough computations\n",
    "\n",
    "In general, a best practice to follow is persisting the dataset(s) you expect to use the most throughout computations."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2b87bd8b-96fe-4e88-9a19-79871ece71a0",
   "metadata": {},
   "source": [
    "# Avoid repeated compute calls\n",
    "\n",
    "When working with related results that share computations between one another, calling `.compute()` on each object individually forces us to discard shared work that could otherwise be used to speed up future computations.\n",
    "\n",
    "For example:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "09785813-271d-4314-8d57-5024c5972a13",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "trip_frac = (ddf.tips > 0).mean()\n",
    "gb_sum = ddf.loc[lambda x: x.tips > 0].groupby(\"service_names\").tips.sum()\n",
    "gb_mean = ddf.loc[lambda x: x.tips > 0].groupby(\"service_names\").tips.mean()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d84a1502",
   "metadata": {},
   "source": [
    "Intuitively, we know that `gb_sum` and `gb_mean` both depend on `ddf.loc[lambda x: x.tips > 0].groupby(\"service_names\")`, but calling `.compute()` on each object forces us to compute this result twice.\n",
    "\n",
    "To compute all of these objects in parallel and compute shared parts of the computation only once, we can use [`dask.compute()`](https://docs.dask.org/en/stable/api.html#dask.compute):"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bfd793ed-6a19-40f0-ac1a-3fd8c93e336c",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "%%time\n",
    "\n",
    "import dask\n",
    "\n",
    "trip_frac, gb_sum, gb_mean = dask.compute(trip_frac, gb_sum, gb_mean)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b8d460cc-d6c3-46c8-983a-9f4a88663db7",
   "metadata": {},
   "source": [
    "# Store data efficiently\n",
    "\n",
    "Up until this point, all of our performance optimizations have taken place after the initial reading of the data.\n",
    "However, as ability to compute increases, data access and I/O become more significant bottlenecks.\n",
    "Additionally, parallel computing will often add new constraints to how your store your data, particularly around providing random access to blocks of your data that are in line with how you plan to compute on it.\n",
    "\n",
    "## File format\n",
    "\n",
    "[Parquet](https://parquet.apache.org) is a popular, columnar file format designed for efficient data storage and retrieval. It handles random access, metadata storage, and binary encoding well. We [recommend using Parquet](https://docs.dask.org/en/stable/dataframe-best-practices.html#use-parquet) when working with tabular data."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "343c83ef-aeb0-460b-9853-3ba026cffb62",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "%%time\n",
    "\n",
    "import dask.dataframe as dd\n",
    "\n",
    "# ddf = dd.read_csv(\n",
    "#     \"s3://coiled-datasets/uber-lyft-tlc-sample/csv-ill/*\", \n",
    "#     dtype={\"wav_match_flag\": \"category\"},\n",
    "# )\n",
    "\n",
    "ddf = dd.read_parquet(\"s3://coiled-datasets/uber-lyft-tlc-sample/parquet-0.2-10/\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "26a2c940-bf12-4f25-b1e8-bedc4a0b7909",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "ddf.dtypes"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6f2e173e",
   "metadata": {},
   "source": [
    "From here, we can see that the same data sanitization as earlier can be done much faster:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "16ca27f6-5d46-451e-bccf-3e2f3b24d0cc",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "%%time\n",
    "\n",
    "# # Convert to datetime\n",
    "# ddf[\"on_scene_datetime\"] = dd.to_datetime(ddf[\"on_scene_datetime\"], format=\"mixed\")\n",
    "# ddf[\"pickup_datetime\"] = dd.to_datetime(ddf[\"pickup_datetime\"], format=\"mixed\")\n",
    "\n",
    "# Unpack columns\n",
    "ddf = ddf.assign(\n",
    "    accessible_vehicle=ddf.on_scene_datetime.isnull(),\n",
    "    pickup_month=ddf.pickup_datetime.dt.month,\n",
    "    pickup_dow=ddf.pickup_datetime.dt.dayofweek,\n",
    "    pickup_hour=ddf.pickup_datetime.dt.hour,\n",
    ")\n",
    "ddf = ddf.drop(columns=[\"on_scene_datetime\", \"pickup_datetime\"])\n",
    "\n",
    "# Format airport_fee\n",
    "ddf[\"airport_fee\"] = ddf[\"airport_fee\"].fillna(0)\n",
    "\n",
    "# Remove outliers\n",
    "lower_bound = 0\n",
    "Q3 = ddf[\"trip_time\"].quantile(0.75)\n",
    "upper_bound = Q3 + (1.5 * (Q3 - lower_bound))\n",
    "ddf = ddf.loc[(ddf[\"trip_time\"] >= lower_bound) & (ddf[\"trip_time\"] <= upper_bound)]\n",
    "\n",
    "service_names = {\n",
    "    \"HV0002\": \"juno\",\n",
    "    \"HV0005\": \"lyft\",\n",
    "    \"HV0003\": \"uber\",\n",
    "    \"HV0004\": \"via\",\n",
    "}\n",
    "\n",
    "ddf[\"service_names\"] = ddf[\"hvfhs_license_num\"].map(service_names)\n",
    "ddf = ddf.drop(columns=[\"hvfhs_license_num\"])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8a6455e3",
   "metadata": {},
   "source": [
    "Following best practices, we will now persist this sanitized dataset, so we no longer need to incur repeated I/O costs:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4f53c0e3-ad21-4c0a-98dd-a138adf0c378",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "ddf = ddf.persist()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bf8bcba0-1510-45eb-a1be-2bf8604c2ecf",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "%%time\n",
    "\n",
    "from distributed import wait\n",
    "wait(ddf);"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "159af1ee",
   "metadata": {},
   "source": [
    "From here, analysis can continue as normally:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "39634c19-079a-46d0-9e0a-ddca38921df2",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "%%time\n",
    "\n",
    "(ddf.tips > 0).mean().compute()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f4278dd9-c037-4b06-a3a6-ff7d41a7d4cd",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "%%time\n",
    "\n",
    "ddf.loc[lambda x: x.tips > 0].groupby(\"service_names\").tips.sum().compute()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4610c9db-f83d-4b62-ad7d-e4a1fff3e788",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "%%time\n",
    "\n",
    "ddf.loc[lambda x: x.tips > 0].groupby(\"service_names\").tips.mean().compute()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "98af8e81-e9fd-4e0d-b0de-b7e6755a3fbf",
   "metadata": {},
   "source": [
    "Note that since we persisted the data, the impact of the improved I/O is gone by the time we get to the analysis.\n",
    "This is because at this point, the data is stored in memory with pandas objects and datatypes; how it was originally stored no longer matters.\n",
    "Put differently, all analysis beyond I/O and sanitization creates an identical task graph to the previous dataset.\n",
    "In the next section, we will see how to troubleshoot and optimize our analysis independent of I/O.\n",
    "\n",
    "## Partition size\n",
    "\n",
    "So far, we've been working with the default partition size which, for this dataset, is pretty small (~10 MB).\n",
    "A small partition size results in very many partition, which in turn results in very many tasks in our computation graphs.\n",
    "\n",
    "When choosing a partition size, the goal is to give Dask enough to do per task that the scheduler overhead isn't taking up a disproportionate amount of time, but not so much that the workers run out of memory.\n",
    "A good rule of thumb for partition sizes is between 100 MB and 1 GB per partition ([excellent blog post on this](https://blog.dask.org/2021/11/02/choosing-dask-chunk-sizes)).\n",
    "\n",
    "So the first step is to see what our partiton size currently is:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "01581f63-3060-4bba-b97f-3e34ed0aae49",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "import dask\n",
    "dask.utils.format_bytes(ddf.partitions[0].compute().memory_usage(deep=True).sum())"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d7fb4d26-c657-417a-b031-0827711c2a72",
   "metadata": {},
   "source": [
    "Let's repartition to a bigger size."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d0654d49-f504-4a5f-9eda-ce1dcd1f0c36",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "%%time\n",
    "\n",
    "ddf = ddf.repartition(\"100MiB\")\n",
    "ddf = ddf.persist()\n",
    "wait(ddf);"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7e08158e-8d36-4b8d-8025-f977270093fb",
   "metadata": {},
   "source": [
    "Note that we persist after we repartition so we don't repeat the repartitioning work every time we compute.\n",
    "\n",
    "As a sanity check, let's check the new partition size:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "280633da-95ac-4f04-ab23-c0e8a132eed9",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "dask.utils.format_bytes(ddf.partitions[0].compute().memory_usage(deep=True).sum())"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5ad83a25-b00a-46f9-9ae5-22e8ad4cb7b8",
   "metadata": {},
   "source": [
    "Nice! Now let's do our analyses again.\n",
    "Remember that this time, the task graph will be much smaller.\n",
    "You can always inspect the graph by calling `.visualize()` rather than `.compute()` or by looking at the \"Graph\" page in the dashboard."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "27254373-f82a-4fe6-b57e-64217b308e64",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "%%time\n",
    "\n",
    "(ddf.tips > 0).mean().compute()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b460febd-ad73-4b2e-8606-bd29a2868e5f",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "%%time\n",
    "\n",
    "ddf.loc[lambda x: x.tips > 0].groupby(\"service_names\").tips.sum().compute()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "99548028-fc5b-46dc-8300-5a1d0c9e69b5",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "%%time\n",
    "\n",
    "ddf.loc[lambda x: x.tips > 0].groupby(\"service_names\").tips.mean().compute()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4dff0b64-e4d3-4f2d-b00e-b950985d9ee0",
   "metadata": {},
   "source": [
    "That was fast 🔥\n",
    "\n",
    "Here we improved on the task graph by increasing the partition size, but we haven't improved the performance of the tasks themselves.\n",
    "In the next section, we'll explore how changing the data type of your columns can make individual tasks more perfomant."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a309e787-5189-41ef-85ed-5aed652054ff",
   "metadata": {},
   "source": [
    "# Use efficient data types\n",
    "\n",
    "Up until this point, we've been using the default data types inferred by Dask for most of our columns. In the case of string data, this means we are using the Python `object` type, which can be slow to process:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5f0820b9-c715-4c4f-b861-88b384c0a1d1",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "ddf.dtypes"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3f58a23d",
   "metadata": {},
   "source": [
    "Recent versions of [Dask and pandas have improved support for PyArrow data types, most notably PyArrow strings](https://medium.com/coiled-hq/pyarrow-strings-in-dask-dataframes-55a0c4871586), which are faster and more memory efficient than Python `objects`.\n",
    "\n",
    "Let's enjoy some of the benefits of PyArrow strings by casting relevant string columns to `string[pyarrow]`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "da37e748-9f2a-4e88-a176-6b48bf82f94e",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "%%time\n",
    "\n",
    "ddf = ddf.astype({\n",
    "    \"service_names\": \"string[pyarrow]\",\n",
    "    \"dispatching_base_num\": \"string[pyarrow]\",\n",
    "    \"originating_base_num\": \"string[pyarrow]\",\n",
    "})\n",
    "\n",
    "ddf = ddf.persist()\n",
    "wait(ddf);"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "644e13ff-3dfa-446e-900a-c47e76317428",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "ddf.dtypes"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6e518106",
   "metadata": {},
   "source": [
    "With that done, let's revisit our partition sizes to see how they've been impacted:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4da5a812-ff2a-4cf0-bfad-b65f7108d3de",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "dask.utils.format_bytes(ddf.partitions[1].compute().memory_usage(deep=True).sum())"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "76399f57",
   "metadata": {},
   "source": [
    "Nice! With PyArrow strings, our partitions are noticeably smaller, and we can once again repartition our data to land at a solid 100 MB partition size:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2da4346f-a50f-4c32-9596-73f6d95c1dd1",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "%%time\n",
    "\n",
    "ddf = ddf.repartition(\"100MB\")\n",
    "ddf = ddf.persist()\n",
    "wait(ddf);"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "29c5a0a3",
   "metadata": {},
   "source": [
    "With these new data types, we can now see that the analyses results in an even smaller task graph; on top of that, the improved performance of the PyArrow strings means that each individual task is more performant:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c8512af8-3a9f-4c85-a3c5-0643fb8b6331",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "%%time\n",
    "\n",
    "(ddf.tips != 0).mean().compute()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e4122917-3f67-42e2-83c6-e9c40d6e954d",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "%%time\n",
    "\n",
    "ddf.loc[lambda x: x.tips > 0].groupby(\"service_names\").tips.sum().compute()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9aedb328-377d-4e35-9c8b-793700dc9a16",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "%%time\n",
    "\n",
    "ddf.loc[lambda x: x.tips > 0].groupby(\"service_names\").tips.mean().compute()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a15d4325",
   "metadata": {},
   "source": [
    "Note that as of `dask=2023.3.1`, we can skip the effort of manually recasting Python object columns to PyArrow strings by modifying the value of `dataframe.convert-string` in our Dask config:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "29eef754-42e6-4251-b7da-acb4468305dd",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# dask.config.set({\"dataframe.convert-string\": True});"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "410f2fbf",
   "metadata": {},
   "source": [
    "The benefits of PyArrow strings aren't just limited to computation. By setting them as the default data type when reading in Parquet data, we can also improve the performance of I/O.\n",
    "\n",
    "# Summary\n",
    "\n",
    "In this notebook, we took a look at a representative Dask DataFrame workload that could benefit from Dask.\n",
    "\n",
    "Starting from a suboptimal place performance-wise, we explored the dashboard to find potentials areas for improvement.\n",
    "We then went through some basic Dask best practices that allowed us to shrink our task graph and improve the performance of individual tasks, which was reflected both in our analyses runtimes and dashboard plots.\n",
    "\n",
    "# Additional Resources\n",
    "\n",
    "- Repositories on GitHub:\n",
    "    - Dask https://github.com/dask/dask\n",
    "    - Distributed https://github.com/dask/distributed\n",
    "\n",
    "- Documentation:\n",
    "    - Dask documentation https://docs.dask.org\n",
    "    - Distributed documentation https://distributed.dask.org\n",
    "\n",
    "- If you have a Dask usage questions, please ask it on the [Dask GitHub discussions board](https://github.com/dask/dask/discussions).\n",
    "\n",
    "- If you run into a bug, feel free to file a report on the [Dask GitHub issue tracker](https://github.com/dask/dask/issues).\n",
    "\n",
    "- If you're interested in getting involved and contributing to Dask. Please check out our [contributing guide](https://docs.dask.org/en/latest/develop.html).\n",
    "\n",
    "# Thank you!"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}