cheatsheet.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "2af8fc07",
   "metadata": {},
   "source": [
    "## Run This Code First! This creates the Spark session and loads data."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4066398c",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "\n",
    "from pyspark.sql import SparkSession, SQLContext\n",
    "from pyspark.sql.functions import col\n",
    "\n",
    "from delta import *\n",
    "\n",
    "# Create our Spark session and SQL Context.\n",
    "warehouse_path = \"file://{}/spark_warehouse\".format(os.getcwd())\n",
    "builder = (\n",
    "    SparkSession.builder.master(\"local[*]\")\n",
    "    .config(\"spark.executor.memory\", \"2G\")\n",
    "    .config(\"spark.driver.memory\", \"2G\")\n",
    "    .config(\"spark.sql.extensions\", \"io.delta.sql.DeltaSparkSessionExtension\")\n",
    "    .config(\n",
    "        \"spark.sql.catalog.spark_catalog\",\n",
    "        \"org.apache.spark.sql.delta.catalog.DeltaCatalog\",\n",
    "    )\n",
    "    .config(\"spark.sql.warehouse.dir\", warehouse_path)\n",
    "    .appName(\"cheatsheet\")\n",
    ")\n",
    "spark = configure_spark_with_delta_pip(builder).getOrCreate()\n",
    "sqlContext = SQLContext(spark)\n",
    "\n",
    "# Unmodified Auto dataset.\n",
    "auto_df = spark.read.format(\"csv\").option(\"header\", True).load(\"data/auto-mpg.csv\")\n",
    "\n",
    "# Fixed Auto dataset.\n",
    "auto_df_fixed = spark.read.format(\"csv\").option(\"header\", True).load(\"data/auto-mpg-fixed.csv\")\n",
    "for (column_name) in (\"mpg cylinders displacement horsepower weight acceleration\".split()):\n",
    "    auto_df_fixed = auto_df_fixed.withColumn(column_name, col(column_name).cast(\"double\"))\n",
    "auto_df_fixed = auto_df_fixed.withColumn(\"modelyear\", col(\"modelyear\").cast(\"int\"))\n",
    "auto_df_fixed = auto_df_fixed.withColumn(\"origin\", col(\"origin\").cast(\"int\"))\n",
    "\n",
    "# Cover type dataset.\n",
    "covtype_df = spark.read.format(\"parquet\").load(\"data/covtype.parquet\")\n",
    "for column_name in covtype_df.columns:\n",
    "    covtype_df = covtype_df.withColumn(column_name, col(column_name).cast(\"int\"))\n",
    "\n",
    "# Customer spend dataset.\n",
    "spend_df = spark.read.format(\"csv\").option(\"header\", True).load(\"data/customer_spend.csv\")\n",
    "\n",
    "# Weblog.\n",
    "weblog_df = spark.read.format(\"csv\").option(\"header\", True).load(\"data/weblog.csv\")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "eee9ed6d",
   "metadata": {},
   "source": [
    "## Loading data stored in filesystems or databases, and saving it."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "be4bf903",
   "metadata": {},
   "source": [
    "**Load a DataFrame from CSV**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "00494ff1",
   "metadata": {},
   "outputs": [],
   "source": [
    "df = spark.read.format(\"csv\").option(\"header\", True).load(\"data/auto-mpg.csv\")\n",
    "df.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "928ca6d4",
   "metadata": {},
   "source": [
    "**Load a DataFrame from a Tab Separated Value (TSV) file**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e197530d",
   "metadata": {},
   "outputs": [],
   "source": [
    "df = (\n",
    "    spark.read.format(\"csv\")\n",
    "    .option(\"header\", True)\n",
    "    .option(\"sep\", \"\\t\")\n",
    "    .load(\"data/auto-mpg.tsv\")\n",
    ")\n",
    "df.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fb5bd3e6",
   "metadata": {},
   "source": [
    "**Save a DataFrame in CSV format**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "85ad41c9",
   "metadata": {},
   "outputs": [],
   "source": [
    "auto_df.write.csv(\"output.csv\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2cdef498",
   "metadata": {},
   "source": [
    "**Load a DataFrame from Parquet**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c63e8ca1",
   "metadata": {},
   "outputs": [],
   "source": [
    "df = spark.read.format(\"parquet\").load(\"data/auto-mpg.parquet\")\n",
    "df.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2b351a78",
   "metadata": {},
   "source": [
    "**Save a DataFrame in Parquet format**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bf83fa40",
   "metadata": {},
   "outputs": [],
   "source": [
    "auto_df.write.parquet(\"output.parquet\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d1070509",
   "metadata": {},
   "source": [
    "**Load a DataFrame from JSON Lines (jsonl) Formatted Data**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4d9342d9",
   "metadata": {},
   "outputs": [],
   "source": [
    "df = spark.read.json(\"data/weblog.jsonl\")\n",
    "df.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3d143f2f",
   "metadata": {},
   "source": [
    "**Save a DataFrame into a Hive catalog table**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "151789cc",
   "metadata": {},
   "outputs": [],
   "source": [
    "auto_df.write.mode(\"overwrite\").saveAsTable(\"autompg\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a246e195",
   "metadata": {},
   "source": [
    "**Load a Hive catalog table into a DataFrame**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "39b11a0a",
   "metadata": {},
   "outputs": [],
   "source": [
    "df = spark.table(\"autompg\")\n",
    "df.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "553656c5",
   "metadata": {},
   "source": [
    "**Load a DataFrame from a SQL query**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7c1bf0df",
   "metadata": {},
   "outputs": [],
   "source": [
    "df = sqlContext.sql(\n",
    "    \"select carname, mpg, horsepower from autompg where horsepower > 100 and mpg > 25\"\n",
    ")\n",
    "df.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f355d231",
   "metadata": {},
   "source": [
    "**Load a CSV file from Amazon S3**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "02358f4b",
   "metadata": {},
   "outputs": [],
   "source": [
    "import configparser\n",
    "import os\n",
    "\n",
    "config = configparser.ConfigParser()\n",
    "config.read(os.path.expanduser(\"~/.aws/credentials\"))\n",
    "access_key = config.get(\"default\", \"aws_access_key_id\").replace('\"', \"\")\n",
    "secret_key = config.get(\"default\", \"aws_secret_access_key\").replace('\"', \"\")\n",
    "\n",
    "# Requires compatible hadoop-aws and aws-java-sdk-bundle JARs.\n",
    "spark.conf.set(\n",
    "    \"fs.s3a.aws.credentials.provider\",\n",
    "    \"org.apache.hadoop.fs.s3a.SimpleAWSCredentialsProvider\",\n",
    ")\n",
    "spark.conf.set(\"fs.s3a.access.key\", access_key)\n",
    "spark.conf.set(\"fs.s3a.secret.key\", secret_key)\n",
    "\n",
    "df = (\n",
    "    spark.read.format(\"csv\")\n",
    "    .option(\"header\", True)\n",
    "    .load(\"s3a://cheatsheet111/auto-mpg.csv\")\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1c84988e",
   "metadata": {},
   "source": [
    "**Load a CSV file from Oracle Cloud Infrastructure (OCI) Object Storage**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d5edcd2d",
   "metadata": {},
   "outputs": [],
   "source": [
    "import oci\n",
    "\n",
    "oci_config = oci.config.from_file()\n",
    "conf = spark.sparkContext.getConf()\n",
    "conf.set(\"fs.oci.client.auth.tenantId\", oci_config[\"tenancy\"])\n",
    "conf.set(\"fs.oci.client.auth.userId\", oci_config[\"user\"])\n",
    "conf.set(\"fs.oci.client.auth.fingerprint\", oci_config[\"fingerprint\"])\n",
    "conf.set(\"fs.oci.client.auth.pemfilepath\", oci_config[\"key_file\"])\n",
    "conf.set(\n",
    "    \"fs.oci.client.hostname\",\n",
    "    \"https://objectstorage.{0}.oraclecloud.com\".format(oci_config[\"region\"]),\n",
    ")\n",
    "PATH = \"oci://<your_bucket>@<your_namespace/<your_path>\"\n",
    "df = spark.read.format(\"csv\").option(\"header\", True).load(PATH)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "223d5750",
   "metadata": {},
   "source": [
    "**Read an Oracle DB table into a DataFrame using a Wallet**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e227d9f9",
   "metadata": {},
   "outputs": [],
   "source": [
    "password = \"my_password\"\n",
    "table = \"source_table\"\n",
    "tnsname = \"my_tns_name\"\n",
    "user = \"ADMIN\"\n",
    "wallet_path = \"/path/to/your/wallet\"\n",
    "\n",
    "properties = {\n",
    "    \"driver\": \"oracle.jdbc.driver.OracleDriver\",\n",
    "    \"oracle.net.tns_admin\": tnsname,\n",
    "    \"password\": password,\n",
    "    \"user\": user,\n",
    "}\n",
    "url = f\"jdbc:oracle:thin:@{tnsname}?TNS_ADMIN={wallet_path}\"\n",
    "df = spark.read.jdbc(url=url, table=table, properties=properties)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "50084e7d",
   "metadata": {},
   "source": [
    "**Write a DataFrame to an Oracle DB table using a Wallet**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2f84d5c3",
   "metadata": {},
   "outputs": [],
   "source": [
    "password = \"my_password\"\n",
    "table = \"target_table\"\n",
    "tnsname = \"my_tns_name\"\n",
    "user = \"ADMIN\"\n",
    "wallet_path = \"/path/to/your/wallet\"\n",
    "\n",
    "properties = {\n",
    "    \"driver\": \"oracle.jdbc.driver.OracleDriver\",\n",
    "    \"oracle.net.tns_admin\": tnsname,\n",
    "    \"password\": password,\n",
    "    \"user\": user,\n",
    "}\n",
    "url = f\"jdbc:oracle:thin:@{tnsname}?TNS_ADMIN={wallet_path}\"\n",
    "\n",
    "# Possible modes are \"Append\", \"Overwrite\", \"Ignore\", \"Error\"\n",
    "df.write.jdbc(url=url, table=table, mode=\"Append\", properties=properties)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8019dc83",
   "metadata": {},
   "source": [
    "**Write a DataFrame to a Postgres table**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "92917172",
   "metadata": {},
   "outputs": [],
   "source": [
    "pg_database = os.environ.get(\"PGDATABASE\") or \"postgres\"\n",
    "pg_host = os.environ.get(\"PGHOST\") or \"localhost\"\n",
    "pg_password = os.environ.get(\"PGPASSWORD\") or \"password\"\n",
    "pg_user = os.environ.get(\"PGUSER\") or \"postgres\"\n",
    "table = \"autompg\"\n",
    "\n",
    "properties = {\n",
    "    \"driver\": \"org.postgresql.Driver\",\n",
    "    \"user\": pg_user,\n",
    "    \"password\": pg_password,\n",
    "}\n",
    "url = f\"jdbc:postgresql://{pg_host}:5432/{pg_database}\"\n",
    "auto_df.write.jdbc(url=url, table=table, mode=\"Append\", properties=properties)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0473f712",
   "metadata": {},
   "source": [
    "**Read a Postgres table into a DataFrame**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d273f5e1",
   "metadata": {},
   "outputs": [],
   "source": [
    "pg_database = os.environ.get(\"PGDATABASE\") or \"postgres\"\n",
    "pg_host = os.environ.get(\"PGHOST\") or \"localhost\"\n",
    "pg_password = os.environ.get(\"PGPASSWORD\") or \"password\"\n",
    "pg_user = os.environ.get(\"PGUSER\") or \"postgres\"\n",
    "table = \"autompg\"\n",
    "\n",
    "properties = {\n",
    "    \"driver\": \"org.postgresql.Driver\",\n",
    "    \"user\": pg_user,\n",
    "    \"password\": pg_password,\n",
    "}\n",
    "url = f\"jdbc:postgresql://{pg_host}:5432/{pg_database}\"\n",
    "df = spark.read.jdbc(url=url, table=table, properties=properties)\n",
    "df.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7c487d0b",
   "metadata": {},
   "source": [
    "## Special data handling scenarios."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c31fbf38",
   "metadata": {},
   "source": [
    "**Provide the schema when loading a DataFrame from CSV**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0da90d00",
   "metadata": {},
   "outputs": [],
   "source": [
    "from pyspark.sql.types import (\n",
    "    DoubleType,\n",
    "    IntegerType,\n",
    "    StringType,\n",
    "    StructField,\n",
    "    StructType,\n",
    ")\n",
    "\n",
    "schema = StructType(\n",
    "    [\n",
    "        StructField(\"mpg\", DoubleType(), True),\n",
    "        StructField(\"cylinders\", IntegerType(), True),\n",
    "        StructField(\"displacement\", DoubleType(), True),\n",
    "        StructField(\"horsepower\", DoubleType(), True),\n",
    "        StructField(\"weight\", DoubleType(), True),\n",
    "        StructField(\"acceleration\", DoubleType(), True),\n",
    "        StructField(\"modelyear\", IntegerType(), True),\n",
    "        StructField(\"origin\", IntegerType(), True),\n",
    "        StructField(\"carname\", StringType(), True),\n",
    "    ]\n",
    ")\n",
    "df = (\n",
    "    spark.read.format(\"csv\")\n",
    "    .option(\"header\", \"true\")\n",
    "    .schema(schema)\n",
    "    .load(\"data/auto-mpg.csv\")\n",
    ")\n",
    "df.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "78df5e40",
   "metadata": {},
   "source": [
    "**Save a DataFrame to CSV, overwriting existing data**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "42922272",
   "metadata": {},
   "outputs": [],
   "source": [
    "auto_df.write.mode(\"overwrite\").csv(\"output.csv\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d54862ca",
   "metadata": {},
   "source": [
    "**Save a DataFrame to CSV with a header**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c5b3c74a",
   "metadata": {},
   "outputs": [],
   "source": [
    "auto_df.coalesce(1).write.csv(\"header.csv\", header=\"true\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "92a985da",
   "metadata": {},
   "source": [
    "**Save a DataFrame in a single CSV file**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "304017c7",
   "metadata": {},
   "outputs": [],
   "source": [
    "auto_df.coalesce(1).write.csv(\"single.csv\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e27fd7ad",
   "metadata": {},
   "source": [
    "**Save DataFrame as a dynamic partitioned table**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9af5631b",
   "metadata": {},
   "outputs": [],
   "source": [
    "spark.conf.set(\"spark.sql.sources.partitionOverwriteMode\", \"dynamic\")\n",
    "auto_df.write.mode(\"append\").partitionBy(\"modelyear\").saveAsTable(\n",
    "    \"autompg_partitioned\"\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0ffc99b6",
   "metadata": {},
   "source": [
    "**Overwrite specific partitions**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fa7fe0eb",
   "metadata": {},
   "outputs": [],
   "source": [
    "spark.conf.set(\"spark.sql.sources.partitionOverwriteMode\", \"dynamic\")\n",
    "your_dataframe.write.mode(\"overwrite\").insertInto(\"your_table\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d672aee1",
   "metadata": {},
   "source": [
    "**Load a CSV file with a money column into a DataFrame**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0f517050",
   "metadata": {},
   "outputs": [],
   "source": [
    "from pyspark.sql.functions import udf\n",
    "from pyspark.sql.types import DecimalType\n",
    "from decimal import Decimal\n",
    "\n",
    "# Load the text file.\n",
    "df = (\n",
    "    spark.read.format(\"csv\")\n",
    "    .option(\"header\", True)\n",
    "    .load(\"data/customer_spend.csv\")\n",
    ")\n",
    "\n",
    "# Convert with a hardcoded custom UDF.\n",
    "money_udf = udf(lambda x: Decimal(x[1:].replace(\",\", \"\")), DecimalType(8, 4))\n",
    "money1 = df.withColumn(\"spend_dollars\", money_udf(df.spend_dollars))\n",
    "\n",
    "# Convert with the money_parser library (much safer).\n",
    "from money_parser import price_str\n",
    "\n",
    "money_convert = udf(\n",
    "    lambda x: Decimal(price_str(x)) if x is not None else None,\n",
    "    DecimalType(8, 4),\n",
    ")\n",
    "df = df.withColumn(\"spend_dollars\", money_convert(df.spend_dollars))\n",
    "df.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "04cc4046",
   "metadata": {},
   "source": [
    "## Adding, removing and modifying DataFrame columns."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4f0e460b",
   "metadata": {},
   "source": [
    "**Add a new column to a DataFrame**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "70876715",
   "metadata": {},
   "outputs": [],
   "source": [
    "from pyspark.sql.functions import upper, lower\n",
    "\n",
    "df = auto_df.withColumn(\"upper\", upper(auto_df.carname)).withColumn(\n",
    "    \"lower\", lower(auto_df.carname)\n",
    ")\n",
    "df.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "53caa355",
   "metadata": {},
   "source": [
    "**Modify a DataFrame column**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "93663a30",
   "metadata": {},
   "outputs": [],
   "source": [
    "from pyspark.sql.functions import col, concat, lit\n",
    "\n",
    "df = auto_df.withColumn(\"modelyear\", concat(lit(\"19\"), col(\"modelyear\")))\n",
    "df.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4598af4a",
   "metadata": {},
   "source": [
    "**Add a column with multiple conditions**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5ff7f7ac",
   "metadata": {},
   "outputs": [],
   "source": [
    "from pyspark.sql.functions import col, when\n",
    "\n",
    "df = auto_df.withColumn(\n",
    "    \"mpg_class\",\n",
    "    when(col(\"mpg\") <= 20, \"low\")\n",
    "    .when(col(\"mpg\") <= 30, \"mid\")\n",
    "    .when(col(\"mpg\") <= 40, \"high\")\n",
    "    .otherwise(\"very high\"),\n",
    ")\n",
    "df.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "696f8117",
   "metadata": {},
   "source": [
    "**Add a constant column**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b5f37b0e",
   "metadata": {},
   "outputs": [],
   "source": [
    "from pyspark.sql.functions import lit\n",
    "\n",
    "df = auto_df.withColumn(\"one\", lit(1))\n",
    "df.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9020391a",
   "metadata": {},
   "source": [
    "**Concatenate columns**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c66f81d1",
   "metadata": {},
   "outputs": [],
   "source": [
    "from pyspark.sql.functions import concat, col, lit\n",
    "\n",
    "df = auto_df.withColumn(\n",
    "    \"concatenated\", concat(col(\"cylinders\"), lit(\"_\"), col(\"mpg\"))\n",
    ")\n",
    "df.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "44209a3a",
   "metadata": {},
   "source": [
    "**Drop a column**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ca5f9717",
   "metadata": {},
   "outputs": [],
   "source": [
    "df = auto_df.drop(\"horsepower\")\n",
    "df.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ab977b19",
   "metadata": {},
   "source": [
    "**Change a column name**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c797801f",
   "metadata": {},
   "outputs": [],
   "source": [
    "df = auto_df.withColumnRenamed(\"horsepower\", \"horses\")\n",
    "df.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c3c3770a",
   "metadata": {},
   "source": [
    "**Change multiple column names**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "45bec5e8",
   "metadata": {},
   "outputs": [],
   "source": [
    "df = auto_df.withColumnRenamed(\"horsepower\", \"horses\").withColumnRenamed(\n",
    "    \"modelyear\", \"year\"\n",
    ")\n",
    "df.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7e854cb6",
   "metadata": {},
   "source": [
    "**Change all column names at once**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c5a6a194",
   "metadata": {},
   "outputs": [],
   "source": [
    "df = auto_df.toDF(*[\"X\" + name for name in auto_df.columns])\n",
    "df.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "efb3f83c",
   "metadata": {},
   "source": [
    "**Convert a DataFrame column to a Python list**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c6b3b9dc",
   "metadata": {},
   "outputs": [],
   "source": [
    "names = auto_df.select(\"carname\").rdd.flatMap(lambda x: x).collect()\n",
    "print(str(names[:10]))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7c5780ff",
   "metadata": {},
   "source": [
    "**Convert a scalar query to a Python value**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e9a0ac99",
   "metadata": {},
   "outputs": [],
   "source": [
    "average = auto_df.agg(dict(mpg=\"avg\")).first()[0]\n",
    "print(str(average))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6f52351c",
   "metadata": {},
   "source": [
    "**Consume a DataFrame row-wise as Python dictionaries**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0de6594c",
   "metadata": {},
   "outputs": [],
   "source": [
    "first_three = auto_df.limit(3)\n",
    "for row in first_three.collect():\n",
    "    my_dict = row.asDict()\n",
    "    print(my_dict)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f70b6853",
   "metadata": {},
   "source": [
    "**Select particular columns from a DataFrame**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8c05ef50",
   "metadata": {},
   "outputs": [],
   "source": [
    "df = auto_df.select([\"mpg\", \"cylinders\", \"displacement\"])\n",
    "df.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2a24245c",
   "metadata": {},
   "source": [
    "**Create an empty dataframe with a specified schema**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bdd2a301",
   "metadata": {},
   "outputs": [],
   "source": [
    "from pyspark.sql.types import StructField, StructType, LongType, StringType\n",
    "\n",
    "schema = StructType(\n",
    "    [\n",
    "        StructField(\"my_id\", LongType(), True),\n",
    "        StructField(\"my_string\", StringType(), True),\n",
    "    ]\n",
    ")\n",
    "df = spark.createDataFrame([], schema)\n",
    "df.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e982254b",
   "metadata": {},
   "source": [
    "**Create a constant dataframe**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3a71b704",
   "metadata": {},
   "outputs": [],
   "source": [
    "import datetime\n",
    "from pyspark.sql.types import (\n",
    "    StructField,\n",
    "    StructType,\n",
    "    LongType,\n",
    "    StringType,\n",
    "    TimestampType,\n",
    ")\n",
    "\n",
    "schema = StructType(\n",
    "    [\n",
    "        StructField(\"my_id\", LongType(), True),\n",
    "        StructField(\"my_string\", StringType(), True),\n",
    "        StructField(\"my_timestamp\", TimestampType(), True),\n",
    "    ]\n",
    ")\n",
    "df = spark.createDataFrame(\n",
    "    [\n",
    "        (1, \"foo\", datetime.datetime.strptime(\"2021-01-01\", \"%Y-%m-%d\")),\n",
    "        (2, \"bar\", datetime.datetime.strptime(\"2021-01-02\", \"%Y-%m-%d\")),\n",
    "    ],\n",
    "    schema,\n",
    ")\n",
    "df.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2358f747",
   "metadata": {},
   "source": [
    "**Convert String to Double**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a71dc647",
   "metadata": {},
   "outputs": [],
   "source": [
    "from pyspark.sql.functions import col\n",
    "\n",
    "df = auto_df.withColumn(\"horsepower\", col(\"horsepower\").cast(\"double\"))\n",
    "df.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a102835a",
   "metadata": {},
   "source": [
    "**Convert String to Integer**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "38209b6b",
   "metadata": {},
   "outputs": [],
   "source": [
    "from pyspark.sql.functions import col\n",
    "\n",
    "df = auto_df.withColumn(\"horsepower\", col(\"horsepower\").cast(\"int\"))\n",
    "df.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a4da6aab",
   "metadata": {},
   "source": [
    "**Get the size of a DataFrame**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "344a0daa",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(\"{} rows\".format(auto_df.count()))\n",
    "print(\"{} columns\".format(len(auto_df.columns)))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1595d2bc",
   "metadata": {},
   "source": [
    "**Get a DataFrame's number of partitions**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "eb68828b",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(\"{} partition(s)\".format(auto_df.rdd.getNumPartitions()))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "43dcb80f",
   "metadata": {},
   "source": [
    "**Get data types of a DataFrame's columns**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b5113ebf",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(auto_df.dtypes)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "354bd9de",
   "metadata": {},
   "source": [
    "**Convert an RDD to Data Frame**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2f6486bc",
   "metadata": {},
   "outputs": [],
   "source": [
    "from pyspark.sql import Row\n",
    "\n",
    "# First, get the RDD from the DataFrame.\n",
    "rdd = auto_df.rdd\n",
    "\n",
    "# This converts it back to an RDD with no changes.\n",
    "df = rdd.map(lambda x: Row(**x.asDict())).toDF()\n",
    "\n",
    "# This changes the rows before creating the DataFrame.\n",
    "df = rdd.map(\n",
    "    lambda x: Row(**{k: v * 2 for (k, v) in x.asDict().items()})\n",
    ").toDF()\n",
    "df.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c404b00a",
   "metadata": {},
   "source": [
    "**Print the contents of an RDD**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "90e77034",
   "metadata": {},
   "outputs": [],
   "source": [
    "rdd = auto_df.rdd\n",
    "print(rdd.take(10))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "31fbda1b",
   "metadata": {},
   "source": [
    "**Print the contents of a DataFrame**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e08221c8",
   "metadata": {},
   "outputs": [],
   "source": [
    "auto_df.show(10)\n",
    "df.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1f085452",
   "metadata": {},
   "source": [
    "**Process each row of a DataFrame**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6ad95bee",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "\n",
    "def foreach_function(row):\n",
    "    if row.horsepower is not None:\n",
    "        os.system(\"echo \" + row.horsepower)\n",
    "\n",
    "auto_df.foreach(foreach_function)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4b40f01c",
   "metadata": {},
   "source": [
    "**DataFrame Map example**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "33e89edb",
   "metadata": {},
   "outputs": [],
   "source": [
    "def map_function(row):\n",
    "    if row.horsepower is not None:\n",
    "        return [float(row.horsepower) * 10]\n",
    "    else:\n",
    "        return [None]\n",
    "\n",
    "df = auto_df.rdd.map(map_function).toDF()\n",
    "df.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "52cf2e31",
   "metadata": {},
   "source": [
    "**DataFrame Flatmap example**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "eac44402",
   "metadata": {},
   "outputs": [],
   "source": [
    "from pyspark.sql.types import Row\n",
    "\n",
    "def flatmap_function(row):\n",
    "    if row.cylinders is not None:\n",
    "        return list(range(int(row.cylinders)))\n",
    "    else:\n",
    "        return [None]\n",
    "\n",
    "rdd = auto_df.rdd.flatMap(flatmap_function)\n",
    "row = Row(\"val\")\n",
    "df = rdd.map(row).toDF()\n",
    "df.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a08b6826",
   "metadata": {},
   "source": [
    "**Create a custom UDF**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6018cc6f",
   "metadata": {},
   "outputs": [],
   "source": [
    "from pyspark.sql.functions import col, udf\n",
    "from pyspark.sql.types import StringType\n",
    "\n",
    "first_word_udf = udf(lambda x: x.split()[0], StringType())\n",
    "df = auto_df.withColumn(\"manufacturer\", first_word_udf(col(\"carname\")))\n",
    "df.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5b068fdf",
   "metadata": {},
   "source": [
    "## Data conversions and other modifications."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "619d3a4f",
   "metadata": {},
   "source": [
    "**Run a SparkSQL Statement on a DataFrame**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3440a3bc",
   "metadata": {},
   "outputs": [],
   "source": [
    "from pyspark.sql.functions import col, regexp_extract\n",
    "\n",
    "auto_df.registerTempTable(\"auto_df\")\n",
    "df = sqlContext.sql(\n",
    "    \"select modelyear, avg(mpg) from auto_df group by modelyear\"\n",
    ")\n",
    "df.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "be4c75a3",
   "metadata": {},
   "source": [
    "**Extract data from a string using a regular expression**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1c5b2e2e",
   "metadata": {},
   "outputs": [],
   "source": [
    "from pyspark.sql.functions import col, regexp_extract\n",
    "\n",
    "group = 0\n",
    "df = (\n",
    "    auto_df.withColumn(\n",
    "        \"identifier\", regexp_extract(col(\"carname\"), \"(\\S?\\d+)\", group)\n",
    "    )\n",
    "    .drop(\"acceleration\")\n",
    "    .drop(\"cylinders\")\n",
    "    .drop(\"displacement\")\n",
    "    .drop(\"modelyear\")\n",
    "    .drop(\"mpg\")\n",
    "    .drop(\"origin\")\n",
    "    .drop(\"horsepower\")\n",
    "    .drop(\"weight\")\n",
    ")\n",
    "df.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b7979d0a",
   "metadata": {},
   "source": [
    "**Fill NULL values in specific columns**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "de505160",
   "metadata": {},
   "outputs": [],
   "source": [
    "df = auto_df.fillna({\"horsepower\": 0})\n",
    "df.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5e95098c",
   "metadata": {},
   "source": [
    "**Fill NULL values with column average**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5bcd520a",
   "metadata": {},
   "outputs": [],
   "source": [
    "from pyspark.sql.functions import avg\n",
    "\n",
    "df = auto_df.fillna({\"horsepower\": auto_df.agg(avg(\"horsepower\")).first()[0]})\n",
    "df.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d710c307",
   "metadata": {},
   "source": [
    "**Fill NULL values with group average**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "41f5deec",
   "metadata": {},
   "outputs": [],
   "source": [
    "from pyspark.sql.functions import coalesce\n",
    "\n",
    "unmodified_columns = auto_df.columns\n",
    "unmodified_columns.remove(\"horsepower\")\n",
    "manufacturer_avg = auto_df.groupBy(\"cylinders\").agg({\"horsepower\": \"avg\"})\n",
    "df = auto_df.join(manufacturer_avg, \"cylinders\").select(\n",
    "    *unmodified_columns,\n",
    "    coalesce(\"horsepower\", \"avg(horsepower)\").alias(\"horsepower\"),\n",
    ")\n",
    "df.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6c1741b7",
   "metadata": {},
   "source": [
    "**Unpack a DataFrame's JSON column to a new DataFrame**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "657c233c",
   "metadata": {},
   "outputs": [],
   "source": [
    "from pyspark.sql.functions import col, json_tuple\n",
    "\n",
    "source = spark.sparkContext.parallelize(\n",
    "    [[\"1\", '{ \"a\" : 10, \"b\" : 11 }'], [\"2\", '{ \"a\" : 20, \"b\" : 21 }']]\n",
    ").toDF([\"id\", \"json\"])\n",
    "df = source.select(\"id\", json_tuple(col(\"json\"), \"a\", \"b\"))\n",
    "df.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "89c601c5",
   "metadata": {},
   "source": [
    "**Query a JSON column**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "44976758",
   "metadata": {},
   "outputs": [],
   "source": [
    "from pyspark.sql.functions import col, json_tuple\n",
    "\n",
    "source = spark.sparkContext.parallelize(\n",
    "    [[\"1\", '{ \"a\" : 10, \"b\" : 11 }'], [\"2\", '{ \"a\" : 20, \"b\" : 21 }']]\n",
    ").toDF([\"id\", \"json\"])\n",
    "df = (\n",
    "    source.select(\"id\", json_tuple(col(\"json\"), \"a\", \"b\"))\n",
    "    .withColumnRenamed(\"c0\", \"a\")\n",
    "    .withColumnRenamed(\"c1\", \"b\")\n",
    "    .where(col(\"b\") > 15)\n",
    ")\n",
    "df.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0e8719b8",
   "metadata": {},
   "source": [
    "## Filtering, sorting, removing duplicates and more."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "46c1db94",
   "metadata": {},
   "source": [
    "**Filter a column using a condition**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a5a3cf22",
   "metadata": {},
   "outputs": [],
   "source": [
    "from pyspark.sql.functions import col\n",
    "\n",
    "df = auto_df.filter(col(\"mpg\") > \"30\")\n",
    "df.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f991878e",
   "metadata": {},
   "source": [
    "**Filter based on a specific column value**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0ed2162e",
   "metadata": {},
   "outputs": [],
   "source": [
    "from pyspark.sql.functions import col\n",
    "\n",
    "df = auto_df.where(col(\"cylinders\") == \"8\")\n",
    "df.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "73f9852f",
   "metadata": {},
   "source": [
    "**Filter based on an IN list**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "16077fb1",
   "metadata": {},
   "outputs": [],
   "source": [
    "from pyspark.sql.functions import col\n",
    "\n",
    "df = auto_df.where(col(\"cylinders\").isin([\"4\", \"6\"]))\n",
    "df.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "579899e4",
   "metadata": {},
   "source": [
    "**Filter based on a NOT IN list**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "554db3ba",
   "metadata": {},
   "outputs": [],
   "source": [
    "from pyspark.sql.functions import col\n",
    "\n",
    "df = auto_df.where(~col(\"cylinders\").isin([\"4\", \"6\"]))\n",
    "df.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "49c9042e",
   "metadata": {},
   "source": [
    "**Filter values based on keys in another DataFrame**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f047c2d6",
   "metadata": {},
   "outputs": [],
   "source": [
    "from pyspark.sql.functions import col\n",
    "\n",
    "# Our DataFrame of keys to exclude.\n",
    "exclude_keys = auto_df.select(\n",
    "    (col(\"modelyear\") + 1).alias(\"adjusted_year\")\n",
    ").distinct()\n",
    "\n",
    "# The anti join returns only keys with no matches.\n",
    "filtered = auto_df.join(\n",
    "    exclude_keys,\n",
    "    how=\"left_anti\",\n",
    "    on=auto_df.modelyear == exclude_keys.adjusted_year,\n",
    ")\n",
    "\n",
    "# Alternatively we can register a temporary table and use a SQL expression.\n",
    "exclude_keys.registerTempTable(\"exclude_keys\")\n",
    "df = auto_df.filter(\n",
    "    \"modelyear not in ( select adjusted_year from exclude_keys )\"\n",
    ")\n",
    "df.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f7bdf568",
   "metadata": {},
   "source": [
    "**Get Dataframe rows that match a substring**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f74b1fab",
   "metadata": {},
   "outputs": [],
   "source": [
    "df = auto_df.where(auto_df.carname.contains(\"custom\"))\n",
    "df.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bb5ecd9f",
   "metadata": {},
   "source": [
    "**Filter a Dataframe based on a custom substring search**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9fddbd57",
   "metadata": {},
   "outputs": [],
   "source": [
    "from pyspark.sql.functions import col\n",
    "\n",
    "df = auto_df.where(col(\"carname\").like(\"%custom%\"))\n",
    "df.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "69540a15",
   "metadata": {},
   "source": [
    "**Filter based on a column's length**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a2dc9b84",
   "metadata": {},
   "outputs": [],
   "source": [
    "from pyspark.sql.functions import col, length\n",
    "\n",
    "df = auto_df.where(length(col(\"carname\")) < 12)\n",
    "df.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e4c02822",
   "metadata": {},
   "source": [
    "**Multiple filter conditions**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5a9b85cb",
   "metadata": {},
   "outputs": [],
   "source": [
    "from pyspark.sql.functions import col\n",
    "\n",
    "# OR\n",
    "df = auto_df.filter((col(\"mpg\") > \"30\") | (col(\"acceleration\") < \"10\"))\n",
    "# AND\n",
    "df = auto_df.filter((col(\"mpg\") > \"30\") & (col(\"acceleration\") < \"13\"))\n",
    "df.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e370fb8e",
   "metadata": {},
   "source": [
    "**Sort DataFrame by a column**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8b7e38e9",
   "metadata": {},
   "outputs": [],
   "source": [
    "from pyspark.sql.functions import col\n",
    "\n",
    "df = auto_df.orderBy(\"carname\")\n",
    "df = auto_df.orderBy(col(\"carname\").desc())\n",
    "df.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "56e050fc",
   "metadata": {},
   "source": [
    "**Take the first N rows of a DataFrame**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c26c843b",
   "metadata": {},
   "outputs": [],
   "source": [
    "n = 10\n",
    "df = auto_df.limit(n)\n",
    "df.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e59e684b",
   "metadata": {},
   "source": [
    "**Get distinct values of a column**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c71584e9",
   "metadata": {},
   "outputs": [],
   "source": [
    "df = auto_df.select(\"cylinders\").distinct()\n",
    "df.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "486bcea0",
   "metadata": {},
   "source": [
    "**Remove duplicates**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a7876216",
   "metadata": {},
   "outputs": [],
   "source": [
    "df = auto_df.dropDuplicates([\"carname\"])\n",
    "df.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fe1e82ad",
   "metadata": {},
   "source": [
    "## Group DataFrame data by key to perform aggregates like counting, sums, averages, etc."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "710bd79d",
   "metadata": {},
   "source": [
    "**count(*) on a particular column**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "46216a82",
   "metadata": {},
   "outputs": [],
   "source": [
    "from pyspark.sql.functions import desc\n",
    "\n",
    "# No sorting.\n",
    "df = auto_df.groupBy(\"cylinders\").count()\n",
    "\n",
    "# With sorting.\n",
    "df = auto_df.groupBy(\"cylinders\").count().orderBy(desc(\"count\"))\n",
    "df.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a759e27f",
   "metadata": {},
   "source": [
    "**Group and sort**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "25296bfe",
   "metadata": {},
   "outputs": [],
   "source": [
    "from pyspark.sql.functions import avg, desc\n",
    "\n",
    "df = (\n",
    "    auto_df.groupBy(\"cylinders\")\n",
    "    .agg(avg(\"horsepower\").alias(\"avg_horsepower\"))\n",
    "    .orderBy(desc(\"avg_horsepower\"))\n",
    ")\n",
    "df.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "26fba7c5",
   "metadata": {},
   "source": [
    "**Filter groups based on an aggregate value, equivalent to SQL HAVING clause**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9b7601b6",
   "metadata": {},
   "outputs": [],
   "source": [
    "from pyspark.sql.functions import col, desc\n",
    "\n",
    "df = (\n",
    "    auto_df.groupBy(\"cylinders\")\n",
    "    .count()\n",
    "    .orderBy(desc(\"count\"))\n",
    "    .filter(col(\"count\") > 100)\n",
    ")\n",
    "df.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6f7feaad",
   "metadata": {},
   "source": [
    "**Group by multiple columns**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "94861105",
   "metadata": {},
   "outputs": [],
   "source": [
    "from pyspark.sql.functions import avg, desc\n",
    "\n",
    "df = (\n",
    "    auto_df.groupBy([\"modelyear\", \"cylinders\"])\n",
    "    .agg(avg(\"horsepower\").alias(\"avg_horsepower\"))\n",
    "    .orderBy(desc(\"avg_horsepower\"))\n",
    ")\n",
    "df.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f78a6f8a",
   "metadata": {},
   "source": [
    "**Aggregate multiple columns**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "43dc1467",
   "metadata": {},
   "outputs": [],
   "source": [
    "expressions = dict(horsepower=\"avg\", weight=\"max\", displacement=\"max\")\n",
    "df = auto_df.groupBy(\"modelyear\").agg(expressions)\n",
    "df.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "74555642",
   "metadata": {},
   "source": [
    "**Aggregate multiple columns with custom orderings**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "66de12db",
   "metadata": {},
   "outputs": [],
   "source": [
    "from pyspark.sql.functions import asc, desc_nulls_last\n",
    "\n",
    "expressions = dict(horsepower=\"avg\", weight=\"max\", displacement=\"max\")\n",
    "orderings = [\n",
    "    desc_nulls_last(\"max(displacement)\"),\n",
    "    desc_nulls_last(\"avg(horsepower)\"),\n",
    "    asc(\"max(weight)\"),\n",
    "]\n",
    "df = auto_df.groupBy(\"modelyear\").agg(expressions).orderBy(*orderings)\n",
    "df.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2955b569",
   "metadata": {},
   "source": [
    "**Get the maximum of a column**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "390defaf",
   "metadata": {},
   "outputs": [],
   "source": [
    "from pyspark.sql.functions import col, max\n",
    "\n",
    "df = auto_df.select(max(col(\"horsepower\")).alias(\"max_horsepower\"))\n",
    "df.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "aa209709",
   "metadata": {},
   "source": [
    "**Sum a list of columns**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5e7d004c",
   "metadata": {},
   "outputs": [],
   "source": [
    "exprs = {x: \"sum\" for x in (\"weight\", \"cylinders\", \"mpg\")}\n",
    "df = auto_df.agg(exprs)\n",
    "df.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "08d43fb9",
   "metadata": {},
   "source": [
    "**Sum a column**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "18ca4f04",
   "metadata": {},
   "outputs": [],
   "source": [
    "from pyspark.sql.functions import sum\n",
    "\n",
    "df = auto_df.groupBy(\"cylinders\").agg(sum(\"weight\").alias(\"total_weight\"))\n",
    "df.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5f1fa319",
   "metadata": {},
   "source": [
    "**Aggregate all numeric columns**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "eaccf481",
   "metadata": {},
   "outputs": [],
   "source": [
    "numerics = set([\"decimal\", \"double\", \"float\", \"integer\", \"long\", \"short\"])\n",
    "exprs = {x[0]: \"sum\" for x in auto_df_fixed.dtypes if x[1] in numerics}\n",
    "df = auto_df_fixed.agg(exprs)\n",
    "df.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d7130827",
   "metadata": {},
   "source": [
    "**Count unique after grouping**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b98955cd",
   "metadata": {},
   "outputs": [],
   "source": [
    "from pyspark.sql.functions import countDistinct\n",
    "\n",
    "df = auto_df.groupBy(\"cylinders\").agg(countDistinct(\"mpg\"))\n",
    "df.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bf2d5672",
   "metadata": {},
   "source": [
    "**Count distinct values on all columns**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f8f02df5",
   "metadata": {},
   "outputs": [],
   "source": [
    "from pyspark.sql.functions import countDistinct\n",
    "\n",
    "df = auto_df.agg(*(countDistinct(c) for c in auto_df.columns))\n",
    "df.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ebb7b6f2",
   "metadata": {},
   "source": [
    "**Group by then filter on the count**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4a27fcbc",
   "metadata": {},
   "outputs": [],
   "source": [
    "from pyspark.sql.functions import col\n",
    "\n",
    "df = auto_df.groupBy(\"cylinders\").count().where(col(\"count\") > 100)\n",
    "df.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "853b589a",
   "metadata": {},
   "source": [
    "**Find the top N per row group (use N=1 for maximum)**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a7d21719",
   "metadata": {},
   "outputs": [],
   "source": [
    "from pyspark.sql.functions import col, row_number\n",
    "from pyspark.sql.window import Window\n",
    "\n",
    "# To get the maximum per group, set n=1.\n",
    "n = 5\n",
    "w = Window().partitionBy(\"cylinders\").orderBy(col(\"horsepower\").desc())\n",
    "df = (\n",
    "    auto_df.withColumn(\"horsepower\", col(\"horsepower\").cast(\"double\"))\n",
    "    .withColumn(\"rn\", row_number().over(w))\n",
    "    .where(col(\"rn\") <= n)\n",
    "    .select(\"*\")\n",
    ")\n",
    "df.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5c8c45c1",
   "metadata": {},
   "source": [
    "**Group key/values into a list**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e7859641",
   "metadata": {},
   "outputs": [],
   "source": [
    "from pyspark.sql.functions import col, collect_list\n",
    "\n",
    "df = auto_df.groupBy(\"cylinders\").agg(\n",
    "    collect_list(col(\"carname\")).alias(\"models\")\n",
    ")\n",
    "df.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "11093356",
   "metadata": {},
   "source": [
    "**Compute a histogram**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "55996d3f",
   "metadata": {},
   "outputs": [],
   "source": [
    "from pyspark.sql.functions import col\n",
    "\n",
    "# Target column must be numeric.\n",
    "df = auto_df.withColumn(\"horsepower\", col(\"horsepower\").cast(\"double\"))\n",
    "\n",
    "# N is the number of bins.\n",
    "N = 11\n",
    "histogram = df.select(\"horsepower\").rdd.flatMap(lambda x: x).histogram(N)\n",
    "print(histogram)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e41bbb49",
   "metadata": {},
   "source": [
    "**Compute global percentiles**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5f104d08",
   "metadata": {},
   "outputs": [],
   "source": [
    "from pyspark.sql.functions import col, ntile\n",
    "from pyspark.sql.window import Window\n",
    "\n",
    "w = Window().orderBy(col(\"mpg\").desc())\n",
    "df = auto_df.withColumn(\"ntile4\", ntile(4).over(w))\n",
    "df.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4acb48a7",
   "metadata": {},
   "source": [
    "**Compute percentiles within a partition**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fff99179",
   "metadata": {},
   "outputs": [],
   "source": [
    "from pyspark.sql.functions import col, ntile\n",
    "from pyspark.sql.window import Window\n",
    "\n",
    "w = Window().partitionBy(\"cylinders\").orderBy(col(\"mpg\").desc())\n",
    "df = auto_df.withColumn(\"ntile4\", ntile(4).over(w))\n",
    "df.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6955acfc",
   "metadata": {},
   "source": [
    "**Compute percentiles after aggregating**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b4f13508",
   "metadata": {},
   "outputs": [],
   "source": [
    "from pyspark.sql.functions import col, ntile\n",
    "from pyspark.sql.window import Window\n",
    "\n",
    "grouped = auto_df.groupBy(\"modelyear\").count()\n",
    "w = Window().orderBy(col(\"count\").desc())\n",
    "df = grouped.withColumn(\"ntile4\", ntile(4).over(w))\n",
    "df.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "95da72f7",
   "metadata": {},
   "source": [
    "**Filter rows with values below a target percentile**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e8a1cfdd",
   "metadata": {},
   "outputs": [],
   "source": [
    "from pyspark.sql.functions import col, lit\n",
    "import pyspark.sql.functions as F\n",
    "\n",
    "target_percentile = auto_df.agg(\n",
    "    F.expr(\"percentile(mpg, 0.9)\").alias(\"target_percentile\")\n",
    ").first()[0]\n",
    "df = auto_df.filter(col(\"mpg\") > lit(target_percentile))\n",
    "df.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2322fa05",
   "metadata": {},
   "source": [
    "**Aggregate and rollup**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "893365bb",
   "metadata": {},
   "outputs": [],
   "source": [
    "from pyspark.sql.functions import avg, col, count, desc\n",
    "\n",
    "subset = auto_df.filter(col(\"modelyear\") > 79)\n",
    "df = (\n",
    "    subset.rollup(\"modelyear\", \"cylinders\")\n",
    "    .agg(\n",
    "        avg(\"horsepower\").alias(\"avg_horsepower\"),\n",
    "        count(\"modelyear\").alias(\"count\"),\n",
    "    )\n",
    "    .orderBy(desc(\"modelyear\"), desc(\"cylinders\"))\n",
    ")\n",
    "df.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "55992054",
   "metadata": {},
   "source": [
    "**Aggregate and cube**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b34a37e6",
   "metadata": {},
   "outputs": [],
   "source": [
    "from pyspark.sql.functions import avg, col, count, desc\n",
    "\n",
    "subset = auto_df.filter(col(\"modelyear\") > 79)\n",
    "df = (\n",
    "    subset.cube(\"modelyear\", \"cylinders\")\n",
    "    .agg(\n",
    "        avg(\"horsepower\").alias(\"avg_horsepower\"),\n",
    "        count(\"modelyear\").alias(\"count\"),\n",
    "    )\n",
    "    .orderBy(desc(\"modelyear\"), desc(\"cylinders\"))\n",
    ")\n",
    "df.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "833d97dd",
   "metadata": {},
   "source": [
    "## Spark allows DataFrames to be joined similarly to how tables are joined in an RDBMS. The diagram below shows join types available in Spark.\n",
    "\n",
    "![Spark Join Types](images/jointypes.webp)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8dea4916",
   "metadata": {},
   "source": [
    "**Join two DataFrames by column name**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a027b4d0",
   "metadata": {},
   "outputs": [],
   "source": [
    "from pyspark.sql.functions import udf\n",
    "from pyspark.sql.types import StringType\n",
    "\n",
    "# Load a list of manufacturer / country pairs.\n",
    "countries = (\n",
    "    spark.read.format(\"csv\")\n",
    "    .option(\"header\", True)\n",
    "    .load(\"data/manufacturers.csv\")\n",
    ")\n",
    "\n",
    "# Add a manufacturers column, to join with the manufacturers list.\n",
    "first_word_udf = udf(lambda x: x.split()[0], StringType())\n",
    "df = auto_df.withColumn(\"manufacturer\", first_word_udf(auto_df.carname))\n",
    "\n",
    "# The actual join.\n",
    "df = df.join(countries, \"manufacturer\")\n",
    "df.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3291ccd3",
   "metadata": {},
   "source": [
    "**Join two DataFrames with an expression**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d63fbac1",
   "metadata": {},
   "outputs": [],
   "source": [
    "from pyspark.sql.functions import udf\n",
    "from pyspark.sql.types import StringType\n",
    "\n",
    "# Load a list of manufacturer / country pairs.\n",
    "countries = (\n",
    "    spark.read.format(\"csv\")\n",
    "    .option(\"header\", True)\n",
    "    .load(\"data/manufacturers.csv\")\n",
    ")\n",
    "\n",
    "# Add a manufacturers column, to join with the manufacturers list.\n",
    "first_word_udf = udf(lambda x: x.split()[0], StringType())\n",
    "df = auto_df.withColumn(\"manufacturer\", first_word_udf(auto_df.carname))\n",
    "\n",
    "# The actual join.\n",
    "df = df.join(countries, df.manufacturer == countries.manufacturer)\n",
    "df.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "073f10ed",
   "metadata": {},
   "source": [
    "**Multiple join conditions**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0b22ec74",
   "metadata": {},
   "outputs": [],
   "source": [
    "from pyspark.sql.functions import udf\n",
    "from pyspark.sql.types import StringType\n",
    "\n",
    "# Load a list of manufacturer / country pairs.\n",
    "countries = (\n",
    "    spark.read.format(\"csv\")\n",
    "    .option(\"header\", True)\n",
    "    .load(\"data/manufacturers.csv\")\n",
    ")\n",
    "\n",
    "# Add a manufacturers column, to join with the manufacturers list.\n",
    "first_word_udf = udf(lambda x: x.split()[0], StringType())\n",
    "df = auto_df.withColumn(\"manufacturer\", first_word_udf(auto_df.carname))\n",
    "\n",
    "# The actual join.\n",
    "df = df.join(\n",
    "    countries,\n",
    "    (df.manufacturer == countries.manufacturer)\n",
    "    | (df.mpg == countries.manufacturer),\n",
    ")\n",
    "df.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e86f916a",
   "metadata": {},
   "source": [
    "**Various Spark join types**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e9254dd2",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Inner join on one column.\n",
    "joined = auto_df.join(auto_df, \"carname\")\n",
    "\n",
    "# Left (outer) join.\n",
    "joined = auto_df.join(auto_df, \"carname\", \"left\")\n",
    "\n",
    "# Left anti (not in) join.\n",
    "joined = auto_df.join(auto_df, \"carname\", \"left_anti\")\n",
    "\n",
    "# Right (outer) join.\n",
    "joined = auto_df.join(auto_df, \"carname\", \"right\")\n",
    "\n",
    "# Full join.\n",
    "joined = auto_df.join(auto_df, \"carname\", \"full\")\n",
    "\n",
    "# Cross join.\n",
    "df = auto_df.crossJoin(auto_df)\n",
    "df.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6d3e7ef9",
   "metadata": {},
   "source": [
    "**Concatenate two DataFrames**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3f3b5272",
   "metadata": {},
   "outputs": [],
   "source": [
    "df1 = spark.read.format(\"csv\").option(\"header\", True).load(\"data/part1.csv\")\n",
    "df2 = spark.read.format(\"csv\").option(\"header\", True).load(\"data/part2.csv\")\n",
    "df = df1.union(df2)\n",
    "df.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "19b8769f",
   "metadata": {},
   "source": [
    "**Load multiple files into a single DataFrame**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a999fe28",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Approach 1: Use a list.\n",
    "df = (\n",
    "    spark.read.format(\"csv\")\n",
    "    .option(\"header\", True)\n",
    "    .load([\"data/part1.csv\", \"data/part2.csv\"])\n",
    ")\n",
    "\n",
    "# Approach 2: Use a wildcard.\n",
    "df = spark.read.format(\"csv\").option(\"header\", True).load(\"data/part*.csv\")\n",
    "df.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "eec90102",
   "metadata": {},
   "source": [
    "**Subtract DataFrames**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "10a16368",
   "metadata": {},
   "outputs": [],
   "source": [
    "from pyspark.sql.functions import col\n",
    "\n",
    "df = auto_df.subtract(auto_df.where(col(\"mpg\") < \"25\"))\n",
    "df.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "88d54e1c",
   "metadata": {},
   "source": [
    "## Loading File Metadata and Processing Files."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d91f7633",
   "metadata": {},
   "source": [
    "**Load Local File Details into a DataFrame**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f457d3bd",
   "metadata": {},
   "outputs": [],
   "source": [
    "from pyspark.sql.types import (\n",
    "    StructField,\n",
    "    StructType,\n",
    "    LongType,\n",
    "    StringType,\n",
    "    TimestampType,\n",
    ")\n",
    "import datetime\n",
    "import glob\n",
    "import os\n",
    "\n",
    "# Simple: Use glob and only file names.\n",
    "files = [[x] for x in glob.glob(\"/etc/*\")]\n",
    "df = spark.createDataFrame(files)\n",
    "\n",
    "# Advanced: Use os.walk and extended attributes.\n",
    "target_path = \"/etc\"\n",
    "entries = []\n",
    "walker = os.walk(target_path)\n",
    "for root, dirs, files in walker:\n",
    "    for file in files:\n",
    "        full_path = os.path.join(root, file)\n",
    "        try:\n",
    "            stat_info = os.stat(full_path)\n",
    "            entries.append(\n",
    "                [\n",
    "                    file,\n",
    "                    full_path,\n",
    "                    stat_info.st_size,\n",
    "                    datetime.datetime.fromtimestamp(stat_info.st_mtime),\n",
    "                ]\n",
    "            )\n",
    "        except:\n",
    "            pass\n",
    "schema = StructType(\n",
    "    [\n",
    "        StructField(\"file\", StringType(), False),\n",
    "        StructField(\"path\", StringType(), False),\n",
    "        StructField(\"size\", LongType(), False),\n",
    "        StructField(\"mtime\", TimestampType(), False),\n",
    "    ]\n",
    ")\n",
    "df = spark.createDataFrame(entries, schema)\n",
    "df.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bc18ef12",
   "metadata": {},
   "source": [
    "**Load Files from Oracle Cloud Infrastructure into a DataFrame**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cd3fe18f",
   "metadata": {},
   "outputs": [],
   "source": [
    "import oci\n",
    "from pyspark.sql.types import (\n",
    "    StructField,\n",
    "    StructType,\n",
    "    LongType,\n",
    "    StringType,\n",
    "    TimestampType,\n",
    ")\n",
    "\n",
    "def get_authenticated_client(client):\n",
    "    config = oci.config.from_file()\n",
    "    authenticated_client = client(config)\n",
    "    return authenticated_client\n",
    "\n",
    "object_store_client = get_authenticated_client(\n",
    "    oci.object_storage.ObjectStorageClient\n",
    ")\n",
    "\n",
    "# Requires an object_store_client object.\n",
    "# See https://oracle-cloud-infrastructure-python-sdk.readthedocs.io/en/latest/api/object_storage/client/oci.object_storage.ObjectStorageClient.html\n",
    "input_bucket = \"oow_2019_dataflow_lab\"\n",
    "raw_inputs = object_store_client.list_objects(\n",
    "    object_store_client.get_namespace().data,\n",
    "    input_bucket,\n",
    "    fields=\"size,md5,timeModified\",\n",
    ")\n",
    "files = [\n",
    "    [x.name, x.size, x.time_modified, x.md5] for x in raw_inputs.data.objects\n",
    "]\n",
    "schema = StructType(\n",
    "    [\n",
    "        StructField(\"name\", StringType(), False),\n",
    "        StructField(\"size\", LongType(), True),\n",
    "        StructField(\"modified\", TimestampType(), True),\n",
    "        StructField(\"md5\", StringType(), True),\n",
    "    ]\n",
    ")\n",
    "df = spark.createDataFrame(files, schema)\n",
    "df.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7e64f198",
   "metadata": {},
   "source": [
    "**Transform Many Images using Pillow**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8bfe8e6c",
   "metadata": {},
   "outputs": [],
   "source": [
    "from PIL import Image\n",
    "import glob\n",
    "\n",
    "def resize_an_image(row):\n",
    "    width, height = 128, 128\n",
    "    file_name = row._1\n",
    "    new_name = file_name.replace(\".png\", \".resized.png\")\n",
    "    img = Image.open(file_name)\n",
    "    img = img.resize((width, height), Image.ANTIALIAS)\n",
    "    img.save(new_name)\n",
    "\n",
    "files = [[x] for x in glob.glob(\"data/resize_image?.png\")]\n",
    "df = spark.createDataFrame(files)\n",
    "df.foreach(resize_an_image)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "302ed111",
   "metadata": {},
   "source": [
    "## Dealing with NULLs and NaNs in DataFrames."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ac839ccf",
   "metadata": {},
   "source": [
    "**Filter rows with None or Null values**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3e318f0c",
   "metadata": {},
   "outputs": [],
   "source": [
    "from pyspark.sql.functions import col\n",
    "\n",
    "df = auto_df.where(col(\"horsepower\").isNull())\n",
    "df = auto_df.where(col(\"horsepower\").isNotNull())\n",
    "df.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ef83fc83",
   "metadata": {},
   "source": [
    "**Drop rows with Null values**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "de2e38a8",
   "metadata": {},
   "outputs": [],
   "source": [
    "df = auto_df.na.drop(thresh=1, subset=(\"horsepower\",))\n",
    "df.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "259caa41",
   "metadata": {},
   "source": [
    "**Count all Null or NaN values in a DataFrame**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b6620316",
   "metadata": {},
   "outputs": [],
   "source": [
    "from pyspark.sql.functions import col, count, isnan, when\n",
    "\n",
    "df = auto_df.select(\n",
    "    [count(when(isnan(c), c)).alias(c) for c in auto_df.columns]\n",
    ")\n",
    "df = auto_df.select(\n",
    "    [count(when(col(c).isNull(), c)).alias(c) for c in auto_df.columns]\n",
    ")\n",
    "df.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4744c44f",
   "metadata": {},
   "source": [
    "## Parsing and processing dates and times."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "521c23ef",
   "metadata": {},
   "source": [
    "**Convert an ISO 8601 formatted date string to date type**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4f4aaa65",
   "metadata": {},
   "outputs": [],
   "source": [
    "from pyspark.sql.functions import col\n",
    "\n",
    "df = spark.sparkContext.parallelize([[\"2021-01-01\"], [\"2022-01-01\"]]).toDF(\n",
    "    [\"date_col\"]\n",
    ")\n",
    "df = df.withColumn(\"date_col\", col(\"date_col\").cast(\"date\"))\n",
    "df.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9c32e1d1",
   "metadata": {},
   "source": [
    "**Convert a custom formatted date string to date type**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5cafe3d4",
   "metadata": {},
   "outputs": [],
   "source": [
    "from pyspark.sql.functions import col, to_date\n",
    "\n",
    "df = spark.sparkContext.parallelize([[\"20210101\"], [\"20220101\"]]).toDF(\n",
    "    [\"date_col\"]\n",
    ")\n",
    "df = df.withColumn(\"date_col\", to_date(col(\"date_col\"), \"yyyyddMM\"))\n",
    "df.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "80badd09",
   "metadata": {},
   "source": [
    "**Get the last day of the current month**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7f7eb3e9",
   "metadata": {},
   "outputs": [],
   "source": [
    "from pyspark.sql.functions import col, last_day\n",
    "\n",
    "df = spark.sparkContext.parallelize([[\"2020-01-01\"], [\"1712-02-10\"]]).toDF(\n",
    "    [\"date_col\"]\n",
    ")\n",
    "df = df.withColumn(\"date_col\", col(\"date_col\").cast(\"date\")).withColumn(\n",
    "    \"last_day\", last_day(col(\"date_col\"))\n",
    ")\n",
    "df.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "00e4a479",
   "metadata": {},
   "source": [
    "**Convert UNIX (seconds since epoch) timestamp to date**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c1081992",
   "metadata": {},
   "outputs": [],
   "source": [
    "from pyspark.sql.functions import col, from_unixtime\n",
    "\n",
    "df = spark.sparkContext.parallelize([[\"1590183026\"], [\"2000000000\"]]).toDF(\n",
    "    [\"ts_col\"]\n",
    ")\n",
    "df = df.withColumn(\"date_col\", from_unixtime(col(\"ts_col\")))\n",
    "df.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "078588b4",
   "metadata": {},
   "source": [
    "**Load a CSV file with complex dates into a DataFrame**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5c9c10e5",
   "metadata": {},
   "outputs": [],
   "source": [
    "from pyspark.sql.functions import udf\n",
    "from pyspark.sql.types import TimestampType\n",
    "import dateparser\n",
    "\n",
    "# Use the dateparser module to convert many formats into timestamps.\n",
    "date_convert = udf(\n",
    "    lambda x: dateparser.parse(x) if x is not None else None, TimestampType()\n",
    ")\n",
    "df = (\n",
    "    spark.read.format(\"csv\")\n",
    "    .option(\"header\", True)\n",
    "    .load(\"data/date_examples.csv\")\n",
    ")\n",
    "df = df.withColumn(\"parsed\", date_convert(df.date))\n",
    "df.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f26e0cc9",
   "metadata": {},
   "source": [
    "## Analyzing unstructured data like [JSON](https://spark.apache.org/docs/latest/sql-data-sources-json.html), XML, etc."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "66533e3a",
   "metadata": {},
   "source": [
    "**Flatten top level text fields in a JSONl document**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "09cf9714",
   "metadata": {},
   "outputs": [],
   "source": [
    "from pyspark.sql.functions import col\n",
    "\n",
    "# Load JSONl into a DataFrame. Schema is inferred automatically.\n",
    "base = spark.read.json(\"data/financial.jsonl\")\n",
    "\n",
    "# Extract interesting fields. Alias keeps columns readable.\n",
    "target_json_fields = [\n",
    "    col(\"symbol\").alias(\"symbol\"),\n",
    "    col(\"quoteType.longName\").alias(\"longName\"),\n",
    "    col(\"price.marketCap.raw\").alias(\"marketCap\"),\n",
    "    col(\"summaryDetail.previousClose.raw\").alias(\"previousClose\"),\n",
    "    col(\"summaryDetail.fiftyTwoWeekHigh.raw\").alias(\"fiftyTwoWeekHigh\"),\n",
    "    col(\"summaryDetail.fiftyTwoWeekLow.raw\").alias(\"fiftyTwoWeekLow\"),\n",
    "    col(\"summaryDetail.trailingPE.raw\").alias(\"trailingPE\"),\n",
    "]\n",
    "df = base.select(target_json_fields)\n",
    "df.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a94f7a14",
   "metadata": {},
   "source": [
    "**Flatten top level text fields from a JSON column**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "30142914",
   "metadata": {},
   "outputs": [],
   "source": [
    "from pyspark.sql.functions import col, from_json, schema_of_json\n",
    "\n",
    "# quote/escape options needed when loading CSV containing JSON.\n",
    "base = (\n",
    "    spark.read.format(\"csv\")\n",
    "    .option(\"header\", True)\n",
    "    .option(\"quote\", '\"')\n",
    "    .option(\"escape\", '\"')\n",
    "    .load(\"data/financial.csv\")\n",
    ")\n",
    "\n",
    "# Infer JSON schema from one entry in the DataFrame.\n",
    "sample_json_document = base.select(\"financial_data\").first()[0]\n",
    "schema = schema_of_json(sample_json_document)\n",
    "\n",
    "# Parse using this schema.\n",
    "parsed = base.withColumn(\"parsed\", from_json(\"financial_data\", schema))\n",
    "\n",
    "# Extract interesting fields.\n",
    "target_json_fields = [\n",
    "    col(\"parsed.symbol\").alias(\"symbol\"),\n",
    "    col(\"parsed.quoteType.longName\").alias(\"longName\"),\n",
    "    col(\"parsed.price.marketCap.raw\").alias(\"marketCap\"),\n",
    "    col(\"parsed.summaryDetail.previousClose.raw\").alias(\"previousClose\"),\n",
    "    col(\"parsed.summaryDetail.fiftyTwoWeekHigh.raw\").alias(\"fiftyTwoWeekHigh\"),\n",
    "    col(\"parsed.summaryDetail.fiftyTwoWeekLow.raw\").alias(\"fiftyTwoWeekLow\"),\n",
    "    col(\"parsed.summaryDetail.trailingPE.raw\").alias(\"trailingPE\"),\n",
    "]\n",
    "df = parsed.select(target_json_fields)\n",
    "df.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7b805343",
   "metadata": {},
   "source": [
    "**Unnest an array of complex structures**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "30b74b58",
   "metadata": {},
   "outputs": [],
   "source": [
    "from pyspark.sql.functions import col, explode\n",
    "\n",
    "base = spark.read.json(\"data/financial.jsonl\")\n",
    "\n",
    "# Analyze balance sheet data, which is held in an array of complex types.\n",
    "target_json_fields = [\n",
    "    col(\"symbol\").alias(\"symbol\"),\n",
    "    col(\"balanceSheetHistoryQuarterly.balanceSheetStatements\").alias(\n",
    "        \"balanceSheetStatements\"\n",
    "    ),\n",
    "]\n",
    "selected = base.select(target_json_fields)\n",
    "\n",
    "# Select a few fields from the balance sheet statement data.\n",
    "target_json_fields = [\n",
    "    col(\"symbol\").alias(\"symbol\"),\n",
    "    col(\"col.endDate.fmt\").alias(\"endDate\"),\n",
    "    col(\"col.cash.raw\").alias(\"cash\"),\n",
    "    col(\"col.totalAssets.raw\").alias(\"totalAssets\"),\n",
    "    col(\"col.totalLiab.raw\").alias(\"totalLiab\"),\n",
    "]\n",
    "\n",
    "# Balance sheet data is in an array, use explode to generate one row per entry.\n",
    "df = selected.select(\"symbol\", explode(\"balanceSheetStatements\")).select(\n",
    "    target_json_fields\n",
    ")\n",
    "df.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "34ec3253",
   "metadata": {},
   "source": [
    "## Using Python's Pandas library to augment Spark. Some operations require the pyarrow library."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a9473b6e",
   "metadata": {},
   "source": [
    "**Convert Spark DataFrame to Pandas DataFrame**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "12e50e63",
   "metadata": {},
   "outputs": [],
   "source": [
    "pandas_df = auto_df.toPandas()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4a4cbdb7",
   "metadata": {},
   "source": [
    "**Convert Pandas DataFrame to Spark DataFrame with Schema Detection**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5c0f212a",
   "metadata": {},
   "outputs": [],
   "source": [
    "df = spark.createDataFrame(pandas_df)\n",
    "df.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4be0178b",
   "metadata": {},
   "source": [
    "**Convert Pandas DataFrame to Spark DataFrame using a Custom Schema**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a63a37f1",
   "metadata": {},
   "outputs": [],
   "source": [
    "# This code converts everything to strings.\n",
    "# If you want to preserve types, see https://gist.github.com/tonyfraser/79a255aa8a9d765bd5cf8bd13597171e\n",
    "from pyspark.sql.types import StructField, StructType, StringType\n",
    "\n",
    "schema = StructType(\n",
    "    [StructField(name, StringType(), True) for name in pandas_df.columns]\n",
    ")\n",
    "df = spark.createDataFrame(pandas_df, schema)\n",
    "df.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1dcb0766",
   "metadata": {},
   "source": [
    "**Convert N rows from a DataFrame to a Pandas DataFrame**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c72d70c4",
   "metadata": {},
   "outputs": [],
   "source": [
    "N = 10\n",
    "pdf = auto_df.limit(N).toPandas()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5ccccf97",
   "metadata": {},
   "source": [
    "**Grouped Aggregation with Pandas**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c7792da1",
   "metadata": {},
   "outputs": [],
   "source": [
    "from pyspark.sql.functions import pandas_udf\n",
    "from pandas import DataFrame\n",
    "\n",
    "@pandas_udf(\"double\")\n",
    "def mean_udaf(pdf: DataFrame) -> float:\n",
    "    return pdf.mean()\n",
    "\n",
    "df = auto_df.groupby(\"cylinders\").agg(mean_udaf(auto_df[\"mpg\"]))\n",
    "df.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "06de5ce0",
   "metadata": {},
   "source": [
    "**Use a Pandas Grouped Map Function via applyInPandas**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e8374234",
   "metadata": {},
   "outputs": [],
   "source": [
    "def rescale(pdf):\n",
    "    minv = pdf.horsepower.min()\n",
    "    maxv = pdf.horsepower.max() - minv\n",
    "    return pdf.assign(horsepower=(pdf.horsepower - minv) / maxv * 100)\n",
    "\n",
    "df = auto_df.groupby(\"cylinders\").applyInPandas(rescale, auto_df.schema)\n",
    "df.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8bc6ca18",
   "metadata": {},
   "source": [
    "## Extracting key statistics out of a body of data."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fc35b502",
   "metadata": {},
   "source": [
    "**Compute the number of NULLs across all columns**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "49a856f4",
   "metadata": {},
   "outputs": [],
   "source": [
    "from pyspark.sql.functions import col, count, when\n",
    "\n",
    "df = auto_df.select(\n",
    "    [count(when(col(c).isNull(), c)).alias(c) for c in auto_df.columns]\n",
    ")\n",
    "df.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ea4a7f02",
   "metadata": {},
   "source": [
    "**Compute average values of all numeric columns**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "dbb32806",
   "metadata": {},
   "outputs": [],
   "source": [
    "numerics = set([\"decimal\", \"double\", \"float\", \"integer\", \"long\", \"short\"])\n",
    "exprs = {x[0]: \"avg\" for x in auto_df_fixed.dtypes if x[1] in numerics}\n",
    "df = auto_df_fixed.agg(exprs)\n",
    "df.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b092c12f",
   "metadata": {},
   "source": [
    "**Compute minimum values of all numeric columns**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "dc45036f",
   "metadata": {},
   "outputs": [],
   "source": [
    "numerics = set([\"decimal\", \"double\", \"float\", \"integer\", \"long\", \"short\"])\n",
    "exprs = {x[0]: \"min\" for x in auto_df_fixed.dtypes if x[1] in numerics}\n",
    "df = auto_df_fixed.agg(exprs)\n",
    "df.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b9d39935",
   "metadata": {},
   "source": [
    "**Compute maximum values of all numeric columns**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8e6ac9d4",
   "metadata": {},
   "outputs": [],
   "source": [
    "numerics = set([\"decimal\", \"double\", \"float\", \"integer\", \"long\", \"short\"])\n",
    "exprs = {x[0]: \"max\" for x in auto_df_fixed.dtypes if x[1] in numerics}\n",
    "df = auto_df_fixed.agg(exprs)\n",
    "df.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b99bada1",
   "metadata": {},
   "source": [
    "**Compute median values of all numeric columns**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "487048ad",
   "metadata": {},
   "outputs": [],
   "source": [
    "import pyspark.sql.functions as F\n",
    "\n",
    "numerics = set([\"decimal\", \"double\", \"float\", \"integer\", \"long\", \"short\"])\n",
    "aggregates = []\n",
    "for name, dtype in auto_df_fixed.dtypes:\n",
    "    if dtype not in numerics:\n",
    "        continue\n",
    "    aggregates.append(\n",
    "        F.expr(\"percentile({}, 0.5)\".format(name)).alias(\n",
    "            \"{}_median\".format(name)\n",
    "        )\n",
    "    )\n",
    "df = auto_df_fixed.agg(*aggregates)\n",
    "\n",
    "df.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f181aa55",
   "metadata": {},
   "source": [
    "**Identify Outliers in a DataFrame**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fd0806c9",
   "metadata": {},
   "outputs": [],
   "source": [
    "from pyspark.sql.functions import col, sqrt\n",
    "\n",
    "target_column = \"mpg\"\n",
    "z_score_threshold = 2\n",
    "\n",
    "# Compute the median of the target column.\n",
    "target_df = auto_df.select(target_column)\n",
    "target_df.registerTempTable(\"target_column\")\n",
    "profiled = sqlContext.sql(\n",
    "    f\"select percentile({target_column}, 0.5) as median from target_column\"\n",
    ")\n",
    "\n",
    "# Compute deviations.\n",
    "deviations = target_df.crossJoin(profiled).withColumn(\n",
    "    \"deviation\", sqrt((target_df[target_column] - profiled[\"median\"]) ** 2)\n",
    ")\n",
    "deviations.registerTempTable(\"deviations\")\n",
    "\n",
    "# The Median Absolute Deviation\n",
    "mad = sqlContext.sql(\"select percentile(deviation, 0.5) as mad from deviations\")\n",
    "\n",
    "# Add a modified z score to the original DataFrame.\n",
    "df = (\n",
    "    auto_df.crossJoin(mad)\n",
    "    .crossJoin(profiled)\n",
    "    .withColumn(\n",
    "        \"zscore\",\n",
    "        0.6745\n",
    "        * sqrt((auto_df[target_column] - profiled[\"median\"]) ** 2)\n",
    "        / mad[\"mad\"],\n",
    "    )\n",
    ")\n",
    "\n",
    "df = df.where(col(\"zscore\") > z_score_threshold)\n",
    "df.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b411cb9c",
   "metadata": {},
   "source": [
    "## Upserts, updates and deletes on data."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3d9a45c8",
   "metadata": {},
   "source": [
    "**Save to a Delta Table**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "25c6c753",
   "metadata": {},
   "outputs": [],
   "source": [
    "auto_df.write.mode(\"overwrite\").format(\"delta\").saveAsTable(\"delta_table\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5ffd2a30",
   "metadata": {},
   "source": [
    "**Update records in a DataFrame using Delta Tables**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1e5fe280",
   "metadata": {},
   "outputs": [],
   "source": [
    "from pyspark.sql.functions import expr\n",
    "\n",
    "output_path = \"delta_tests\"\n",
    "\n",
    "# Currently you have to save/reload to convert from table to DataFrame.\n",
    "auto_df.write.mode(\"overwrite\").format(\"delta\").save(output_path)\n",
    "dt = DeltaTable.forPath(spark, output_path)\n",
    "\n",
    "# Run a SQL update operation.\n",
    "dt.update(\n",
    "    condition=expr(\"carname like 'Volks%'\"), set={\"carname\": expr(\"carname\")}\n",
    ")\n",
    "\n",
    "# Convert back to a DataFrame.\n",
    "df = dt.toDF()\n",
    "df.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f36f260d",
   "metadata": {},
   "source": [
    "**Merge into a Delta table**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0d32fb8b",
   "metadata": {},
   "outputs": [],
   "source": [
    "from pyspark.sql.functions import col, expr\n",
    "\n",
    "# Save the original data.\n",
    "output_path = \"delta_tests\"\n",
    "auto_df.write.mode(\"overwrite\").format(\"delta\").save(output_path)\n",
    "\n",
    "# Load data that corrects some car names.\n",
    "corrected_df = (\n",
    "    spark.read.format(\"csv\")\n",
    "    .option(\"header\", True)\n",
    "    .load(\"data/auto-mpg-fixed.csv\")\n",
    ")\n",
    "\n",
    "# Merge the corrected data in.\n",
    "dt = DeltaTable.forPath(spark, output_path)\n",
    "ret = (\n",
    "    dt.alias(\"original\")\n",
    "    .merge(\n",
    "        corrected_df.alias(\"corrected\"),\n",
    "        \"original.modelyear = corrected.modelyear and original.weight = corrected.weight and original.acceleration = corrected.acceleration\",\n",
    "    )\n",
    "    .whenMatchedUpdate(\n",
    "        condition=expr(\"original.carname <> corrected.carname\"),\n",
    "        set={\"carname\": col(\"corrected.carname\")},\n",
    "    )\n",
    "    .whenNotMatchedInsertAll()\n",
    "    .execute()\n",
    ")\n",
    "\n",
    "# Show select table history.\n",
    "df = dt.history().select(\"version operation operationMetrics\".split())\n",
    "\n",
    "df.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "de18a446",
   "metadata": {},
   "source": [
    "**Show Table Version History**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3785523e",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Load our table.\n",
    "output_path = \"delta_tests\"\n",
    "dt = DeltaTable.forPath(spark, output_path)\n",
    "\n",
    "# Show select table history.\n",
    "df = dt.history().select(\"version timestamp operation\".split())\n",
    "df.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "85ea11f6",
   "metadata": {},
   "source": [
    "**Load a Delta Table by Version ID (Time Travel Query)**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f382661e",
   "metadata": {},
   "outputs": [],
   "source": [
    "from pyspark.sql.functions import desc\n",
    "\n",
    "# Get versions.\n",
    "output_path = \"delta_tests\"\n",
    "dt = DeltaTable.forPath(spark, output_path)\n",
    "versions = (\n",
    "    dt.history().select(\"version timestamp\".split()).orderBy(desc(\"version\"))\n",
    ")\n",
    "most_recent_version = versions.first()[0]\n",
    "print(\"Most recent version is\", most_recent_version)\n",
    "\n",
    "# Load the most recent data.\n",
    "df = (\n",
    "    spark.read.format(\"delta\")\n",
    "    .option(\"versionAsOf\", most_recent_version)\n",
    "    .load(output_path)\n",
    ")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "61a96ee2",
   "metadata": {},
   "source": [
    "**Load a Delta Table by Timestamp (Time Travel Query)**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "30bcf495",
   "metadata": {},
   "outputs": [],
   "source": [
    "from pyspark.sql.functions import desc\n",
    "\n",
    "# Get versions.\n",
    "output_path = \"delta_tests\"\n",
    "dt = DeltaTable.forPath(spark, output_path)\n",
    "versions = dt.history().select(\"version timestamp\".split()).orderBy(\"timestamp\")\n",
    "most_recent_timestamp = versions.first()[1]\n",
    "print(\"Most recent timestamp is\", most_recent_timestamp)\n",
    "\n",
    "# Load the oldest version by timestamp.\n",
    "df = (\n",
    "    spark.read.format(\"delta\")\n",
    "    .option(\"timestampAsOf\", most_recent_timestamp)\n",
    "    .load(output_path)\n",
    ")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3eb6ae32",
   "metadata": {},
   "source": [
    "**Compact a Delta Table**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "37bbe406",
   "metadata": {},
   "outputs": [],
   "source": [
    "output_path = \"delta_tests\"\n",
    "\n",
    "# Load table.\n",
    "dt = DeltaTable.forPath(spark, output_path)\n",
    "\n",
    "# Clean up data older than the given window.\n",
    "retention_window_hours = 168\n",
    "dt.vacuum(retention_window_hours)\n",
    "\n",
    "# Show the new versions.\n",
    "df = dt.history().select(\"version timestamp\".split()).orderBy(\"version\")\n",
    "df.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c0a1fd3a",
   "metadata": {},
   "source": [
    "**Add custom metadata to a Delta table write**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b6cfd7fe",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import time\n",
    "\n",
    "extra_properties = dict(\n",
    "    user=os.environ.get(\"USER\"),\n",
    "    write_timestamp=time.time(),\n",
    ")\n",
    "auto_df.write.mode(\"append\").option(\"userMetadata\", extra_properties).format(\n",
    "    \"delta\"\n",
    ").save(\"delta_table_metadata\")\n",
    "df.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "eec22fb3",
   "metadata": {},
   "source": [
    "**Read custom Delta table metadata**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5d508b90",
   "metadata": {},
   "outputs": [],
   "source": [
    "dt = DeltaTable.forPath(spark, \"delta_table_metadata\")\n",
    "df = dt.history().select(\"version timestamp userMetadata\".split())\n",
    "df.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b8820f84",
   "metadata": {},
   "source": [
    "## Spark Streaming (Focuses on Structured Streaming)."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3112a808",
   "metadata": {},
   "source": [
    "**Connect to Kafka using SASL PLAIN authentication**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e9070f7c",
   "metadata": {},
   "outputs": [],
   "source": [
    "options = {\n",
    "    \"kafka.sasl.jaas.config\": 'org.apache.kafka.common.security.plain.PlainLoginModule required username=\"USERNAME\" password=\"PASSWORD\";',\n",
    "    \"kafka.sasl.mechanism\": \"PLAIN\",\n",
    "    \"kafka.security.protocol\": \"SASL_SSL\",\n",
    "    \"kafka.bootstrap.servers\": \"server:9092\",\n",
    "    \"group.id\": \"my_group\",\n",
    "    \"subscribe\": \"my_topic\",\n",
    "}\n",
    "df = spark.readStream.format(\"kafka\").options(**options).load()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d7174c13",
   "metadata": {},
   "source": [
    "**Create a windowed Structured Stream over input CSV files**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "89263c95",
   "metadata": {},
   "outputs": [],
   "source": [
    "from pyspark.sql.functions import avg, count, current_timestamp, window\n",
    "from pyspark.sql.types import (\n",
    "    StructField,\n",
    "    StructType,\n",
    "    DoubleType,\n",
    "    IntegerType,\n",
    "    StringType,\n",
    ")\n",
    "\n",
    "input_location = \"streaming/input\"\n",
    "schema = StructType(\n",
    "    [\n",
    "        StructField(\"mpg\", DoubleType(), True),\n",
    "        StructField(\"cylinders\", IntegerType(), True),\n",
    "        StructField(\"displacement\", DoubleType(), True),\n",
    "        StructField(\"horsepower\", DoubleType(), True),\n",
    "        StructField(\"weight\", DoubleType(), True),\n",
    "        StructField(\"acceleration\", DoubleType(), True),\n",
    "        StructField(\"modelyear\", IntegerType(), True),\n",
    "        StructField(\"origin\", IntegerType(), True),\n",
    "        StructField(\"carname\", StringType(), True),\n",
    "        StructField(\"manufacturer\", StringType(), True),\n",
    "    ]\n",
    ")\n",
    "df = spark.readStream.csv(path=input_location, schema=schema).withColumn(\n",
    "    \"timestamp\", current_timestamp()\n",
    ")\n",
    "\n",
    "aggregated = (\n",
    "    df.groupBy(window(df.timestamp, \"1 minute\"), \"manufacturer\")\n",
    "    .agg(\n",
    "        avg(\"horsepower\").alias(\"avg_horsepower\"),\n",
    "        avg(\"timestamp\").alias(\"avg_timestamp\"),\n",
    "        count(\"modelyear\").alias(\"count\"),\n",
    "    )\n",
    "    .coalesce(10)\n",
    ")\n",
    "summary = aggregated.orderBy(\"window\", \"manufacturer\")\n",
    "query = (\n",
    "    summary.writeStream.outputMode(\"complete\")\n",
    "    .format(\"console\")\n",
    "    .option(\"truncate\", False)\n",
    "    .start()\n",
    ")\n",
    "query.awaitTermination()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "caa86145",
   "metadata": {},
   "source": [
    "**Create an unwindowed Structured Stream over input CSV files**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9563ae9d",
   "metadata": {},
   "outputs": [],
   "source": [
    "from pyspark.sql.functions import avg, count, desc\n",
    "from pyspark.sql.types import (\n",
    "    StructField,\n",
    "    StructType,\n",
    "    DoubleType,\n",
    "    IntegerType,\n",
    "    StringType,\n",
    ")\n",
    "\n",
    "input_location = \"streaming/input\"\n",
    "schema = StructType(\n",
    "    [\n",
    "        StructField(\"mpg\", DoubleType(), True),\n",
    "        StructField(\"cylinders\", IntegerType(), True),\n",
    "        StructField(\"displacement\", DoubleType(), True),\n",
    "        StructField(\"horsepower\", DoubleType(), True),\n",
    "        StructField(\"weight\", DoubleType(), True),\n",
    "        StructField(\"acceleration\", DoubleType(), True),\n",
    "        StructField(\"modelyear\", IntegerType(), True),\n",
    "        StructField(\"origin\", IntegerType(), True),\n",
    "        StructField(\"carname\", StringType(), True),\n",
    "        StructField(\"manufacturer\", StringType(), True),\n",
    "    ]\n",
    ")\n",
    "\n",
    "df = spark.readStream.csv(path=input_location, schema=schema)\n",
    "summary = (\n",
    "    df.groupBy(\"modelyear\")\n",
    "    .agg(\n",
    "        avg(\"horsepower\").alias(\"avg_horsepower\"),\n",
    "        count(\"modelyear\").alias(\"count\"),\n",
    "    )\n",
    "    .orderBy(desc(\"modelyear\"))\n",
    "    .coalesce(10)\n",
    ")\n",
    "query = summary.writeStream.outputMode(\"complete\").format(\"console\").start()\n",
    "query.awaitTermination()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b3e1aebe",
   "metadata": {},
   "source": [
    "**Add the current timestamp to a DataFrame**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0e5e30e9",
   "metadata": {},
   "outputs": [],
   "source": [
    "from pyspark.sql.functions import current_timestamp\n",
    "\n",
    "df = auto_df.withColumn(\"timestamp\", current_timestamp())\n",
    "df.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "48565009",
   "metadata": {},
   "source": [
    "**Session analytics on a DataFrame**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c07059d3",
   "metadata": {},
   "outputs": [],
   "source": [
    "from pyspark.sql.functions import hash, session_window\n",
    "\n",
    "hits_per_session = (\n",
    "    weblog_df.groupBy(\"ip\", session_window(\"time\", \"5 minutes\"))\n",
    "    .count()\n",
    "    .withColumn(\"session\", hash(\"ip\", \"session_window\"))\n",
    ")\n",
    "df.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d25a3f0d",
   "metadata": {},
   "source": [
    "**Call a UDF only when a threshold is reached**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0c90463e",
   "metadata": {},
   "outputs": [],
   "source": [
    "from pyspark.sql.types import BooleanType\n",
    "from pyspark.sql.functions import udf\n",
    "\n",
    "@udf(returnType=BooleanType())\n",
    "def myudf(short_circuit, state, value):\n",
    "    if short_circuit == True:\n",
    "        return True\n",
    "\n",
    "    # Log, send an alert, etc.\n",
    "    return False\n",
    "\n",
    "df = (\n",
    "    spark.readStream.format(\"socket\")\n",
    "    .option(\"host\", \"localhost\")\n",
    "    .option(\"port\", \"9090\")\n",
    "    .load()\n",
    ")\n",
    "parsed = (\n",
    "    df.selectExpr(\n",
    "        \"split(value,',')[0] as state\",\n",
    "        \"split(value,',')[1] as zipcode\",\n",
    "        \"split(value,',')[2] as spend\",\n",
    "    )\n",
    "    .groupBy(\"state\")\n",
    "    .agg({\"spend\": \"avg\"})\n",
    "    .orderBy(desc(\"avg(spend)\"))\n",
    ")\n",
    "tagged = parsed.withColumn(\n",
    "    \"below\", myudf(col(\"avg(spend)\") < 100, col(\"state\"), col(\"avg(spend)\"))\n",
    ")\n",
    "\n",
    "tagged.writeStream.outputMode(\"complete\").format(\n",
    "    \"console\"\n",
    ").start().awaitTermination()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cb8ef79e",
   "metadata": {},
   "source": [
    "**Streaming Machine Learning**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7df91adf",
   "metadata": {},
   "outputs": [],
   "source": [
    "from pyspark.ml import PipelineModel\n",
    "\n",
    "pipeline_model = PipelineModel.load(\"path/to/pipeline\")\n",
    "df = pipeline.transform(input_df)\n",
    "df.writeStream.format(\"console\").start().awaitTermination()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ec19d0d9",
   "metadata": {},
   "source": [
    "**Control stream processing frequency**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "39c04e1b",
   "metadata": {},
   "outputs": [],
   "source": [
    "df.writeStream.outputMode(\"complete\").format(\"console\").trigger(\n",
    "    processingTime=\"10 seconds\"\n",
    ").start().awaitTermination()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0939cc28",
   "metadata": {},
   "source": [
    "**Write a streaming DataFrame to a database**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d8227c57",
   "metadata": {},
   "outputs": [],
   "source": [
    "import time\n",
    "\n",
    "pg_database = os.environ.get(\"PGDATABASE\") or \"postgres\"\n",
    "pg_host = os.environ.get(\"PGHOST\") or \"localhost\"\n",
    "pg_password = os.environ.get(\"PGPASSWORD\") or \"password\"\n",
    "pg_user = os.environ.get(\"PGUSER\") or \"postgres\"\n",
    "url = f\"jdbc:postgresql://{pg_host}:5432/{pg_database}\"\n",
    "table = \"streaming\"\n",
    "properties = {\n",
    "    \"driver\": \"org.postgresql.Driver\",\n",
    "    \"user\": pg_user,\n",
    "    \"password\": pg_password,\n",
    "}\n",
    "\n",
    "def foreach_batch_function(my_df, epoch_id):\n",
    "    my_df.write.jdbc(url=url, table=table, mode=\"Append\", properties=properties)\n",
    "\n",
    "df = (\n",
    "    spark.readStream.format(\"rate\")\n",
    "    .option(\"rowPerSecond\", 100)\n",
    "    .option(\"numPartitions\", 2)\n",
    "    .load()\n",
    ")\n",
    "query = df.writeStream.foreachBatch(foreach_batch_function).start()\n",
    "\n",
    "# Wait for some data to be processed and exit.\n",
    "for i in range(10):\n",
    "    time.sleep(5)\n",
    "    if len(query.recentProgress) > 0:\n",
    "        query.stop()\n",
    "        break\n",
    "\n",
    "df = spark.read.jdbc(url=url, table=table, properties=properties)\n",
    "result = \"{} rows written to database\".format(df.count())\n",
    "print(result)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0d6e89be",
   "metadata": {},
   "source": [
    "## Techniques for dealing with time series data."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a56271e7",
   "metadata": {},
   "source": [
    "**Zero fill missing values in a timeseries**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a769222e",
   "metadata": {},
   "outputs": [],
   "source": [
    "from pyspark.sql.functions import coalesce, lit\n",
    "\n",
    "# Use distinct values of customer and date from the dataset itself.\n",
    "# In general it's safer to use known reference tables for IDs and dates.\n",
    "df = spend_df.join(\n",
    "    spend_df.select(\"customer_id\")\n",
    "    .distinct()\n",
    "    .crossJoin(spend_df.select(\"date\").distinct()),\n",
    "    [\"date\", \"customer_id\"],\n",
    "    \"right\",\n",
    ").select(\"date\", \"customer_id\", coalesce(\"spend_dollars\", lit(0)))\n",
    "df.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fbc4fe8e",
   "metadata": {},
   "source": [
    "**First Time an ID is Seen**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e2d261f8",
   "metadata": {},
   "outputs": [],
   "source": [
    "from pyspark.sql.functions import first\n",
    "from pyspark.sql.window import Window\n",
    "\n",
    "w = Window().partitionBy(\"customer_id\").orderBy(\"date\")\n",
    "df = spend_df.withColumn(\"first_seen\", first(\"date\").over(w))\n",
    "df.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "55fe0508",
   "metadata": {},
   "source": [
    "**Cumulative Sum**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b17e9d1f",
   "metadata": {},
   "outputs": [],
   "source": [
    "from pyspark.sql.functions import sum\n",
    "from pyspark.sql.window import Window\n",
    "\n",
    "w = (\n",
    "    Window()\n",
    "    .partitionBy(\"customer_id\")\n",
    "    .orderBy(\"date\")\n",
    "    .rangeBetween(Window.unboundedPreceding, 0)\n",
    ")\n",
    "df = spend_df.withColumn(\"running_sum\", sum(\"spend_dollars\").over(w))\n",
    "df.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cc743ae0",
   "metadata": {},
   "source": [
    "**Cumulative Sum in a Period**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0177b033",
   "metadata": {},
   "outputs": [],
   "source": [
    "from pyspark.sql.functions import sum, year\n",
    "from pyspark.sql.window import Window\n",
    "\n",
    "# Add an additional partition clause for the sub-period.\n",
    "w = (\n",
    "    Window()\n",
    "    .partitionBy([\"customer_id\", year(\"date\")])\n",
    "    .orderBy(\"date\")\n",
    "    .rangeBetween(Window.unboundedPreceding, 0)\n",
    ")\n",
    "df = spend_df.withColumn(\"running_sum\", sum(\"spend_dollars\").over(w))\n",
    "df.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "88ccb73f",
   "metadata": {},
   "source": [
    "**Cumulative Average**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "967c0682",
   "metadata": {},
   "outputs": [],
   "source": [
    "from pyspark.sql.functions import avg\n",
    "from pyspark.sql.window import Window\n",
    "\n",
    "w = (\n",
    "    Window()\n",
    "    .partitionBy(\"customer_id\")\n",
    "    .orderBy(\"date\")\n",
    "    .rangeBetween(Window.unboundedPreceding, 0)\n",
    ")\n",
    "df = spend_df.withColumn(\"running_avg\", avg(\"spend_dollars\").over(w))\n",
    "df.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "251c0f81",
   "metadata": {},
   "source": [
    "**Cumulative Average in a Period**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "76d4ae01",
   "metadata": {},
   "outputs": [],
   "source": [
    "from pyspark.sql.functions import avg, year\n",
    "from pyspark.sql.window import Window\n",
    "\n",
    "# Add an additional partition clause for the sub-period.\n",
    "w = (\n",
    "    Window()\n",
    "    .partitionBy([\"customer_id\", year(\"date\")])\n",
    "    .orderBy(\"date\")\n",
    "    .rangeBetween(Window.unboundedPreceding, 0)\n",
    ")\n",
    "df = spend_df.withColumn(\"running_avg\", avg(\"spend_dollars\").over(w))\n",
    "df.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f0fe315a",
   "metadata": {},
   "source": [
    "## Machine Learning is a deep subject, too much to cover in this cheatsheet which is intended for code you can easily paste into your apps. The examples below will show basics of ML in Spark. It is helpful to understand the terminology of ML like Features, Estimators and Models. If you want some background on these things consider courses like \"Google crash course in ML\" or Udemy's \"Machine Learning Course with Python\".\n",
    "\n",
    "A brief introduction to Spark ML terms:\n",
    "* Feature: A Feature is an individual measurement. For example if you want to predict height based on age and sex, a combination of age and sex is a Feature.\n",
    "* Vector: A Vector is a special Spark data type similar to an array of numbers. Spark ML algorithms require Features to be loaded into Vectors for training and predictions.\n",
    "* Vector Column: Model training requires considering many Features at the same time. Spark ML operates on `DataFrame`s. Before training can happen you need to construct a `DataFrame` column of type Vector. See examples below.\n",
    "* Label: Supervised ML algorithms like regression and classification require a label when training. In Spark you will put labels in a column in a `DataFrame` such that each row has both a Feature and its associated Label.\n",
    "* Model: A Model is an algorithm capable of turning Feature vectors into values, usually thought of as predictions.\n",
    "* Estimator: An Estimator builds a mathematical model that transforms input values into outputs. Estimators do double duty in Spark, some Estimators like regression and classification build statistical models. Some Estimators are purely for data preparation like the StringIndexer which builds a Model containing a dictionary that maps strings to numbers in a deterministic way.\n",
    "* Fitting: Fitting is the process of building a Model using an Estimator and an input DataFrame you provide.\n",
    "* Transformer: Transformers create new DataFrames using the `transform` API, which applies algorithms to the input DataFrame and outputs a DataFrame with additional columns. The nature of the `transform` could be statistical or it could be a simple algorithm, depending on the type of Estimator that created the Model.\n",
    "* Pipelines: Pipelines are a series of Estimators that apply a series of `transform`s to a `DataFrame` before calling `fit` on the final Estimator in the Pipeline. The Pipeline is itself an Estimator. When you `fit` a Pipeline, Spark `fit`s the first Estimator in the Pipeline using an input `DataFrame` you provide. This produces a Model. If there are additional Estimators in the Pipeline, the newly created Model's `transform` method is called against the input `DataFrame` to create a new `DataFrame`. The process then begins again with the newly created `DataFrame` being passed to the next Estimator's `fit` method. Fitting a Pipeline produces a `PipelineModel`.\n",
    "\n",
    "This image helps visualize the relationship between Spark ML classes.\n",
    "![Hierarchy of Spark ML Classes](images/mlhierarchy.webp)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "dc3cb777",
   "metadata": {},
   "source": [
    "**Prepare data for training with a VectorAssembler**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f5babdb2",
   "metadata": {},
   "outputs": [],
   "source": [
    "from pyspark.ml.feature import VectorAssembler\n",
    "\n",
    "vectorAssembler = VectorAssembler(\n",
    "    inputCols=[\n",
    "        \"cylinders\",\n",
    "        \"displacement\",\n",
    "        \"acceleration\",\n",
    "    ],\n",
    "    outputCol=\"features\",\n",
    "    handleInvalid=\"skip\",\n",
    ")\n",
    "assembled = vectorAssembler.transform(auto_df_fixed)\n",
    "assembled = assembled.select(\n",
    "    [\"cylinders\", \"displacement\", \"acceleration\", \"features\"]\n",
    ")\n",
    "print(\"Data type for features column is:\", assembled.dtypes[-1][1])\n",
    "df.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "91147546",
   "metadata": {},
   "source": [
    "**A basic Random Forest Regression model**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "42e02249",
   "metadata": {},
   "outputs": [],
   "source": [
    "from pyspark.ml.feature import VectorAssembler\n",
    "from pyspark.ml.regression import RandomForestRegressor\n",
    "\n",
    "vectorAssembler = VectorAssembler(\n",
    "    inputCols=[\n",
    "        \"cylinders\",\n",
    "        \"displacement\",\n",
    "        \"horsepower\",\n",
    "    ],\n",
    "    outputCol=\"features\",\n",
    "    handleInvalid=\"skip\",\n",
    ")\n",
    "assembled = vectorAssembler.transform(auto_df_fixed)\n",
    "assembled = assembled.select([\"features\", \"mpg\", \"carname\"])\n",
    "\n",
    "# Define the estimator.\n",
    "rf = RandomForestRegressor(\n",
    "    numTrees=20,\n",
    "    featuresCol=\"features\",\n",
    "    labelCol=\"mpg\",\n",
    ")\n",
    "\n",
    "# Fit the model.\n",
    "rf_model = rf.fit(assembled)\n",
    "\n",
    "# Save the model.\n",
    "rf_model.write().overwrite().save(\"rf_regression_simple.model\")\n",
    "\n",
    "df.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f4654c8a",
   "metadata": {},
   "source": [
    "**Hyperparameter tuning**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e1c68903",
   "metadata": {},
   "outputs": [],
   "source": [
    "from pyspark.ml import Pipeline\n",
    "from pyspark.ml.evaluation import RegressionEvaluator\n",
    "from pyspark.ml.feature import StringIndexer, VectorAssembler\n",
    "from pyspark.ml.regression import RandomForestRegressor\n",
    "from pyspark.ml.tuning import CrossValidator, ParamGridBuilder\n",
    "from pyspark.sql.functions import udf\n",
    "from pyspark.sql.types import StringType\n",
    "\n",
    "# Set up our main ML pipeline.\n",
    "columns_to_assemble = [\n",
    "    \"cylinders\",\n",
    "    \"displacement\",\n",
    "    \"acceleration\",\n",
    "]\n",
    "vector_assembler = VectorAssembler(\n",
    "    inputCols=columns_to_assemble,\n",
    "    outputCol=\"features\",\n",
    "    handleInvalid=\"skip\",\n",
    ")\n",
    "\n",
    "# Define the model.\n",
    "rf = RandomForestRegressor(\n",
    "    featuresCol=\"features\",\n",
    "    labelCol=\"mpg\",\n",
    ")\n",
    "\n",
    "# Run the pipeline.\n",
    "pipeline = Pipeline(stages=[vector_assembler, rf])\n",
    "\n",
    "# Hyperparameter search.\n",
    "target_metric = \"rmse\"\n",
    "paramGrid = (\n",
    "    ParamGridBuilder().addGrid(rf.numTrees, list(range(20, 100, 10))).build()\n",
    ")\n",
    "cross_validator = CrossValidator(\n",
    "    estimator=pipeline,\n",
    "    estimatorParamMaps=paramGrid,\n",
    "    evaluator=RegressionEvaluator(\n",
    "        labelCol=\"mpg\", predictionCol=\"prediction\", metricName=target_metric\n",
    "    ),\n",
    "    numFolds=2,\n",
    "    parallelism=4,\n",
    ")\n",
    "\n",
    "# Run cross-validation to get the best parameters.\n",
    "fit_cross_validator = cross_validator.fit(auto_df_fixed)\n",
    "best_pipeline_model = fit_cross_validator.bestModel\n",
    "best_regressor = best_pipeline_model.stages[1]\n",
    "print(\"Best model has {} trees.\".format(best_regressor.getNumTrees))\n",
    "\n",
    "# Save the Cross Validator, to capture everything including stats.\n",
    "fit_cross_validator.write().overwrite().save(\"rf_regression_optimized.model\")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1e992f37",
   "metadata": {},
   "source": [
    "**Encode string variables as numbers**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0c87117b",
   "metadata": {},
   "outputs": [],
   "source": [
    "from pyspark.ml.feature import StringIndexer\n",
    "from pyspark.sql.functions import udf\n",
    "from pyspark.sql.types import StringType\n",
    "\n",
    "# Add manufacturer name we will use as a string column.\n",
    "first_word_udf = udf(lambda x: x.split()[0], StringType())\n",
    "df = auto_df_fixed.withColumn(\n",
    "    \"manufacturer\", first_word_udf(auto_df_fixed.carname)\n",
    ")\n",
    "\n",
    "# Encode the manufactor name into numbers.\n",
    "indexer = StringIndexer(\n",
    "    inputCol=\"manufacturer\", outputCol=\"manufacturer_encoded\"\n",
    ")\n",
    "encoded = (\n",
    "    indexer.fit(df)\n",
    "    .transform(df)\n",
    "    .select([\"manufacturer\", \"manufacturer_encoded\"])\n",
    ")\n",
    "df.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d7fae768",
   "metadata": {},
   "source": [
    "**One-hot encode a categorical variable**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4be2fec9",
   "metadata": {},
   "outputs": [],
   "source": [
    "from pyspark.ml.feature import OneHotEncoder\n",
    "\n",
    "# Turn the model year into categories.\n",
    "year_encoder = OneHotEncoder(\n",
    "    inputCol=\"modelyear\", outputCol=\"modelyear_encoded\"\n",
    ")\n",
    "encoded = (\n",
    "    year_encoder.fit(auto_df_fixed)\n",
    "    .transform(auto_df_fixed)\n",
    "    .select([\"modelyear\", \"modelyear_encoded\"])\n",
    ")\n",
    "df.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "eaae016e",
   "metadata": {},
   "source": [
    "**Optimize a model after a data preparation pipeline**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "52615f0f",
   "metadata": {},
   "outputs": [],
   "source": [
    "from pyspark.ml import Pipeline\n",
    "from pyspark.ml.evaluation import RegressionEvaluator\n",
    "from pyspark.ml.feature import OneHotEncoder, StringIndexer, VectorAssembler\n",
    "from pyspark.ml.regression import RandomForestRegressor\n",
    "from pyspark.ml.tuning import CrossValidator, ParamGridBuilder\n",
    "from pyspark.sql.functions import udf\n",
    "from pyspark.sql.types import StringType\n",
    "\n",
    "### Phase 1.\n",
    "# Add manufacturer name we will use as a string column.\n",
    "first_word_udf = udf(lambda x: x.split()[0], StringType())\n",
    "df = auto_df_fixed.withColumn(\n",
    "    \"manufacturer\", first_word_udf(auto_df_fixed.carname)\n",
    ")\n",
    "\n",
    "# Encode the manufactor name into numbers.\n",
    "manufacturer_encoder = StringIndexer(\n",
    "    inputCol=\"manufacturer\", outputCol=\"manufacturer_encoded\"\n",
    ")\n",
    "# Turn the model year into categories.\n",
    "year_encoder = OneHotEncoder(\n",
    "    inputCol=\"modelyear\", outputCol=\"modelyear_encoded\"\n",
    ")\n",
    "\n",
    "# Run data preparation as a pipeline.\n",
    "data_prep_pipeline = Pipeline(stages=[manufacturer_encoder, year_encoder])\n",
    "prepared = data_prep_pipeline.fit(df).transform(df)\n",
    "\n",
    "### Phase 2.\n",
    "# Assemble vectors.\n",
    "columns_to_assemble = [\n",
    "    \"cylinders\",\n",
    "    \"displacement\",\n",
    "    \"horsepower\",\n",
    "    \"weight\",\n",
    "    \"acceleration\",\n",
    "    \"manufacturer_encoded\",\n",
    "    \"modelyear_encoded\",\n",
    "]\n",
    "vector_assembler = VectorAssembler(\n",
    "    inputCols=columns_to_assemble,\n",
    "    outputCol=\"features\",\n",
    "    handleInvalid=\"skip\",\n",
    ")\n",
    "\n",
    "# Define the model.\n",
    "rf = RandomForestRegressor(\n",
    "    featuresCol=\"features\",\n",
    "    labelCol=\"mpg\",\n",
    ")\n",
    "\n",
    "# Define the Pipeline.\n",
    "pipeline = Pipeline(stages=[vector_assembler, rf])\n",
    "\n",
    "# Run cross-validation to get the best parameters.\n",
    "paramGrid = (\n",
    "    ParamGridBuilder().addGrid(rf.numTrees, list(range(20, 100, 10))).build()\n",
    ")\n",
    "cross_validator = CrossValidator(\n",
    "    estimator=pipeline,\n",
    "    estimatorParamMaps=paramGrid,\n",
    "    evaluator=RegressionEvaluator(\n",
    "        labelCol=\"mpg\", predictionCol=\"prediction\", metricName=\"rmse\"\n",
    "    ),\n",
    "    numFolds=2,\n",
    "    parallelism=4,\n",
    ")\n",
    "fit_cross_validator = cross_validator.fit(prepared)\n",
    "\n",
    "# Save the Cross Validator, to capture everything including stats.\n",
    "fit_cross_validator.write().overwrite().save(\"rf_regression_full.model\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f0554fb2",
   "metadata": {},
   "source": [
    "**Evaluate Model Performance**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "59c21629",
   "metadata": {},
   "outputs": [],
   "source": [
    "from pyspark.ml import Pipeline\n",
    "from pyspark.ml.evaluation import RegressionEvaluator\n",
    "from pyspark.ml.feature import OneHotEncoder, StringIndexer, VectorAssembler\n",
    "from pyspark.ml.regression import RandomForestRegressionModel\n",
    "from pyspark.ml.tuning import CrossValidatorModel\n",
    "from pyspark.sql.functions import udf\n",
    "from pyspark.sql.types import StringType\n",
    "\n",
    "# Metrics supported by RegressionEvaluator.\n",
    "metrics = \"rmse mse r2 mae\".split()\n",
    "\n",
    "# Load the simple model and compute its predictions.\n",
    "simple_assembler = VectorAssembler(\n",
    "    inputCols=[\n",
    "        \"cylinders\",\n",
    "        \"displacement\",\n",
    "        \"horsepower\",\n",
    "    ],\n",
    "    outputCol=\"features\",\n",
    "    handleInvalid=\"skip\",\n",
    ")\n",
    "simple_input = simple_assembler.transform(auto_df_fixed).select(\n",
    "    [\"features\", \"mpg\"]\n",
    ")\n",
    "rf_simple_model = RandomForestRegressionModel.load(\"rf_regression_simple.model\")\n",
    "simple_predictions = rf_simple_model.transform(simple_input)\n",
    "\n",
    "# Load the complex model and compute its predictions.\n",
    "first_word_udf = udf(lambda x: x.split()[0], StringType())\n",
    "df = auto_df_fixed.withColumn(\n",
    "    \"manufacturer\", first_word_udf(auto_df_fixed.carname)\n",
    ")\n",
    "manufacturer_encoder = StringIndexer(\n",
    "    inputCol=\"manufacturer\", outputCol=\"manufacturer_encoded\"\n",
    ")\n",
    "year_encoder = OneHotEncoder(\n",
    "    inputCol=\"modelyear\", outputCol=\"modelyear_encoded\"\n",
    ")\n",
    "data_prep_pipeline = Pipeline(stages=[manufacturer_encoder, year_encoder])\n",
    "prepared = data_prep_pipeline.fit(df).transform(df)\n",
    "columns_to_assemble = [\n",
    "    \"cylinders\",\n",
    "    \"displacement\",\n",
    "    \"horsepower\",\n",
    "    \"weight\",\n",
    "    \"acceleration\",\n",
    "    \"manufacturer_encoded\",\n",
    "    \"modelyear_encoded\",\n",
    "]\n",
    "complex_assembler = VectorAssembler(\n",
    "    inputCols=columns_to_assemble,\n",
    "    outputCol=\"features\",\n",
    "    handleInvalid=\"skip\",\n",
    ")\n",
    "complex_input = complex_assembler.transform(prepared).select(\n",
    "    [\"features\", \"mpg\"]\n",
    ")\n",
    "cv_model = CrossValidatorModel.load(\"rf_regression_full.model\")\n",
    "best_pipeline = cv_model.bestModel\n",
    "rf_complex_model = best_pipeline.stages[-1]\n",
    "complex_predictions = rf_complex_model.transform(complex_input)\n",
    "\n",
    "# Evaluate performances.\n",
    "evaluator = RegressionEvaluator(predictionCol=\"prediction\", labelCol=\"mpg\")\n",
    "performances = [\n",
    "    [\"simple\", simple_predictions, dict()],\n",
    "    [\"complex\", complex_predictions, dict()],\n",
    "]\n",
    "for label, predictions, tracker in performances:\n",
    "    for metric in metrics:\n",
    "        tracker[metric] = evaluator.evaluate(\n",
    "            predictions, {evaluator.metricName: metric}\n",
    "        )\n",
    "print(performances)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3d6fd073",
   "metadata": {},
   "source": [
    "**Get feature importances of a trained model**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f63162ed",
   "metadata": {},
   "outputs": [],
   "source": [
    "from pyspark.ml.tuning import CrossValidatorModel\n",
    "\n",
    "# Load the model we fit earlier.\n",
    "model = CrossValidatorModel.load(\"rf_regression_full.model\")\n",
    "\n",
    "# Get the best model.\n",
    "best_pipeline = model.bestModel\n",
    "\n",
    "# Get the names of assembled columns.\n",
    "assembler = best_pipeline.stages[0]\n",
    "original_columns = assembler.getInputCols()\n",
    "\n",
    "# Get feature importances.\n",
    "real_model = best_pipeline.stages[1]\n",
    "for feature, importance in zip(original_columns, real_model.featureImportances):\n",
    "    print(\"{} contributes {:0.3f}%\".format(feature, importance * 100))\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0b02b61b",
   "metadata": {},
   "source": [
    "**Plot Hyperparameter tuning metrics**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1bdca5d9",
   "metadata": {},
   "outputs": [],
   "source": [
    "from pyspark.ml import Pipeline\n",
    "from pyspark.ml.evaluation import RegressionEvaluator\n",
    "from pyspark.ml.feature import StringIndexer, VectorAssembler\n",
    "from pyspark.ml.regression import RandomForestRegressor\n",
    "from pyspark.ml.tuning import CrossValidator, ParamGridBuilder\n",
    "from pyspark.sql.functions import udf\n",
    "from pyspark.sql.types import StringType\n",
    "\n",
    "# Add manufacturer name we will use as a string column.\n",
    "first_word_udf = udf(lambda x: x.split()[0], StringType())\n",
    "df = auto_df_fixed.withColumn(\n",
    "    \"manufacturer\", first_word_udf(auto_df_fixed.carname)\n",
    ")\n",
    "manufacturer_encoded = StringIndexer(\n",
    "    inputCol=\"manufacturer\", outputCol=\"manufacturer_encoded\"\n",
    ")\n",
    "encoded_df = manufacturer_encoded.fit(df).transform(df)\n",
    "\n",
    "# Set up our main ML pipeline.\n",
    "columns_to_assemble = [\n",
    "    \"manufacturer_encoded\",\n",
    "    \"cylinders\",\n",
    "    \"displacement\",\n",
    "    \"horsepower\",\n",
    "    \"weight\",\n",
    "    \"acceleration\",\n",
    "]\n",
    "vector_assembler = VectorAssembler(\n",
    "    inputCols=columns_to_assemble,\n",
    "    outputCol=\"features\",\n",
    "    handleInvalid=\"skip\",\n",
    ")\n",
    "\n",
    "# Define the model.\n",
    "rf = RandomForestRegressor(\n",
    "    numTrees=20,\n",
    "    featuresCol=\"features\",\n",
    "    labelCol=\"mpg\",\n",
    ")\n",
    "\n",
    "# Run the pipeline.\n",
    "pipeline = Pipeline(stages=[vector_assembler, rf])\n",
    "\n",
    "# Hyperparameter search.\n",
    "target_metric = \"rmse\"\n",
    "paramGrid = (\n",
    "    ParamGridBuilder().addGrid(rf.numTrees, list(range(20, 100, 10))).build()\n",
    ")\n",
    "crossval = CrossValidator(\n",
    "    estimator=pipeline,\n",
    "    estimatorParamMaps=paramGrid,\n",
    "    evaluator=RegressionEvaluator(\n",
    "        labelCol=\"mpg\", predictionCol=\"prediction\", metricName=target_metric\n",
    "    ),\n",
    "    numFolds=2,\n",
    "    parallelism=4,\n",
    ")\n",
    "\n",
    "# Run cross-validation, get metrics for each parameter.\n",
    "model = crossval.fit(encoded_df)\n",
    "\n",
    "# Plot results using matplotlib.\n",
    "import pandas\n",
    "import matplotlib\n",
    "\n",
    "parameter_grid = [\n",
    "    {k.name: v for k, v in p.items()} for p in model.getEstimatorParamMaps()\n",
    "]\n",
    "pdf = pandas.DataFrame(\n",
    "    model.avgMetrics,\n",
    "    index=[x[\"numTrees\"] for x in parameter_grid],\n",
    "    columns=[target_metric],\n",
    ")\n",
    "ax = pdf.plot(style=\"*-\")\n",
    "ax.figure.suptitle(\"Hyperparameter Search: RMSE by Number of Trees\")\n",
    "ax.figure.savefig(\"hyperparameters.png\")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fa9d58bc",
   "metadata": {},
   "source": [
    "**Compute correlation matrix**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "76520b11",
   "metadata": {},
   "outputs": [],
   "source": [
    "from pyspark.ml.feature import VectorAssembler\n",
    "from pyspark.ml.stat import Correlation\n",
    "\n",
    "# Remove non-numeric columns.\n",
    "df = auto_df_fixed.drop(\"carname\")\n",
    "\n",
    "# Assemble all columns except mpg into a vector.\n",
    "feature_columns = list(df.columns)\n",
    "feature_columns.remove(\"mpg\")\n",
    "vector_col = \"features\"\n",
    "vector_assembler = VectorAssembler(\n",
    "    inputCols=feature_columns,\n",
    "    outputCol=vector_col,\n",
    "    handleInvalid=\"skip\",\n",
    ")\n",
    "df_vector = vector_assembler.transform(df).select(vector_col)\n",
    "\n",
    "# Compute the correlation matrix.\n",
    "matrix = Correlation.corr(df_vector, vector_col)\n",
    "corr_array = matrix.collect()[0][\"pearson({})\".format(vector_col)].toArray()\n",
    "\n",
    "# This part is just for pretty-printing.\n",
    "pdf = pandas.DataFrame(\n",
    "    corr_array, index=feature_columns, columns=feature_columns\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c1a3c97d",
   "metadata": {},
   "source": [
    "**Save a model**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "58dc5fe4",
   "metadata": {},
   "outputs": [],
   "source": [
    "from pyspark.ml.feature import VectorAssembler\n",
    "from pyspark.ml.regression import RandomForestRegressor\n",
    "from pyspark.ml.classification import RandomForestClassifier\n",
    "\n",
    "vectorAssembler = VectorAssembler(\n",
    "    inputCols=[\n",
    "        \"cylinders\",\n",
    "        \"displacement\",\n",
    "        \"horsepower\",\n",
    "    ],\n",
    "    outputCol=\"features\",\n",
    "    handleInvalid=\"skip\",\n",
    ")\n",
    "assembled = vectorAssembler.transform(auto_df_fixed)\n",
    "\n",
    "# Random test/train split.\n",
    "train_df, test_df = assembled.randomSplit([0.7, 0.3])\n",
    "\n",
    "# A regression model.\n",
    "rf_regressor = RandomForestRegressor(\n",
    "    numTrees=50,\n",
    "    featuresCol=\"features\",\n",
    "    labelCol=\"mpg\",\n",
    ")\n",
    "\n",
    "# A classification model.\n",
    "rf_classifier = RandomForestClassifier(\n",
    "    numTrees=50,\n",
    "    featuresCol=\"features\",\n",
    "    labelCol=\"origin\",\n",
    ")\n",
    "\n",
    "# Fit the models.\n",
    "rf_regressor_model = rf_regressor.fit(train_df)\n",
    "rf_regressor_model.write().overwrite().save(\"rf_regressor_saveload.model\")\n",
    "rf_classifier_model = rf_classifier.fit(train_df)\n",
    "rf_classifier_model.write().overwrite().save(\"rf_classifier_saveload.model\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e2f57589",
   "metadata": {},
   "source": [
    "**Load a model and use it for transformations**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d40017f9",
   "metadata": {},
   "outputs": [],
   "source": [
    "from pyspark.ml.feature import VectorAssembler\n",
    "from pyspark.ml.regression import RandomForestRegressionModel\n",
    "\n",
    "# Model type and assembled features need to agree with the trained model.\n",
    "rf_model = RandomForestRegressionModel.load(\"rf_regressor_saveload.model\")\n",
    "\n",
    "# The input DataFrame needs the same structure we used when we fit.\n",
    "vectorAssembler = VectorAssembler(\n",
    "    inputCols=[\n",
    "        \"cylinders\",\n",
    "        \"displacement\",\n",
    "        \"horsepower\",\n",
    "    ],\n",
    "    outputCol=\"features\",\n",
    "    handleInvalid=\"skip\",\n",
    ")\n",
    "assembled = vectorAssembler.transform(auto_df_fixed)\n",
    "predictions = rf_model.transform(assembled).select(\n",
    "    \"carname\", \"mpg\", \"prediction\"\n",
    ")\n",
    "df.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "27ce7edf",
   "metadata": {},
   "source": [
    "**Load a model and use it for predictions**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "75c781a0",
   "metadata": {},
   "outputs": [],
   "source": [
    "from pyspark.ml.linalg import Vectors\n",
    "from pyspark.ml.regression import RandomForestRegressionModel\n",
    "\n",
    "# Model type and assembled features need to agree with the trained model.\n",
    "rf_model = RandomForestRegressionModel.load(\"rf_regressor_saveload.model\")\n",
    "\n",
    "input_vector = Vectors.dense([8, 307.0, 130.0])\n",
    "prediction = rf_model.predict(input_vector)\n",
    "print(\"Prediction is\", prediction)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "06169a54",
   "metadata": {},
   "source": [
    "**Load a classification model and use it to compute confidences for output labels**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b9c6f9fe",
   "metadata": {},
   "outputs": [],
   "source": [
    "from pyspark.ml.linalg import Vectors\n",
    "from pyspark.ml.classification import RandomForestClassificationModel\n",
    "\n",
    "# Model type and assembled features need to agree with the trained model.\n",
    "rf_model = RandomForestClassificationModel.load(\"rf_classifier_saveload.model\")\n",
    "\n",
    "input_vector = Vectors.dense([8, 307.0, 130.0])\n",
    "prediction = rf_model.predictProbability(input_vector)\n",
    "print(\"Predictions are\", prediction)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "97cc96ae",
   "metadata": {},
   "source": [
    "## A few performance tips and tricks."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "48f0e462",
   "metadata": {},
   "source": [
    "**Get the Spark version**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "adb25113",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(spark.sparkContext.version)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "94dcae5a",
   "metadata": {},
   "source": [
    "**Log messages using Spark's Log4J**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "98d9da0f",
   "metadata": {},
   "outputs": [],
   "source": [
    "logger = spark.sparkContext._jvm.org.apache.log4j.Logger.getRootLogger()\n",
    "logger.warn(\"WARNING LEVEL LOG MESSAGE\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "91fe2ea1",
   "metadata": {},
   "source": [
    "**Cache a DataFrame**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f5003915",
   "metadata": {},
   "outputs": [],
   "source": [
    "from pyspark import StorageLevel\n",
    "from pyspark.sql.functions import lit\n",
    "\n",
    "# Make some copies of the DataFrame.\n",
    "df1 = auto_df.where(lit(1) > lit(0))\n",
    "df2 = auto_df.where(lit(2) > lit(0))\n",
    "df3 = auto_df.where(lit(3) > lit(0))\n",
    "\n",
    "print(\"Show the default storage level (NONE).\")\n",
    "print(auto_df.storageLevel)\n",
    "\n",
    "print(\"\\nChange storage level to Memory/Disk via the cache shortcut.\")\n",
    "df1.cache()\n",
    "print(df1.storageLevel)\n",
    "\n",
    "print(\n",
    "    \"\\nChange storage level to the equivalent of cache using an explicit StorageLevel.\"\n",
    ")\n",
    "df2.persist(storageLevel=StorageLevel(True, True, False, True, 1))\n",
    "print(df2.storageLevel)\n",
    "\n",
    "print(\"\\nSet storage level to NONE using an explicit StorageLevel.\")\n",
    "df3.persist(storageLevel=StorageLevel(False, False, False, False, 1))\n",
    "print(df3.storageLevel)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4416a4ef",
   "metadata": {},
   "source": [
    "**Show the execution plan, with costs**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ac50f8e8",
   "metadata": {},
   "outputs": [],
   "source": [
    "df = auto_df.groupBy(\"cylinders\").count()\n",
    "execution_plan = str(df.explain(mode=\"cost\"))\n",
    "print(execution_plan)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6265383f",
   "metadata": {},
   "source": [
    "**Partition by a Column Value**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "52745c3f",
   "metadata": {},
   "outputs": [],
   "source": [
    "# rows is an iterable, e.g. itertools.chain\n",
    "def number_in_partition(rows):\n",
    "    try:\n",
    "        first_row = next(rows)\n",
    "        partition_size = sum(1 for x in rows) + 1\n",
    "        partition_value = first_row.modelyear\n",
    "        print(f\"Partition {partition_value} has {partition_size} records\")\n",
    "    except StopIteration:\n",
    "        print(\"Empty partition\")\n",
    "\n",
    "df = auto_df.repartition(20, \"modelyear\")\n",
    "df.foreachPartition(number_in_partition)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "455bc098",
   "metadata": {},
   "source": [
    "**Range Partition a DataFrame**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "079578b5",
   "metadata": {},
   "outputs": [],
   "source": [
    "from pyspark.sql.functions import col\n",
    "\n",
    "# rows is an iterable, e.g. itertools.chain\n",
    "def count_in_partition(rows):\n",
    "    my_years = set()\n",
    "    number_in_partition = 0\n",
    "    for row in rows:\n",
    "        my_years.add(row.modelyear)\n",
    "        number_in_partition += 1\n",
    "    seen_years = sorted(list(my_years))\n",
    "    if len(seen_years) > 0:\n",
    "        seen_values = \",\".join(seen_years)\n",
    "        print(\n",
    "            f\"This partition has {number_in_partition} records with years {seen_values}\"\n",
    "        )\n",
    "    else:\n",
    "        print(\"Empty partition\")\n",
    "\n",
    "number_of_partitions = 5\n",
    "df = auto_df.repartitionByRange(number_of_partitions, col(\"modelyear\"))\n",
    "df.foreachPartition(count_in_partition)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fb7ea8de",
   "metadata": {},
   "source": [
    "**Change Number of DataFrame Partitions**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f648ebae",
   "metadata": {},
   "outputs": [],
   "source": [
    "from pyspark.sql.functions import col\n",
    "\n",
    "df = auto_df.repartition(col(\"modelyear\"))\n",
    "number_of_partitions = 5\n",
    "df = auto_df.repartition(number_of_partitions, col(\"mpg\"))\n",
    "df.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ed3b4f4f",
   "metadata": {},
   "source": [
    "**Coalesce DataFrame partitions**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c8c83c99",
   "metadata": {},
   "outputs": [],
   "source": [
    "import math\n",
    "\n",
    "target_partitions = math.ceil(auto_df.rdd.getNumPartitions() / 2)\n",
    "df = auto_df.coalesce(target_partitions)\n",
    "df.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4cc5ef5d",
   "metadata": {},
   "source": [
    "**Set the number of shuffle partitions**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9c211529",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Default shuffle partitions is usually 200.\n",
    "grouped1 = auto_df.groupBy(\"cylinders\").count()\n",
    "print(\"{} partition(s)\".format(grouped1.rdd.getNumPartitions()))\n",
    "\n",
    "# Set the shuffle partitions to 20.\n",
    "# This can reduce the number of files generated when saving DataFrames.\n",
    "spark.conf.set(\"spark.sql.shuffle.partitions\", 20)\n",
    "\n",
    "grouped2 = auto_df.groupBy(\"cylinders\").count()\n",
    "print(\"{} partition(s)\".format(grouped2.rdd.getNumPartitions()))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f341eba4",
   "metadata": {},
   "source": [
    "**Sample a subset of a DataFrame**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bc723821",
   "metadata": {},
   "outputs": [],
   "source": [
    "df = (\n",
    "    spark.read.format(\"csv\")\n",
    "    .option(\"header\", True)\n",
    "    .load(\"data/auto-mpg.csv\")\n",
    "    .sample(0.1)\n",
    ")\n",
    "df.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d2793529",
   "metadata": {},
   "source": [
    "**Run multiple concurrent jobs in different pools**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "73c811dc",
   "metadata": {},
   "outputs": [],
   "source": [
    "import concurrent.futures\n",
    "import datetime\n",
    "import time\n",
    "\n",
    "other_df = auto_df.toDF(*(\"_\" + c for c in auto_df.columns))\n",
    "target_frames = [\n",
    "    auto_df.crossJoin(other_df.limit((11 - i) * 20))\n",
    "    .groupBy(column_name)\n",
    "    .count()\n",
    "    for i, column_name in enumerate(auto_df.columns)\n",
    "]\n",
    "\n",
    "def launch(i, target):\n",
    "    spark.sparkContext.setLocalProperty(\"spark.scheduler.pool\", f\"pool{i}\")\n",
    "    print(\"Job\", i, \"returns\", target.first(), \"at\", datetime.datetime.now())\n",
    "\n",
    "with concurrent.futures.ThreadPoolExecutor() as executor:\n",
    "    futures = []\n",
    "    for i, target in enumerate(target_frames):\n",
    "        print(\"Starting job\", i, \"at\", datetime.datetime.now())\n",
    "        futures.append(executor.submit(launch, i, target))\n",
    "        time.sleep(0.2)\n",
    "    concurrent.futures.wait(futures)\n",
    "    for future in futures:\n",
    "        print(future.result())"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9b9396e6",
   "metadata": {},
   "source": [
    "**Print Spark configuration properties**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "72106d16",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(spark.sparkContext.getConf().getAll())"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a7ee3ba5",
   "metadata": {},
   "source": [
    "**Set Spark configuration properties**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f2afd8ec",
   "metadata": {},
   "outputs": [],
   "source": [
    "key = \"spark.hadoop.mapreduce.fileoutputcommitter.algorithm.version\"\n",
    "value = 2\n",
    "\n",
    "# Wrong! Settings cannot be changed this way.\n",
    "# spark.sparkContext.getConf().set(key, value)\n",
    "\n",
    "# Correct.\n",
    "spark.conf.set(key, value)\n",
    "\n",
    "# Alternatively: Set at build time.\n",
    "# Some settings can only be made at build time.\n",
    "spark_builder = SparkSession.builder.appName(\"My App\")\n",
    "spark_builder.config(key, value)\n",
    "spark = spark_builder.getOrCreate()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9340cc6b",
   "metadata": {},
   "source": [
    "**Publish Metrics to Graphite**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1539dcdb",
   "metadata": {},
   "outputs": [],
   "source": [
    "pass"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1da49ec2",
   "metadata": {},
   "source": [
    "**Increase Spark driver/executor heap space**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "abbd57ed",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Memory configuration depends entirely on your runtime.\n",
    "# In OCI Data Flow you control memory by selecting a larger or smaller VM.\n",
    "# No other configuration is needed.\n",
    "#\n",
    "# For other environments see the Spark \"Cluster Mode Overview\" to get started.\n",
    "# https://spark.apache.org/docs/latest/cluster-overview.html\n",
    "# And good luck!"
   ]
  }
 ],
 "metadata": {},
 "nbformat": 4,
 "nbformat_minor": 5
}