ch15.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Chapter 15: Reshaping"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import polars as pl\n",
    "pl.__version__  # The book is built with Polars version 1.20.0"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Wide Versus Long DataFrames"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "grades_wide = pl.DataFrame(\n",
    "    {\n",
    "        \"student\": [\"Jeroen\", \"Thijs\", \"Ritchie\"],\n",
    "        \"math\": [85, 78, 92],\n",
    "        \"science\": [90, 82, 85],\n",
    "        \"history\": [88, 80, 87],\n",
    "    }\n",
    ")\n",
    "\n",
    "grades_wide"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "grades_long = pl.DataFrame(\n",
    "    {\n",
    "        \"student\": [\n",
    "            \"Jeroen\",\n",
    "            \"Jeroen\",\n",
    "            \"Jeroen\",\n",
    "            \"Thijs\",\n",
    "            \"Thijs\",\n",
    "            \"Thijs\",\n",
    "            \"Ritchie\",\n",
    "            \"Ritchie\",\n",
    "            \"Ritchie\",\n",
    "        ],\n",
    "        \"subject\": [\n",
    "            \"Math\",\n",
    "            \"Science\",\n",
    "            \"History\",\n",
    "            \"Math\",\n",
    "            \"Science\",\n",
    "            \"History\",\n",
    "            \"Math\",\n",
    "            \"Science\",\n",
    "            \"History\",\n",
    "        ],\n",
    "        \"grade\": [85, 90, 88, 78, 82, 80, 92, 85, 87],\n",
    "    }\n",
    ")\n",
    "\n",
    "grades_long"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Pivot to a Wider DataFrame"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "grades = pl.DataFrame(\n",
    "    {\n",
    "        \"student\": [\n",
    "            \"Jeroen\",\n",
    "            \"Jeroen\",\n",
    "            \"Jeroen\",\n",
    "            \"Thijs\",\n",
    "            \"Thijs\",\n",
    "            \"Thijs\",\n",
    "            \"Ritchie\",\n",
    "            \"Ritchie\",\n",
    "            \"Ritchie\",\n",
    "        ],\n",
    "        \"subject\": [\n",
    "            \"Math\",\n",
    "            \"Science\",\n",
    "            \"History\",\n",
    "            \"Math\",\n",
    "            \"Science\",\n",
    "            \"History\",\n",
    "            \"Math\",\n",
    "            \"Science\",\n",
    "            \"History\",\n",
    "        ],\n",
    "        \"grade\": [85, 90, 88, 78, 82, 80, 92, 85, 87],\n",
    "    }\n",
    ")\n",
    "\n",
    "grades"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "grades.pivot(index=\"student\", on=\"subject\", values=\"grade\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "multiple_grades = pl.DataFrame(\n",
    "    {\n",
    "        \"student\": [\n",
    "            \"Jeroen\",\n",
    "            \"Jeroen\",\n",
    "            \"Jeroen\",\n",
    "            \"Jeroen\",\n",
    "            \"Jeroen\",\n",
    "            \"Jeroen\",\n",
    "            \"Thijs\",\n",
    "            \"Thijs\",\n",
    "            \"Thijs\",\n",
    "            \"Thijs\",\n",
    "            \"Thijs\",\n",
    "            \"Thijs\",\n",
    "        ],\n",
    "        \"subject\": [\n",
    "            \"Math\",\n",
    "            \"Math\",\n",
    "            \"Math\",\n",
    "            \"Science\",\n",
    "            \"Science\",\n",
    "            \"Science\",\n",
    "            \"Math\",\n",
    "            \"Math\",\n",
    "            \"Math\",\n",
    "            \"Science\",\n",
    "            \"Science\",\n",
    "            \"Science\",\n",
    "        ],\n",
    "        \"grade\": [85, 88, 85, 60, 66, 63, 51, 79, 62, 82, 85, 82],\n",
    "    }\n",
    ")\n",
    "\n",
    "multiple_grades"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "multiple_grades.pivot(\n",
    "    index=\"student\", on=\"subject\", values=\"grade\", aggregate_function=\"mean\"\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "multiple_grades.pivot(\n",
    "    index=\"student\",\n",
    "    on=\"subject\",\n",
    "    values=\"grade\",\n",
    "    aggregate_function=pl.element().max() - pl.element().min(),\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "lf = pl.LazyFrame(\n",
    "    {\n",
    "        \"col1\": [\"a\", \"a\", \"a\", \"b\", \"b\", \"b\"],\n",
    "        \"col2\": [\"x\", \"x\", \"x\", \"x\", \"y\", \"y\"],\n",
    "        \"col3\": [6, 7, 3, 2, 5, 7],\n",
    "    }\n",
    ")\n",
    "\n",
    "index = pl.col(\"col1\")\n",
    "on = pl.col(\"col2\")\n",
    "values = pl.col(\"col3\")\n",
    "unique_column_values = [\"x\", \"y\"]\n",
    "aggregate_function = lambda col: col.tanh().mean()\n",
    "\n",
    "lf.group_by(index).agg(\n",
    "    aggregate_function(values.filter(on == value)).alias(value)\n",
    "    for value in unique_column_values\n",
    ").collect()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Unpivot to a Longer DataFrame"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "grades_wide = pl.DataFrame(\n",
    "    {\n",
    "        \"student\": [\"Jeroen\", \"Thijs\", \"Ritchie\"],\n",
    "        \"math\": [85, 78, 92],\n",
    "        \"science\": [90, 82, 85],\n",
    "        \"history\": [88, 80, 87],\n",
    "    }\n",
    ")\n",
    "\n",
    "grades_wide"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "grades_wide.unpivot(\n",
    "    index=[\"student\"],\n",
    "    on=[\"math\", \"science\", \"history\"],\n",
    "    variable_name=\"subject\",\n",
    "    value_name=\"grade\",\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "df = pl.DataFrame(\n",
    "    {\n",
    "        \"student\": [\"Jeroen\", \"Thijs\", \"Ritchie\", \"Jeroen\", \"Thijs\", \"Ritchie\"],\n",
    "        \"class\": [\n",
    "            \"Math101\",\n",
    "            \"Math101\",\n",
    "            \"Math101\",\n",
    "            \"Math102\",\n",
    "            \"Math102\",\n",
    "            \"Math102\",\n",
    "        ],\n",
    "        \"age\": [20, 21, 22, 20, 21, 22],\n",
    "        \"semester\": [\"Fall\", \"Fall\", \"Fall\", \"Spring\", \"Spring\", \"Spring\"],\n",
    "        \"math\": [85, 78, 92, 88, 79, 95],\n",
    "        \"science\": [90, 82, 85, 92, 81, 87],\n",
    "        \"history\": [88, 80, 87, 85, 82, 89],\n",
    "    }\n",
    ")\n",
    "df"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "df.unpivot(\n",
    "    index=[\"student\", \"class\", \"age\", \"semester\"],\n",
    "    on=[\"math\", \"science\", \"history\"],\n",
    "    variable_name=\"subject\",\n",
    "    value_name=\"grade\",\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Transposing"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "grades_wide = pl.DataFrame(\n",
    "    {\n",
    "        \"student\": [\"Jeroen\", \"Thijs\", \"Ritchie\"],\n",
    "        \"math\": [85, 78, 92],\n",
    "        \"science\": [90, 82, 85],\n",
    "        \"history\": [88, 80, 87],\n",
    "    }\n",
    ")\n",
    "\n",
    "grades_wide"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "report_columns = (f\"report_{i + 1}\" for i, _ in enumerate(df.columns))  \n",
    "\n",
    "grades_wide.transpose(\n",
    "    include_header=True,\n",
    "    header_name=\"original_headers\",\n",
    "    column_names=report_columns,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Exploding"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "grades_nested = pl.DataFrame(\n",
    "    {\n",
    "        \"student\": [\"Jeroen\", \"Thijs\", \"Ritchie\"],\n",
    "        \"math\": [[85, 90, 88], [78, 82, 80], [92, 85, 87]],\n",
    "    }\n",
    ")\n",
    "\n",
    "grades_nested"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "grades_nested.explode(\"math\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "grades_nested = pl.DataFrame(\n",
    "    {\n",
    "        \"student\": [\"Jeroen\", \"Thijs\", \"Ritchie\"],\n",
    "        \"math\": [[85, 90, 88], [78, 82, 80], [92, 85, 87]],\n",
    "        \"science\": [[85, 90, 88], [78, 82], [92, 85, 87]],\n",
    "        \"history\": [[85, 90, 88], [78, 82], [92, 85, 87]],\n",
    "    }\n",
    ")\n",
    "\n",
    "grades_nested"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# This raises a ShapeError:\n",
    "# grades_nested.explode(\"math\", \"science\", \"history\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "grades_nested_long = grades_nested.unpivot(\n",
    "    index=\"student\", variable_name=\"subject\", value_name=\"grade\"\n",
    ")\n",
    "\n",
    "grades_nested_long"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "grades_nested_long.explode(\"grade\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "nested_lists = pl.DataFrame(\n",
    "    {\n",
    "        \"id\": [1, 2],\n",
    "        \"nested_value\": [[[\"a\", \"b\"]], [[\"c\"], [\"d\", \"e\"]]],\n",
    "    },\n",
    "    strict=False,\n",
    ")\n",
    "nested_lists"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "nested_lists.explode(\"nested_value\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "nested_lists.explode(\"nested_value\").explode(\"nested_value\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Partition into Multiple DataFrames"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sales = pl.DataFrame(\n",
    "    {\n",
    "        \"OrderID\": [1, 2, 3, 4, 5, 6],\n",
    "        \"Product\": [\"A\", \"B\", \"A\", \"C\", \"B\", \"A\"],\n",
    "        \"Quantity\": [10, 5, 8, 7, 3, 12],\n",
    "        \"Region\": [\"North\", \"South\", \"North\", \"West\", \"South\", \"West\"],\n",
    "    }\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sales.partition_by(\"Region\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sales.partition_by(\"Region\", include_key=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sales_dict = sales.partition_by([\"Region\"], as_dict=True)\n",
    "\n",
    "sales_dict"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sales_dict[(\"North\",)]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Takeaways"
   ]
  }
 ],
 "metadata": {},
 "nbformat": 4,
 "nbformat_minor": 5
}