04_processing_column_data/19_using_case_and_when.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Using CASE and WHEN\n",
    "\n",
    "Let us understand how to perform conditional operations using `CASE` and `WHEN` in Spark."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "tags": [
     "remove-cell"
    ]
   },
   "outputs": [
    {
     "data": {
      "text/html": [
       "<iframe width=\"560\" height=\"315\" src=\"https://www.youtube.com/embed/ABMRqy_YO2w?rel=0&amp;controls=1&amp;showinfo=0\" frameborder=\"0\" allowfullscreen></iframe>\n"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "%%HTML\n",
    "<iframe width=\"560\" height=\"315\" src=\"https://www.youtube.com/embed/ABMRqy_YO2w?rel=0&amp;controls=1&amp;showinfo=0\" frameborder=\"0\" allowfullscreen></iframe>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "* `CASE` and `WHEN` is typically used to apply transformations based up on conditions. We can use `CASE` and `WHEN` similar to SQL using `expr` or `selectExpr`.\n",
    "* If we want to use APIs, Spark provides functions such as `when` and `otherwise`. `when` is available as part of `pyspark.sql.functions`. On top of column type that is generated using `when` we should be able to invoke `otherwise`."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Let us start spark context for this Notebook so that we can execute the code provided. You can sign up for our [10 node state of the art cluster/labs](https://labs.itversity.com/plans) to learn Spark SQL using our unique integrated LMS."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "from pyspark.sql import SparkSession\n",
    "\n",
    "import getpass\n",
    "username = getpass.getuser()\n",
    "\n",
    "spark = SparkSession. \\\n",
    "    builder. \\\n",
    "    config('spark.ui.port', '0'). \\\n",
    "    config(\"spark.sql.warehouse.dir\", f\"/user/{username}/warehouse\"). \\\n",
    "    enableHiveSupport(). \\\n",
    "    appName(f'{username} | Python - Processing Column Data'). \\\n",
    "    master('yarn'). \\\n",
    "    getOrCreate()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "If you are going to use CLIs, you can use Spark SQL using one of the 3 approaches.\n",
    "\n",
    "**Using Spark SQL**\n",
    "\n",
    "```\n",
    "spark2-sql \\\n",
    "    --master yarn \\\n",
    "    --conf spark.ui.port=0 \\\n",
    "    --conf spark.sql.warehouse.dir=/user/${USER}/warehouse\n",
    "```\n",
    "\n",
    "**Using Scala**\n",
    "\n",
    "```\n",
    "spark2-shell \\\n",
    "    --master yarn \\\n",
    "    --conf spark.ui.port=0 \\\n",
    "    --conf spark.sql.warehouse.dir=/user/${USER}/warehouse\n",
    "```\n",
    "\n",
    "**Using Pyspark**\n",
    "\n",
    "```\n",
    "pyspark2 \\\n",
    "    --master yarn \\\n",
    "    --conf spark.ui.port=0 \\\n",
    "    --conf spark.sql.warehouse.dir=/user/${USER}/warehouse\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "employees = [(1, \"Scott\", \"Tiger\", 1000.0, 10,\n",
    "                      \"united states\", \"+1 123 456 7890\", \"123 45 6789\"\n",
    "                     ),\n",
    "                     (2, \"Henry\", \"Ford\", 1250.0, None,\n",
    "                      \"India\", \"+91 234 567 8901\", \"456 78 9123\"\n",
    "                     ),\n",
    "                     (3, \"Nick\", \"Junior\", 750.0, '',\n",
    "                      \"united KINGDOM\", \"+44 111 111 1111\", \"222 33 4444\"\n",
    "                     ),\n",
    "                     (4, \"Bill\", \"Gomes\", 1500.0, 10,\n",
    "                      \"AUSTRALIA\", \"+61 987 654 3210\", \"789 12 6118\"\n",
    "                     )\n",
    "                ]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "employeesDF = spark. \\\n",
    "    createDataFrame(employees,\n",
    "                    schema=\"\"\"employee_id INT, first_name STRING, \n",
    "                    last_name STRING, salary FLOAT, bonus STRING, nationality STRING,\n",
    "                    phone_number STRING, ssn STRING\"\"\"\n",
    "                   )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "+-----------+----------+---------+------+-----+--------------+----------------+-----------+\n",
      "|employee_id|first_name|last_name|salary|bonus|   nationality|    phone_number|        ssn|\n",
      "+-----------+----------+---------+------+-----+--------------+----------------+-----------+\n",
      "|          1|     Scott|    Tiger|1000.0|   10| united states| +1 123 456 7890|123 45 6789|\n",
      "|          2|     Henry|     Ford|1250.0| null|         India|+91 234 567 8901|456 78 9123|\n",
      "|          3|      Nick|   Junior| 750.0|     |united KINGDOM|+44 111 111 1111|222 33 4444|\n",
      "|          4|      Bill|    Gomes|1500.0|   10|     AUSTRALIA|+61 987 654 3210|789 12 6118|\n",
      "+-----------+----------+---------+------+-----+--------------+----------------+-----------+\n",
      "\n"
     ]
    }
   ],
   "source": [
    "employeesDF.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "* Let us transform bonus to 0 in case of null or empty, otherwise return the bonus amount."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "from pyspark.sql.functions import coalesce, lit, col"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "+-----------+----------+---------+------+-----+--------------+----------------+-----------+------+\n",
      "|employee_id|first_name|last_name|salary|bonus|   nationality|    phone_number|        ssn|bonus1|\n",
      "+-----------+----------+---------+------+-----+--------------+----------------+-----------+------+\n",
      "|          1|     Scott|    Tiger|1000.0|   10| united states| +1 123 456 7890|123 45 6789|    10|\n",
      "|          2|     Henry|     Ford|1250.0| null|         India|+91 234 567 8901|456 78 9123|     0|\n",
      "|          3|      Nick|   Junior| 750.0|     |united KINGDOM|+44 111 111 1111|222 33 4444|     0|\n",
      "|          4|      Bill|    Gomes|1500.0|   10|     AUSTRALIA|+61 987 654 3210|789 12 6118|    10|\n",
      "+-----------+----------+---------+------+-----+--------------+----------------+-----------+------+\n",
      "\n"
     ]
    }
   ],
   "source": [
    "employeesDF. \\\n",
    "    withColumn('bonus1', coalesce(col('bonus').cast('int'), lit(0))). \\\n",
    "    show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "from pyspark.sql.functions import expr"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "+-----------+----------+---------+------+-----+--------------+----------------+-----------+\n",
      "|employee_id|first_name|last_name|salary|bonus|   nationality|    phone_number|        ssn|\n",
      "+-----------+----------+---------+------+-----+--------------+----------------+-----------+\n",
      "|          1|     Scott|    Tiger|1000.0|   10| united states| +1 123 456 7890|123 45 6789|\n",
      "|          2|     Henry|     Ford|1250.0|    0|         India|+91 234 567 8901|456 78 9123|\n",
      "|          3|      Nick|   Junior| 750.0|    0|united KINGDOM|+44 111 111 1111|222 33 4444|\n",
      "|          4|      Bill|    Gomes|1500.0|   10|     AUSTRALIA|+61 987 654 3210|789 12 6118|\n",
      "+-----------+----------+---------+------+-----+--------------+----------------+-----------+\n",
      "\n"
     ]
    }
   ],
   "source": [
    "employeesDF. \\\n",
    "    withColumn(\n",
    "        'bonus', \n",
    "        expr(\"\"\"\n",
    "            CASE WHEN bonus IS NULL OR bonus = '' THEN 0\n",
    "            ELSE bonus\n",
    "            END\n",
    "            \"\"\")\n",
    "    ). \\\n",
    "    show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "from pyspark.sql.functions import when"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "\u001b[0;31mSignature:\u001b[0m \u001b[0mwhen\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mcondition\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mvalue\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
       "\u001b[0;31mDocstring:\u001b[0m\n",
       "Evaluates a list of conditions and returns one of multiple possible result expressions.\n",
       "If :func:`Column.otherwise` is not invoked, None is returned for unmatched conditions.\n",
       "\n",
       ":param condition: a boolean :class:`Column` expression.\n",
       ":param value: a literal value, or a :class:`Column` expression.\n",
       "\n",
       ">>> df.select(when(df['age'] == 2, 3).otherwise(4).alias(\"age\")).collect()\n",
       "[Row(age=3), Row(age=4)]\n",
       "\n",
       ">>> df.select(when(df.age == 2, df.age + 1).alias(\"age\")).collect()\n",
       "[Row(age=3), Row(age=None)]\n",
       "\n",
       ".. versionadded:: 1.4\n",
       "\u001b[0;31mFile:\u001b[0m      /usr/hdp/current/spark2-client/python/pyspark/sql/functions.py\n",
       "\u001b[0;31mType:\u001b[0m      function\n"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "when?"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "+-----------+----------+---------+------+-----+--------------+----------------+-----------+\n",
      "|employee_id|first_name|last_name|salary|bonus|   nationality|    phone_number|        ssn|\n",
      "+-----------+----------+---------+------+-----+--------------+----------------+-----------+\n",
      "|          1|     Scott|    Tiger|1000.0|   10| united states| +1 123 456 7890|123 45 6789|\n",
      "|          2|     Henry|     Ford|1250.0|    0|         India|+91 234 567 8901|456 78 9123|\n",
      "|          3|      Nick|   Junior| 750.0|    0|united KINGDOM|+44 111 111 1111|222 33 4444|\n",
      "|          4|      Bill|    Gomes|1500.0|   10|     AUSTRALIA|+61 987 654 3210|789 12 6118|\n",
      "+-----------+----------+---------+------+-----+--------------+----------------+-----------+\n",
      "\n"
     ]
    }
   ],
   "source": [
    "employeesDF. \\\n",
    "    withColumn(\n",
    "        'bonus',\n",
    "        when((col('bonus').isNull()) | (col('bonus') == lit('')), 0).otherwise(col('bonus'))\n",
    "    ). \\\n",
    "    show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "* Create a dataframe using list called as persons and categorize them based up on following rules.\n",
    "\n",
    "|Age range|Category|\n",
    "|---------|--------|\n",
    "|0 to 2 Months|New Born|\n",
    "|2+ Months to 12 Months|Infant|\n",
    "|12+ Months to 48 Months|Toddler|\n",
    "|48+ Months to 144 Months|Kids|\n",
    "|144+ Months|Teenager or Adult|"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "persons = [\n",
    "    (1, 1),\n",
    "    (2, 13),\n",
    "    (3, 18),\n",
    "    (4, 60),\n",
    "    (5, 120),\n",
    "    (6, 0),\n",
    "    (7, 12),\n",
    "    (8, 160)\n",
    "]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "personsDF = spark.createDataFrame(persons, schema='id INT, age INT')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "+---+---+\n",
      "| id|age|\n",
      "+---+---+\n",
      "|  1|  1|\n",
      "|  2| 13|\n",
      "|  3| 18|\n",
      "|  4| 60|\n",
      "|  5|120|\n",
      "|  6|  0|\n",
      "|  7| 12|\n",
      "|  8|160|\n",
      "+---+---+\n",
      "\n"
     ]
    }
   ],
   "source": [
    "personsDF.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "+---+---+-----------------+\n",
      "| id|age|         category|\n",
      "+---+---+-----------------+\n",
      "|  1|  1|         New Born|\n",
      "|  2| 13|          Toddler|\n",
      "|  3| 18|          Toddler|\n",
      "|  4| 60|              Kid|\n",
      "|  5|120|              Kid|\n",
      "|  6|  0|         New Born|\n",
      "|  7| 12|           Infant|\n",
      "|  8|160|Teenager or Adult|\n",
      "+---+---+-----------------+\n",
      "\n"
     ]
    }
   ],
   "source": [
    "personsDF. \\\n",
    "    withColumn(\n",
    "        'category',\n",
    "        expr(\"\"\"\n",
    "            CASE\n",
    "            WHEN age BETWEEN 0 AND 2 THEN 'New Born'\n",
    "            WHEN age > 2 AND age <= 12 THEN 'Infant'\n",
    "            WHEN age > 12 AND age <= 48 THEN 'Toddler'\n",
    "            WHEN age > 48 AND age <= 144 THEN 'Kid'\n",
    "            ELSE 'Teenager or Adult'\n",
    "            END\n",
    "        \"\"\")\n",
    "    ). \\\n",
    "    show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "+---+---+-----------------+\n",
      "| id|age|         category|\n",
      "+---+---+-----------------+\n",
      "|  1|  1|         New Born|\n",
      "|  2| 13|          Toddler|\n",
      "|  3| 18|          Toddler|\n",
      "|  4| 60|              Kid|\n",
      "|  5|120|              Kid|\n",
      "|  6|  0|         New Born|\n",
      "|  7| 12|           Infant|\n",
      "|  8|160|Teenager or Adult|\n",
      "+---+---+-----------------+\n",
      "\n"
     ]
    }
   ],
   "source": [
    "personsDF. \\\n",
    "    withColumn(\n",
    "        'category',\n",
    "        when(col('age').between(0, 2), 'New Born').\n",
    "        when((col('age') > 2) & (col('age') <= 12), 'Infant').\n",
    "        when((col('age') > 12) & (col('age') <= 48), 'Toddler').\n",
    "        when((col('age') > 48) & (col('age') <= 144), 'Kid').\n",
    "        otherwise('Teenager or Adult')\n",
    "    ). \\\n",
    "    show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Pyspark 2",
   "language": "python",
   "name": "pyspark2"
  },
  "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.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}