{
 "nbformat": 4,
 "nbformat_minor": 5,
 "metadata": {
  "colab": {
   "provenance": [],
   "toc_visible": true
  },
  "kernelspec": {
   "name": "python3",
   "display_name": "Python 3"
  },
  "language_info": {
   "name": "python"
  }
 },
 "cells": [
  {
   "id": "m01-l2-0",
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Lesson 1.2 · Relational Joins and SQL over DataFrames with DuckDB\n",
    "\n",
    "*Accelerated Foundation in Applied Causal Inference*\n",
    "\n",
    "**By the end of this lesson you can:**\n",
    "\n",
    "- Join tidy tables with DuckDB SQL directly on Polars DataFrames\n",
    "- Choose between INNER, LEFT and ANTI joins and audit what each one drops\n",
    "- Detect missing periods before they bias a panel estimate\n",
    "- Align covariates in time so that they are measured before the outcome\n",
    "\n",
    "> Read the full explanations on the course website, then run each cell here in order (**Runtime → Run all** works on a fresh runtime)."
   ]
  },
  {
   "id": "m01-l2-1",
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Rebuild the tidy tables"
   ]
  },
  {
   "id": "m01-l2-2",
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "!pip install -q polars duckdb pandas pyarrow\n",
    "import polars as pl, duckdb, pandas as pd, numpy as np\n",
    "\n",
    "np.random.seed(42)\n",
    "dates = pd.date_range(\"2020-01-01\", periods=48, freq=\"ME\")\n",
    "\n",
    "inflation = (\n",
    "    pl.DataFrame({\n",
    "        \"date\": dates,\n",
    "        \"Ontario\": np.random.normal(0.02, 0.005, 48).cumsum() + 2.0,\n",
    "        \"Quebec\":  np.random.normal(0.02, 0.006, 48).cumsum() + 1.8,\n",
    "    })\n",
    "    .unpivot(index=\"date\", variable_name=\"province\", value_name=\"inflation_rate\")\n",
    ")\n",
    "\n",
    "unemployment = pl.DataFrame({\n",
    "    \"record_date\": np.tile(dates, 2),\n",
    "    \"province\": [\"Ontario\"] * 48 + [\"Quebec\"] * 48,\n",
    "    \"unemp_rate\": np.random.normal(5.5, 0.5, 96),\n",
    "})\n",
    "# Drop three Quebec months to mimic a gap in the feed\n",
    "gap = [dates[10], dates[11], dates[30]]\n",
    "unemployment = unemployment.filter(\n",
    "    ~((pl.col(\"province\") == \"Quebec\") & pl.col(\"record_date\").is_in(gap))\n",
    ")\n",
    "print(inflation.shape, unemployment.shape)   # (96, 3) (93, 3)"
   ]
  },
  {
   "id": "m01-l2-3",
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## SQL joins on DataFrames"
   ]
  },
  {
   "id": "m01-l2-4",
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "> **Question it answers:** What were inflation and unemployment in each province and month, side by side?"
   ]
  },
  {
   "id": "m01-l2-5",
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "query = \"\"\"\n",
    "    SELECT i.date, i.province, i.inflation_rate, u.unemp_rate\n",
    "    FROM inflation AS i\n",
    "    LEFT JOIN unemployment AS u\n",
    "      ON i.date = u.record_date\n",
    "     AND i.province = u.province\n",
    "    ORDER BY i.province, i.date\n",
    "\"\"\"\n",
    "panel = duckdb.sql(query).pl()\n",
    "print(panel.head(6))\n",
    "print(\"Rows:\", panel.height)"
   ]
  },
  {
   "id": "m01-l2-6",
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Auditing a join: what got lost?"
   ]
  },
  {
   "id": "m01-l2-7",
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "> **Question it answers:** Which province-months are missing from the unemployment feed?"
   ]
  },
  {
   "id": "m01-l2-8",
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "missing = duckdb.sql(\"\"\"\n",
    "    SELECT i.province, i.date\n",
    "    FROM inflation i\n",
    "    ANTI JOIN unemployment u\n",
    "      ON i.date = u.record_date AND i.province = u.province\n",
    "\"\"\").pl()\n",
    "print(missing)\n",
    "\n",
    "audit = duckdb.sql(\"\"\"\n",
    "    SELECT province,\n",
    "           COUNT(*)                         AS months,\n",
    "           COUNT(unemp_rate)                AS months_with_unemployment,\n",
    "           COUNT(*) - COUNT(unemp_rate)     AS gaps\n",
    "    FROM panel\n",
    "    GROUP BY province ORDER BY province\n",
    "\"\"\").pl()\n",
    "print(audit)"
   ]
  },
  {
   "id": "m01-l2-9",
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Timing: lags, leads and publication delays"
   ]
  },
  {
   "id": "m01-l2-10",
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "> **Question it answers:** What was unemployment in the previous month, for each province?"
   ]
  },
  {
   "id": "m01-l2-11",
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "panel_lagged = duckdb.sql(\"\"\"\n",
    "    SELECT *,\n",
    "           LAG(unemp_rate, 1) OVER (PARTITION BY province ORDER BY date) AS unemp_lag1,\n",
    "           LAG(unemp_rate, 3) OVER (PARTITION BY province ORDER BY date) AS unemp_lag3\n",
    "    FROM panel\n",
    "    ORDER BY province, date\n",
    "\"\"\").pl()\n",
    "print(panel_lagged.head(5))"
   ]
  },
  {
   "id": "m01-l2-12",
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "panel_lagged_pl = panel.sort(\"province\", \"date\").with_columns(\n",
    "    pl.col(\"unemp_rate\").shift(1).over(\"province\").alias(\"unemp_lag1\")\n",
    ")\n",
    "assert panel_lagged_pl[\"unemp_lag1\"].equals(panel_lagged[\"unemp_lag1\"], check_names=False)\n",
    "print(\"Polars and DuckDB agree.\")"
   ]
  },
  {
   "id": "m01-l2-13",
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Exercise: an ASOF join"
   ]
  },
  {
   "id": "m01-l2-14",
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### ✍️ Exercise\n",
    "\n",
    "Statistics are published with a delay. Suppose a policy committee meets on the dates below, and the\n",
    "unemployment figure for a month is only **published 35 days after month-end**. For each meeting, find the\n",
    "latest unemployment rate for Ontario that was actually *known* on the meeting day.\n",
    "\n",
    "Hint: create a `published` column, then use DuckDB's `ASOF LEFT JOIN ... ON m.meeting >= u.published`."
   ]
  },
  {
   "id": "m01-l2-15",
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": [
    "# Your answer here"
   ]
  },
  {
   "id": "m01-l2-16",
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Checklist and recap"
   ]
  }
 ]
}