dask-bag-and-dataframe-example.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "1bdf748a-1807-4094-a5d0-b9f1e6cc62ef",
   "metadata": {},
   "source": [
    "# dask.bag and dask.dataframe\n",
    "\n",
    "The standard imports for `dask.bag` and `dask.dataframe` are `db` and `dd`, respectively."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d28b2625-95d5-4669-9f25-2fbf5517b9a9",
   "metadata": {},
   "outputs": [],
   "source": [
    "from distributed import Client\n",
    "client = Client(\"127.0.0.1:61329\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ff87f121-82a6-456e-82eb-ee6060cb92bb",
   "metadata": {},
   "outputs": [],
   "source": [
    "import dask.bag as db\n",
    "import dask.dataframe as dd"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "25d5edcb-4731-45e9-8405-ddc72a50c020",
   "metadata": {},
   "source": [
    "The simplest bag example possible:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "85a22c4b-9a4a-4e88-a8bd-0d200dc88edc",
   "metadata": {},
   "outputs": [],
   "source": [
    "bag = db.from_sequence([1, 2, 3, 4, 5, 6], npartitions=2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "64e93556-8406-4b53-9f85-8904fec7ab3b",
   "metadata": {},
   "outputs": [],
   "source": [
    "evens_squared = bag.filter(lambda x: x % 2 == 0).map(lambda x: x ** 2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6f8a8348-4e01-4f6f-891d-63613dde56fa",
   "metadata": {},
   "outputs": [],
   "source": [
    "evens_squared.compute()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "183d8716-47c2-44d4-8211-052869398360",
   "metadata": {
    "tags": []
   },
   "source": [
    "Let's look at something a bit more interesting, some US Federal Election Commission political campaign contribution data stored in a number of plain text files."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fa1cc229-dc3c-47d2-9dbe-98b49dbef6c3",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "from pathlib import Path\n",
    "columns = Path(\"data/fec/columns.csv\").read_text().strip().split(\",\")\n",
    "columns"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "44559c92-1a6d-4823-888e-f84453187648",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "!head -n 3 data/fec/contributions/d100.txt"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "05049eb9-9b7d-4950-9cd5-24e4e4fca7af",
   "metadata": {},
   "source": [
    "We'll create a bag by:\n",
    "- reading the text files (with `read_text`),\n",
    "- doing a bit of per-line processing (with `map`),\n",
    "- and limit the dataset to Atlanta, GA (with `filter`)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9f8a5df0-6e6f-4f3f-9de0-eff8c878ab4c",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "bag = (db.read_text(\"data/fec/contributions/*\")\n",
    "         .map(lambda x: {k: v for k, v in zip(columns, x.split(\"|\"))})\n",
    "         .filter(lambda x: x[\"STATE\"] == \"GA\" and x[\"CITY\"] == \"ATLANTA\"))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f5e973c9-6226-435e-a848-f7871e9e0e7b",
   "metadata": {},
   "source": [
    "A look at the most frequent donation:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "31819007-364f-464e-ae1d-7e1edc9f4777",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "freqs = bag.map(lambda x: x[\"TRANSACTION_AMT\"]).frequencies(sort=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8e9158ec-13e2-4c04-923c-83ecc9235454",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "freqs.take(10)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "57b07110-db29-4353-84a7-cdf9a3f58745",
   "metadata": {},
   "source": [
    "We can convert the existing bag to a dataframe so we can take advantage of both the pandas-like API and the pandas-backend performance."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ba2257fe-325f-47fe-a84f-716776db34a2",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "df = bag.to_dataframe(columns=columns)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7508e1f2-00e9-4778-8f4b-a0b1c5364f28",
   "metadata": {},
   "source": [
    "We can see what the average contribution is in Atlanta:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ed34458a-eb66-4355-b37d-bc7a7efc7057",
   "metadata": {},
   "outputs": [],
   "source": [
    "df.TRANSACTION_AMT.mean()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "891ad1cf-d17d-4f25-b006-946c48743384",
   "metadata": {},
   "source": [
    "Ah, well that didn't work. Let's use our pandas-like API for some help:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "932e76e2-7a7e-46c6-ad5a-9691c179fb6c",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "ta = dd.to_numeric(df.TRANSACTION_AMT, errors=\"coerce\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "63b125a0-a521-4acf-8a2b-627e0fb00dfe",
   "metadata": {},
   "source": [
    "Using `persist` instead of `compute` will force the result of the computation to persist in memory, allowing for more performant repeated use."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f4b39e2e-55ab-4832-92ea-7db8e00037ed",
   "metadata": {},
   "outputs": [],
   "source": [
    "ta = ta.persist()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e9214031-ac81-4e59-b9d9-e644bee156a6",
   "metadata": {},
   "outputs": [],
   "source": [
    "ta = ta[(ta > 0) & (ta < 20000)]\n",
    "ta.mean().compute()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4feabe97-4489-47b8-8c14-8e65e848026c",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "ta[ta > 150].mean().compute()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "48d49e99-e106-44fa-953b-f34f41425566",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "n, bins, pathes = plt.hist(ta[ta < 200].to_dask_array().compute(), bins=50)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5c58b0cd-e5dd-448a-a483-c09b3f1a9815",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "ta.value_counts().nlargest(10).compute()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "39ba33e5-6cb2-4ee7-910a-3e0ebcfec73e",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}