docs/videos/ide.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# IDE demo\n",
    "\n",
    "This notebook contains the code accompanying the IDE demo. You can find the video [here](\n",
    "https://github.com/kaiko-ai/typedspark/assets/47976799/e6f7fa9c-6d14-4f68-baba-fe3c22f75b67)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "from pyspark.sql.types import DateType, LongType, StringType\n",
    "from typedspark import DataSet, Column, Schema\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]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_dogs(pets: DataSet[Pets]) -> DataSet[Pets]:\n",
    "    return pets.filter(Pets.species == \"dog\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "from chispa.dataframe_comparer import assert_df_equality\n",
    "from pyspark.sql import SparkSession\n",
    "from typedspark import create_partially_filled_dataset\n",
    "\n",
    "\n",
    "def test_get_dogs(spark: SparkSession):\n",
    "    pets = create_partially_filled_dataset(\n",
    "        spark,\n",
    "        Pets,\n",
    "        {\n",
    "            Pets.pet_id: [1, 2, 3],\n",
    "            Pets.species: [\"dog\", \"cat\", \"dog\"],\n",
    "        },\n",
    "    )\n",
    "\n",
    "    observed = get_dogs(pets)\n",
    "    expected = create_partially_filled_dataset(\n",
    "        spark,\n",
    "        Pets,\n",
    "        {\n",
    "            Pets.pet_id: [1, 3],\n",
    "            Pets.species: [\"dog\", \"dog\"],\n",
    "        },\n",
    "    )\n",
    "\n",
    "    assert_df_equality(\n",
    "        observed,\n",
    "        expected,\n",
    "        ignore_row_order=True,\n",
    "        ignore_nullable=True,\n",
    "    )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "from pyspark.sql.functions import concat_ws\n",
    "from typedspark import (\n",
    "    register_schema_to_dataset,\n",
    "    transform_to_schema,\n",
    ")\n",
    "\n",
    "\n",
    "class Reminder(Schema):\n",
    "    owner_id: Column[LongType]\n",
    "    pet_id: Column[LongType]\n",
    "    vaccination_id: Column[LongType]\n",
    "    full_name: Column[StringType]\n",
    "    email_address: Column[StringType]\n",
    "    pet_name: Column[StringType]\n",
    "    vaccine: Column[StringType]\n",
    "    due: Column[DateType]\n",
    "\n",
    "\n",
    "def find_owners_who_need_to_renew_their_pets_vaccinations(\n",
    "    owners: DataSet[Owners],\n",
    "    pets: DataSet[Pets],\n",
    "    vaccinations: DataSet[Vaccinations],\n",
    ") -> DataSet[Reminder]:\n",
    "    _owners = register_schema_to_dataset(owners, Owners)\n",
    "    _pets = register_schema_to_dataset(pets, Pets)\n",
    "    _vaccinations = register_schema_to_dataset(vaccinations, Vaccinations)\n",
    "\n",
    "    return transform_to_schema(\n",
    "        owners.join(\n",
    "            pets,\n",
    "            _owners.owner_id == _pets.owner_id,\n",
    "            \"inner\",\n",
    "        ).join(\n",
    "            vaccinations,\n",
    "            _pets.pet_id == _vaccinations.pet_id,\n",
    "            \"inner\",\n",
    "        ),\n",
    "        Reminder,\n",
    "        {\n",
    "            Reminder.owner_id: _owners.owner_id,\n",
    "            Reminder.pet_id: _pets.pet_id,\n",
    "            Reminder.vaccination_id: _vaccinations.vaccination_id,\n",
    "            Reminder.full_name: concat_ws(\n",
    "                \" \",\n",
    "                _owners.first_name,\n",
    "                _owners.last_name,\n",
    "            ),\n",
    "            Reminder.email_address: _owners.email,\n",
    "            Reminder.pet_name: _pets.pet_name,\n",
    "            Reminder.vaccine: _vaccinations.vaccine_name,\n",
    "            Reminder.due: _vaccinations.next_due_date,\n",
    "        },\n",
    "    )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "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
}