spark/1_data_skew/3_solving_data_skew_aqe_broadcast.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "application/vnd.databricks.v1+cell": {
     "cellMetadata": {},
     "inputWidgets": {},
     "nuid": "12fa052b-65ab-42fe-92e9-866a91ab0479",
     "showTitle": false,
     "title": ""
    }
   },
   "source": [
    "<h2> Imports & Configuration </h2>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "from IPython.core.interactiveshell import InteractiveShell\n",
    "InteractiveShell.ast_node_interactivity = 'all'\n",
    "\n",
    "import warnings\n",
    "warnings.filterwarnings('ignore')\n",
    "warnings.simplefilter('ignore')\n",
    "import time"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "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 IntegerType\n",
    "import pyspark.sql.functions as F\n",
    "from pyspark.sql import SparkSession"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "23/07/09 09:02:54 WARN Utils: Your hostname, Afaques-MacBook-Pro.local resolves to a loopback address: 127.0.0.1; using 10.0.0.4 instead (on interface en0)\n",
      "23/07/09 09:02:54 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/09 09:02:54 WARN NativeCodeLoader: Unable to load native-hadoop library for your platform... using builtin-java classes where applicable\n"
     ]
    }
   ],
   "source": [
    "spark = SparkSession.builder.master(\"local[4]\").getOrCreate()\n",
    "sc = spark.sparkContext\n",
    "sc.setLogLevel(\"ERROR\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "# spark.conf.set(\"spark.sql.shuffle.partitions\", \"3\")\n",
    "spark.conf.set(\"spark.sql.adaptive.enabled\", \"false\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Join Skews"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "                                                                                \r"
     ]
    }
   ],
   "source": [
    "transactions_file = \"../../data/data_skew/transactions.parquet\"\n",
    "customer_file = \"../../data/data_skew/customers.parquet\"\n",
    "\n",
    "df_transactions = spark.read.parquet(transactions_file)\n",
    "df_customers = spark.read.parquet(customer_file)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "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"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\r",
      "[Stage 2:>                                                          (0 + 1) / 1]\r"
     ]
    },
    {
     "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",
      "                                                                                \r"
     ]
    }
   ],
   "source": [
    "df_transactions.printSchema()\n",
    "df_transactions.show(5, False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "root\n",
      " |-- cust_id: string (nullable = true)\n",
      " |-- name: string (nullable = true)\n",
      " |-- age: string (nullable = true)\n",
      " |-- gender: string (nullable = true)\n",
      " |-- birthday: string (nullable = true)\n",
      " |-- zip: string (nullable = true)\n",
      " |-- city: string (nullable = true)\n",
      "\n",
      "+----------+-------------+---+------+----------+-----+-----------+\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.printSchema()\n",
    "df_customers.show(5, False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "[Stage 6:=======================================>               (142 + 6) / 200]\r"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "+----------+--------+\n",
      "|cust_id   |ct      |\n",
      "+----------+--------+\n",
      "|C0YDPQWPBJ|17539732|\n",
      "|CBW3FMEAU7|7999    |\n",
      "|C3KUDEN3KO|7999    |\n",
      "|C89FCEGPJP|7999    |\n",
      "|CHNFNR89ZV|7998    |\n",
      "+----------+--------+\n",
      "only showing top 5 rows\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\r",
      "[Stage 6:================================================>      (177 + 5) / 200]\r",
      "\r",
      "                                                                                \r"
     ]
    }
   ],
   "source": [
    "(\n",
    "    df_transactions\n",
    "    .groupBy(\"cust_id\")\n",
    "    .agg(F.countDistinct(\"txn_id\").alias(\"ct\"))\n",
    "    .orderBy(F.desc(\"ct\"))\n",
    "    .show(5, False)\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "spark.conf.set(\"spark.sql.autoBroadcastJoinThreshold\", -1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "df_txn_details = (\n",
    "    df_transactions.join(\n",
    "        df_customers,\n",
    "        on=\"cust_id\",\n",
    "        how=\"inner\"\n",
    "    )\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "                                                                                \r"
     ]
    },
    {
     "data": {
      "text/plain": [
       "39790092"
      ]
     },
     "execution_count": 11,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "time taken: 11.315901279449463\n"
     ]
    }
   ],
   "source": [
    "start_time = time.time()\n",
    "df_txn_details.count()\n",
    "print(f\"time taken: {time.time() - start_time}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Using AQE"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "spark.conf.set(\"spark.sql.adaptive.enabled\", \"true\")\n",
    "spark.conf.set(\"spark.sql.adaptive.skewedJoin.enabled\", \"true\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "df_txn_details = (\n",
    "    df_transactions.join(\n",
    "        df_customers,\n",
    "        on=\"cust_id\",\n",
    "        how=\"inner\"\n",
    "    )\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "                                                                                \r"
     ]
    },
    {
     "data": {
      "text/plain": [
       "39790092"
      ]
     },
     "execution_count": 14,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "time taken: 9.6234769821167\n"
     ]
    }
   ],
   "source": [
    "start_time = time.time()\n",
    "df_txn_details.count()\n",
    "print(f\"time taken: {time.time() - start_time}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Using Broadcast Joins"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 10MB = 10485760 Bytes\n",
    "spark.conf.set(\"spark.sql.autoBroadcastJoinThreshold\", 10485760)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [],
   "source": [
    "df_txn_details = (\n",
    "    df_transactions.join(\n",
    "        F.broadcast(df_customers),\n",
    "        on=\"cust_id\",\n",
    "        how=\"inner\"\n",
    "    )\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "                                                                                \r"
     ]
    },
    {
     "data": {
      "text/plain": [
       "39790092"
      ]
     },
     "execution_count": 17,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "time taken: 2.7766311168670654\n"
     ]
    }
   ],
   "source": [
    "start_time = time.time()\n",
    "df_txn_details.count()\n",
    "print(f\"time taken: {time.time() - start_time}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [],
   "source": [
    "# spark.stop()"
   ]
  }
 ],
 "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
}