spark/2_reading_query_plans.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {},
     "inputWidgets": {},
     "nuid": "12fa052b-65ab-42fe-92e9-866a91ab0479",
     "showTitle": false,
     "title": ""
    }
   },
   "source": [
    "# Imports & Configuration"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import warnings\n",
    "warnings.filterwarnings(\"ignore\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "from IPython.core.interactiveshell import InteractiveShell\n",
    "InteractiveShell.ast_node_interactivity = 'all'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {
      "byteLimit": 2048000,
      "rowLimit": 10000
     },
     "inputWidgets": {},
     "nuid": "e02171b5-8b8d-4b81-8826-2fb3ee1a8a24",
     "showTitle": false,
     "title": ""
    }
   },
   "outputs": [],
   "source": [
    "from pyspark.sql.types import *\n",
    "import pyspark.sql.functions as F\n",
    "from pyspark.sql import SparkSession"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "23/07/21 18:42:50 WARN Utils: Your hostname, Afaques-MacBook-Pro.local resolves to a loopback address: 127.0.0.1; using 10.0.0.6 instead (on interface en0)\n",
      "23/07/21 18:42:50 WARN Utils: Set SPARK_LOCAL_IP if you need to bind to another address\n",
      "Using Spark's default log4j profile: org/apache/spark/log4j-defaults.properties\n",
      "Setting default log level to \"WARN\".\n",
      "To adjust logging level use sc.setLogLevel(newLevel). For SparkR, use setLogLevel(newLevel).\n",
      "23/07/21 18:42:51 WARN NativeCodeLoader: Unable to load native-hadoop library for your platform... using builtin-java classes where applicable\n"
     ]
    }
   ],
   "source": [
    "spark = SparkSession.builder.master(\"local[*]\").getOrCreate()\n",
    "sc = spark.sparkContext\n",
    "sc.setLogLevel(\"ERROR\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Reading File"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "                                                                                \r"
     ]
    }
   ],
   "source": [
    "transactions_file = \"../data/data_skew/transactions.parquet\"\n",
    "df_transactions = spark.read.parquet(transactions_file)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "+----------+----------+----------+---------------+----------+----+-----+---+-------------+------+-----------+\n",
      "|cust_id   |start_date|end_date  |txn_id         |date      |year|month|day|expense_type |amt   |city       |\n",
      "+----------+----------+----------+---------------+----------+----+-----+---+-------------+------+-----------+\n",
      "|C0YDPQWPBJ|2010-07-01|2018-12-01|TZ5SMKZY9S03OQJ|2018-10-07|2018|10   |7  |Entertainment|10.42 |boston     |\n",
      "|C0YDPQWPBJ|2010-07-01|2018-12-01|TYIAPPNU066CJ5R|2016-03-27|2016|3    |27 |Motor/Travel |44.34 |portland   |\n",
      "|C0YDPQWPBJ|2010-07-01|2018-12-01|TETSXIK4BLXHJ6W|2011-04-11|2011|4    |11 |Entertainment|3.18  |chicago    |\n",
      "|C0YDPQWPBJ|2010-07-01|2018-12-01|TQKL1QFJY3EM8LO|2018-02-22|2018|2    |22 |Groceries    |268.97|los_angeles|\n",
      "|C0YDPQWPBJ|2010-07-01|2018-12-01|TYL6DFP09PPXMVB|2010-10-16|2010|10   |16 |Entertainment|2.66  |chicago    |\n",
      "+----------+----------+----------+---------------+----------+----+-----+---+-------------+------+-----------+\n",
      "only showing top 5 rows\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\r",
      "[Stage 1:>                                                          (0 + 1) / 1]\r",
      "\r",
      "                                                                                \r"
     ]
    }
   ],
   "source": [
    "df_transactions.show(5, False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "customers_file = \"../data/data_skew/customers.parquet\"\n",
    "df_customers = spark.read.parquet(customers_file)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "+----------+-------------+---+------+----------+-----+-----------+\n",
      "|cust_id   |name         |age|gender|birthday  |zip  |city       |\n",
      "+----------+-------------+---+------+----------+-----+-----------+\n",
      "|C007YEYTX9|Aaron Abbott |34 |Female|7/13/1991 |97823|boston     |\n",
      "|C00B971T1J|Aaron Austin |37 |Female|12/16/2004|30332|chicago    |\n",
      "|C00WRSJF1Q|Aaron Barnes |29 |Female|3/11/1977 |23451|denver     |\n",
      "|C01AZWQMF3|Aaron Barrett|31 |Male  |7/9/1998  |46613|los_angeles|\n",
      "|C01BKUFRHA|Aaron Becker |54 |Male  |11/24/1979|40284|san_diego  |\n",
      "+----------+-------------+---+------+----------+-----+-----------+\n",
      "only showing top 5 rows\n",
      "\n"
     ]
    }
   ],
   "source": [
    "df_customers.show(5, False)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Spark's Query Plans\n",
    "\n",
    "Image Credits: `Databricks`"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "![Spark Execution](../data/spark-execution.png)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Narrow Transformations\n",
    "- `filter` rows where `city='boston'`\n",
    "- `add` a new column: adding `first_name` and `last_name`\n",
    "- `alter` an exisitng column: adding 5 to `age` column\n",
    "- `select` relevant columns"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "+----------+----------+---------+----+------+---------+\n",
      "|cust_id   |first_name|last_name|age |gender|birthday |\n",
      "+----------+----------+---------+----+------+---------+\n",
      "|C007YEYTX9|Aaron     |Abbott   |39.0|Female|7/13/1991|\n",
      "|C08XAQUY73|Aaron     |Lambert  |59.0|Female|11/5/1966|\n",
      "|C094P1VXF9|Aaron     |Lindsey  |29.0|Male  |9/21/1990|\n",
      "|C097SHE1EF|Aaron     |Lopez    |27.0|Female|4/18/2001|\n",
      "|C0DTC6436T|Aaron     |Schwartz |57.0|Female|7/9/1962 |\n",
      "+----------+----------+---------+----+------+---------+\n",
      "only showing top 5 rows\n",
      "\n",
      "== Parsed Logical Plan ==\n",
      "'Project ['cust_id, 'first_name, 'last_name, 'age, 'gender, 'birthday]\n",
      "+- Project [cust_id#67, name#68, (cast(age#69 as double) + cast(5 as double)) AS age#129, gender#70, birthday#71, zip#72, city#73, first_name#110, last_name#119]\n",
      "   +- Project [cust_id#67, name#68, age#69, gender#70, birthday#71, zip#72, city#73, first_name#110, split(name#68,  , -1)[1] AS last_name#119]\n",
      "      +- Project [cust_id#67, name#68, age#69, gender#70, birthday#71, zip#72, city#73, split(name#68,  , -1)[0] AS first_name#110]\n",
      "         +- Filter (city#73 = boston)\n",
      "            +- Relation [cust_id#67,name#68,age#69,gender#70,birthday#71,zip#72,city#73] parquet\n",
      "\n",
      "== Analyzed Logical Plan ==\n",
      "cust_id: string, first_name: string, last_name: string, age: double, gender: string, birthday: string\n",
      "Project [cust_id#67, first_name#110, last_name#119, age#129, gender#70, birthday#71]\n",
      "+- Project [cust_id#67, name#68, (cast(age#69 as double) + cast(5 as double)) AS age#129, gender#70, birthday#71, zip#72, city#73, first_name#110, last_name#119]\n",
      "   +- Project [cust_id#67, name#68, age#69, gender#70, birthday#71, zip#72, city#73, first_name#110, split(name#68,  , -1)[1] AS last_name#119]\n",
      "      +- Project [cust_id#67, name#68, age#69, gender#70, birthday#71, zip#72, city#73, split(name#68,  , -1)[0] AS first_name#110]\n",
      "         +- Filter (city#73 = boston)\n",
      "            +- Relation [cust_id#67,name#68,age#69,gender#70,birthday#71,zip#72,city#73] parquet\n",
      "\n",
      "== Optimized Logical Plan ==\n",
      "Project [cust_id#67, split(name#68,  , -1)[0] AS first_name#110, split(name#68,  , -1)[1] AS last_name#119, (cast(age#69 as double) + 5.0) AS age#129, gender#70, birthday#71]\n",
      "+- Filter (isnotnull(city#73) AND (city#73 = boston))\n",
      "   +- Relation [cust_id#67,name#68,age#69,gender#70,birthday#71,zip#72,city#73] parquet\n",
      "\n",
      "== Physical Plan ==\n",
      "*(1) Project [cust_id#67, split(name#68,  , -1)[0] AS first_name#110, split(name#68,  , -1)[1] AS last_name#119, (cast(age#69 as double) + 5.0) AS age#129, gender#70, birthday#71]\n",
      "+- *(1) Filter (isnotnull(city#73) AND (city#73 = boston))\n",
      "   +- *(1) ColumnarToRow\n",
      "      +- FileScan parquet [cust_id#67,name#68,age#69,gender#70,birthday#71,city#73] Batched: true, DataFilters: [isnotnull(city#73), (city#73 = boston)], Format: Parquet, Location: InMemoryFileIndex(1 paths)[file:/Users/afaqueahmad/Documents/youtube/spark-experiments/data/data_..., PartitionFilters: [], PushedFilters: [IsNotNull(city), EqualTo(city,boston)], ReadSchema: struct<cust_id:string,name:string,age:string,gender:string,birthday:string,city:string>\n",
      "\n"
     ]
    }
   ],
   "source": [
    "df_narrow_transform = (\n",
    "    df_customers\n",
    "    .filter(F.col(\"city\") == \"boston\")\n",
    "    .withColumn(\"first_name\", F.split(\"name\", \" \").getItem(0))\n",
    "    .withColumn(\"last_name\", F.split(\"name\", \" \").getItem(1))\n",
    "    .withColumn(\"age\", F.col(\"age\") + F.lit(5))\n",
    "    .select(\"cust_id\", \"first_name\", \"last_name\", \"age\", \"gender\", \"birthday\")\n",
    ")\n",
    "\n",
    "df_narrow_transform.show(5, False)\n",
    "df_narrow_transform.explain(True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Wide Transformations\n",
    "1. Repartition\n",
    "2. Coalesce\n",
    "3. Joins\n",
    "4. GroupBy\n",
    "   - `count`\n",
    "   - `countDistinct`\n",
    "   - `sum`"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 1. Repartition"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "13"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "df_transactions.rdd.getNumPartitions()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "== Parsed Logical Plan ==\n",
      "Repartition 24, true\n",
      "+- Relation [cust_id#0,start_date#1,end_date#2,txn_id#3,date#4,year#5,month#6,day#7,expense_type#8,amt#9,city#10] parquet\n",
      "\n",
      "== Analyzed Logical Plan ==\n",
      "cust_id: string, start_date: string, end_date: string, txn_id: string, date: string, year: string, month: string, day: string, expense_type: string, amt: string, city: string\n",
      "Repartition 24, true\n",
      "+- Relation [cust_id#0,start_date#1,end_date#2,txn_id#3,date#4,year#5,month#6,day#7,expense_type#8,amt#9,city#10] parquet\n",
      "\n",
      "== Optimized Logical Plan ==\n",
      "Repartition 24, true\n",
      "+- Relation [cust_id#0,start_date#1,end_date#2,txn_id#3,date#4,year#5,month#6,day#7,expense_type#8,amt#9,city#10] parquet\n",
      "\n",
      "== Physical Plan ==\n",
      "AdaptiveSparkPlan isFinalPlan=false\n",
      "+- Exchange RoundRobinPartitioning(24), REPARTITION_BY_NUM, [id=#72]\n",
      "   +- FileScan parquet [cust_id#0,start_date#1,end_date#2,txn_id#3,date#4,year#5,month#6,day#7,expense_type#8,amt#9,city#10] Batched: true, DataFilters: [], Format: Parquet, Location: InMemoryFileIndex(1 paths)[file:/Users/afaqueahmad/Documents/youtube/spark-experiments/data/data_..., PartitionFilters: [], PushedFilters: [], ReadSchema: struct<cust_id:string,start_date:string,end_date:string,txn_id:string,date:string,year:string,mon...\n",
      "\n"
     ]
    }
   ],
   "source": [
    "df_transactions.repartition(24).explain(True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 2. Coalesce"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "== Parsed Logical Plan ==\n",
      "Repartition 5, false\n",
      "+- Relation [cust_id#0,start_date#1,end_date#2,txn_id#3,date#4,year#5,month#6,day#7,expense_type#8,amt#9,city#10] parquet\n",
      "\n",
      "== Analyzed Logical Plan ==\n",
      "cust_id: string, start_date: string, end_date: string, txn_id: string, date: string, year: string, month: string, day: string, expense_type: string, amt: string, city: string\n",
      "Repartition 5, false\n",
      "+- Relation [cust_id#0,start_date#1,end_date#2,txn_id#3,date#4,year#5,month#6,day#7,expense_type#8,amt#9,city#10] parquet\n",
      "\n",
      "== Optimized Logical Plan ==\n",
      "Repartition 5, false\n",
      "+- Relation [cust_id#0,start_date#1,end_date#2,txn_id#3,date#4,year#5,month#6,day#7,expense_type#8,amt#9,city#10] parquet\n",
      "\n",
      "== Physical Plan ==\n",
      "Coalesce 5\n",
      "+- *(1) ColumnarToRow\n",
      "   +- FileScan parquet [cust_id#0,start_date#1,end_date#2,txn_id#3,date#4,year#5,month#6,day#7,expense_type#8,amt#9,city#10] Batched: true, DataFilters: [], Format: Parquet, Location: InMemoryFileIndex(1 paths)[file:/Users/afaqueahmad/Documents/youtube/spark-experiments/data/data_..., PartitionFilters: [], PushedFilters: [], ReadSchema: struct<cust_id:string,start_date:string,end_date:string,txn_id:string,date:string,year:string,mon...\n",
      "\n"
     ]
    }
   ],
   "source": [
    "df_transactions.coalesce(5).explain(True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Why doesn't `.coalesce()` explicitly show the partitioning scheme?\n",
    "\n",
    "`.coalesce` doesn't show the partitioning scheme e.g. `RoundRobinPartitioning` because: \n",
    "- The operation only minimizes data movement by merging into fewer partitions, it doesn't do any shuffling.\n",
    "- Because no shuffling is done, the partitioning scheme remains the same as the original DataFrame and Spark doesn't include it explicitly in it's plan as the partitioning scheme is unaffected by `.coalesce`"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 3. Joins"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "spark.conf.set(\"spark.sql.autoBroadcastJoinThreshold\", -1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "df_joined = (\n",
    "    df_transactions.join(\n",
    "        df_customers,\n",
    "        how=\"inner\",\n",
    "        on=\"cust_id\"\n",
    "    )\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "== Parsed Logical Plan ==\n",
      "'Join UsingJoin(Inner,Buffer(cust_id))\n",
      ":- Relation [cust_id#0,start_date#1,end_date#2,txn_id#3,date#4,year#5,month#6,day#7,expense_type#8,amt#9,city#10] parquet\n",
      "+- Relation [cust_id#67,name#68,age#69,gender#70,birthday#71,zip#72,city#73] parquet\n",
      "\n",
      "== Analyzed Logical Plan ==\n",
      "cust_id: string, start_date: string, end_date: string, txn_id: string, date: string, year: string, month: string, day: string, expense_type: string, amt: string, city: string, name: string, age: string, gender: string, birthday: string, zip: string, city: string\n",
      "Project [cust_id#0, start_date#1, end_date#2, txn_id#3, date#4, year#5, month#6, day#7, expense_type#8, amt#9, city#10, name#68, age#69, gender#70, birthday#71, zip#72, city#73]\n",
      "+- Join Inner, (cust_id#0 = cust_id#67)\n",
      "   :- Relation [cust_id#0,start_date#1,end_date#2,txn_id#3,date#4,year#5,month#6,day#7,expense_type#8,amt#9,city#10] parquet\n",
      "   +- Relation [cust_id#67,name#68,age#69,gender#70,birthday#71,zip#72,city#73] parquet\n",
      "\n",
      "== Optimized Logical Plan ==\n",
      "Project [cust_id#0, start_date#1, end_date#2, txn_id#3, date#4, year#5, month#6, day#7, expense_type#8, amt#9, city#10, name#68, age#69, gender#70, birthday#71, zip#72, city#73]\n",
      "+- Join Inner, (cust_id#0 = cust_id#67)\n",
      "   :- Filter isnotnull(cust_id#0)\n",
      "   :  +- Relation [cust_id#0,start_date#1,end_date#2,txn_id#3,date#4,year#5,month#6,day#7,expense_type#8,amt#9,city#10] parquet\n",
      "   +- Filter isnotnull(cust_id#67)\n",
      "      +- Relation [cust_id#67,name#68,age#69,gender#70,birthday#71,zip#72,city#73] parquet\n",
      "\n",
      "== Physical Plan ==\n",
      "AdaptiveSparkPlan isFinalPlan=false\n",
      "+- Project [cust_id#0, start_date#1, end_date#2, txn_id#3, date#4, year#5, month#6, day#7, expense_type#8, amt#9, city#10, name#68, age#69, gender#70, birthday#71, zip#72, city#73]\n",
      "   +- SortMergeJoin [cust_id#0], [cust_id#67], Inner\n",
      "      :- Sort [cust_id#0 ASC NULLS FIRST], false, 0\n",
      "      :  +- Exchange hashpartitioning(cust_id#0, 200), ENSURE_REQUIREMENTS, [id=#114]\n",
      "      :     +- Filter isnotnull(cust_id#0)\n",
      "      :        +- FileScan parquet [cust_id#0,start_date#1,end_date#2,txn_id#3,date#4,year#5,month#6,day#7,expense_type#8,amt#9,city#10] Batched: true, DataFilters: [isnotnull(cust_id#0)], Format: Parquet, Location: InMemoryFileIndex(1 paths)[file:/Users/afaqueahmad/Documents/youtube/spark-experiments/data/data_..., PartitionFilters: [], PushedFilters: [IsNotNull(cust_id)], ReadSchema: struct<cust_id:string,start_date:string,end_date:string,txn_id:string,date:string,year:string,mon...\n",
      "      +- Sort [cust_id#67 ASC NULLS FIRST], false, 0\n",
      "         +- Exchange hashpartitioning(cust_id#67, 200), ENSURE_REQUIREMENTS, [id=#115]\n",
      "            +- Filter isnotnull(cust_id#67)\n",
      "               +- FileScan parquet [cust_id#67,name#68,age#69,gender#70,birthday#71,zip#72,city#73] Batched: true, DataFilters: [isnotnull(cust_id#67)], Format: Parquet, Location: InMemoryFileIndex(1 paths)[file:/Users/afaqueahmad/Documents/youtube/spark-experiments/data/data_..., PartitionFilters: [], PushedFilters: [IsNotNull(cust_id)], ReadSchema: struct<cust_id:string,name:string,age:string,gender:string,birthday:string,zip:string,city:string>\n",
      "\n"
     ]
    }
   ],
   "source": [
    "df_joined.explain(True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 4. GroupBy"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "root\n",
      " |-- cust_id: string (nullable = true)\n",
      " |-- start_date: string (nullable = true)\n",
      " |-- end_date: string (nullable = true)\n",
      " |-- txn_id: string (nullable = true)\n",
      " |-- date: string (nullable = true)\n",
      " |-- year: string (nullable = true)\n",
      " |-- month: string (nullable = true)\n",
      " |-- day: string (nullable = true)\n",
      " |-- expense_type: string (nullable = true)\n",
      " |-- amt: string (nullable = true)\n",
      " |-- city: string (nullable = true)\n",
      "\n"
     ]
    }
   ],
   "source": [
    "df_transactions.printSchema()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### GroupBy Count"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [],
   "source": [
    "df_city_counts = (\n",
    "    df_transactions\n",
    "    .groupBy(\"city\")\n",
    "    .count()\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "== Parsed Logical Plan ==\n",
      "'Aggregate ['city], ['city, count(1) AS count#199L]\n",
      "+- Relation [cust_id#0,start_date#1,end_date#2,txn_id#3,date#4,year#5,month#6,day#7,expense_type#8,amt#9,city#10] parquet\n",
      "\n",
      "== Analyzed Logical Plan ==\n",
      "city: string, count: bigint\n",
      "Aggregate [city#10], [city#10, count(1) AS count#199L]\n",
      "+- Relation [cust_id#0,start_date#1,end_date#2,txn_id#3,date#4,year#5,month#6,day#7,expense_type#8,amt#9,city#10] parquet\n",
      "\n",
      "== Optimized Logical Plan ==\n",
      "Aggregate [city#10], [city#10, count(1) AS count#199L]\n",
      "+- Project [city#10]\n",
      "   +- Relation [cust_id#0,start_date#1,end_date#2,txn_id#3,date#4,year#5,month#6,day#7,expense_type#8,amt#9,city#10] parquet\n",
      "\n",
      "== Physical Plan ==\n",
      "AdaptiveSparkPlan isFinalPlan=false\n",
      "+- HashAggregate(keys=[city#10], functions=[count(1)], output=[city#10, count#199L])\n",
      "   +- Exchange hashpartitioning(city#10, 200), ENSURE_REQUIREMENTS, [id=#131]\n",
      "      +- HashAggregate(keys=[city#10], functions=[partial_count(1)], output=[city#10, count#203L])\n",
      "         +- FileScan parquet [city#10] Batched: true, DataFilters: [], Format: Parquet, Location: InMemoryFileIndex(1 paths)[file:/Users/afaqueahmad/Documents/youtube/spark-experiments/data/data_..., PartitionFilters: [], PushedFilters: [], ReadSchema: struct<city:string>\n",
      "\n"
     ]
    }
   ],
   "source": [
    "df_city_counts.explain(True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [],
   "source": [
    "df_txn_amt_city = (\n",
    "    df_transactions\n",
    "    .groupBy(\"city\")\n",
    "    .agg(F.sum(\"amt\").alias(\"txn_amt\"))\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "== Parsed Logical Plan ==\n",
      "'Aggregate ['city], ['city, sum('amt) AS txn_amt#216]\n",
      "+- Relation [cust_id#0,start_date#1,end_date#2,txn_id#3,date#4,year#5,month#6,day#7,expense_type#8,amt#9,city#10] parquet\n",
      "\n",
      "== Analyzed Logical Plan ==\n",
      "city: string, txn_amt: double\n",
      "Aggregate [city#10], [city#10, sum(cast(amt#9 as double)) AS txn_amt#216]\n",
      "+- Relation [cust_id#0,start_date#1,end_date#2,txn_id#3,date#4,year#5,month#6,day#7,expense_type#8,amt#9,city#10] parquet\n",
      "\n",
      "== Optimized Logical Plan ==\n",
      "Aggregate [city#10], [city#10, sum(cast(amt#9 as double)) AS txn_amt#216]\n",
      "+- Project [amt#9, city#10]\n",
      "   +- Relation [cust_id#0,start_date#1,end_date#2,txn_id#3,date#4,year#5,month#6,day#7,expense_type#8,amt#9,city#10] parquet\n",
      "\n",
      "== Physical Plan ==\n",
      "AdaptiveSparkPlan isFinalPlan=false\n",
      "+- HashAggregate(keys=[city#10], functions=[sum(cast(amt#9 as double))], output=[city#10, txn_amt#216])\n",
      "   +- Exchange hashpartitioning(city#10, 200), ENSURE_REQUIREMENTS, [id=#144]\n",
      "      +- HashAggregate(keys=[city#10], functions=[partial_sum(cast(amt#9 as double))], output=[city#10, sum#220])\n",
      "         +- FileScan parquet [amt#9,city#10] Batched: true, DataFilters: [], Format: Parquet, Location: InMemoryFileIndex(1 paths)[file:/Users/afaqueahmad/Documents/youtube/spark-experiments/data/data_..., PartitionFilters: [], PushedFilters: [], ReadSchema: struct<amt:string,city:string>\n",
      "\n"
     ]
    }
   ],
   "source": [
    "df_txn_amt_city.explain(True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### GroupBy Count Distinct "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [],
   "source": [
    "df_txn_per_city = (\n",
    "    df_transactions\n",
    "    .groupBy(\"cust_id\")\n",
    "    .agg(F.countDistinct(\"city\").alias(\"city_count\"))\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "[Stage 5:====================================================>    (12 + 1) / 13]\r"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "+----------+----------+\n",
      "|cust_id   |city_count|\n",
      "+----------+----------+\n",
      "|CPP8BY8U93|10        |\n",
      "|CYB8BX9LU1|10        |\n",
      "|CFRT841CCD|10        |\n",
      "|CA0TSNMYDK|10        |\n",
      "|COZ8NONEVZ|10        |\n",
      "+----------+----------+\n",
      "only showing top 5 rows\n",
      "\n",
      "== Parsed Logical Plan ==\n",
      "'Aggregate ['cust_id], ['cust_id, 'count(distinct 'city) AS city_count#232]\n",
      "+- Relation [cust_id#0,start_date#1,end_date#2,txn_id#3,date#4,year#5,month#6,day#7,expense_type#8,amt#9,city#10] parquet\n",
      "\n",
      "== Analyzed Logical Plan ==\n",
      "cust_id: string, city_count: bigint\n",
      "Aggregate [cust_id#0], [cust_id#0, count(distinct city#10) AS city_count#232L]\n",
      "+- Relation [cust_id#0,start_date#1,end_date#2,txn_id#3,date#4,year#5,month#6,day#7,expense_type#8,amt#9,city#10] parquet\n",
      "\n",
      "== Optimized Logical Plan ==\n",
      "Aggregate [cust_id#0], [cust_id#0, count(distinct city#10) AS city_count#232L]\n",
      "+- Project [cust_id#0, city#10]\n",
      "   +- Relation [cust_id#0,start_date#1,end_date#2,txn_id#3,date#4,year#5,month#6,day#7,expense_type#8,amt#9,city#10] parquet\n",
      "\n",
      "== Physical Plan ==\n",
      "AdaptiveSparkPlan isFinalPlan=false\n",
      "+- HashAggregate(keys=[cust_id#0], functions=[count(distinct city#10)], output=[cust_id#0, city_count#232L])\n",
      "   +- Exchange hashpartitioning(cust_id#0, 200), ENSURE_REQUIREMENTS, [id=#266]\n",
      "      +- HashAggregate(keys=[cust_id#0], functions=[partial_count(distinct city#10)], output=[cust_id#0, count#244L])\n",
      "         +- HashAggregate(keys=[cust_id#0, city#10], functions=[], output=[cust_id#0, city#10])\n",
      "            +- Exchange hashpartitioning(cust_id#0, city#10, 200), ENSURE_REQUIREMENTS, [id=#262]\n",
      "               +- HashAggregate(keys=[cust_id#0, city#10], functions=[], output=[cust_id#0, city#10])\n",
      "                  +- FileScan parquet [cust_id#0,city#10] Batched: true, DataFilters: [], Format: Parquet, Location: InMemoryFileIndex(1 paths)[file:/Users/afaqueahmad/Documents/youtube/spark-experiments/data/data_..., PartitionFilters: [], PushedFilters: [], ReadSchema: struct<cust_id:string,city:string>\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\r",
      "                                                                                \r"
     ]
    }
   ],
   "source": [
    "df_txn_per_city.show(5, False)\n",
    "df_txn_per_city.explain(True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 5. Interesting Observations"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Why is a filter step present despite predicate pushdown? \n",
    "\n",
    "This is largely due to the way `Spark's Catalyst Optimizer` works. Specifically, due to two separate stages of the query optimization process: Physical Planning and Logical Planning.\n",
    "\n",
    "- **Logical Planning**: Catalyst optimizer simplifies the unresolved logical plan (which represents the user's query) by applying various rule-based optimizations. This includes `predicate pushdown`, `projection pushdown` where filter conditions and column projections are moved as close to the data source as possible.\n",
    "\n",
    "- **Physical Planning** phase is where the logical plan is translated into one or more physical plans, which can actually be executed on the cluster. This includes operations like file `scans`, `filters`, `projections`, etc.\n",
    "\n",
    "In this case, during the logical planning phase, the predicate (`F.col(\"city\") == \"boston\"`) has been pushed down and will be applied during the scan of the Parquet file (`PushedFilters: [IsNotNull(city), EqualTo(city,boston)]`), thus improving performance.\n",
    "\n",
    "Now, during the physical planning phase, the same filter condition (`+- *(1) Filter (isnotnull(city#73) AND (city#73 = boston))`) is applied again to the data that's been loaded into memory. This is because of the following reasons:\n",
    "\n",
    "1. **Guaranteed Correctness:** It might seem **redundant**, but remember that not all data sources can handle pushed-down predicates, and not all predicates can be pushed down. Therefore, **even if a predicate is pushed down to the data source, Spark still includes the predicate in the physical plan** to cover cases where the data source might not have been able to fully apply the predicate. This is Spark's way of making sure the correct data is always returned, no matter the capabilities of the data source.\n",
    "\n",
    "2. **No Assumptions**: Spark's Catalyst optimizer doesn't make assumptions about the data source's ability to handle pushed-down predicates. The optimizer aims to generate plans that return correct results across a wide range of scenarios. Even if the filter is pushed down, Spark does not have the feedback from data source whether the pushdown was successful or not, so it includes the filter operation in the physical plan as well.\n",
    "\n",
    "It is more of a **fail-safe mechanism** to ensure data **integrity** and **correctness**.\n",
    "\n",
    "---"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### In what cases will predicate pushdown not work?\n",
    "\n",
    "2 Examples where **filter pushdown** will not work:\n",
    "\n",
    "1. **Complex Data Types**: Spark's Parquet data source does not push down filters that involve **complex types**, such as **arrays**, **maps**, and **structs**. This is because these complex data types can have complicated nested structures that the Parquet reader cannot easily filter on.\n",
    "\n",
    "Here's an example:\n",
    "\n",
    "```\n",
    "root\n",
    " |-- Name: string (nullable = true)\n",
    " |-- properties: map (nullable = true)\n",
    " |    |-- key: string\n",
    " |    |-- value: string (valueContainsNull = true)\n",
    "\n",
    "+----------+-----------------------------+\n",
    "|Name      |properties                   |\n",
    "+----------+-----------------------------+\n",
    "|Afaque    |[eye -> black, hair -> black]|\n",
    "|Naved     |[eye ->, hair -> brown]      |\n",
    "|Ali       |[eye -> black, hair -> red]  |\n",
    "|Amaan     |[eye -> grey, hair -> grey]  |\n",
    "|Omaira    |[eye -> , hair -> brown]     |\n",
    "+----------+-----------------------------+\n",
    "```\n",
    "\n",
    "```python\n",
    "df.filter(df.properties.getItem(\"eye\") == \"brown\").show()\n",
    "```\n",
    "\n",
    "```\n",
    "== Physical Plan ==\n",
    "*(1) Filter (metadata#123[key] = value)\n",
    "+- *(1) ColumnarToRow\n",
    "   +- FileScan parquet [id#122,metadata#123] Batched: true, DataFilters: [(metadata#123[key] = value)], Format: Parquet, ...\n",
    "```\n",
    "\n",
    "------------------------------------------------\n",
    "\n",
    "3. Unsupported Expressions: \n",
    "\n",
    "In Spark, `Parquet` data source does not support pushdown for filters involving a `.cast` operation. The reason for this behaviour is as follows:\n",
    "- `.cast` changes the datatype of the column, and the Parquet data source may not be able to perform the filter operation correctly on the cast data.\n",
    "\n",
    "**Note**: This behavior may vary based on the data source. For example, if you're working with a JDBC data source connected to a database that supports SQL-like operations, the `.cast` filter could potentially be pushed down to the database."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Example of operation where filter pushdown doesn't work"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "+----------+--------------+---+------+----------+-----+------------+\n",
      "|cust_id   |name          |age|gender|birthday  |zip  |city        |\n",
      "+----------+--------------+---+------+----------+-----+------------+\n",
      "|C01BKUFRHA|Aaron Becker  |54 |Male  |11/24/1979|40284|san_diego   |\n",
      "|C01WMZQ7PN|Aaron Brady   |51 |Female|8/20/1994 |52204|philadelphia|\n",
      "|C021567NJZ|Aaron Briggs  |57 |Male  |3/10/1990 |22008|philadelphia|\n",
      "|C02JNTM46B|Aaron Chambers|51 |Male  |1/6/2001  |63337|new_york    |\n",
      "|C030A69V1L|Aaron Clarke  |55 |Male  |4/28/1999 |77176|philadelphia|\n",
      "+----------+--------------+---+------+----------+-----+------------+\n",
      "only showing top 5 rows\n",
      "\n",
      "== Parsed Logical Plan ==\n",
      "'Filter (cast('age as int) > 50)\n",
      "+- Relation [cust_id#67,name#68,age#69,gender#70,birthday#71,zip#72,city#73] parquet\n",
      "\n",
      "== Analyzed Logical Plan ==\n",
      "cust_id: string, name: string, age: string, gender: string, birthday: string, zip: string, city: string\n",
      "Filter (cast(age#69 as int) > 50)\n",
      "+- Relation [cust_id#67,name#68,age#69,gender#70,birthday#71,zip#72,city#73] parquet\n",
      "\n",
      "== Optimized Logical Plan ==\n",
      "Filter (isnotnull(age#69) AND (cast(age#69 as int) > 50))\n",
      "+- Relation [cust_id#67,name#68,age#69,gender#70,birthday#71,zip#72,city#73] parquet\n",
      "\n",
      "== Physical Plan ==\n",
      "*(1) Filter (isnotnull(age#69) AND (cast(age#69 as int) > 50))\n",
      "+- *(1) ColumnarToRow\n",
      "   +- FileScan parquet [cust_id#67,name#68,age#69,gender#70,birthday#71,zip#72,city#73] Batched: true, DataFilters: [isnotnull(age#69), (cast(age#69 as int) > 50)], Format: Parquet, Location: InMemoryFileIndex(1 paths)[file:/Users/afaqueahmad/Documents/youtube/spark-experiments/data/data_..., PartitionFilters: [], PushedFilters: [IsNotNull(age)], ReadSchema: struct<cust_id:string,name:string,age:string,gender:string,birthday:string,zip:string,city:string>\n",
      "\n"
     ]
    }
   ],
   "source": [
    "df_customer_gt_50 = (\n",
    "    df_customers\n",
    "    .filter(F.col(\"age\").cast(\"int\") > 50)\n",
    ")\n",
    "df_customer_gt_50.show(5, False)\n",
    "df_customer_gt_50.explain(True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "application/vnd.databricks.v1+notebook": {
   "dashboards": [],
   "language": "python",
   "notebookMetadata": {
    "pythonIndentUnit": 4
   },
   "notebookName": "data-skew-explanation",
   "notebookOrigID": 4018749166498458,
   "widgets": {}
  },
  "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.8.16"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}