3. Basic Aggregations.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "4f18842e",
   "metadata": {},
   "source": [
    "In this notebook, we will:\n",
    "    \n",
    "**1.** Analyze the physical plans of basic aggregations\n",
    "\n",
    "**2.** Demonstrate the actual computation happening under the hood at each executor."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "e7086be5",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/root/miniconda/envs/py36/lib/python3.6/site-packages/pyspark/context.py:238: FutureWarning: Python 3.6 support is deprecated in Spark 3.2.\n",
      "  FutureWarning\n"
     ]
    }
   ],
   "source": [
    "from pyspark.sql import SparkSession, functions as F\n",
    "\n",
    "spark = SparkSession.builder.config(\"spark.sql.shuffle.partitions\", 16).getOrCreate()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a0144de9",
   "metadata": {},
   "source": [
    "Following cell creates the demo dataset. It is the same function used in the first notebook. You can skip it."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "1cc09de5",
   "metadata": {},
   "outputs": [],
   "source": [
    "import pandas as pd\n",
    "import numpy as np\n",
    "\n",
    "def create_demo_data(n_products, n_stores, start_date=\"2021-01-01\", end_date=\"2022-01-01\"):\n",
    "    \"\"\"Creates demo data, writes it as parquet partitioned by date, reads it and returns the dataframe\"\"\"\n",
    "    dates = pd.date_range(start_date, end_date)\n",
    "    dates = [str(date)[:10] for date in dates]\n",
    "\n",
    "    day_index = np.arange(len(dates))\n",
    "    result = []\n",
    "    for product in range(n_products):\n",
    "        for store in range(n_stores):\n",
    "            sales = np.random.poisson(10, size=len(dates))\n",
    "            partial_df = (\n",
    "                pd.DataFrame(dates, columns=[\"date\"])\n",
    "                .assign(product_id=product)\n",
    "                .assign(store_id=store)\n",
    "                .assign(day_index=day_index)\n",
    "                .assign(sales_quantity=sales)\n",
    "            )\n",
    "            result.append(partial_df)\n",
    "    pdf = pd.concat(result)\n",
    "    result = spark.createDataFrame(pdf)\n",
    "    result.repartition(\"date\").write.partitionBy(\"date\").parquet(\"demo-data\", mode=\"overwrite\")\n",
    "    return spark.read.parquet(\"demo-data\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "f06d32a9",
   "metadata": {},
   "outputs": [],
   "source": [
    "df = create_demo_data(n_products=10, n_stores=10, start_date=\"2021-01-01\", end_date=\"2021-01-31\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "518a42d3",
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "16"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "df.rdd.getNumPartitions()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d65cc737",
   "metadata": {},
   "source": [
    "## Basic Sum:\n",
    "\n",
    "- We all know this. Let's see the physical plan."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "633e4cce",
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "== Physical Plan ==\n",
      "AdaptiveSparkPlan isFinalPlan=false\n",
      "+- HashAggregate(keys=[], functions=[sum(sales_quantity#19L)])\n",
      "   +- Exchange SinglePartition, ENSURE_REQUIREMENTS, [id=#46]\n",
      "      +- HashAggregate(keys=[], functions=[partial_sum(sales_quantity#19L)])\n",
      "         +- Project [sales_quantity#19L]\n",
      "            +- FileScan parquet [sales_quantity#19L,date#20] Batched: true, DataFilters: [], Format: Parquet, Location: InMemoryFileIndex(1 paths)[file:/workspaces/rocks/Untitled Folder/demo-data], PartitionFilters: [], PushedFilters: [], ReadSchema: struct<sales_quantity:bigint>\n",
      "\n",
      "\n"
     ]
    }
   ],
   "source": [
    "df.select(F.sum(\"sales_quantity\")).explain()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "16246e6e",
   "metadata": {},
   "source": [
    "This is a typical map reduce pattern. The plan translates to:\n",
    "\n",
    "**1.** read parquet: ``+- FileScan parquet``\n",
    "\n",
    "**2.** select sales_quantity column: ``+- Project [sales_quantity#19L]``\n",
    "\n",
    "**3.** Compute the sum at each partition (partial_sum):  ``+- HashAggregate(keys=[], functions=[partial_sum(sales_quantity#19L)])``\n",
    "\n",
    "**4.** Move the partial results into a single partition ``+- Exchange SinglePartition, ENSURE_REQUIREMENTS, [id=#80]``\n",
    "\n",
    "**5.** Sum up all the partial results to get final result.: ``+- HashAggregate(keys=[], functions=[sum(sales_quantity#50L)])``\n",
    "\n",
    "In this case we have 16 partitions in the dataframe. At the end of **Step 3**, we will have **a single row at each partition** which only contains the sum of sales_quantity at that partition. Then these results are moved into a single partition.\n",
    "\n",
    "**Takeaway**:\n",
    "\n",
    "- Sum requires the shuffle of partial results. However, the amount of data is reduced greatly before this shuffle (to 16 rows in this case). Therefore, this is not a shuffle to worry about."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "eeaaaae2",
   "metadata": {},
   "source": [
    "## Groupby aggregations\n",
    "\n",
    "\n",
    "### Groupby Sum"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "6aff6dcb",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "== Physical Plan ==\n",
      "AdaptiveSparkPlan isFinalPlan=false\n",
      "+- HashAggregate(keys=[product_id#16L], functions=[sum(sales_quantity#19L)])\n",
      "   +- Exchange hashpartitioning(product_id#16L, 16), ENSURE_REQUIREMENTS, [id=#61]\n",
      "      +- HashAggregate(keys=[product_id#16L], functions=[partial_sum(sales_quantity#19L)])\n",
      "         +- Project [product_id#16L, sales_quantity#19L]\n",
      "            +- FileScan parquet [product_id#16L,sales_quantity#19L,date#20] Batched: true, DataFilters: [], Format: Parquet, Location: InMemoryFileIndex(1 paths)[file:/workspaces/rocks/Untitled Folder/demo-data], PartitionFilters: [], PushedFilters: [], ReadSchema: struct<product_id:bigint,sales_quantity:bigint>\n",
      "\n",
      "\n"
     ]
    }
   ],
   "source": [
    "product_sum = df.groupBy(\"product_id\").agg(F.sum(\"sales_quantity\").alias(\"product_sum\"))\n",
    "product_sum.explain()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e928c5f4",
   "metadata": {},
   "source": [
    "- Again, full shuffle is not necessary for groupby aggregations. For ``groupby.sum`` we can find the sum for each partition and then shuffle just the results of each partition and combine them.\n",
    "\n",
    "- Observe that there are two HashAggregate operations in the physical plan. 1st computes the partial sum, 2nd computes the final sum.\n",
    "\n",
    "**Note:** We read the physical plan from bottom to top.\n",
    "\n",
    "The plan is basically:\n",
    "\n",
    "        Read parquet -> Select necessary columns -> Partial sum -> Shuffle partial results -> Final sum\n",
    "\n",
    "Spark explains this as follows.\n",
    "\n",
    "        FileScan parquet -> Project -> HashAggregate -> Exchange hashpartitioning -> HashAggregate\n",
    " \n",
    "``HashAggregate``: Aggregates each partition locally. Decreases the number of rows in a partition.\n",
    "\n",
    "``Exchange hashpartitioning``: Some cool name for shuffle. Moves the data between executors. Expensive.\n",
    "\n",
    "\n",
    "**Demonstration:**\n",
    "\n",
    "\n",
    "**Initial state:**\n",
    "\n",
    "Executor 1\n",
    "\n",
    "```\n",
    "+-------------------------------------------------+\n",
    "|product_id|store_id|  date      | sales_quantity |\n",
    "|          |        |            |                |\n",
    "|    1     |    1   |  2000-01-01|       10       |\n",
    "|    1     |    1   |  2001-01-02|       12       |\n",
    "|    2     |    1   |  2000-01-01|       9        |\n",
    "|    2     |    1   |  2000-01-02|       5        |\n",
    "|    2     |    2   |  2000-01-01|       3        |\n",
    "+--------------------------------+----------------+\n",
    "```\n",
    "\n",
    "Executor 2\n",
    "\n",
    "```\n",
    "+-------------------------------------------------+\n",
    "|product_id|store_id|  date      | sales_quantity |\n",
    "|          |        |            |                |\n",
    "|    1     |    1   |  1999-01-02|       1        |\n",
    "|    2     |    1   |  1999-01-01|       7        |\n",
    "|    2     |    1   |  1999-01-02|       13       |\n",
    "|    3     |    1   |  1999-01-02|       100      |\n",
    "+--------------------------------+----------------+\n",
    "```\n",
    "\n",
    "**1. Aggregate each partition locally:** HashAggregate(keys=[product_id#67], functions=[partial_sum(sales_quantity#71L)])\n",
    "\n",
    "Executor 1\n",
    "\n",
    "```\n",
    "+----------------------+\n",
    "|product_id|partial_sum|\n",
    "|          |           |\n",
    "|    1     |    22     |\n",
    "|    2     |    17     |\n",
    "+----------------------+\n",
    "```\n",
    "\n",
    "Executor 2\n",
    "\n",
    "```\n",
    "+----------------------+\n",
    "|product_id|partial_sum|\n",
    "|          |           |\n",
    "|    1     |    1      |\n",
    "|    2     |    20     |\n",
    "|    3     |    100    |\n",
    "+----------------------+\n",
    "```\n",
    "\n",
    "**2. Shuffle partial results**: Exchange hashpartitioning(product_id#67, 16), ENSURE_REQUIREMENTS, [id=#364]\n",
    "\n",
    "Executor 1\n",
    "\n",
    "```\n",
    "+----------------------+\n",
    "|product_id|partial_sum|\n",
    "|          |           |\n",
    "|    1     |    22     |\n",
    "|    1     |    1      |\n",
    "+----------------------+\n",
    "```\n",
    "\n",
    "Executor 2\n",
    "\n",
    "```\n",
    "+----------------------+\n",
    "|product_id|partial_sum|\n",
    "|          |           |\n",
    "|    2     |    20     |\n",
    "|    2     |    17     |\n",
    "|    3     |    100    |\n",
    "+----------------------+\n",
    "```\n",
    "\n",
    "**3. Compute final sums:**  HashAggregate(keys=[product_id#67], functions=[sum(sales_quantity#71L)])\n",
    "\n",
    "Executor 1\n",
    "\n",
    "```\n",
    "+----------------------+\n",
    "|product_id|partial_sum|\n",
    "|          |           |\n",
    "|    1     |    23     |\n",
    "+----------------------+\n",
    "```\n",
    "\n",
    "Executor 2\n",
    "\n",
    "```\n",
    "+----------------------+\n",
    "|product_id|partial_sum|\n",
    "|          |           |\n",
    "|    2     |    37     |\n",
    "|    3     |    100    |\n",
    "+----------------------+\n",
    "```\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a2155356",
   "metadata": {},
   "source": [
    "### Groupby stddev\n",
    "\n",
    "Again, full shuffle is not necessary. There is a partial_stddev implementation.\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "d0397645",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "== Physical Plan ==\n",
      "AdaptiveSparkPlan isFinalPlan=false\n",
      "+- HashAggregate(keys=[product_id#16L], functions=[stddev_samp(cast(sales_quantity#19L as double))])\n",
      "   +- Exchange hashpartitioning(product_id#16L, 16), ENSURE_REQUIREMENTS, [id=#76]\n",
      "      +- HashAggregate(keys=[product_id#16L], functions=[partial_stddev_samp(cast(sales_quantity#19L as double))])\n",
      "         +- Project [product_id#16L, sales_quantity#19L]\n",
      "            +- FileScan parquet [product_id#16L,sales_quantity#19L,date#20] Batched: true, DataFilters: [], Format: Parquet, Location: InMemoryFileIndex(1 paths)[file:/workspaces/rocks/Untitled Folder/demo-data], PartitionFilters: [], PushedFilters: [], ReadSchema: struct<product_id:bigint,sales_quantity:bigint>\n",
      "\n",
      "\n"
     ]
    }
   ],
   "source": [
    "df.groupBy(\"product_id\").agg(F.stddev(\"sales_quantity\")).explain()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "225c07fd",
   "metadata": {},
   "source": [
    "### Groupby countDistinct"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "40182f32",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "== Physical Plan ==\n",
      "AdaptiveSparkPlan isFinalPlan=false\n",
      "+- HashAggregate(keys=[product_id#16L], functions=[count(distinct store_id#17L)])\n",
      "   +- Exchange hashpartitioning(product_id#16L, 16), ENSURE_REQUIREMENTS, [id=#101]\n",
      "      +- HashAggregate(keys=[product_id#16L], functions=[partial_count(distinct store_id#17L)])\n",
      "         +- HashAggregate(keys=[product_id#16L, store_id#17L], functions=[])\n",
      "            +- Exchange hashpartitioning(product_id#16L, store_id#17L, 16), ENSURE_REQUIREMENTS, [id=#97]\n",
      "               +- HashAggregate(keys=[product_id#16L, store_id#17L], functions=[])\n",
      "                  +- Project [product_id#16L, store_id#17L]\n",
      "                     +- FileScan parquet [product_id#16L,store_id#17L,date#20] Batched: true, DataFilters: [], Format: Parquet, Location: InMemoryFileIndex(1 paths)[file:/workspaces/rocks/Untitled Folder/demo-data], PartitionFilters: [], PushedFilters: [], ReadSchema: struct<product_id:bigint,store_id:bigint>\n",
      "\n",
      "\n"
     ]
    }
   ],
   "source": [
    "df.groupBy(\"product_id\").agg(F.countDistinct(\"store_id\")).explain()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e273626f",
   "metadata": {},
   "source": [
    "\n",
    "**1.** Get distinct store/products at each executor\n",
    "\n",
    "**2.** Shuffle to move store/product pairs to the same partition.\n",
    "\n",
    "**3.** Again, get distinct store/products at each executor.\n",
    "\n",
    "**4.** Compute partial counts.\n",
    "\n",
    "**5.** Shuffle the partial results.\n",
    "\n",
    "**6.** Compute the final results.\n",
    "\n",
    "Following demonstration will make it crystal clear.\n",
    "\n",
    "**Demonstration**\n",
    "\n",
    "\n",
    "- Assume parquet is partitioned by date. So, all rows for a date is in the same partition\n",
    "- Assume we have two executors.\n",
    "\n",
    "**Initial state:**\n",
    "\n",
    "Executor 1\n",
    "\n",
    "```\n",
    "+--------------------------------+\n",
    "|product_id|store_id|  date      |\n",
    "|          |        |            |\n",
    "|    1     |    1   |  2000-01-01|\n",
    "|    1     |    1   |  2001-01-02|\n",
    "|    2     |    1   |  2000-01-01|\n",
    "|    2     |    1   |  2000-01-02|\n",
    "|    2     |    2   |  2000-01-01|\n",
    "+--------------------------------+\n",
    "```\n",
    "Executor 2\n",
    "\n",
    "```\n",
    "+--------------------------------+\n",
    "|product_id|store_id|  date      |\n",
    "|          |        |            |\n",
    "|    1     |    1   |  1999-01-01|\n",
    "|    1     |    1   |  1999-01-02|\n",
    "|    2     |    1   |  1999-01-01|\n",
    "|    2     |    1   |  1999-01-02|\n",
    "|    2     |    2   |  1999-01-01|\n",
    "|    3     |    3   |  1999-01-01|\n",
    "+--------------------------------+\n",
    "```\n",
    "\n",
    "\n",
    "**1: Get distinct store/product locally:** HashAggregate(keys=[product_id#67, store_id#68], functions=[])\n",
    "\n",
    "Executor 1\n",
    "\n",
    "```\n",
    "+-------------------+\n",
    "|product_id|store_id|\n",
    "|          |        |\n",
    "|    1     |    1   |\n",
    "|    2     |    1   |\n",
    "|    2     |    2   |\n",
    "+-------------------+\n",
    "```\n",
    "\n",
    "\n",
    "Executor 2\n",
    "\n",
    "```\n",
    "+-------------------+\n",
    "|product_id|store_id|\n",
    "|          |        |\n",
    "|    1     |    1   |\n",
    "|    2     |    1   |\n",
    "|    2     |    2   |\n",
    "|    3     |    3   |\n",
    "+-------------------+\n",
    "```\n",
    "\n",
    "**2. Shuffle**: Exchange hashpartitioning(product_id#67, store_id#68, 1200), ENSURE_REQUIREMENTS, [id=#370]\n",
    "\n",
    "- Moves all rows for a store/product pair to the same partition\n",
    "- Shuffling does not change the total number of rows of the dataframe.\n",
    "- Note there are some duplicates after this step.\n",
    "\n",
    "Executor 1:\n",
    "\n",
    "```\n",
    "+-------------------+\n",
    "|product_id|store_id|\n",
    "|          |        |\n",
    "|    1     |    1   |\n",
    "|    1     |    1   |\n",
    "|    2     |    1   |\n",
    "|    2     |    1   |\n",
    "+-------------------+\n",
    "```\n",
    "\n",
    "\n",
    "Executor 2\n",
    "\n",
    "```\n",
    "+-------------------+\n",
    "|product_id|store_id|\n",
    "|          |        |\n",
    "|    2     |    2   |\n",
    "|    2     |    2   |\n",
    "|    3     |    3   |\n",
    "+-------------------+\n",
    "```\n",
    "\n",
    "**3. Again, get distinct store/products locally**: HashAggregate(keys=[product_id#67, store_id#68], functions=[])\n",
    "\n",
    "Executor 1:\n",
    "\n",
    "```\n",
    "+-------------------+\n",
    "|product_id|store_id|\n",
    "|          |        |\n",
    "|    1     |    1   |\n",
    "|    2     |    1   |\n",
    "+-------------------+\n",
    "```\n",
    "\n",
    "\n",
    "Executor 2\n",
    "\n",
    "```\n",
    "+-------------------+\n",
    "|product_id|store_id|\n",
    "|          |        |\n",
    "|    2     |    2   |\n",
    "|    3     |    3   |\n",
    "+-------------------+\n",
    "```\n",
    "\n",
    "**4. Compute partial counts at each executor**: HashAggregate(keys=[product_id#67], functions=[partial_count(distinct store_id#68)])\n",
    "\n",
    "We have previously shuffled for product/store we can get partial counts and then add them up later.\n",
    "\n",
    "```\n",
    "+------------------------+\n",
    "|product_id|partial_count|\n",
    "|          |             |\n",
    "|    1     |      1      |\n",
    "|    2     |      1      |\n",
    "+------------------------+\n",
    "```\n",
    "\n",
    "Executor 2\n",
    "\n",
    "```\n",
    "+------------------------+\n",
    "|product_id|partial_count|\n",
    "|          |             |\n",
    "|    2     |      1      |\n",
    "|    3     |      1      |\n",
    "+------------------------+\n",
    "```\n",
    "\n",
    "**5. Shuffle for products:** Exchange hashpartitioning(product_id#67, 1200), ENSURE_REQUIREMENTS, [id=#374]\n",
    "\n",
    "Executor 1:\n",
    "\n",
    "```\n",
    "+------------------------+\n",
    "|product_id|partial_count|\n",
    "|          |             |\n",
    "|    2     |      1      |\n",
    "|    2     |      1      |\n",
    "+------------------------+\n",
    "```\n",
    "\n",
    "Executor 2\n",
    "\n",
    "```\n",
    "+------------------------+\n",
    "|product_id|partial_count|\n",
    "|          |             |\n",
    "|    1     |      1      |\n",
    "|    3     |      1      |\n",
    "+------------------------+\n",
    "```\n",
    "\n",
    "**6. Sum up the partial counts to get final result.**: HashAggregate(keys=[product_id#67], functions=[count(distinct store_id#68)])\n",
    "\n",
    "Executor 1\n",
    "\n",
    "```\n",
    "+----------------------------+\n",
    "|product_id|distinct_store_id|\n",
    "|          |                 |\n",
    "|    2     |      2          |\n",
    "+----------------------------+\n",
    "```\n",
    "\n",
    "Executor 2\n",
    "\n",
    "```\n",
    "+----------------------------+\n",
    "|product_id|distinct_store_id|\n",
    "|          |                 |\n",
    "|    1     |      1          |\n",
    "|    3     |      1          |\n",
    "+----------------------------+\n",
    "```\n",
    "\n",
    "**Takeaways:**\n",
    "\n",
    "- ``HashAggregate`` -> ``Exchange`` -> ``HashAggregate`` is a common pattern.\n",
    "\n",
    "- If it is possible, Spark will do local aggregation first to decrease the amount of data that needs to be shuffled.\n",
    "\n",
    "- Groupby aggregations are your friend since they avoid the full shuffle."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "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.6.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}