From 9fad4304bd9102161b6bcdb2a0c535a966e3c938 Mon Sep 17 00:00:00 2001 From: Your Name Date: Fri, 17 Jul 2026 09:24:51 +0100 Subject: [PATCH] migrate: convert numbers.ipynb from pandas to polars --- numbers.ipynb | 194 +++++++++++++++++++++++++++++++++----------------- 1 file changed, 130 insertions(+), 64 deletions(-) diff --git a/numbers.ipynb b/numbers.ipynb index 0b96fbb..e73d803 100644 --- a/numbers.ipynb +++ b/numbers.ipynb @@ -13,7 +13,7 @@ "\n", "### Prerequisites\n", "\n", - "This chapter mostly uses functions from **pandas**, which you are likely to already have installed but you can install using `uv add pandas` in the terminal. We'll use real examples from nycflights13, as well as toy examples made with fake data.\n", + "This chapter mostly uses functions from **pandas**, which you are likely to already have installed but you can install using `uv add polars` in the terminal. We'll use real examples from nycflights13, as well as toy examples made with fake data.\n", "\n", "Let's first load up the NYC flights data\n" ] @@ -44,10 +44,11 @@ "metadata": {}, "outputs": [], "source": [ - "import pandas as pd\n", + "import polars as pl\n", "\n", "url = \"https://raw.githubusercontent.com/byuidatascience/data4python4ds/master/data-raw/flights/flights.csv\"\n", - "flights = pd.read_csv(url)" + "\n", + "flights = pl.read_csv(url, use_pyarrow=True)" ] }, { @@ -57,7 +58,7 @@ "source": [ "### Counts\n", "\n", - "It's surprising how much data science you can do with just counts and a little basic arithmetic, so **pandas** strives to make counting as easy as possible with `.count()` and `.value_counts()`. The former just provides a straight count of all the non NA items:" + "It's surprising how much data science you can do with just counts and a little basic arithmetic, so **polars** strives to make counting as easy as possible with `.count()` and `.value_counts()`. The former provides a straight count of all the non-null items:" ] }, { @@ -104,12 +105,14 @@ "outputs": [], "source": [ "(\n", - " flights.groupby([\"dest\"])\n", + " flights.group_by(\"dest\")\n", " .agg(\n", - " mean_delay=(\"dep_delay\", \"mean\"),\n", - " count_flights=(\"dest\", \"count\"),\n", + " [\n", + " pl.col(\"dep_delay\").mean().alias(\"mean_delay\"),\n", + " pl.col(\"dest\").count().alias(\"count_flights\"),\n", + " ]\n", " )\n", - " .sort_values(by=\"count_flights\", ascending=False)\n", + " .sort(\"count_flights\", descending=True)\n", ")" ] }, @@ -128,7 +131,7 @@ "metadata": {}, "outputs": [], "source": [ - "(flights.groupby(\"tailnum\").agg(miles=(\"distance\", \"sum\")))" + "flights.group_by(\"tailnum\").agg(pl.col(\"distance\").sum().alias(\"miles\"))" ] }, { @@ -136,7 +139,7 @@ "id": "1b489765", "metadata": {}, "source": [ - "You can count missing values by combining `sum()` and `isnull()`. In the flights dataset this represents flights that are cancelled. Note that because there isn't a simple string name for applying `.isnull()` followed by `.sum()` (unlike just running `sum()`, which would be given by the string \"sum\"), we need to use a lambda function in the below:" + "You can count missing values by combining sum() and is_null(). In the flights dataset this represents flights that are cancelled:" ] }, { @@ -146,7 +149,7 @@ "metadata": {}, "outputs": [], "source": [ - "(flights.groupby(\"dest\").agg(n_cancelled=(\"dep_time\", lambda x: x.isnull().sum())))" + "flights.group_by(\"dest\").agg(pl.col(\"dep_time\").is_null().sum().alias(\"n_cancelled\"))" ] }, { @@ -160,11 +163,11 @@ "\n", "Basic number arithmatic is achieved by `+` (addition), `-` (subtraction), `*` (multiplication), `/` (division), `**` (powers), `%` (modulo), and `@` (tensor product). Most of these functions don't need a huge amount of explanation because you'll be familiar with them already (and you can look up the others when you do need them).\n", "\n", - "When you have two numeric columns of equal length and you add or subtract them, it's pretty obvious what's going to happen. But we do need to talk about what happens when there is a variable involved that is *not* as long as the column. This is important for operations like `flights.assign(air_time = air_time / 60)` because there are 336,776 numbers on the left of `/` but only one on the right. In this case, **pandas** will understand that you'd like to divide *all* values of air time by 60. This is sometimes called 'broadcasting'. Below is a digram that tries to explain what's going on:\n", + "When you have two numeric columns of equal length and you add or subtract them, it's pretty obvious what's going to happen. But we do need to talk about what happens when there is a variable involved that is *not* as long as the column. This is important for operations like `flights.assign(air_time = air_time / 60)` because there are 336,776 numbers on the left of `/` but only one on the right. In this case, **polars** will understand that you'd like to divide *all* values of air time by 60. This is sometimes called 'broadcasting'. Below is a digram that tries to explain what's going on:\n", "\n", "![](https://numpy.org/doc/stable/_images/broadcasting_1.png)\n", "\n", - "You can find out much more about [broadcasting on the **numpy** documentation](https://numpy.org/doc/stable/user/basics.broadcasting.html). **pandas** is built on top of **numpy** and inherits some of its functionality. \n" + "You can find out much more about [broadcasting on the **numpy** documentation](https://numpy.org/doc/stable/user/basics.broadcasting.html). **polars** is built on top of **numpy** and inherits some of its functionality. \n" ] }, { @@ -172,7 +175,7 @@ "id": "adb10890", "metadata": {}, "source": [ - "When operating on two columns, **pandas** compares their shapes element-wise. Two columns are compatible when they are equal, or one of them is a scalar. If these conditions are not met, you will get an error.\n", + "When operating on two columns, **polars** compares their shapes element-wise. Two columns are compatible when they are equal, or one of them is a scalar. If these conditions are not met, you will get an error.\n", "\n" ] }, @@ -201,7 +204,7 @@ "id": "dd24faa5", "metadata": {}, "source": [ - "Sometimes, you'd like to look at the maximum or minimum value across rows *or* columns. As often is the case with **pandas**, you can specify rows or columns to apply functions to by passing `axis=0` (index) or `axis=1` (columns) to that function. The axis designation can be confusing: remember that you are asking which dimension you wish to aggregate over, leaving you with the other dimension. So if we wish to find the minimum in each row, we aggregate / collapse columns, so we need to pass `axis=1`." + "Sometimes, you'd like to look at the maximum or minimum value across rows or columns. To find the minimum in each row, you can use the horizontal functions:" ] }, { @@ -211,7 +214,7 @@ "metadata": {}, "outputs": [], "source": [ - "df = pd.DataFrame({\"x\": [1, 5, 7], \"y\": [3, 2, pd.NA]})\n", + "df = pl.DataFrame({\"x\": [1, 5, 7], \"y\": [3, 2, None]})\n", "df" ] }, @@ -230,7 +233,7 @@ "metadata": {}, "outputs": [], "source": [ - "df.min(axis=1)" + "df.select(pl.min_horizontal(\"x\", \"y\"))" ] }, { @@ -273,9 +276,11 @@ "metadata": {}, "outputs": [], "source": [ - "flights.assign(\n", - " hour=lambda x: x[\"sched_dep_time\"] // 100,\n", - " minute=lambda x: x[\"sched_dep_time\"] % 100,\n", + "flights.with_columns(\n", + " [\n", + " (pl.col(\"sched_dep_time\") // 100).alias(\"hour\"),\n", + " (pl.col(\"sched_dep_time\") % 100).alias(\"minute\"),\n", + " ]\n", ")" ] }, @@ -300,8 +305,11 @@ "\n", "starting = 100\n", "interest = 1.05\n", - "money = pd.DataFrame(\n", - " {\"year\": 2000 + np.arange(1, 51), \"money\": starting * interest ** np.arange(1, 51)}\n", + "money = pl.DataFrame(\n", + " {\n", + " \"year\": 2000 + np.arange(1, 51),\n", + " \"money\": starting * interest ** np.arange(1, 51),\n", + " }\n", ")\n", "money.head()" ] @@ -321,7 +329,10 @@ "metadata": {}, "outputs": [], "source": [ - "money.plot(x=\"year\", y=\"money\");" + "# Convert to pandas for plotting\n", + "money_pd = money.to_pandas()\n", + "money_pd.plot(x=\"year\", y=\"money\")\n", + "plt.show()" ] }, { @@ -339,7 +350,8 @@ "metadata": {}, "outputs": [], "source": [ - "money.plot(x=\"year\", y=\"money\", logy=True);" + "money_pd.plot(x=\"year\", y=\"money\", logy=True)\n", + "plt.show()" ] }, { @@ -371,7 +383,7 @@ "metadata": {}, "outputs": [], "source": [ - "money.head().round(2)" + "money.select(pl.all().round(2)).head()" ] }, { @@ -379,7 +391,7 @@ "id": "2b26e633", "metadata": {}, "source": [ - "This can also be applied to individual columns or differentially to columns via a dictionary:" + "This can also be applied to individual columns:" ] }, { @@ -389,7 +401,12 @@ "metadata": {}, "outputs": [], "source": [ - "money.head().round({\"year\": 0, \"money\": 1})" + "money.with_columns(\n", + " [\n", + " pl.col(\"year\").round(0),\n", + " pl.col(\"money\").round(1),\n", + " ]\n", + ").head()" ] }, { @@ -407,7 +424,12 @@ "metadata": {}, "outputs": [], "source": [ - "money.tail().round({\"year\": 0, \"money\": -2})" + "money.with_columns(\n", + " [\n", + " pl.col(\"year\").round(),\n", + " ((pl.col(\"money\") / 100).round() * 100).alias(\"money\"),\n", + " ]\n", + ").head()" ] }, { @@ -415,7 +437,7 @@ "id": "6544b94e", "metadata": {}, "source": [ - "Sometimes, you'll want to round according to significant figures rather than decimal places. There's not a really easy way to do this but you can define a custom function to do it. Here's an example of rounding to 2 significant figures (change the `2` in the below to round to a different number of significant figures):" + "Sometimes, you'll want to round according to significant figures rather than decimal places. There's not a really easy way to do this but you can define a custom function to do it. Here's an example of rounding to 2 significant figures (change the `2` in the code block below to round to a different number of significant figures):" ] }, { @@ -425,7 +447,9 @@ "metadata": {}, "outputs": [], "source": [ - "money[\"money\"].head().apply(lambda x: float(f'{float(f\"{x:.2g}\"):g}'))" + "money.with_columns(\n", + " pl.col(\"money\").map_elements(lambda x: float(f'{float(f\"{x:.2g}\"):g}'))\n", + ")" ] }, { @@ -490,7 +514,7 @@ "id": "c22b0eda", "metadata": {}, "source": [ - "Remember that you can always apply **numpy** functions to **pandas** data frame columns like so:" + "Remember that you can always apply numpy functions to **polars** data frame columns using `map_elements()`:" ] }, { @@ -500,7 +524,7 @@ "metadata": {}, "outputs": [], "source": [ - "money[\"money\"].head().apply(np.ceil)" + "money.with_columns(pl.col(\"money\").map_elements(np.ceil)).head()" ] }, { @@ -510,7 +534,7 @@ "source": [ "### Cumulative and Rolling Aggregates\n", "\n", - "**pandas** has several cumulative functions, including `.cumsum()`, `.cummax()` and `.cummin()`, and `.cumprod()`. " + "**polars** has several cumulative functions, including `.cumsum()`, `.cummax()` and `.cummin()`, and `.cumprod()`. " ] }, { @@ -520,15 +544,24 @@ "metadata": {}, "outputs": [], "source": [ - "money[\"money\"].tail().cumsum()" - ] - }, - { - "cell_type": "markdown", - "id": "c4b569cb", - "metadata": {}, - "source": [ - "As ever, this can be applied across rows too by passing `axis=1`." + "import numpy as np\n", + "import polars as pl\n", + "\n", + "# Create the money DataFrame\n", + "starting = 100\n", + "interest = 1.05\n", + "money = pl.DataFrame(\n", + " {\n", + " \"year\": 2000 + np.arange(1, 51),\n", + " \"money\": starting * interest ** np.arange(1, 51),\n", + " }\n", + ")\n", + "\n", + "# Show the data\n", + "money.head()\n", + "\n", + "# Cumulative sum\n", + "money[\"money\"].cum_sum().tail()" ] }, { @@ -542,7 +575,7 @@ "\n", "### Ranking\n", "\n", - "**pandas**' rank function is `.rank()`. Let's look back at the data we made earlier and rank it:" + "**polars**' rank function is `.rank()`. Let's look back at the data we made earlier and rank it:" ] }, { @@ -562,7 +595,7 @@ "metadata": {}, "outputs": [], "source": [ - "df.rank()" + "df.select(pl.all().rank())" ] }, { @@ -570,7 +603,7 @@ "id": "85345760", "metadata": {}, "source": [ - "Of course, there's no change here because the items were already ranked! We can also do a pct rank by passing the keyword argument `pct=True`." + "We can also do a percent rank by passing the keyword argument `method=\"average\"` and then scaling:" ] }, { @@ -580,7 +613,7 @@ "metadata": {}, "outputs": [], "source": [ - "df.rank(pct=True)" + "df.select(pl.all().rank(method=\"average\") / df.height)" ] }, { @@ -602,9 +635,15 @@ "metadata": {}, "outputs": [], "source": [ - "money[\"money_lag_5\"] = money[\"money\"].shift(5)\n", - "money[\"money_lead_10\"] = money[\"money\"].shift(-10)\n", - "money.set_index(\"year\").plot();" + "money_with_shifts = money.with_columns(\n", + " [\n", + " pl.col(\"money\").shift(5).alias(\"money_lag_5\"),\n", + " pl.col(\"money\").shift(-10).alias(\"money_lead_10\"),\n", + " ]\n", + ")\n", + "money_with_shifts_pd = money_with_shifts.to_pandas()\n", + "money_with_shifts_pd.set_index(\"year\").plot()\n", + "plt.show()" ] }, { @@ -636,7 +675,7 @@ "source": [ "## More Useful Summary Statistics\n", "\n", - "We've seen how useful `.mean()`, `.count()`, and `.value_counts()` can be for analysis. **pandas** has a great many more built-in summary statistics functions, however. These include `.median()` (you may find it interesting to compare the mean vs the median when looking at the hourly departure delay in the flights data), `.mode()`, `.min()`, and `.max()`.\n", + "We've seen how useful `.mean()`, `.count()`, and `.value_counts()` can be for analysis. **polars** has a great many more built-in summary statistics functions, however. These include `.median()`, `.mode()`, `.min()`, and `.max()`.\n", "\n", "A class of useful summary statistics are provided by the `.quantile` function, which is the same as `median` for `.quantile(0.5)`. The quantile at x% is the value that x% of values are below. (Note that under this definition, `.quantile(1)` will be the same as `.max()`.) Let's see an example with the 25th percentile." ] @@ -656,7 +695,7 @@ "id": "02c212f4", "metadata": {}, "source": [ - "Sometimes you don't just want one percentile, but a bunch of them. **pandas** makes this very easy by allowing you to pass a list of quantiles:" + "Sometimes you don't just want one percentile, but a bunch of them. **polars** makes this very easy by allowing you to pass a list of quantiles:" ] }, { @@ -666,7 +705,30 @@ "metadata": {}, "outputs": [], "source": [ - "money[\"money\"].quantile([0, 0.25, 0.5, 0.75])" + "# Create a struct with quantile values\n", + "money.select(\n", + " [\n", + " pl.struct(\n", + " [\n", + " pl.lit(0).alias(\"0%\"),\n", + " pl.lit(0.25).alias(\"25%\"),\n", + " pl.lit(0.5).alias(\"50%\"),\n", + " pl.lit(0.75).alias(\"75%\"),\n", + " ]\n", + " ).alias(\"quantiles\")\n", + " ]\n", + ")\n", + "\n", + "# Create a DataFrame with the quantile results\n", + "quantile_df = money.select(\n", + " [\n", + " pl.lit(money[\"money\"].quantile(0)).alias(\"0%\"),\n", + " pl.lit(money[\"money\"].quantile(0.25)).alias(\"25%\"),\n", + " pl.lit(money[\"money\"].quantile(0.5)).alias(\"50%\"),\n", + " pl.lit(money[\"money\"].quantile(0.75)).alias(\"75%\"),\n", + " ]\n", + ")\n", + "quantile_df" ] }, { @@ -690,12 +752,16 @@ "outputs": [], "source": [ "(\n", - " flights.groupby([\"origin\", \"dest\"])\n", + " flights.group_by([\"origin\", \"dest\"])\n", " .agg(\n", - " distance_sd=(\"distance\", lambda x: x.quantile(0.75) - x.quantile(0.25)),\n", - " count=(\"distance\", \"count\"),\n", + " [\n", + " (\n", + " pl.col(\"distance\").quantile(0.75) - pl.col(\"distance\").quantile(0.25)\n", + " ).alias(\"distance_sd\"),\n", + " pl.col(\"distance\").count().alias(\"count\"),\n", + " ]\n", " )\n", - " .query(\"distance_sd > 0\")\n", + " .filter(pl.col(\"distance_sd\") > 0)\n", ")" ] }, @@ -718,7 +784,9 @@ "metadata": {}, "outputs": [], "source": [ - "flights[\"dep_delay\"].plot.hist(bins=50, title=\" Distribution: length of delay\");" + "flights_pd = flights.to_pandas()\n", + "flights_pd[\"dep_delay\"].plot.hist(bins=50, title=\"Distribution: length of delay\")\n", + "plt.show()" ] }, { @@ -728,9 +796,10 @@ "metadata": {}, "outputs": [], "source": [ - "flights.query(\"dep_delay <= 120\")[\"dep_delay\"].plot.hist(\n", - " bins=50, title=\" Distribution: length of delay\"\n", - ");" + "flights_pd.query(\"dep_delay <= 120\")[\"dep_delay\"].plot.hist(\n", + " bins=50, title=\"Distribution: length of delay\"\n", + ")\n", + "plt.show()" ] }, { @@ -768,9 +837,6 @@ } ], "metadata": { - "interpreter": { - "hash": "9d7534ecd9fbc7d385378f8400cf4d6cb9c6175408a574f1c99c5269f08771cc" - }, "jupytext": { "cell_metadata_filter": "-all", "encoding": "# -*- coding: utf-8 -*-", @@ -778,7 +844,7 @@ "main_language": "python" }, "kernelspec": { - "display_name": "Python 3 (ipykernel)", + "display_name": ".venv", "language": "python", "name": "python3" },