{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# DataFrames: Groupby\n",
    "\n",
    "This notebook uses the Pandas groupby-aggregate and groupby-apply on scalable Dask dataframes.  It will discuss both common use and best practices."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Start Dask Client for Dashboard\n",
    "\n",
    "Starting the Dask Client is optional.  It will provide a dashboard which \n",
    "is useful to gain insight on the computation.  \n",
    "\n",
    "The link to the dashboard will become visible when you create the client below.  We recommend having it open on one side of your screen while using your notebook on the other side.  This can take some effort to arrange your windows, but seeing them both at the same is very useful when learning."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from dask.distributed import Client\n",
    "client = Client(n_workers=1, threads_per_worker=4, processes=True, memory_limit='2GB')\n",
    "client"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Artifical dataset\n",
    "\n",
    "We create an artificial timeseries dataset to help us work with groupby operations"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import dask\n",
    "df = dask.datasets.timeseries()\n",
    "df"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This dataset is small enough to fit in the cluster's memory, so we persist it now.\n",
    "\n",
    "You would skip this step if your dataset becomes too large to fit into memory."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "df = df.persist()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Groupby Aggregations\n",
    "\n",
    "Dask dataframes implement a commonly used subset of the Pandas groupby API (see [Pandas Groupby Documentation](https://pandas.pydata.org/pandas-docs/stable/groupby.html).\n",
    "\n",
    "We start with groupby aggregations.  These are generally fairly efficient, assuming that the number of groups is small (less than a million)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "df.groupby('name').x.mean().compute()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Performance will depend on the aggregation you do (mean vs std), the key on which you group (name vs id), and the number of total groups"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%time _ = df.groupby('id').x.mean().compute()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%time _ = df.groupby('name').x.mean().compute()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%time df.groupby('name').agg({'x': ['mean', 'std'], 'y': ['mean', 'count']}).compute().head()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This is the same as with Pandas.  Generally speaking, Dask.dataframe groupby-aggregations are roughly same performance as Pandas groupby-aggregations, just more scalable.\n",
    "\n",
    "You can read more about Pandas' common aggregations in [the Pandas documentation](https://pandas.pydata.org/pandas-docs/stable/user_guide/groupby.html#aggregation).\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Custom Aggregations\n",
    "\n",
    "Dask dataframe Aggregate is available for custom aggregations (See [Dask dataframe Aggregate Documentation](https://docs.dask.org/en/latest/dataframe-groupby.html#aggregate))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Many groups\n",
    "\n",
    "By default groupby-aggregations (like groupby-mean or groupby-sum) return the result as a single-partition Dask dataframe.  Their results are *usually* quite small, so this is *usually* a good choice.\n",
    "\n",
    "However, sometimes people want to do groupby aggregations on *many* groups (millions or more).  In these cases the full result may not fit into a single Pandas dataframe output, and you may need to split your output into multiple partitions.  You can control this with the `split_out=` parameter"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Computational graph of a single output aggregation (for a small number of groups, like 1000)\n",
    "df.groupby('name').x.mean().visualize(node_attr={'penwidth': '6'})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Computational graph of an aggregation to four outputs (for a larger number of groups, like 1000000)\n",
    "df.groupby('id').x.mean(split_out=4).visualize(node_attr={'penwidth': '6'})"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Groupby Apply\n",
    "\n",
    "Groupby-aggregations are generally quite fast because they can be broken down easily into well known operations.  The data doesn't have to move around too much and we can just pass around small intermediate values across the network.\n",
    "\n",
    "For some operations however the function to be applied requires *all* data from a given group (like every record of someone named \"Alice\").  This will force a great deal of communication and be more expensive, but is still possible with the Groupby-apply method.  This should be avoided if a groupby-aggregation works.\n",
    "\n",
    "In the following example we train a simple Scikit-Learn machine learning model on every person's name."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from sklearn.linear_model import LinearRegression\n",
    "\n",
    "def train(partition):\n",
    "    if partition.empty:\n",
    "        return\n",
    "    est = LinearRegression()\n",
    "    est.fit(partition[['x', 'id']].values, partition.y.values)\n",
    "    return est"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%time df.groupby('name').apply(train, meta=object).compute().sort_index()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "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.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}