docs/videos/notebook.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Notebook demo\n",
    "\n",
    "This notebook contains the code accompanying the notebook demo. You can find the video [here](https://github.com/kaiko-ai/typedspark/assets/47976799/39e157c3-6db0-436a-9e72-44b2062df808).\n",
    "\n",
    "## Set-up tables"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "from pyspark.sql import SparkSession\n",
    "from pyspark.sql.types import DateType, LongType, StringType\n",
    "from typedspark import Column, Schema, create_empty_dataset, create_partially_filled_dataset\n",
    "from datetime import datetime, timedelta"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "DataFrame[]"
      ]
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "spark = SparkSession.builder.getOrCreate()\n",
    "spark.sql(\"CREATE DATABASE IF NOT EXISTS vet\")\n",
    "spark.sql(\"CREATE DATABASE IF NOT EXISTS library\")\n",
    "spark.sql(\"CREATE DATABASE IF NOT EXISTS store\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "                                                                                \r"
     ]
    }
   ],
   "source": [
    "class Appointments(Schema):\n",
    "    appointment_id: Column[LongType]\n",
    "    pet_id: Column[LongType]\n",
    "    appointment_date: Column[DateType]\n",
    "    appointment_reason: Column[StringType]\n",
    "    veterinarian_name: Column[StringType]\n",
    "    notes: Column[StringType]\n",
    "\n",
    "\n",
    "class Pets(Schema):\n",
    "    pet_id: Column[LongType]\n",
    "    owner_id: Column[LongType]\n",
    "    pet_name: Column[StringType]\n",
    "    species: Column[StringType]\n",
    "    breed: Column[StringType]\n",
    "    age: Column[LongType]\n",
    "    birthdate: Column[DateType]\n",
    "    gender: Column[StringType]\n",
    "\n",
    "\n",
    "class Vaccinations(Schema):\n",
    "    vaccination_id: Column[LongType]\n",
    "    pet_id: Column[LongType]\n",
    "    vaccine_name: Column[StringType]\n",
    "    vaccine_date: Column[DateType]\n",
    "    next_due_date: Column[DateType]\n",
    "\n",
    "\n",
    "class Owners(Schema):\n",
    "    owner_id: Column[LongType]\n",
    "    first_name: Column[StringType]\n",
    "    last_name: Column[StringType]\n",
    "    email: Column[StringType]\n",
    "    phone_number: Column[StringType]\n",
    "    address: Column[StringType]\n",
    "\n",
    "\n",
    "create_empty_dataset(spark, Owners).write.saveAsTable(\n",
    "    \"vet.owners\", format=\"parquet\", mode=\"overwrite\"\n",
    ")\n",
    "create_empty_dataset(spark, Pets).write.saveAsTable(\"vet.pets\", format=\"parquet\", mode=\"overwrite\")\n",
    "create_empty_dataset(spark, Appointments).write.saveAsTable(\n",
    "    \"vet.appointments\", format=\"parquet\", mode=\"overwrite\"\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "date = datetime(2023, 10, 2)\n",
    "\n",
    "create_partially_filled_dataset(\n",
    "    spark,\n",
    "    Vaccinations,\n",
    "    {\n",
    "        Vaccinations.vaccination_id: [1, 2, 3, 4, 5, 6, 7],\n",
    "        Vaccinations.pet_id: [1, 2, 3, 1, 3, 2, 3],\n",
    "        Vaccinations.vaccine_name: [\n",
    "            \"rabies\",\n",
    "            \"rabies\",\n",
    "            \"rabies\",\n",
    "            \"lyme\",\n",
    "            \"lyme\",\n",
    "            \"influenza\",\n",
    "            \"influenza\",\n",
    "        ],\n",
    "        Vaccinations.next_due_date: [\n",
    "            date + timedelta(days=32),\n",
    "            date + timedelta(days=6),\n",
    "            date + timedelta(days=12),\n",
    "            date + timedelta(days=15),\n",
    "            date + timedelta(days=2),\n",
    "            date + timedelta(days=1),\n",
    "            date + timedelta(days=3),\n",
    "        ],\n",
    "        Vaccinations.vaccine_date: [\n",
    "            date + timedelta(days=32) - timedelta(days=365),\n",
    "            date + timedelta(days=6) - timedelta(days=365),\n",
    "            date + timedelta(days=12) - timedelta(days=365),\n",
    "            date + timedelta(days=15) - timedelta(days=365),\n",
    "            date + timedelta(days=2) - timedelta(days=365),\n",
    "            date + timedelta(days=1) - timedelta(days=365),\n",
    "            date + timedelta(days=3) - timedelta(days=365),\n",
    "        ],\n",
    "    },\n",
    ").write.saveAsTable(\"vet.vaccinations\", format=\"parquet\", mode=\"overwrite\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## The demo"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typedspark import Catalogs, create_schema\n",
    "from pyspark.sql import SparkSession\n",
    "import pyspark.sql.functions as F\n",
    "\n",
    "spark = SparkSession.builder.getOrCreate()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "db = Catalogs(spark)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "vaccinations, Vaccinations = db.spark_catalog.vet.vaccinations()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "+--------------+------+------------+------------+-------------+\n",
      "|vaccination_id|pet_id|vaccine_name|vaccine_date|next_due_date|\n",
      "+--------------+------+------------+------------+-------------+\n",
      "|             6|     2|   influenza|  2022-10-03|   2023-10-03|\n",
      "|             7|     3|   influenza|  2022-10-05|   2023-10-05|\n",
      "|             1|     1|      rabies|  2022-11-03|   2023-11-03|\n",
      "|             3|     3|      rabies|  2022-10-14|   2023-10-14|\n",
      "|             2|     2|      rabies|  2022-10-08|   2023-10-08|\n",
      "|             5|     3|        lyme|  2022-10-04|   2023-10-04|\n",
      "|             4|     1|        lyme|  2022-10-17|   2023-10-17|\n",
      "+--------------+------+------------+------------+-------------+\n",
      "\n"
     ]
    }
   ],
   "source": [
    "vaccinations.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "\n",
       "from pyspark.sql.types import DateType, LongType, StringType\n",
       "\n",
       "from typedspark import Column, Schema\n",
       "\n",
       "\n",
       "class Vaccinations(Schema):\n",
       "    vaccination_id: Column[LongType]\n",
       "    pet_id: Column[LongType]\n",
       "    vaccine_name: Column[StringType]\n",
       "    vaccine_date: Column[DateType]\n",
       "    next_due_date: Column[DateType]"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "Vaccinations"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "+--------------+------+------------+------------+-------------+\n",
      "|vaccination_id|pet_id|vaccine_name|vaccine_date|next_due_date|\n",
      "+--------------+------+------------+------------+-------------+\n",
      "|             1|     1|      rabies|  2022-11-03|   2023-11-03|\n",
      "|             3|     3|      rabies|  2022-10-14|   2023-10-14|\n",
      "|             2|     2|      rabies|  2022-10-08|   2023-10-08|\n",
      "+--------------+------+------------+------------+-------------+\n",
      "\n"
     ]
    }
   ],
   "source": [
    "vaccinations.filter(Vaccinations.vaccine_name == \"rabies\").show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "pivot = (\n",
    "    vaccinations.groupby(Vaccinations.pet_id)\n",
    "    .pivot(Vaccinations.vaccine_name.str)\n",
    "    .agg(F.first(Vaccinations.next_due_date))\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "+------+----------+----------+----------+\n",
      "|pet_id| influenza|      lyme|    rabies|\n",
      "+------+----------+----------+----------+\n",
      "|     2|2023-10-03|      NULL|2023-10-08|\n",
      "|     3|2023-10-05|2023-10-04|2023-10-14|\n",
      "|     1|      NULL|2023-10-17|2023-11-03|\n",
      "+------+----------+----------+----------+\n",
      "\n"
     ]
    }
   ],
   "source": [
    "pivot.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "pivot, Pivot = create_schema(pivot)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "\n",
       "from pyspark.sql.types import DateType, LongType\n",
       "\n",
       "from typedspark import Column, Schema\n",
       "\n",
       "\n",
       "class DynamicallyLoadedSchema(Schema):\n",
       "    pet_id: Column[LongType]\n",
       "    influenza: Column[DateType]\n",
       "    lyme: Column[DateType]\n",
       "    rabies: Column[DateType]"
      ]
     },
     "execution_count": 14,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "Pivot"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "+------+----------+----------+----------+\n",
      "|pet_id| influenza|      lyme|    rabies|\n",
      "+------+----------+----------+----------+\n",
      "|     2|2023-10-03|      NULL|2023-10-08|\n",
      "|     3|2023-10-05|2023-10-04|2023-10-14|\n",
      "+------+----------+----------+----------+\n",
      "\n"
     ]
    }
   ],
   "source": [
    "pivot.filter(Pivot.influenza.isNotNull()).show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Post-demo tear-down"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "DataFrame[]"
      ]
     },
     "execution_count": 16,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "spark.sql(\"DROP DATABASE vet CASCADE\")\n",
    "spark.sql(\"DROP DATABASE library CASCADE\")\n",
    "spark.sql(\"DROP DATABASE store CASCADE\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "typedspark",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}