src/Chapter 7 - High Performance Code/Section 1.2 - Joins on Skewed Data (Null Keys).ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "A `skewed dataset` is defined by a dataset that has a class imbalance, this leads to poor or failing spark jobs that often get a `OOM` (out of memory) error.\n",
    "\n",
    "When performing a `join` onto a `skewed dataset` it's usually the case where there is an imbalance on the `key`(s) on which the join is performed on. This results in a majority of the data falls onto a single partition, which will take longer to complete than the other partitions.\n",
    "\n",
    "Some hints to detect skewness is:\n",
    "1. The `key`(s) consist mainly of `null` values which fall onto a single partition.\n",
    "2. There is a subset of values for the `key`(s) that makeup the high percentage of the total keys which fall onto a single partition.\n",
    "\n",
    "We go through both these cases and see how we can combat it."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Library Imports"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "from pyspark.sql import SparkSession\n",
    "from pyspark.sql import functions as F"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Template"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "spark = (\n",
    "    SparkSession.builder\n",
    "    .master(\"local\")\n",
    "    .appName(\"Exploring Joins\")\n",
    "    .config(\"spark.some.config.option\", \"some-value\")\n",
    "    .getOrCreate()\n",
    ")\n",
    "\n",
    "sc = spark.sparkContext"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Situation 1: Null Keys"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Inital Datasets"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "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>id</th>\n",
       "      <th>card_id</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>1</td>\n",
       "      <td>NaN</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>2</td>\n",
       "      <td>NaN</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>3</td>\n",
       "      <td>1.0</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "   id  card_id\n",
       "0   1      NaN\n",
       "1   2      NaN\n",
       "2   3      1.0"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "customers = spark.createDataFrame([\n",
    "    (1, None), \n",
    "    (2, None), \n",
    "    (3, 1),\n",
    "], [\"id\", \"card_id\"])\n",
    "\n",
    "customers.toPandas()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "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>card_id</th>\n",
       "      <th>first_name</th>\n",
       "      <th>last_name</th>\n",
       "      <th>age</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>1</td>\n",
       "      <td>john</td>\n",
       "      <td>doe</td>\n",
       "      <td>21</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>2</td>\n",
       "      <td>rick</td>\n",
       "      <td>roll</td>\n",
       "      <td>10</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>3</td>\n",
       "      <td>bob</td>\n",
       "      <td>brown</td>\n",
       "      <td>2</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "   card_id first_name last_name  age\n",
       "0        1       john       doe   21\n",
       "1        2       rick      roll   10\n",
       "2        3        bob     brown    2"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "cards = spark.createDataFrame([\n",
    "    (1, \"john\", \"doe\", 21), \n",
    "    (2, \"rick\", \"roll\", 10), \n",
    "    (3, \"bob\", \"brown\", 2)\n",
    "], [\"card_id\", \"first_name\", \"last_name\", \"age\"])\n",
    "\n",
    "cards.toPandas()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Option #1: Join Regularly"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "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>card_id</th>\n",
       "      <th>id</th>\n",
       "      <th>first_name</th>\n",
       "      <th>last_name</th>\n",
       "      <th>age</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>NaN</td>\n",
       "      <td>1</td>\n",
       "      <td>None</td>\n",
       "      <td>None</td>\n",
       "      <td>NaN</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>NaN</td>\n",
       "      <td>2</td>\n",
       "      <td>None</td>\n",
       "      <td>None</td>\n",
       "      <td>NaN</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>1.0</td>\n",
       "      <td>3</td>\n",
       "      <td>john</td>\n",
       "      <td>doe</td>\n",
       "      <td>21.0</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "   card_id  id first_name last_name   age\n",
       "0      NaN   1       None      None   NaN\n",
       "1      NaN   2       None      None   NaN\n",
       "2      1.0   3       john       doe  21.0"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "df = customers.join(cards, \"card_id\", \"left\")\n",
    "\n",
    "df.toPandas()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "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>card_id</th>\n",
       "      <th>id</th>\n",
       "      <th>first_name</th>\n",
       "      <th>last_name</th>\n",
       "      <th>age</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>1</td>\n",
       "      <td>3</td>\n",
       "      <td>john</td>\n",
       "      <td>doe</td>\n",
       "      <td>21</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "   card_id  id first_name last_name  age\n",
       "0        1   3       john       doe   21"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "df = customers.join(cards, \"card_id\")\n",
    "\n",
    "df.toPandas()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**What Happened**:\n",
    "* Rows that didn't join up were brought to the join.\n",
    "\n",
    "* For a `left join`, they will get `Null` values for the right side columns, what's the point of being them in?\n",
    "* For a `inner join`, they rows will get dropped, so again what's the point of being them in?\n",
    "\n",
    "**Results**:\n",
    "* We brought more rows to the join than we had to. These rows get normally get put onto a single partition. \n",
    "* If the data is large enough and the percentage of keys that are null is high. The program could OOM out."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Option #2: Filter Null Keys First, then Join, then Union"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "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>id</th>\n",
       "      <th>card_id</th>\n",
       "      <th>first_name</th>\n",
       "      <th>last_name</th>\n",
       "      <th>age</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>1</td>\n",
       "      <td>NaN</td>\n",
       "      <td>None</td>\n",
       "      <td>None</td>\n",
       "      <td>NaN</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>2</td>\n",
       "      <td>NaN</td>\n",
       "      <td>None</td>\n",
       "      <td>None</td>\n",
       "      <td>NaN</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>3</td>\n",
       "      <td>1.0</td>\n",
       "      <td>john</td>\n",
       "      <td>doe</td>\n",
       "      <td>21.0</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "   id  card_id first_name last_name   age\n",
       "0   1      NaN       None      None   NaN\n",
       "1   2      NaN       None      None   NaN\n",
       "2   3      1.0       john       doe  21.0"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "def null_skew_helper(left, right, key):\n",
    "    \"\"\"\n",
    "    Steps:\n",
    "        1. Filter out the null rows.\n",
    "        2. Create the columns you would get from the join.\n",
    "        3. Join the tables.\n",
    "        4. Union the null rows to joined table.\n",
    "    \"\"\"\n",
    "    df1 = left.where(F.col(key).isNull())\n",
    "    for f in right.schema.fields:\n",
    "            df1 = df1.withColumn(f.name, F.lit(None).cast(f.dataType))\n",
    "    \n",
    "    df2 = left.where(F.col(key).isNotNull())\n",
    "    df2 = df2.join(right, key, \"left\")\n",
    "    \n",
    "    return df1.union(df2.select(df1.columns))\n",
    "    \n",
    "    \n",
    "df = null_skew_helper(customers, cards, \"card_id\")\n",
    "\n",
    "df.toPandas()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "== Physical Plan ==\n",
      "Union\n",
      ":- *(1) Project [id#0L, null AS card_id#23L, null AS first_name#26, null AS last_name#30, null AS age#35L]\n",
      ":  +- *(1) Filter isnull(card_id#1L)\n",
      ":     +- Scan ExistingRDD[id#0L,card_id#1L]\n",
      "+- *(5) Project [id#0L, card_id#1L, first_name#5, last_name#6, age#7L]\n",
      "   +- SortMergeJoin [card_id#1L], [card_id#4L], LeftOuter\n",
      "      :- *(3) Sort [card_id#1L ASC NULLS FIRST], false, 0\n",
      "      :  +- Exchange hashpartitioning(card_id#1L, 200)\n",
      "      :     +- *(2) Filter isnotnull(card_id#1L)\n",
      "      :        +- Scan ExistingRDD[id#0L,card_id#1L]\n",
      "      +- *(4) Sort [card_id#4L ASC NULLS FIRST], false, 0\n",
      "         +- Exchange hashpartitioning(card_id#4L, 200)\n",
      "            +- Scan ExistingRDD[card_id#4L,first_name#5,last_name#6,age#7L]\n"
     ]
    }
   ],
   "source": [
    "df.explain()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**What Happened**:\n",
    "* We seperated the data into 2 sets:\n",
    "  * one where the `key`s are not `null`.\n",
    "  * one where the `key`s are `null`.\n",
    "* We perform the join on the set where the keys are not null, then union it back with the set where the keys are null. (This step is not necessary when doing an inner join).\n",
    "\n",
    "**Results**:\n",
    "* We brought less data to the join.\n",
    "* We read the data twice; more time was spent on reading data from disk."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Option #3: Cache the Table, Filter Null Keys First, then Join, then Union"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Helper Function**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "def null_skew_helper(left, right, key):\n",
    "    \"\"\"\n",
    "    Steps:\n",
    "        1. Cache table.\n",
    "        2. Filter out the null rows.\n",
    "        3. Create the columns you would get from the join.\n",
    "        4. Join the tables.\n",
    "        5. Union the null rows to joined table.\n",
    "    \"\"\"\n",
    "    left = left.cache()\n",
    "    \n",
    "    df1 = left.where(F.col(key).isNull())\n",
    "    for f in right.schema.fields:\n",
    "            df1 = df1.withColumn(f.name, F.lit(None).cast(f.dataType))\n",
    "    \n",
    "    df2 = left.where(F.col(key).isNotNull())\n",
    "    df2 = df2.join(right, key, \"left\")\n",
    "    \n",
    "    return df1.union(df2.select(df1.columns))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "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>id</th>\n",
       "      <th>card_id</th>\n",
       "      <th>first_name</th>\n",
       "      <th>last_name</th>\n",
       "      <th>age</th>\n",
       "    </tr>\n",
       "  </thead>\n",
       "  <tbody>\n",
       "    <tr>\n",
       "      <th>0</th>\n",
       "      <td>1</td>\n",
       "      <td>NaN</td>\n",
       "      <td>None</td>\n",
       "      <td>None</td>\n",
       "      <td>NaN</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>1</th>\n",
       "      <td>2</td>\n",
       "      <td>NaN</td>\n",
       "      <td>None</td>\n",
       "      <td>None</td>\n",
       "      <td>NaN</td>\n",
       "    </tr>\n",
       "    <tr>\n",
       "      <th>2</th>\n",
       "      <td>3</td>\n",
       "      <td>1.0</td>\n",
       "      <td>john</td>\n",
       "      <td>doe</td>\n",
       "      <td>21.0</td>\n",
       "    </tr>\n",
       "  </tbody>\n",
       "</table>\n",
       "</div>"
      ],
      "text/plain": [
       "   id  card_id first_name last_name   age\n",
       "0   1      NaN       None      None   NaN\n",
       "1   2      NaN       None      None   NaN\n",
       "2   3      1.0       john       doe  21.0"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "df = null_skew_helper(customers, cards, \"card_id\")\n",
    "\n",
    "df.toPandas()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "== Physical Plan ==\n",
      "Union\n",
      ":- *(1) Project [id#0L, null AS card_id#68L, null AS first_name#71, null AS last_name#75, null AS age#80L]\n",
      ":  +- *(1) Filter isnull(card_id#1L)\n",
      ":     +- *(1) InMemoryTableScan [card_id#1L, id#0L], [isnull(card_id#1L)]\n",
      ":           +- InMemoryRelation [id#0L, card_id#1L], true, 10000, StorageLevel(disk, memory, deserialized, 1 replicas)\n",
      ":                 +- Scan ExistingRDD[id#0L,card_id#1L]\n",
      "+- *(5) Project [id#0L, card_id#1L, first_name#5, last_name#6, age#7L]\n",
      "   +- SortMergeJoin [card_id#1L], [card_id#4L], LeftOuter\n",
      "      :- *(3) Sort [card_id#1L ASC NULLS FIRST], false, 0\n",
      "      :  +- Exchange hashpartitioning(card_id#1L, 200)\n",
      "      :     +- *(2) Filter isnotnull(card_id#1L)\n",
      "      :        +- *(2) InMemoryTableScan [id#0L, card_id#1L], [isnotnull(card_id#1L)]\n",
      "      :              +- InMemoryRelation [id#0L, card_id#1L], true, 10000, StorageLevel(disk, memory, deserialized, 1 replicas)\n",
      "      :                    +- Scan ExistingRDD[id#0L,card_id#1L]\n",
      "      +- *(4) Sort [card_id#4L ASC NULLS FIRST], false, 0\n",
      "         +- Exchange hashpartitioning(card_id#4L, 200)\n",
      "            +- Scan ExistingRDD[card_id#4L,first_name#5,last_name#6,age#7L]\n"
     ]
    }
   ],
   "source": [
    "df.explain()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**What Happened**:\n",
    "* Similar to option #2, but we did a `InMemoryTableScan` instead of two reads of the data.\n",
    "\n",
    "**Results**:\n",
    "* We brought less data to the join.\n",
    "* We did 1 less read, but we used more memory."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Summary\n",
    "\n",
    "All to say:\n",
    "* It's definitely better to bring less data to a join, so performing a filter for `null keys` before the join is definitely suggested.\n",
    "* For `left join`s:\n",
    "    * By doing a union, this will result in an extra read of data or memory usage.\n",
    "    * Decide what you can afford; the extra read vs memory usage and `cache` the table before the `filter`.\n",
    "\n",
    "Always check the spread the values for the `join key`, to detect if there's any skew and pre filters that can be performed."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 2",
   "language": "python",
   "name": "python2"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 2
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython2",
   "version": "2.7.15"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}