python_adv_group_joining.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "11.5 ms ± 201 μs per loop (mean ± std. dev. of 7 runs, 100 loops each)\n"
     ]
    }
   ],
   "source": [
    "%%timeit\n",
    "# Take the eight CSV files and turn them into a data frame.\n",
    "\n",
    "# We are only interested in the first three columns from each CSV file: the date\n",
    "# and time, the max temperature, and the min temperature.\n",
    "\n",
    "# Add city and state columns that contain the city and state from the filename\n",
    "# and allow us to distinguish between rows.\n",
    "\n",
    "\n",
    "import pandas as pd\n",
    "from pathlib import Path\n",
    "path = \"data/weather\"\n",
    "files = Path(path).glob('*.csv')\n",
    "\n",
    "df_parts = []\n",
    "names = [\"date_time\", \"max_temp\", \"min_temp\"]\n",
    "for f in files:\n",
    "    # take city and state from filename\n",
    "    file_name = str(f).split('/')\n",
    "    city, state = file_name[-1].split('.')[0].replace(\"+\", \" \").split(',')\n",
    "    \n",
    "    # read csv and rename columns\n",
    "    df_part = pd.read_csv(f)\n",
    "    cols = df_part.columns.to_list()\n",
    "    cols[0:3] = names\n",
    "    df_part.columns = cols\n",
    "    \n",
    "    # drop other unused colums and add city, state\n",
    "    (df_parts\n",
    "     .append(\n",
    "         df_part\n",
    "         .drop(df_part.columns[3:], axis=1)\n",
    "         .assign(city=city, state=state)))\n",
    "\n",
    "# concat tables to create df\n",
    "df = pd.concat(df_parts, ignore_index=True, axis=0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "1.9 ms ± 21.2 μs per loop (mean ± std. dev. of 7 runs, 1,000 loops each)\n"
     ]
    }
   ],
   "source": [
    "%%timeit\n",
    "# Does the data for each city and state start and end at (roughly) the same \n",
    "# time? How do you know?\n",
    "df.groupby([\"state\", \"city\"]).agg(min=(\"date_time\", \"min\"), max=(\"date_time\", \"max\"))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "1.11 ms ± 19.8 μs per loop (mean ± std. dev. of 7 runs, 1,000 loops each)\n"
     ]
    }
   ],
   "source": [
    "%%timeit\n",
    "# What is the lowest minimum temperature recorded for each city in the data set?\n",
    "df.groupby([\"state\", \"city\"]).agg(lowest_temp=(\"min_temp\", \"min\"))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "833 μs ± 12.7 μs per loop (mean ± std. dev. of 7 runs, 1,000 loops each)\n"
     ]
    }
   ],
   "source": [
    "%%timeit\n",
    "# What is the highest maximum temperature recorded in each state in the data set?\n",
    "df.groupby(\"state\").agg(highest_temp=(\"max_temp\", \"max\"))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "1.53 ms ± 15.8 μs per loop (mean ± std. dev. of 7 runs, 1,000 loops each)\n"
     ]
    }
   ],
   "source": [
    "%%timeit\n",
    "# What is the average difference in temperature (i.e., max – min) for each of the\n",
    "# cities in our data set\n",
    "(\n",
    "    df\n",
    "    .assign(delta=(df[\"max_temp\"] - df[\"min_temp\"]))\n",
    "    .groupby([\"state\", \"city\"])\n",
    "    .agg(mean_delta=(\"delta\", \"mean\"))\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "31.5 ms ± 1.1 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)\n"
     ]
    }
   ],
   "source": [
    "%%timeit\n",
    "import polars as pl\n",
    "from pathlib import Path\n",
    "path = \"data/weather\"\n",
    "files = Path(path).glob('*.csv')\n",
    "\n",
    "df_parts = []\n",
    "for file in files:\n",
    "    file_name = str(file).split('/')\n",
    "    city, state = file_name[-1].split('.')[0].replace(\"+\", \" \").split(',')\n",
    "    df_part = (\n",
    "            pl\n",
    "            .scan_csv(file, try_parse_dates=True)\n",
    "            .select(\n",
    "                pl.col(\"date_time\"), \n",
    "                pl.col(\"^[a-z,+]+_maxtempC$\").alias(\"max_temp\"), \n",
    "                pl.col(\"^[a-z,+]+_mintempC$\").alias(\"min_temp\"))\n",
    "            .with_columns(\n",
    "                pl.lit(city).alias(\"city\"), \n",
    "                pl.lit(state).alias(\"state\"))\n",
    "        )\n",
    "    df_parts.append(df_part)\n",
    "\n",
    "df = pl.concat(df_parts).collect()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "855 μs ± 12.2 μs per loop (mean ± std. dev. of 7 runs, 1,000 loops each)\n"
     ]
    }
   ],
   "source": [
    "%%timeit\n",
    "# Does the data for each city and state start and end at (roughly) the same \n",
    "# time? How do you know?\n",
    "(\n",
    "    df\n",
    "    .group_by(pl.col(\"state\"), pl.col(\"city\"))\n",
    "    .agg(\n",
    "        pl.min(\"date_time\").alias(\"min\"), \n",
    "        pl.max(\"date_time\").alias(\"max\")\n",
    "    )\n",
    ")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 37,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "826 μs ± 11.9 μs per loop (mean ± std. dev. of 7 runs, 1,000 loops each)\n"
     ]
    }
   ],
   "source": [
    "%%timeit\n",
    "# What is the lowest minimum temperature recorded for each city in the data set?\n",
    "df.group_by(pl.col(\"state\"), pl.col(\"city\")).agg(pl.min(\"min_temp\").alias(\"min\"))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 38,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "825 μs ± 10.5 μs per loop (mean ± std. dev. of 7 runs, 1,000 loops each)\n"
     ]
    }
   ],
   "source": [
    "%%timeit\n",
    "# What is the highest maximum temperature recorded in each state in the data set?\n",
    "df.group_by(pl.col(\"state\"), pl.col(\"city\")).agg(pl.max(\"max_temp\").alias(\"max\"))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 39,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "1.42 ms ± 13.6 μs per loop (mean ± std. dev. of 7 runs, 1,000 loops each)\n"
     ]
    }
   ],
   "source": [
    "%%timeit\n",
    "# What is the average difference in temperature (i.e., max – min) for each of the\n",
    "# cities in our data set\n",
    "(\n",
    "    df\n",
    "    .with_columns((pl.col(\"max_temp\") - pl.col(\"min_temp\")).alias(\"delta\"))\n",
    "    .group_by(pl.col(\"state\"), pl.col(\"city\"))\n",
    "    .agg(pl.mean(\"delta\").alias(\"mean\"))\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 41,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "6.7 ms ± 72 μs per loop (mean ± std. dev. of 7 runs, 100 loops each)\n"
     ]
    }
   ],
   "source": [
    "%%timeit\n",
    "import pyarrow as pa\n",
    "import pyarrow.csv\n",
    "\n",
    "from pathlib import Path\n",
    "path = \"data/weather\"\n",
    "files = list(Path(path).glob('*.csv'))\n",
    "\n",
    "df_parts = []\n",
    "for file in files:\n",
    "    # take city and state from filename\n",
    "    file_name = str(file).split('/')\n",
    "    city, state = file_name[-1].split('.')[0].replace(\"+\", \" \").split(',')\n",
    "\n",
    "    table = pyarrow.csv.read_csv(file)\n",
    "    df_parts.append(\n",
    "        table\n",
    "        .drop_columns(table.column_names[3:])\n",
    "        .rename_columns([\"date_time\", \"max_temp\", \"min_temp\"])\n",
    "        .append_column(\"city\", [[pa.scalar(city)]*table.num_rows])\n",
    "        .append_column(\"state\", [[pa.scalar(state)]*table.num_rows])\n",
    "    )\n",
    "df = pa.concat_tables(df_parts)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# import pyarrow.parquet as pq\n",
    "# pq.write_table(table, 'example.parquet')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 42,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "206 μs ± 2.34 μs per loop (mean ± std. dev. of 7 runs, 1,000 loops each)\n"
     ]
    }
   ],
   "source": [
    "%%timeit\n",
    "# Does the data for each city and state start and end at (roughly) the same \n",
    "# time? How do you know?\n",
    "(\n",
    "    df\n",
    "    .group_by([\"state\",\"city\"])\n",
    "    .aggregate([(\"date_time\", \"min\"), (\"date_time\", \"max\")]\n",
    "    )\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "190 μs ± 910 ns per loop (mean ± std. dev. of 7 runs, 10,000 loops each)\n"
     ]
    }
   ],
   "source": [
    "%%timeit\n",
    "# What is the lowest minimum temperature recorded for each city in the data set?\n",
    "df.group_by([\"city\"]).aggregate([(\"min_temp\", \"min\")])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "191 μs ± 3.78 μs per loop (mean ± std. dev. of 7 runs, 10,000 loops each)\n"
     ]
    }
   ],
   "source": [
    "%%timeit\n",
    "# What is the highest maximum temperature recorded in each state in the data set?\n",
    "df.group_by([\"state\"]).aggregate([(\"max_temp\", \"max\")])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 45,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "216 μs ± 1.37 μs per loop (mean ± std. dev. of 7 runs, 1,000 loops each)\n"
     ]
    }
   ],
   "source": [
    "%%timeit\n",
    "# What is the average difference in temperature (i.e., max – min) for each of the\n",
    "# cities in our data set\n",
    "import pyarrow.compute as pc\n",
    "pc.subtract(df[\"max_temp\"], df[\"min_temp\"])\n",
    "(\n",
    "    df\n",
    "    .append_column(\n",
    "        pa.field(\"delta\", pa.int64()), \n",
    "        pc.subtract(df[\"max_temp\"], df[\"min_temp\"])\n",
    "    )\n",
    "    .group_by([\"state\",\"city\"])\n",
    "    .aggregate([(\"delta\", \"mean\")])\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "SAT scores, revisited"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "1.65 ms ± 11.6 μs per loop (mean ± std. dev. of 7 runs, 1,000 loops each)\n"
     ]
    }
   ],
   "source": [
    "%%timeit\n",
    "import pandas as pd\n",
    "# Read in the scores file (sat-scores.csv). This time, you want the following \n",
    "# columns\n",
    "\n",
    "columns = [\"Year\", \"State.Code\", \"Total.Math\", \"Family Income.Less than 20k.Math\",\n",
    "\"Family Income.Between 20-40k.Math\", \"Family Income.Between 40-60k.Math\",\n",
    "\"Family Income.Between 60-80k.Math\", \"Family Income.Between 80-100k.Math\",\n",
    "\"Family Income.More than 100k.Math\"]\n",
    "\n",
    "# Rename the income-related column names to something shorter\n",
    "\n",
    "new_names = [\"Year\", \"State.Code\", \"Total.Math\", \"income<20k\", \"20k<income<40k\", \n",
    "        \"40k<income<60k\", \"60k<income<80k\", \"80k<income<100k\", \"income>100k\"]\n",
    "cols = dict(zip(columns, new_names))\n",
    "\n",
    "df = (\n",
    "    pd\n",
    "    .read_csv(\"data/sat-scores.csv\", usecols=columns)\n",
    "    .rename(columns=cols)[new_names]\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "239 μs ± 3.31 μs per loop (mean ± std. dev. of 7 runs, 1,000 loops each)\n"
     ]
    }
   ],
   "source": [
    "%%timeit\n",
    "# Find the average SAT math score for each income level, grouped and then\n",
    "# sorted by year\n",
    "df1 = (\n",
    "    df\n",
    "    .drop(columns=[\"State.Code\", \"Total.Math\"])\n",
    "    .groupby(\"Year\").agg(\"mean\")\n",
    "    .sort_index(ascending=True)\n",
    ")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "327 μs ± 5.33 μs per loop (mean ± std. dev. of 7 runs, 1,000 loops each)\n"
     ]
    }
   ],
   "source": [
    "%%timeit\n",
    "# For each year in the data set, determine how much better each income group\n",
    "# did, on average, than the next-poorer group of students. Do you see\n",
    "# any income group that did worse, in any year, than the next-poorer students\n",
    "\n",
    "df2 = df1.shift(-1, axis=1)\n",
    "new_names = [\"(20, 40) vs (0, 20)\", \"(40, 60) vs (20, 40)\", \n",
    "             \"(60, 80) vs (40, 60)\", \"(80, 100) vs (60, 80)\", \n",
    "             \"(100, ) vs (80, 100)\"]\n",
    "diff = (df2 - df1).dropna(axis=1)\n",
    "diff.columns = new_names\n",
    "diff < 0"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "47.5 μs ± 700 ns per loop (mean ± std. dev. of 7 runs, 10,000 loops each)\n"
     ]
    }
   ],
   "source": [
    "%%timeit\n",
    "# Which income bracket, on average, had the greatest advantage over the next-\n",
    "# poorer income bracket?\n",
    "diff.mean(axis=0).idxmax()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "241 μs ± 2.27 μs per loop (mean ± std. dev. of 7 runs, 1,000 loops each)\n"
     ]
    }
   ],
   "source": [
    "%%timeit\n",
    "# Which five states have the greatest gap in SAT math scores between the richest\n",
    "# and poorest students?\n",
    "(\n",
    "    df\n",
    "    .assign(delta=df[\"income>100k\"] - df[\"income<20k\"])\n",
    "    .groupby(\"State.Code\")[\"delta\"]\n",
    "    .mean()\n",
    "    .sort_values(ascending=False)[:5]\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "1 ms ± 2.45 μs per loop (mean ± std. dev. of 7 runs, 1,000 loops each)\n"
     ]
    }
   ],
   "source": [
    "%%timeit\n",
    "import polars as pl\n",
    "# Read in the scores file (sat-scores.csv). This time, you want the following \n",
    "# columns\n",
    "\n",
    "columns = [\"Year\", \"State.Code\", \"Total.Math\", \"Family Income.Less than 20k.Math\",\n",
    "\"Family Income.Between 20-40k.Math\", \"Family Income.Between 40-60k.Math\",\n",
    "\"Family Income.Between 60-80k.Math\", \"Family Income.Between 80-100k.Math\",\n",
    "\"Family Income.More than 100k.Math\"]\n",
    "\n",
    "# Rename the income-related column names to something shorter\n",
    "\n",
    "new_names = [\"Year\", \"State.Code\", \"Total.Math\", \"income<20k\", \"20k<income<40k\", \n",
    "        \"40k<income<60k\", \"60k<income<80k\", \"80k<income<100k\", \"income>100k\"]\n",
    "cols = dict(zip(columns, new_names))\n",
    "\n",
    "df = pl.read_csv(\"data/sat-scores.csv\", columns=columns).rename(cols)[new_names]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "240 μs ± 2.06 μs per loop (mean ± std. dev. of 7 runs, 1,000 loops each)\n"
     ]
    }
   ],
   "source": [
    "%%timeit\n",
    "# Find the average SAT math score for each income level, grouped and then\n",
    "# sorted by year\n",
    "df1 = (\n",
    "    df\n",
    "    .drop([\"State.Code\", \"Total.Math\"])\n",
    "    .group_by(\"Year\").mean()\n",
    "    .sort(by=\"Year\", descending=False)\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "399 μs ± 7.28 μs per loop (mean ± std. dev. of 7 runs, 1,000 loops each)\n"
     ]
    }
   ],
   "source": [
    "%%timeit\n",
    "# For each year in the data set, determine how much better each income group\n",
    "# did, on average, than the next-poorer group of students. Do you see\n",
    "# any income group that did worse, in any year, than the next-poorer students\n",
    "\n",
    "import polars.selectors as cs\n",
    "new_names = [\"(20, 40) vs (0, 20)\", \"(40, 60) vs (20, 40)\", \n",
    "             \"(60, 80) vs (40, 60)\", \"(80, 100) vs (60, 80)\", \n",
    "             \"(100, ) vs (80, 100)\"]\n",
    "\n",
    "tmp = df1.drop(\"Year\").transpose()\n",
    "diff = (tmp.shift(-1) - tmp).transpose().drop(cs.by_index(-1))\n",
    "diff.columns = new_names\n",
    "diff < 0"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "138 μs ± 1.18 μs per loop (mean ± std. dev. of 7 runs, 10,000 loops each)\n"
     ]
    }
   ],
   "source": [
    "%%timeit\n",
    "# Which income bracket, on average, had the greatest advantage over the next-\n",
    "# poorer income bracket?\n",
    "(\n",
    "    diff\n",
    "    .mean()\n",
    "    .transpose(include_header=True)\n",
    "    .sort(by=\"column_0\", descending=True)[:1]\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "776 μs ± 12.1 μs per loop (mean ± std. dev. of 7 runs, 1,000 loops each)\n"
     ]
    }
   ],
   "source": [
    "%%timeit\n",
    "# Which five states have the greatest gap in SAT math scores between the richest\n",
    "# and poorest students?\n",
    "(\n",
    "    df\n",
    "    .with_columns((pl.col(\"income>100k\") - pl.col(\"income<20k\")).alias(\"delta\"))\n",
    "    .group_by(\"State.Code\")\n",
    "    .agg(pl.col(\"delta\").mean())\n",
    "    .sort(by=\"delta\", descending=True)[:5]\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "2.95 ms ± 113 μs per loop (mean ± std. dev. of 7 runs, 100 loops each)\n"
     ]
    }
   ],
   "source": [
    "%%timeit\n",
    "import pyarrow as pa \n",
    "import pyarrow.dataset as ds\n",
    "\n",
    "columns = [\"Year\", \"State.Code\", \"Total.Math\", \"Family Income.Less than 20k.Math\",\n",
    "\"Family Income.Between 20-40k.Math\", \"Family Income.Between 40-60k.Math\",\n",
    "\"Family Income.Between 60-80k.Math\", \"Family Income.Between 80-100k.Math\",\n",
    "\"Family Income.More than 100k.Math\"]\n",
    "fields = [ds.field(name) for name in columns]\n",
    "new_names = [\"Year\", \"State.Code\", \"Total.Math\", \"income<20k\", \"20k<income<40k\", \n",
    "        \"40k<income<60k\", \"60k<income<80k\", \"80k<income<100k\", \"income>100k\"]\n",
    "cols = dict(zip(new_names, fields))\n",
    "df = ds.dataset(\"data/sat-scores.csv\",format=\"csv\").scanner(columns=cols).to_table()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "124 μs ± 2.04 μs per loop (mean ± std. dev. of 7 runs, 10,000 loops each)\n"
     ]
    }
   ],
   "source": [
    "%%timeit\n",
    "# Find the average SAT math score for each income level, grouped and then\n",
    "# sorted by year\n",
    "df1 = (\n",
    "    df\n",
    "    .drop_columns([\"State.Code\", \"Total.Math\"])\n",
    "    .group_by([\"Year\"]).aggregate([(x, \"mean\") for x in new_names[3:]])\n",
    "    .sort_by([(\"Year\", \"ascending\")])\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "82.2 μs ± 757 ns per loop (mean ± std. dev. of 7 runs, 10,000 loops each)\n"
     ]
    }
   ],
   "source": [
    "%%timeit\n",
    "# For each year in the data set, determine how much better each income group\n",
    "# did, on average, than the next-poorer group of students. Do you see\n",
    "# any income group that did worse, in any year, than the next-poorer students\n",
    "import pyarrow.compute as pc\n",
    "\n",
    "cols = df1.column_names[1:]\n",
    "new_names = [\"(20, 40) vs (0, 20)\", \"(40, 60) vs (20, 40)\", \n",
    "             \"(60, 80) vs (40, 60)\", \"(80, 100) vs (60, 80)\", \n",
    "             \"(100, ) vs (80, 100)\"]\n",
    "\n",
    "diff = pa.Table.from_arrays(\n",
    "    [pc.subtract(df1.column(cols[i]), df1.column(cols[i-1])) \n",
    "        for i in range(1, len(cols))], \n",
    "    names=new_names)\n",
    "\n",
    "pa.Table.from_arrays(\n",
    "    [pc.less_equal(diff.column(n), 0) for n in diff.column_names], \n",
    "    names=new_names)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 34,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "20.8 μs ± 676 ns per loop (mean ± std. dev. of 7 runs, 10,000 loops each)\n"
     ]
    }
   ],
   "source": [
    "%%timeit\n",
    "# Which income bracket, on average, had the greatest advantage over the next-\n",
    "# poorer income bracket?\n",
    "sorted(\n",
    "    zip(diff.column_names, \n",
    "        [pc.mean(diff.column(col)) for col in diff.column_names]), \n",
    "        key=lambda x: x[1].as_py()\n",
    ")[-1]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 36,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "86.8 μs ± 1.09 μs per loop (mean ± std. dev. of 7 runs, 10,000 loops each)\n"
     ]
    }
   ],
   "source": [
    "%%timeit\n",
    "# Which five states have the greatest gap in SAT math scores between the richest\n",
    "# and poorest students?\n",
    "(\n",
    "    df\n",
    "    .append_column(pa.field(\"delta\", pa.int64()), \n",
    "                   pc.subtract(df[\"income>100k\"], df[\"income<20k\"]))\n",
    "    .group_by([\"State.Code\"]).aggregate([(\"delta\", \"mean\")])\n",
    "    .sort_by([(\"delta_mean\", \"descending\")])[:5]\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Snowy, rainy cities"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 135,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pandas as pd\n",
    "from pathlib import Path\n",
    "path = \"data/weather\"\n",
    "files = Path(path).glob('*.csv')\n",
    "\n",
    "df_parts = []\n",
    "names = [\"precipMM\", \"max_temp\", \"min_temp\"]\n",
    "for f in files:\n",
    "    # take city and state from filename\n",
    "    file_name = str(f).split('/')\n",
    "    city, state = file_name[-1].split('.')[0].replace(\"+\", \" \").split(',')\n",
    "    \n",
    "    # read csv and rename columns\n",
    "    df_part = pd.read_csv(f)\n",
    "    df_part.columns = [x.split('_')[1] for x in  df_part.columns]\n",
    "    df_part.rename(columns={\"maxtempC\": \"max_temp\", \"mintempC\":\"min_temp\"}, inplace=True)\n",
    "    \n",
    "    # drop other unused colums and add city, state\n",
    "    (df_parts\n",
    "     .append(\n",
    "         df_part[names]\n",
    "         .assign(city=city, state=state)))\n",
    "\n",
    "# concat tables to create df\n",
    "df = pd.concat(df_parts, ignore_index=True, axis=0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "1.56 ms ± 14 μs per loop (mean ± std. dev. of 7 runs, 1,000 loops each)\n"
     ]
    }
   ],
   "source": [
    "%%timeit\n",
    "# Determine which cities had, on at least three occasions, precipitation of \n",
    "# 15 mm or more (book impl)\n",
    "\n",
    "def has_multiple_readings_at_least(mini_df, min_mm, times):\n",
    "    return mini_df.loc[mini_df['precipMM'] >= min_mm, 'precipMM'].count() >= times\n",
    "(\n",
    "    df\n",
    "    .groupby(['city', 'state'])\n",
    "    .filter(has_multiple_readings_at_least, min_mm=15, times=3)[\"city\"].unique()\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 69,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "342 μs ± 12 μs per loop (mean ± std. dev. of 7 runs, 1,000 loops each)\n"
     ]
    }
   ],
   "source": [
    "%%timeit\n",
    "# Determine which cities had, on at least three occasions, precipitation of \n",
    "# 15 mm or more\n",
    "(\n",
    "    df.loc[df[\"precipMM\"] >= 15]\n",
    "    .groupby([\"state\", \"city\"])\n",
    "    .filter(lambda x: len(x) >= 3)[\"city\"]\n",
    "    .unique()\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "1.65 ms ± 16.2 μs per loop (mean ± std. dev. of 7 runs, 100 loops each)\n"
     ]
    }
   ],
   "source": [
    "%%timeit\n",
    "# Find cities that had at least three measurements of 10 mm of precipitation or\n",
    "# more when the temperature was at or below 0° Celsius.(book impl)\n",
    "\n",
    "def has_multiple_readings_at_least(mini_df, min_mm, temp, times):\n",
    "    return (mini_df\n",
    "            .loc[((mini_df['precipMM'] >= min_mm) & \n",
    "                  (mini_df['min_temp'] <= temp)), 'precipMM'].count() >= times)\n",
    "(\n",
    "    df\n",
    "    .groupby(['city', 'state'])\n",
    "    .filter(has_multiple_readings_at_least, min_mm=10, temp=0, times=3)[\"city\"]\n",
    "    .unique()\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 74,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "374 μs ± 4.1 μs per loop (mean ± std. dev. of 7 runs, 1,000 loops each)\n"
     ]
    }
   ],
   "source": [
    "%%timeit\n",
    "# Find cities that had at least three measurements of 10 mm of precipitation or\n",
    "# more when the temperature was at or below 0° Celsius.\n",
    "\n",
    "(\n",
    "    df.loc[(df[\"precipMM\"] >= 10) & (df[\"min_temp\"] <= 0)]\n",
    "    .groupby([\"state\", \"city\"])\n",
    "    .filter(lambda x: len(x) >= 3)[\"city\"]\n",
    "    .unique()\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "886 μs ± 12.8 μs per loop (mean ± std. dev. of 7 runs, 1,000 loops each)\n"
     ]
    }
   ],
   "source": [
    "%%timeit\n",
    "# For each precipitation measurement, calculate the proportion of that city’s\n",
    "# total precipitation(book impl)\n",
    "def proportion_of_city_precip(s):\n",
    "    return s / s.sum()\n",
    "\n",
    "df['precip_pct'] = df.groupby('city')['precipMM'].transform(proportion_of_city_precip)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 143,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "1.53 ms ± 41.9 μs per loop (mean ± std. dev. of 7 runs, 1,000 loops each)\n"
     ]
    }
   ],
   "source": [
    "%%timeit\n",
    "# For each precipitation measurement, calculate the proportion of that city’s\n",
    "# total precipitation\n",
    "df1 = df.join(\n",
    "    df.groupby([\"state\", \"city\"])[\"precipMM\"].sum(), \n",
    "        on=[\"state\", \"city\"], lsuffix='_l', rsuffix='_r')\n",
    "df.assign(frac=df1[\"precipMM_l\"]/df1[\"precipMM_r\"])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 142,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "437 μs ± 8.68 μs per loop (mean ± std. dev. of 7 runs, 1,000 loops each)\n"
     ]
    }
   ],
   "source": [
    "%%timeit\n",
    "# For each city, determine the greatest proportion of that city’s total \n",
    "# precipitation to fall in a given period.\n",
    "df.groupby(['state', 'city'])['frac'].max()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 103,
   "metadata": {},
   "outputs": [],
   "source": [
    "import polars as pl\n",
    "from pathlib import Path\n",
    "path = \"data/weather\"\n",
    "files = Path(path).glob('*.csv')\n",
    "\n",
    "df_parts = []\n",
    "for file in files:\n",
    "    file_name = str(file).split('/')\n",
    "    city, state = file_name[-1].split('.')[0].replace(\"+\", \" \").split(',')\n",
    "    df_part = (\n",
    "            pl\n",
    "            .scan_csv(file, try_parse_dates=True)\n",
    "            .select(\n",
    "                pl.col(\"^[a-z,+]+_precipMM$\").alias(\"precipMM\"), \n",
    "                pl.col(\"^[a-z,+]+_maxtempC$\").alias(\"max_temp\"), \n",
    "                pl.col(\"^[a-z,+]+_mintempC$\").alias(\"min_temp\"))\n",
    "            .with_columns(\n",
    "                pl.lit(city).alias(\"city\"), \n",
    "                pl.lit(state).alias(\"state\"))\n",
    "        )\n",
    "    df_parts.append(df_part)\n",
    "\n",
    "df = pl.concat(df_parts).collect()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "955 μs ± 12.4 μs per loop (mean ± std. dev. of 7 runs, 1,000 loops each)\n"
     ]
    }
   ],
   "source": [
    "%%timeit\n",
    "# Determine which cities had, on at least three occasions, precipitation of \n",
    "# 15 mm or more\n",
    "\n",
    "(\n",
    "    df\n",
    "    .filter(pl.col(\"precipMM\") >= 15)\n",
    "    .group_by(\"state\", \"city\")\n",
    "    .len()\n",
    "    .filter(pl.col(\"len\") >= 3)\n",
    ")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "1.58 ms ± 8.18 μs per loop (mean ± std. dev. of 7 runs, 1,000 loops each)\n"
     ]
    }
   ],
   "source": [
    "%%timeit\n",
    "# Find cities that had at least three measurements of 10 mm of precipitation or\n",
    "# more when the temperature was at or below 0° Celsius.\n",
    "\n",
    "(\n",
    "    df\n",
    "    .filter((pl.col(\"precipMM\") >= 10) & pl.col(\"min_temp\") <= 0)\n",
    "    .group_by(\"state\", \"city\")\n",
    "    .len()\n",
    "    .filter(pl.col(\"len\") >= 3)\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 130,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "890 μs ± 8.77 μs per loop (mean ± std. dev. of 7 runs, 1,000 loops each)\n"
     ]
    }
   ],
   "source": [
    "%%timeit\n",
    "# For each precipitation measurement, calculate the proportion of that city’s\n",
    "# total precipitation(book impl)\n",
    "(\n",
    "    df\n",
    "    .join(\n",
    "        df.group_by(\"state\", \"city\").agg(pl.col(\"precipMM\").sum().alias(\"sum\")),\n",
    "        on=[\"state\", \"city\"], suffix='right'\n",
    "    )\n",
    "    .with_columns((pl.col(\"precipMM\")/pl.col(\"sum\")).alias(\"frac\"))\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 132,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "487 μs ± 7.19 μs per loop (mean ± std. dev. of 7 runs, 1,000 loops each)\n"
     ]
    }
   ],
   "source": [
    "%%timeit\n",
    "# For each city, determine the greatest proportion of that city’s total \n",
    "# precipitation to fall in a given period.\n",
    "df.group_by('state', 'city').agg(pl.col(\"frac\").max())"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "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.12.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}