cheatsheet.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "3b2b8b7e",
   "metadata": {},
   "source": [
    "## Run This Code First! This creates the Spark session and loads data."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "014393a4",
   "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"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6abb889a",
   "metadata": {},
   "source": [
    "## Loading data stored in filesystems or databases, and saving it."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a4e104c1",
   "metadata": {},
   "source": [
    "**Load a DataFrame from CSV**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2c0f8f50",
   "metadata": {},
   "outputs": [],
   "source": [
    "# See https://spark.apache.org/docs/latest/api/java/org/apache/spark/sql/DataFrameReader.html\n",
    "# for a list of supported options.\n",
    "df = spark.read.format(\"csv\").option(\"header\", True).load(\"data/auto-mpg.csv\")\n",
    "df.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "54b3dac2",
   "metadata": {},
   "source": [
    "**Load a DataFrame from a Tab Separated Value (TSV) file**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "53135def",
   "metadata": {},
   "outputs": [],
   "source": [
    "# See https://spark.apache.org/docs/latest/api/java/org/apache/spark/sql/DataFrameReader.html\n",
    "# for a list of supported options.\n",
    "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": "02de4f96",
   "metadata": {},
   "source": [
    "**Save a DataFrame in CSV format**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e4083d26",
   "metadata": {},
   "outputs": [],
   "source": [
    "# See https://spark.apache.org/docs/latest/api/java/org/apache/spark/sql/DataFrameWriter.html\n",
    "# for a list of supported options.\n",
    "auto_df.write.csv(\"output.csv\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ddb2ab0e",
   "metadata": {},
   "source": [
    "**Load a DataFrame from Parquet**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c37defb0",
   "metadata": {},
   "outputs": [],
   "source": [
    "df = spark.read.format(\"parquet\").load(\"data/auto-mpg.parquet\")\n",
    "df.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "de220610",
   "metadata": {},
   "source": [
    "**Save a DataFrame in Parquet format**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "276995fd",
   "metadata": {},
   "outputs": [],
   "source": [
    "auto_df.write.parquet(\"output.parquet\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7a5b576a",
   "metadata": {},
   "source": [
    "**Load a DataFrame from JSON Lines (jsonl) Formatted Data**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "27be2158",
   "metadata": {},
   "outputs": [],
   "source": [
    "# JSON Lines / jsonl format uses one JSON document per line.\n",
    "# If you have data with mostly regular structure this is better than nesting it in an array.\n",
    "# See https://jsonlines.org/\n",
    "df = spark.read.json(\"data/weblog.jsonl\")\n",
    "df.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0e11408d",
   "metadata": {},
   "source": [
    "**Save a DataFrame into a Hive catalog table**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "70defa67",
   "metadata": {},
   "outputs": [],
   "source": [
    "auto_df.write.mode(\"overwrite\").saveAsTable(\"autompg\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c4e36056",
   "metadata": {},
   "source": [
    "**Load a Hive catalog table into a DataFrame**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cecdaa1e",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Load the table previously saved.\n",
    "df = spark.table(\"autompg\")\n",
    "df.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b0493f2b",
   "metadata": {},
   "source": [
    "## Special data handling scenarios."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2c8d8e85",
   "metadata": {},
   "source": [
    "**Provide the schema when loading a DataFrame from CSV**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "393e6378",
   "metadata": {},
   "outputs": [],
   "source": [
    "# See https://spark.apache.org/docs/latest/api/python/_modules/pyspark/sql/types.html\n",
    "# for a list of types.\n",
    "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": "158e6e5a",
   "metadata": {},
   "source": [
    "**Save a DataFrame to CSV, overwriting existing data**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6778774a",
   "metadata": {},
   "outputs": [],
   "source": [
    "auto_df.write.mode(\"overwrite\").csv(\"output.csv\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c5eb0144",
   "metadata": {},
   "source": [
    "**Save a DataFrame to CSV with a header**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c70dc09b",
   "metadata": {},
   "outputs": [],
   "source": [
    "# See https://spark.apache.org/docs/latest/api/java/org/apache/spark/sql/DataFrameWriter.html\n",
    "# for a list of supported options.\n",
    "auto_df.coalesce(1).write.csv(\"header.csv\", header=\"true\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "dccc2b9d",
   "metadata": {},
   "source": [
    "**Save a DataFrame in a single CSV file**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e7cf5cbd",
   "metadata": {},
   "outputs": [],
   "source": [
    "auto_df.coalesce(1).write.csv(\"single.csv\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a2610d21",
   "metadata": {},
   "source": [
    "**Save DataFrame as a dynamic partitioned table**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4b50422a",
   "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": "85306834",
   "metadata": {},
   "source": [
    "**Load a CSV file with a money column into a DataFrame**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c0fabeee",
   "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": "e24b3fba",
   "metadata": {},
   "source": [
    "## Adding, removing and modifying DataFrame columns."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "40b00477",
   "metadata": {},
   "source": [
    "**Add a new column to a DataFrame**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a8137da9",
   "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": "44be86a6",
   "metadata": {},
   "source": [
    "**Modify a DataFrame column**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e9fe5d4b",
   "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": "66773693",
   "metadata": {},
   "source": [
    "**Add a column with multiple conditions**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c13e2e2f",
   "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": "1133c06e",
   "metadata": {},
   "source": [
    "**Add a constant column**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3c06eb14",
   "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": "0f0e9494",
   "metadata": {},
   "source": [
    "**Concatenate columns**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ea3bb979",
   "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": "59a5d7b2",
   "metadata": {},
   "source": [
    "**Drop a column**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ba26fee6",
   "metadata": {},
   "outputs": [],
   "source": [
    "df = auto_df.drop(\"horsepower\")\n",
    "df.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "76f00f40",
   "metadata": {},
   "source": [
    "**Change a column name**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "32e71152",
   "metadata": {},
   "outputs": [],
   "source": [
    "df = auto_df.withColumnRenamed(\"horsepower\", \"horses\")\n",
    "df.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0eae7276",
   "metadata": {},
   "source": [
    "**Change multiple column names**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2cf90d74",
   "metadata": {},
   "outputs": [],
   "source": [
    "df = auto_df.withColumnRenamed(\"horsepower\", \"horses\").withColumnRenamed(\n",
    "    \"modelyear\", \"year\"\n",
    ")\n",
    "df.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "78b09a37",
   "metadata": {},
   "source": [
    "**Convert a DataFrame column to a Python list**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9813cbcc",
   "metadata": {},
   "outputs": [],
   "source": [
    "names = auto_df.select(\"carname\").rdd.flatMap(lambda x: x).collect()\n",
    "print(str(names[:10]))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ed223052",
   "metadata": {},
   "source": [
    "**Convert a scalar query to a Python value**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e609413d",
   "metadata": {},
   "outputs": [],
   "source": [
    "average = auto_df.agg(dict(mpg=\"avg\")).first()[0]\n",
    "print(str(average))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "310a4b97",
   "metadata": {},
   "source": [
    "**Consume a DataFrame row-wise as Python dictionaries**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b9188f11",
   "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": "6d1b6cea",
   "metadata": {},
   "source": [
    "**Select particular columns from a DataFrame**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "67f3d541",
   "metadata": {},
   "outputs": [],
   "source": [
    "df = auto_df.select([\"mpg\", \"cylinders\", \"displacement\"])\n",
    "df.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "43bb00d0",
   "metadata": {},
   "source": [
    "**Create an empty dataframe with a specified schema**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "692bcddb",
   "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": "031d0365",
   "metadata": {},
   "source": [
    "**Create a constant dataframe**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8b019f33",
   "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": "18675a30",
   "metadata": {},
   "source": [
    "**Convert String to Double**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3a7b1247",
   "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": "edb5ed5b",
   "metadata": {},
   "source": [
    "**Convert String to Integer**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6b770d7e",
   "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": "59d6ddb7",
   "metadata": {},
   "source": [
    "**Get the size of a DataFrame**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "59ca0ef9",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(\"{} rows\".format(auto_df.count()))\n",
    "print(\"{} columns\".format(len(auto_df.columns)))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "eb501b19",
   "metadata": {},
   "source": [
    "**Get a DataFrame's number of partitions**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "431a8e78",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(\"{} partition(s)\".format(auto_df.rdd.getNumPartitions()))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "33103290",
   "metadata": {},
   "source": [
    "**Get data types of a DataFrame's columns**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4386635f",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(auto_df.dtypes)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "60a131ce",
   "metadata": {},
   "source": [
    "**Convert an RDD to Data Frame**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6f863566",
   "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": "6df78a9a",
   "metadata": {},
   "source": [
    "**Print the contents of an RDD**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e0791fc2",
   "metadata": {},
   "outputs": [],
   "source": [
    "rdd = auto_df.rdd\n",
    "print(rdd.take(10))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0b8f4377",
   "metadata": {},
   "source": [
    "**Print the contents of a DataFrame**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bbfcee0c",
   "metadata": {},
   "outputs": [],
   "source": [
    "auto_df.show(10)\n",
    "df.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a216a2d0",
   "metadata": {},
   "source": [
    "**Process each row of a DataFrame**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "63b0f670",
   "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": "fe7d7b9e",
   "metadata": {},
   "source": [
    "**DataFrame Map example**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "af67ba67",
   "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": "52b69ca2",
   "metadata": {},
   "source": [
    "**DataFrame Flatmap example**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1b40ef5a",
   "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": "e4317ab1",
   "metadata": {},
   "source": [
    "**Create a custom UDF**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "424d5ab4",
   "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": "9a1998cb",
   "metadata": {},
   "source": [
    "## Data conversions and other modifications."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cfba8916",
   "metadata": {},
   "source": [
    "**Extract data from a string using a regular expression**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "332637d9",
   "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": "02770184",
   "metadata": {},
   "source": [
    "**Fill NULL values in specific columns**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f8e83598",
   "metadata": {},
   "outputs": [],
   "source": [
    "df = auto_df.fillna({\"horsepower\": 0})\n",
    "df.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "14cd1769",
   "metadata": {},
   "source": [
    "**Fill NULL values with column average**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c2b18a7a",
   "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": "b28e94eb",
   "metadata": {},
   "source": [
    "**Fill NULL values with group average**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6524fdf8",
   "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": "ab6ab298",
   "metadata": {},
   "source": [
    "**Unpack a DataFrame's JSON column to a new DataFrame**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d13bd96e",
   "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": "302c1542",
   "metadata": {},
   "source": [
    "**Query a JSON column**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f60d29d5",
   "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": "98bfcfa8",
   "metadata": {},
   "source": [
    "## Filtering, sorting, removing duplicates and more."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "53d22730",
   "metadata": {},
   "source": [
    "**Filter a column using a condition**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d2070907",
   "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": "178ae2c5",
   "metadata": {},
   "source": [
    "**Filter based on a specific column value**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a24cf66e",
   "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": "1ce4f8e3",
   "metadata": {},
   "source": [
    "**Filter based on an IN list**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b4a0499f",
   "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": "544ba006",
   "metadata": {},
   "source": [
    "**Filter based on a NOT IN list**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "64b2955a",
   "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": "8bfc3e50",
   "metadata": {},
   "source": [
    "**Filter values based on keys in another DataFrame**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6bc693ec",
   "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": "403e8304",
   "metadata": {},
   "source": [
    "**Get Dataframe rows that match a substring**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9ee2364a",
   "metadata": {},
   "outputs": [],
   "source": [
    "df = auto_df.where(auto_df.carname.contains(\"custom\"))\n",
    "df.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5f7f337a",
   "metadata": {},
   "source": [
    "**Filter a Dataframe based on a custom substring search**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4ba666f3",
   "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": "91ba7ba8",
   "metadata": {},
   "source": [
    "**Filter based on a column's length**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "924776ae",
   "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": "b2ea46c8",
   "metadata": {},
   "source": [
    "**Multiple filter conditions**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3e24ffab",
   "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": "761adb14",
   "metadata": {},
   "source": [
    "**Sort DataFrame by a column**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "085d0751",
   "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": "9557b1df",
   "metadata": {},
   "source": [
    "**Take the first N rows of a DataFrame**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2202d60d",
   "metadata": {},
   "outputs": [],
   "source": [
    "n = 10\n",
    "df = auto_df.limit(n)\n",
    "df.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7acd77d6",
   "metadata": {},
   "source": [
    "**Get distinct values of a column**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b2f23119",
   "metadata": {},
   "outputs": [],
   "source": [
    "df = auto_df.select(\"cylinders\").distinct()\n",
    "df.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "baf44533",
   "metadata": {},
   "source": [
    "**Remove duplicates**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c2aa6bfd",
   "metadata": {},
   "outputs": [],
   "source": [
    "df = auto_df.dropDuplicates([\"carname\"])\n",
    "df.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6c72aad0",
   "metadata": {},
   "source": [
    "## Group DataFrame data by key to perform aggregates like counting, sums, averages, etc."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c83dfe03",
   "metadata": {},
   "source": [
    "**count(*) on a particular column**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2de59edb",
   "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": "7dadb5b4",
   "metadata": {},
   "source": [
    "**Group and sort**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "19394d57",
   "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": "b009b60a",
   "metadata": {},
   "source": [
    "**Filter groups based on an aggregate value, equivalent to SQL HAVING clause**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3c7294c4",
   "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": "847d4cfc",
   "metadata": {},
   "source": [
    "**Group by multiple columns**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "670462a4",
   "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": "e3b3e4d6",
   "metadata": {},
   "source": [
    "**Aggregate multiple columns**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6bcb8c51",
   "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": "9db7dddd",
   "metadata": {},
   "source": [
    "**Aggregate multiple columns with custom orderings**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9170b66c",
   "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": "181f4988",
   "metadata": {},
   "source": [
    "**Get the maximum of a column**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8f8b2d8d",
   "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": "6ec4ac81",
   "metadata": {},
   "source": [
    "**Sum a list of columns**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "93693dab",
   "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": "e26b0750",
   "metadata": {},
   "source": [
    "**Sum a column**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "002863d5",
   "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": "f20f70f8",
   "metadata": {},
   "source": [
    "**Aggregate all numeric columns**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0051c33e",
   "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": "b07889d7",
   "metadata": {},
   "source": [
    "**Count unique after grouping**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d70db6af",
   "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": "2c1c29d5",
   "metadata": {},
   "source": [
    "**Count distinct values on all columns**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2049ecc3",
   "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": "800741f4",
   "metadata": {},
   "source": [
    "**Group by then filter on the count**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "259a2253",
   "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": "f8c5751d",
   "metadata": {},
   "source": [
    "**Find the top N per row group (use N=1 for maximum)**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1f90c14b",
   "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": "dfe71d3b",
   "metadata": {},
   "source": [
    "**Group key/values into a list**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f6e086ae",
   "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": "1b7fe788",
   "metadata": {},
   "source": [
    "**Compute a histogram**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "41ed5dd0",
   "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": "ae15f665",
   "metadata": {},
   "source": [
    "**Compute global percentiles**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "05a2f0ac",
   "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": "e95bf384",
   "metadata": {},
   "source": [
    "**Compute percentiles within a partition**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7a3ab5da",
   "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": "d8a027af",
   "metadata": {},
   "source": [
    "**Compute percentiles after aggregating**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e621e033",
   "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": "e47759b0",
   "metadata": {},
   "source": [
    "**Filter rows with values below a target percentile**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "97bed4a8",
   "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": "ff2ac541",
   "metadata": {},
   "source": [
    "**Aggregate and rollup**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bab91b8d",
   "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": "71686e09",
   "metadata": {},
   "source": [
    "**Aggregate and cube**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "422e183e",
   "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": "5e05d1cd",
   "metadata": {},
   "source": [
    "## Joining and stacking DataFrames."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "251f1c6f",
   "metadata": {},
   "source": [
    "**Join two DataFrames by column name**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "04b76a76",
   "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": "629fb4c9",
   "metadata": {},
   "source": [
    "**Join two DataFrames with an expression**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4132a9a0",
   "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": "c7fd576e",
   "metadata": {},
   "source": [
    "**Multiple join conditions**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8114eec2",
   "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": "8725e19e",
   "metadata": {},
   "source": [
    "**Various Spark join types**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "36754bb0",
   "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": "456cc09c",
   "metadata": {},
   "source": [
    "**Concatenate two DataFrames**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2db2130c",
   "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": "8e9fa06d",
   "metadata": {},
   "source": [
    "**Load multiple files into a single DataFrame**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a49b9866",
   "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": "d99dd58a",
   "metadata": {},
   "source": [
    "**Subtract DataFrames**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "728b6ff2",
   "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": "59edabf7",
   "metadata": {},
   "source": [
    "## Loading File Metadata and Processing Files."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a88f8c35",
   "metadata": {},
   "source": [
    "**Load Local File Details into a DataFrame**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "af6ee106",
   "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": "305fa4bf",
   "metadata": {},
   "source": [
    "**Load Files from Oracle Cloud Infrastructure into a DataFrame**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d0e71009",
   "metadata": {},
   "outputs": [],
   "source": [
    "from pyspark.sql.types import (\n",
    "    StructField,\n",
    "    StructType,\n",
    "    LongType,\n",
    "    StringType,\n",
    "    TimestampType,\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": "f94ed33b",
   "metadata": {},
   "source": [
    "**Transform Many Images using Pillow**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "78717918",
   "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": "f6116375",
   "metadata": {},
   "source": [
    "## Dealing with NULLs and NaNs in DataFrames."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "facd686d",
   "metadata": {},
   "source": [
    "**Filter rows with None or Null values**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "91a9c238",
   "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": "a567d7dd",
   "metadata": {},
   "source": [
    "**Drop rows with Null values**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2d2654a1",
   "metadata": {},
   "outputs": [],
   "source": [
    "# thresh controls the number of nulls before the row gets dropped.\n",
    "# subset controls the columns to consider.\n",
    "df = auto_df.na.drop(thresh=2, subset=(\"horsepower\",))\n",
    "df.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0cbeacec",
   "metadata": {},
   "source": [
    "**Count all Null or NaN values in a DataFrame**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e6497f17",
   "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": "b3897e38",
   "metadata": {},
   "source": [
    "## Parsing and processing dates and times."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4257afa0",
   "metadata": {},
   "source": [
    "**Convert an ISO 8601 formatted date string to date type**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d94b33a2",
   "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": "8414fc1d",
   "metadata": {},
   "source": [
    "**Convert a custom formatted date string to date type**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "25861597",
   "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": "1176dead",
   "metadata": {},
   "source": [
    "**Get the last day of the current month**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "dac65285",
   "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": "54eea546",
   "metadata": {},
   "source": [
    "**Convert UNIX (seconds since epoch) timestamp to date**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2f617d8c",
   "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": "92883144",
   "metadata": {},
   "source": [
    "**Load a CSV file with complex dates into a DataFrame**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "302f3f2f",
   "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": "5852bdaa",
   "metadata": {},
   "source": [
    "## Analyzing unstructured data like JSON, XML, etc."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "29b3c16a",
   "metadata": {},
   "source": [
    "**Flatten top level text fields in a JSONl document**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b2673e6d",
   "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": "498b58a9",
   "metadata": {},
   "source": [
    "**Flatten top level text fields from a JSON column**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4c487897",
   "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": "a07231c3",
   "metadata": {},
   "source": [
    "**Unnest an array of complex structures**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0431f4de",
   "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": "dc1640b2",
   "metadata": {},
   "source": [
    "## Using Python's Pandas library to augment Spark. Some operations require the pyarrow library."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a4291ed3",
   "metadata": {},
   "source": [
    "**Convert Spark DataFrame to Pandas DataFrame**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e4dac82f",
   "metadata": {},
   "outputs": [],
   "source": [
    "pandas_df = auto_df.toPandas()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "593b6b23",
   "metadata": {},
   "source": [
    "**Convert Pandas DataFrame to Spark DataFrame**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "354180dd",
   "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": "63b67169",
   "metadata": {},
   "source": [
    "**Convert N rows from a DataFrame to a Pandas DataFrame**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b87902f5",
   "metadata": {},
   "outputs": [],
   "source": [
    "N = 10\n",
    "pdf = auto_df.limit(N).toPandas()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "02e8da92",
   "metadata": {},
   "source": [
    "**Grouped Aggregation with Pandas**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3a984d14",
   "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": "359f878b",
   "metadata": {},
   "source": [
    "**Use a Pandas Grouped Map Function via applyInPandas**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fdea6d0b",
   "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": "48fa879a",
   "metadata": {},
   "source": [
    "## Extracting key statistics out of a body of data."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3bec4947",
   "metadata": {},
   "source": [
    "**Compute the number of NULLs across all columns**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3af0af70",
   "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": "10dd3537",
   "metadata": {},
   "source": [
    "**Compute average values of all numeric columns**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5b7541d2",
   "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": "719fec0f",
   "metadata": {},
   "source": [
    "**Compute minimum values of all numeric columns**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4c5c2ec7",
   "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": "45f38e49",
   "metadata": {},
   "source": [
    "**Compute maximum values of all numeric columns**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "aeb7f2f2",
   "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": "586cfc10",
   "metadata": {},
   "source": [
    "**Compute median values of all numeric columns**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d9b9762f",
   "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": "59cb738f",
   "metadata": {},
   "source": [
    "**Identify Outliers in a DataFrame**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c996c09f",
   "metadata": {},
   "outputs": [],
   "source": [
    "# This approach uses the Median Absolute Deviation.\n",
    "# Outliers are based on variances in a single numeric column.\n",
    "# Tune outlier sensitivity using z_score_threshold.\n",
    "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": "cabafe9c",
   "metadata": {},
   "source": [
    "## Upserts, updates and deletes on data."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a021a361",
   "metadata": {},
   "source": [
    "**Save to a Delta Table**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "08fda79c",
   "metadata": {},
   "outputs": [],
   "source": [
    "auto_df.write.mode(\"overwrite\").format(\"delta\").saveAsTable(\"delta_table\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "505c16b0",
   "metadata": {},
   "source": [
    "**Update records in a DataFrame using Delta Tables**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6e52288d",
   "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": "b765c7cd",
   "metadata": {},
   "source": [
    "**Merge into a Delta table**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "87c8941c",
   "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": "9bcfac07",
   "metadata": {},
   "source": [
    "**Show Table Version History**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "68434b6b",
   "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": "d619bab4",
   "metadata": {},
   "source": [
    "**Load a Delta Table by Version ID (Time Travel Query)**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "864214c0",
   "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": "5f15c839",
   "metadata": {},
   "source": [
    "**Load a Delta Table by Timestamp (Time Travel Query)**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "75d25390",
   "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": "61fa187f",
   "metadata": {},
   "source": [
    "**Compact a Delta Table**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e35e0c32",
   "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": "7cc7ec8d",
   "metadata": {},
   "source": [
    "## Spark Streaming (Focuses on Structured Streaming)."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0c480036",
   "metadata": {},
   "source": [
    "**Add the current timestamp to a DataFrame**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "587e5b2c",
   "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": "624b6290",
   "metadata": {},
   "source": [
    "## Techniques for dealing with time series data."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6cef4dfd",
   "metadata": {},
   "source": [
    "**Zero fill missing values in a timeseries**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e49745bf",
   "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": "b760f771",
   "metadata": {},
   "source": [
    "**First Time an ID is Seen**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e61ed068",
   "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": "21fed075",
   "metadata": {},
   "source": [
    "**Cumulative Sum**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2a31a57d",
   "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": "ab6fea07",
   "metadata": {},
   "source": [
    "**Cumulative Sum in a Period**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "28a2df13",
   "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": "549a9c36",
   "metadata": {},
   "source": [
    "**Cumulative Average**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "826a5f11",
   "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": "ce610f43",
   "metadata": {},
   "source": [
    "**Cumulative Average in a Period**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ace18fcb",
   "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": "a7f71fb9",
   "metadata": {},
   "source": [
    "## Machine Learning"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2848d458",
   "metadata": {},
   "source": [
    "**Save a model**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0aff1097",
   "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",
    "        \"weight\",\n",
    "        \"acceleration\",\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",
    "# Define the model.\n",
    "rf = RandomForestRegressor(\n",
    "    numTrees=50,\n",
    "    featuresCol=\"features\",\n",
    "    labelCol=\"mpg\",\n",
    ")\n",
    "\n",
    "# Train the model.\n",
    "rf_model = rf.fit(train_df)\n",
    "rf_model.write().overwrite().save(\"rf_regression.model\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e158983d",
   "metadata": {},
   "source": [
    "**Load a model and use it for predictions**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "812a8183",
   "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_regression.model\")\n",
    "vectorAssembler = VectorAssembler(\n",
    "    inputCols=[\n",
    "        \"cylinders\",\n",
    "        \"displacement\",\n",
    "        \"horsepower\",\n",
    "        \"weight\",\n",
    "        \"acceleration\",\n",
    "    ],\n",
    "    outputCol=\"features\",\n",
    "    handleInvalid=\"skip\",\n",
    ")\n",
    "assembled = vectorAssembler.transform(auto_df_fixed)\n",
    "\n",
    "predictions = rf_model.transform(assembled).select(\n",
    "    \"carname\", \"mpg\", \"prediction\"\n",
    ")\n",
    "df.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "763cbb93",
   "metadata": {},
   "source": [
    "**A basic Linear Regression model**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1b4267cb",
   "metadata": {},
   "outputs": [],
   "source": [
    "from pyspark.ml.feature import VectorAssembler\n",
    "from pyspark.ml.regression import LinearRegression\n",
    "\n",
    "vectorAssembler = VectorAssembler(\n",
    "    inputCols=[\n",
    "        \"cylinders\",\n",
    "        \"displacement\",\n",
    "        \"horsepower\",\n",
    "        \"weight\",\n",
    "        \"acceleration\",\n",
    "    ],\n",
    "    outputCol=\"features\",\n",
    "    handleInvalid=\"skip\",\n",
    ")\n",
    "assembled = vectorAssembler.transform(auto_df_fixed)\n",
    "assembled = assembled.select([\"features\", \"mpg\", \"carname\"])\n",
    "\n",
    "# Random test/train split.\n",
    "train_df, test_df = assembled.randomSplit([0.7, 0.3])\n",
    "\n",
    "# Define the model.\n",
    "lr = LinearRegression(\n",
    "    featuresCol=\"features\",\n",
    "    labelCol=\"mpg\",\n",
    "    maxIter=10,\n",
    "    regParam=0.3,\n",
    "    elasticNetParam=0.8,\n",
    ")\n",
    "\n",
    "# Train the model.\n",
    "lr_model = lr.fit(train_df)\n",
    "\n",
    "# Stats for training.\n",
    "print(\n",
    "    \"RMSE={} r2={}\".format(\n",
    "        lr_model.summary.rootMeanSquaredError, lr_model.summary.r2\n",
    "    )\n",
    ")\n",
    "\n",
    "# Make predictions.\n",
    "df = lr_model.transform(test_df)\n",
    "df.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0915edd9",
   "metadata": {},
   "source": [
    "**A basic Random Forest Regression model**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "923e488b",
   "metadata": {},
   "outputs": [],
   "source": [
    "from pyspark.ml.feature import VectorAssembler\n",
    "from pyspark.ml.regression import RandomForestRegressor\n",
    "from pyspark.ml.evaluation import RegressionEvaluator\n",
    "\n",
    "vectorAssembler = VectorAssembler(\n",
    "    inputCols=[\n",
    "        \"cylinders\",\n",
    "        \"displacement\",\n",
    "        \"horsepower\",\n",
    "        \"weight\",\n",
    "        \"acceleration\",\n",
    "    ],\n",
    "    outputCol=\"features\",\n",
    "    handleInvalid=\"skip\",\n",
    ")\n",
    "assembled = vectorAssembler.transform(auto_df_fixed)\n",
    "assembled = assembled.select([\"features\", \"mpg\", \"carname\"])\n",
    "\n",
    "# Random test/train split.\n",
    "train_df, test_df = assembled.randomSplit([0.7, 0.3])\n",
    "\n",
    "# Define the model.\n",
    "rf = RandomForestRegressor(\n",
    "    numTrees=20,\n",
    "    featuresCol=\"features\",\n",
    "    labelCol=\"mpg\",\n",
    ")\n",
    "\n",
    "# Train the model.\n",
    "rf_model = rf.fit(train_df)\n",
    "\n",
    "# Make predictions.\n",
    "df = rf_model.transform(test_df)\n",
    "\n",
    "# Evaluate the model.\n",
    "r2 = RegressionEvaluator(\n",
    "    labelCol=\"mpg\", predictionCol=\"prediction\", metricName=\"r2\"\n",
    ").evaluate(df)\n",
    "rmse = RegressionEvaluator(\n",
    "    labelCol=\"mpg\", predictionCol=\"prediction\", metricName=\"rmse\"\n",
    ").evaluate(df)\n",
    "print(\"RMSE={} r2={}\".format(rmse, r2))\n",
    "\n",
    "df.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4c9c2ee9",
   "metadata": {},
   "source": [
    "**A basic Random Forest Classification model**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c7a6e2b6",
   "metadata": {},
   "outputs": [],
   "source": [
    "from pyspark.ml.classification import RandomForestClassifier\n",
    "from pyspark.ml.evaluation import MulticlassClassificationEvaluator\n",
    "from pyspark.ml.feature import VectorAssembler\n",
    "\n",
    "label_column = \"cover_type\"\n",
    "vectorAssembler = VectorAssembler(\n",
    "    inputCols=covtype_df.columns,\n",
    "    outputCol=\"features\",\n",
    "    handleInvalid=\"skip\",\n",
    ")\n",
    "assembled = vectorAssembler.transform(covtype_df)\n",
    "\n",
    "# Random test/train split.\n",
    "train_df, test_df = assembled.randomSplit([0.7, 0.3])\n",
    "\n",
    "# Define the model.\n",
    "rf = RandomForestClassifier(\n",
    "    numTrees=50,\n",
    "    featuresCol=\"features\",\n",
    "    labelCol=label_column,\n",
    ")\n",
    "\n",
    "# Train the model.\n",
    "rf_model = rf.fit(train_df)\n",
    "\n",
    "# Make predictions.\n",
    "predictions = rf_model.transform(test_df)\n",
    "\n",
    "# Select (prediction, true label) and compute test error\n",
    "evaluator = MulticlassClassificationEvaluator(\n",
    "    labelCol=label_column, predictionCol=\"prediction\", metricName=\"accuracy\"\n",
    ")\n",
    "accuracy = evaluator.evaluate(predictions)\n",
    "print(\"Test Error = %g\" % (1.0 - accuracy))\n",
    "df = predictions.select([label_column, \"prediction\"])\n",
    "df.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "775d4473",
   "metadata": {},
   "source": [
    "**Encode string variables before using a VectorAssembler**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f2cffb84",
   "metadata": {},
   "outputs": [],
   "source": [
    "from pyspark.ml import Pipeline\n",
    "from pyspark.ml.feature import StringIndexer, VectorAssembler\n",
    "from pyspark.ml.regression import RandomForestRegressor\n",
    "from pyspark.ml.evaluation import RegressionEvaluator\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",
    "# Strings must be indexed or we will get:\n",
    "# pyspark.sql.utils.IllegalArgumentException: Data type string of column manufacturer is not supported.\n",
    "#\n",
    "# We also encode outside of the main pipeline or else we risk getting:\n",
    "#  Caused by: org.apache.spark.SparkException: Unseen label: XXX. To handle unseen labels, set Param handleInvalid to keep.\n",
    "#\n",
    "# This is because training data is selected randomly and may not have all possible categories.\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",
    "# Random test/train split.\n",
    "train_df, test_df = encoded_df.randomSplit([0.7, 0.3])\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",
    "model = pipeline.fit(train_df)\n",
    "\n",
    "# Make predictions.\n",
    "df = model.transform(test_df).select(\"carname\", \"mpg\", \"prediction\")\n",
    "\n",
    "# Select (prediction, true label) and compute test error\n",
    "rmse = RegressionEvaluator(\n",
    "    labelCol=\"mpg\", predictionCol=\"prediction\", metricName=\"rmse\"\n",
    ").evaluate(df)\n",
    "print(\"RMSE={}\".format(rmse))\n",
    "\n",
    "df.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b016b392",
   "metadata": {},
   "source": [
    "**Get feature importances of a trained model**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d1af3fa7",
   "metadata": {},
   "outputs": [],
   "source": [
    "from pyspark.ml import Pipeline\n",
    "from pyspark.ml.feature import StringIndexer, VectorAssembler\n",
    "from pyspark.ml.regression import RandomForestRegressor\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",
    "# Random test/train split.\n",
    "train_df, test_df = encoded_df.randomSplit([0.7, 0.3])\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",
    "model = pipeline.fit(train_df)\n",
    "\n",
    "# Make predictions.\n",
    "predictions = model.transform(test_df).select(\"carname\", \"mpg\", \"prediction\")\n",
    "\n",
    "# Get feature importances.\n",
    "real_model = model.stages[1]\n",
    "for feature, importance in zip(\n",
    "    columns_to_assemble, real_model.featureImportances\n",
    "):\n",
    "    print(\"{} contributes {:0.3f}%\".format(feature, importance * 100))\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9337f57a",
   "metadata": {},
   "source": [
    "**Automatically encode categorical variables**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "344e4272",
   "metadata": {},
   "outputs": [],
   "source": [
    "from pyspark.ml.feature import VectorAssembler, VectorIndexer\n",
    "from pyspark.ml.regression import RandomForestRegressor\n",
    "from pyspark.sql.functions import countDistinct\n",
    "\n",
    "# Remove non-numeric columns.\n",
    "df = auto_df_fixed.drop(\"carname\")\n",
    "\n",
    "# Profile this DataFrame to get a good value for maxCategories.\n",
    "grouped = df.agg(*(countDistinct(c) for c in df.columns))\n",
    "grouped.show()\n",
    "\n",
    "# Assemble all columns except mpg into a vector.\n",
    "feature_columns = list(df.columns)\n",
    "feature_columns.remove(\"mpg\")\n",
    "vector_assembler = VectorAssembler(\n",
    "    inputCols=feature_columns,\n",
    "    outputCol=\"features\",\n",
    "    handleInvalid=\"skip\",\n",
    ")\n",
    "assembled = vector_assembler.transform(df)\n",
    "\n",
    "# From profiling the dataset, 15 is a good value for max categories.\n",
    "indexer = VectorIndexer(\n",
    "    inputCol=\"features\", outputCol=\"indexed\", maxCategories=15\n",
    ")\n",
    "indexed = indexer.fit(assembled).transform(assembled)\n",
    "\n",
    "# Build and train the model.\n",
    "train_df, test_df = indexed.randomSplit([0.7, 0.3])\n",
    "rf = RandomForestRegressor(\n",
    "    numTrees=50,\n",
    "    featuresCol=\"features\",\n",
    "    labelCol=\"mpg\",\n",
    ")\n",
    "rf_model = rf.fit(train_df)\n",
    "\n",
    "# Get feature importances.\n",
    "for feature, importance in zip(feature_columns, rf_model.featureImportances):\n",
    "    print(\"{} contributes {:0.3f}%\".format(feature, importance * 100))\n",
    "\n",
    "# Make predictions.\n",
    "df = rf_model.transform(test_df).select(\"mpg\", \"prediction\")\n",
    "df.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c20566ca",
   "metadata": {},
   "source": [
    "**Hyperparameter tuning**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b8bbfeee",
   "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",
    "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(encoded_df)\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(\"fit_cross_validator.model\")\n",
    "\n",
    "# Or, just save the best model.\n",
    "best_pipeline_model.write().overwrite().save(\"best_pipeline_model.model\")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a6538117",
   "metadata": {},
   "source": [
    "**Plot Hyperparameter tuning metrics**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8219f002",
   "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": "1255e59b",
   "metadata": {},
   "source": [
    "**A Random Forest Classification model with Hyperparameter Tuning**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "dd2f81c0",
   "metadata": {},
   "outputs": [],
   "source": [
    "from pyspark.ml.classification import RandomForestClassifier\n",
    "from pyspark.ml.evaluation import MulticlassClassificationEvaluator\n",
    "from pyspark.ml.feature import VectorAssembler\n",
    "from pyspark.ml import Pipeline\n",
    "from pyspark.ml.tuning import CrossValidator, ParamGridBuilder\n",
    "\n",
    "label_column = \"cover_type\"\n",
    "vector_assembler = VectorAssembler(\n",
    "    inputCols=covtype_df.columns,\n",
    "    outputCol=\"features\",\n",
    "    handleInvalid=\"skip\",\n",
    ")\n",
    "\n",
    "# Define the model.\n",
    "rf = RandomForestClassifier(\n",
    "    numTrees=50,\n",
    "    featuresCol=\"features\",\n",
    "    labelCol=label_column,\n",
    ")\n",
    "\n",
    "# Run the pipeline.\n",
    "pipeline = Pipeline(stages=[vector_assembler, rf])\n",
    "\n",
    "# Hyperparameter search.\n",
    "paramGrid = (\n",
    "    ParamGridBuilder().addGrid(rf.numTrees, list(range(50, 80, 10))).build()\n",
    ")\n",
    "crossval = CrossValidator(\n",
    "    estimator=pipeline,\n",
    "    estimatorParamMaps=paramGrid,\n",
    "    evaluator=MulticlassClassificationEvaluator(\n",
    "        labelCol=label_column, predictionCol=\"prediction\"\n",
    "    ),\n",
    "    numFolds=2,\n",
    "    parallelism=4,\n",
    ")\n",
    "\n",
    "# Run cross-validation and choose the best set of parameters.\n",
    "model = crossval.fit(covtype_df)\n",
    "\n",
    "# Identify the best hyperparameters.\n",
    "real_model = model.bestModel.stages[1]\n",
    "print(\"Best model has {} trees.\".format(real_model.getNumTrees))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "943af15e",
   "metadata": {},
   "source": [
    "**Compute correlation matrix**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a50cf8ce",
   "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": "6808d358",
   "metadata": {},
   "source": [
    "## A few performance tips and tricks."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "728317e8",
   "metadata": {},
   "source": [
    "**Get the Spark version**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2719fdba",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(spark.sparkContext.version)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "65812491",
   "metadata": {},
   "source": [
    "**Log messages using Spark's Log4J**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f4fdc141",
   "metadata": {},
   "outputs": [],
   "source": [
    "logger = spark.sparkContext._jvm.org.apache.log4j.Logger.getRootLogger()\n",
    "logger.warn(\"WARNING LEVEL LOG MESSAGE\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e9d1c344",
   "metadata": {},
   "source": [
    "**Cache a DataFrame**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7788da2d",
   "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": "27ce73d3",
   "metadata": {},
   "source": [
    "**Show the execution plan, with costs**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d8824998",
   "metadata": {},
   "outputs": [],
   "source": [
    "df = auto_df.groupBy(\"cylinders\").count()\n",
    "execution_plan = df.explain(mode=\"cost\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "763147f9",
   "metadata": {},
   "source": [
    "**Partition by a Column Value**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4fb8324a",
   "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": "7e1f4e47",
   "metadata": {},
   "source": [
    "**Range Partition a DataFrame**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "00a6990c",
   "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": "83a65f07",
   "metadata": {},
   "source": [
    "**Change Number of DataFrame Partitions**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d997d084",
   "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.repartitionByRange(number_of_partitions, col(\"mpg\"))\n",
    "df.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "814a93af",
   "metadata": {},
   "source": [
    "**Coalesce DataFrame partitions**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "98802e9e",
   "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": "38538036",
   "metadata": {},
   "source": [
    "**Set the number of shuffle partitions**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a2acc96c",
   "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": "63e52c1b",
   "metadata": {},
   "source": [
    "**Sample a subset of a DataFrame**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5c3c32c3",
   "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": "371698f1",
   "metadata": {},
   "source": [
    "**Print Spark configuration properties**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c30b0e13",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(spark.sparkContext.getConf().getAll())"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7097fe1f",
   "metadata": {},
   "source": [
    "**Increase Spark driver/executor heap space**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "222495c4",
   "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
}