{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# PySpark Data Engineering Assessment (Extended)\n",
    "\n",
    "Welcome! In this notebook, you'll practice:\n",
    "\n",
    "1. Reading the **Titanic CSV** in **Pandas** and **PySpark**.\n",
    "2. **Splitting** a single dataset into two DataFrames and **merging** them back together in both Pandas and Spark.\n",
    "3. Data cleaning and aggregations in Pandas and Spark.\n",
    "4. Writing and reading **Parquet** files.\n",
    "5. Creating a **PySpark UDF** that leverages a **lightweight transformer model** to compute embeddings for passenger names.\n",
    "\n",
    "---\n",
    "\n",
    "## Dataset\n",
    "\n",
    "- **`titanic.csv`**: This file is in the `../data/` directory, containing columns such as:\n",
    "  - `PassengerId`, `Name`, `Sex`, `Age`, `Fare`, `Survived`, etc.\n",
    "\n",
    "We will:\n",
    "1. Read `titanic.csv` into Pandas and Spark.\n",
    "2. Split the original DataFrame into two subsets (simulating two “tables”).\n",
    "3. Demonstrate merges/joins in Pandas and Spark using these subsets.\n",
    "4. Perform data cleaning and transformations.\n",
    "5. Write to Parquet.\n",
    "6. Implement a Spark UDF to generate embeddings for passenger names.\n",
    "\n",
    "---\n",
    "\n",
    "## Instructions\n",
    "\n",
    "Throughout the notebook, you'll see `TODO` sections. Please fill in the required code. Feel free to add extra cells or explanations as needed.\n",
    "\n",
    "When finished, please save or export this notebook and submit according to your instructions.\n",
    "\n",
    "Let's begin!\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 1. Imports and Spark Setup\n",
    "\n",
    "import os\n",
    "import pandas as pd\n",
    "\n",
    "# PySpark imports\n",
    "from pyspark.sql import SparkSession\n",
    "from pyspark.sql import functions as F\n",
    "from pyspark.sql.types import *\n",
    "\n",
    "# Create/initialize Spark session\n",
    "spark = SparkSession.builder \\\n",
    "    .appName(\"TitanicAssessmentExtended\") \\\n",
    "    .getOrCreate()\n",
    "\n",
    "print(\"Spark version:\", spark.version)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 2. Read the Titanic CSV (Pandas & Spark)\n",
    "# ========================================\n",
    "\n",
    "# Path to the CSV file\n",
    "titanic_csv_path = os.path.join(\"..\", \"data\", \"titanic.csv\")\n",
    "\n",
    "# 2.1 TODO: Read 'titanic.csv' into a Pandas DataFrame (pd_df)\n",
    "# pd_df = ?\n",
    "\n",
    "# Inspect the shape and first few rows\n",
    "# print(\"pd_df shape:\", pd_df.shape)\n",
    "# display(pd_df.head())\n",
    "\n",
    "# 2.2 TODO: Read 'titanic.csv' into a Spark DataFrame (spark_df)\n",
    "# spark_df = ?\n",
    "\n",
    "# Check schema and row count\n",
    "# spark_df. ...\n",
    "# print(\"spark_df count:\", spark_df. ...)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 3. Split Data into Two Subsets for Merging/Joining\n",
    "# ==================================================\n",
    "# Split the dataset into two df's by column, then merge them \n",
    "# back together\n",
    "#   df_part1: subset of columns -> PassengerId, Name, Sex, Age\n",
    "#   df_part2: subset of columns -> PassengerId, Fare, Survived, Pclass\n",
    "#\n",
    "# \n",
    "\n",
    "# 3.1 Pandas Split\n",
    "# ----------------\n",
    "\n",
    "# TODO: Create two new DataFrames from pd_df:\n",
    "#    pd_part1 = pd_df[[\"PassengerId\", \"Name\", \"Sex\", \"Age\"]]\n",
    "#    pd_part2 = pd_df[...]\n",
    "\n",
    "# pd_part1 = ?\n",
    "# pd_part2 = ?\n",
    "\n",
    "# display(pd_part1.head())\n",
    "# display(pd_part2.head())\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 3.2 Spark Split\n",
    "# ---------------\n",
    "# TODO: Create two new DataFrames from spark_df:\n",
    "#    spark_part1 = spark_df. ...\n",
    "#    spark_part2 = spark_df. ...\n",
    "\n",
    "# spark_part1 = ?\n",
    "# spark_part2 = ?\n",
    "\n",
    "# spark_part1.show(5)\n",
    "# spark_part2.show(5)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 4. Merging / Joining the Split DataFrames\n",
    "# =========================================\n",
    "\n",
    "# 4.1 Merge in Pandas\n",
    "# -------------------\n",
    "# TODO: Merge pd_part1 and pd_part2 on \"PassengerId\"\n",
    "# We'll call the merged DataFrame \"pd_merged\".\n",
    "#\n",
    "\n",
    "# pd_merged = ?\n",
    "# print(\"pd_merged shape:\", pd_merged.shape)\n",
    "# display(pd_merged.head())\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 4.2 Join in Spark\n",
    "# -----------------\n",
    "# TODO: Join spark_part1 with spark_part2 on \"PassengerId\"\n",
    "# We'll call the joined DataFrame \"spark_merged\".\n",
    "#\n",
    "\n",
    "\n",
    "#Uncomment below\n",
    "# spark_merged = ?\n",
    "# print(\"spark_merged count:\", spark_merged.count())\n",
    "# spark_merged.show(5)\n",
    "# spark_merged.printSchema()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 5. Data Cleaning\n",
    "# ================\n",
    "# We'll focus on the merged DataFrames. For instance, drop rows that have missing\n",
    "# values in certain columns like 'Age' or 'Fare'.\n",
    "\n",
    "# 5.1 TODO: Pandas DataFrame cleaning\n",
    "# Create a cleaned version, 'pd_merged_clean',\n",
    "# dropping nulls in [\"Age\", \"Fare\"].\n",
    "\n",
    "# pd_merged_clean = ?\n",
    "\n",
    "# print(\"Before dropna:\", pd_merged.shape)\n",
    "# print(\"After dropna:\", pd_merged_clean.shape)\n",
    "# pd_merged_clean.head()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 5.2 TODO: Spark DataFrame cleaning\n",
    "# Create a cleaned version, 'spark_merged_clean',\n",
    "# dropping nulls in [\"Age\", \"Fare\"].\n",
    "\n",
    "# spark_merged_clean = ?\n",
    "\n",
    "# print(\"spark_merged count BEFORE dropna:\", spark_merged.count())\n",
    "# print(\"spark_merged_clean count AFTER dropna:\", spark_merged_clean.count())\n",
    "# spark_merged_clean.show(5)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 6. Basic Aggregations\n",
    "# =====================\n",
    "# Let's do a couple of group-by queries to glean insights.\n",
    "\n",
    "# 6.1 TODO: Pandas - Average fare by Pclass\n",
    "# e.g. group by 'Pclass' and compute mean fare in pd_merged_clean\n",
    "\n",
    "# pd_avg_fare = ?\n",
    "# pd_avg_fare\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 6.2 TODO: Spark - Survival rate by Sex and Pclass\n",
    "# Average survival rate by Sex and Pclass\n",
    "#\n",
    "# spark_survival_rate = ?\n",
    "# spark_survival_rate.show()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 7. Writing to Parquet\n",
    "# =====================\n",
    "# We'll write the cleaned Spark DataFrame to a Parquet file (e.g. \"../titanic_merged_clean.parquet\").\n",
    "\n",
    "# 7.1 TODO: Write spark_merged_clean to Parquet\n",
    "# e.g., spark_merged_clean.write. ...\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 7.2 TODO: Read it back into a new Spark DataFrame called 'spark_parquet_df'\n",
    "# spark_parquet_df = ?\n",
    "\n",
    "# print(\"spark_parquet_df count:\", spark_parquet_df.count())\n",
    "# spark_parquet_df.show(5)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 8. Create a Temp View and Query\n",
    "# ========================================\n",
    "# 8.1 TODO: Create a temp view with 'spark_merged_clean' (e.g. \"titanic_merged\")\n",
    "# spark_merged_clean.createOrReplaceTempView(\"titanic_merged\")\n",
    "\n",
    "# 8.2 TODO: Spark SQL query examples\n",
    "\n",
    "#Get the average passenger age grouped by PClass\n",
    "# result_df = spark.sql(\"SELECT ... FROM titanic_merged GROUP BY ...\")\n",
    "# result_df.show()\n",
    "\n",
    "# Calculate the Pearson correlation between passenger Fare and Survival\n",
    "# using either SQL or another method\n",
    "# Corr.(X, Y) = cov(X,Y)/(std(X)*std(Y))\n",
    "# corr = ..."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 9. Bonus 2: Transformer Embeddings UDF\n",
    "# ======================================\n",
    "\n",
    "from sentence_transformers import SentenceTransformer\n",
    "from pyspark.sql.functions import udf\n",
    "from pyspark.sql.types import ArrayType, FloatType\n",
    "\n",
    "# Load the pre-trained MiniLM sentence transformer model\n",
    "model = SentenceTransformer('sentence-transformers/all-MiniLM-L6-v2')\n",
    "\n",
    "# Define a UDF to compute the embeddings\n",
    "def compute_embedding(text):\n",
    "    '''\n",
    "    Your function goes here\n",
    "    '''\n",
    "    pass\n",
    "\n",
    "# Register the UDF in Spark\n",
    "embedding_udf = None #Replace with your udf\n",
    "\n",
    "# Apply the UDF to compute embeddings for each document\n",
    "df_with_embeddings = spark_merged_clean.withColumn('mini-lm-vectors', '...')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": []
  }
 ],
 "metadata": {
  "language_info": {
   "name": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}