From 2998f9b60dc90edadcfe35aa3ecdca106dabdced Mon Sep 17 00:00:00 2001 From: Your Name Date: Thu, 16 Jul 2026 14:47:21 +0100 Subject: [PATCH] migrate: convert joins.ipynb from pandas to polars --- joins.ipynb | 268 ++++++++++++++++++++++++++++++++++++++++++---------- 1 file changed, 217 insertions(+), 51 deletions(-) diff --git a/joins.ipynb b/joins.ipynb index d494d84..8ec4902 100644 --- a/joins.ipynb +++ b/joins.ipynb @@ -11,7 +11,7 @@ "\n", "It's rare that a data analysis involves only a single data frame. Typically you have many data frames, and you must *join* them together to answer the questions that you're interested in.\n", "\n", - "**pandas** has a really rich set of options for combining one or more data frames, with the two most important being concatenate and merge. Some of the examples in this chapter show you how to join a pair of data frames. Fortunately this is enough, since you can combine three data frames by combining two pairs." + "**polars** has a really rich set of options for combining one or more data frames, with the two most important being concatenate and joins. Some of the examples in this chapter show you how to join a pair of data frames. Fortunately this is enough, since you can combine three data frames by combining two pairs." ] }, { @@ -41,7 +41,7 @@ "source": [ "### Prerequisites\n", "\n", - "This chapter will use the **pandas** data analysis package." + "This chapter will use the **polars** data analysis package." ] }, { @@ -51,15 +51,13 @@ "source": [ "## Concatenate\n", "\n", - "If you have two or more data frames with the same index or the same columns, you can glue them together into a single data frame using `pd.concat()`.\n", + "If you have two or more data frames with the same columns, you can glue them together into a single data frame using pl.concat().\n", "\n", "![](https://pandas.pydata.org/docs/_images/08_concat_row.svg)\n", "\n", - "For the same columns, pass `axis=0` to glue the index together; for the same index, pass `axis=1` to glue the columns together. The concatenate function will typically be used on a list of data frames.\n", + "For the same columns, use pl.concat() to stack rows vertically. For the same index, you can use pl.concat() with how=\"horizontal\".\n", "\n", - "If you want to track where the original data came from in the final data frame, use the `keys` keyword.\n", - "\n", - "Here's an example using data on two different states' populations that also makes uses of the `keys` option:" + "Here's an example using data on two different states' populations:" ] }, { @@ -69,7 +67,7 @@ "metadata": {}, "outputs": [], "source": [ - "import pandas as pd\n", + "import polars as pl\n", "\n", "base_url = (\n", " \"https://github.com/aeturrell/coding-for-economists/raw/refs/heads/main/data/\"\n", @@ -77,32 +75,66 @@ "state_codes = [\"ca\", \"il\"]\n", "end_url = \"pop.dta\"\n", "\n", - "# This grabs the two dataframes, one for each state\n", - "list_of_state_dfs = [pd.read_stata(base_url + state + end_url) for state in state_codes]\n", - "# Show example of first entry in list of dataframes\n", + "# Read the Stata files and convert to polars\n", + "list_of_state_dfs = []\n", + "for state in state_codes:\n", + " import pandas as pd\n", + "\n", + " temp_df = pd.read_stata(base_url + state + end_url)\n", + " list_of_state_dfs.append(pl.from_pandas(temp_df))\n", + "\n", + "# shows example of the 2 dataframes\n", "print(list_of_state_dfs[0])\n", + "print(list_of_state_dfs[1])\n", "\n", - "# Concatenate the list of dataframes\n", - "df = pd.concat(list_of_state_dfs, keys=state_codes, axis=0)\n", + "# concatenate the list of dataframes\n", + "df = pl.concat(list_of_state_dfs)\n", "df" ] }, { "cell_type": "markdown", - "id": "e0c343b0", + "id": "79bbd002", "metadata": {}, "source": [ - "Note that the `keys` argument is optional, but is useful for keeping track of origin data frames within the merged data frame.\n", + "Note that when concatenating, you can track the origin of rows by adding a column before concatenation:" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "2adb960a", + "metadata": {}, + "outputs": [], + "source": [ + "# Add a state identifier column to each dataframe\n", + "list_with_state = []\n", + "for state, df_state in zip(state_codes, list_of_state_dfs):\n", + " df_state = df_state.with_columns(pl.lit(state).alias(\"state\"))\n", + " list_with_state.append(df_state)\n", "\n", + "# Concatenate with state identifiers\n", + "df_with_state = pl.concat(list_with_state)\n", + "df_with_state" + ] + }, + { + "cell_type": "markdown", + "id": "e0c343b0", + "metadata": {}, + "source": [ "::: {.callout-tip title=\"Exercise\"}\n", "Concatenate the follow two data frames:\n", "\n", "```python\n", - "df1 = pd.DataFrame([['a', 1], ['b', 2]],\n", - " columns=['letter', 'number'])\n", + "df1 = pl.DataFrame({\n", + " 'letter': ['a', 'b'],\n", + " 'number': [1, 2]\n", + "})\n", "\n", - "df2 = pd.DataFrame([['c', 3], ['d', 4]],\n", - " columns=['letter', 'number'])\n", + "df2 = pl.DataFrame({\n", + " 'letter': ['c', 'd'],\n", + " 'number': [3, 4]\n", "```\n", "\n", ":::" @@ -113,17 +145,25 @@ "id": "f6e91777", "metadata": {}, "source": [ - "## Merge\n", + "## Joins\n", "\n", - "There are so many options for merging data frames using `pd.merge(left, right, on=..., how=...` that we won't be able to cover them all here. The most important features are: the two data frames to be merged, what variables (aka keys) to merge on (and these can be indexes) via `on=`, and *how* to do the merge (eg left, right, outer, inner) via `how=`. This diagram shows an example of a merge using keys from the left-hand data frame:\n", + "There are many options for joining data frames using `pl.DataFrame.join()`. The most important features are: the two data frames to be joined, what variables (aka keys) to join on via `on=`, and *how* to do the join (eg left, right, outer, inner) via `how=`. This diagram shows an example of a left join:\n", "\n", "![](https://pandas.pydata.org/docs/_images/08_merge_left.svg)\n", "\n", "The `how=` keyword works in the following ways:\n", - "- `how='left'` uses keys from the left data frame only to merge.\n", - "- `how='right'` uses keys from the right data frame only to merge.\n", - "- `how='inner'` uses keys that appear in both data frames to merge.\n", - "- `how='outer'` uses the cartesian product of keys in both data frames to merge on.\n", + "\n", + "- how=\"left\" uses keys from the left data frame only to join.\n", + "\n", + "- how=\"right\" uses keys from the right data frame only to join.\n", + "\n", + "- how=\"inner\" uses keys that appear in both data frames to join.\n", + "\n", + "- how=\"outer\" uses all keys from both data frames to join on.\n", + "\n", + "- how=\"semi\" returns rows from the left where keys appear in the right.\n", + "\n", + "- how=\"anti\" returns rows from the left where keys do NOT appear in the right.\n", "\n", "Let's see examples of some of these:" ] @@ -135,7 +175,7 @@ "metadata": {}, "outputs": [], "source": [ - "left = pd.DataFrame(\n", + "left = pl.DataFrame(\n", " {\n", " \"key1\": [\"K0\", \"K0\", \"K1\", \"K2\"],\n", " \"key2\": [\"K0\", \"K1\", \"K0\", \"K1\"],\n", @@ -143,16 +183,32 @@ " \"B\": [\"B0\", \"B1\", \"B2\", \"B3\"],\n", " }\n", ")\n", - "right = pd.DataFrame(\n", + "right = pl.DataFrame(\n", " {\n", " \"key1\": [\"K0\", \"K1\", \"K1\", \"K2\"],\n", " \"key2\": [\"K0\", \"K0\", \"K0\", \"K0\"],\n", " \"C\": [\"C0\", \"C1\", \"C2\", \"C3\"],\n", " \"D\": [\"D0\", \"D1\", \"D2\", \"D3\"],\n", " }\n", - ")\n", - "# Right merge\n", - "pd.merge(left, right, on=[\"key1\", \"key2\"], how=\"right\")" + ")" + ] + }, + { + "cell_type": "markdown", + "id": "cfe5aa30", + "metadata": {}, + "source": [ + "Right join:" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "595f41e6", + "metadata": {}, + "outputs": [], + "source": [ + "left.join(right, on=[\"key1\", \"key2\"], how=\"right\")" ] }, { @@ -160,7 +216,7 @@ "id": "8af0241e", "metadata": {}, "source": [ - "Note that the key combination of K2 and K0 did not exist in the left-hand data frame, and so its entries in the final data frame are NaNs. But it *does* have entries because we chose the keys from the right-hand data frame.\n", + "Note that the key combination of K2 and K0 did not exist in the left-hand data frame, and so its entries in the final data frame are null values. But it does have entries because we chose the keys from the right-hand data frame.\n", "\n", "What about an inner merge?" ] @@ -172,7 +228,7 @@ "metadata": {}, "outputs": [], "source": [ - "pd.merge(left, right, on=[\"key1\", \"key2\"], how=\"inner\")" + "left.join(right, on=[\"key1\", \"key2\"], how=\"inner\")" ] }, { @@ -182,7 +238,7 @@ "source": [ "Now we see that the combination K2 and K0 are excluded because they didn't exist in the overlap of keys in both data frames.\n", "\n", - "Finally, let's take a look at an outer merge that comes with some extra info via the `indicator` keyword:" + "Finally, let's take a look at an outer join:" ] }, { @@ -192,48 +248,158 @@ "metadata": {}, "outputs": [], "source": [ - "pd.merge(left, right, on=[\"key1\", \"key2\"], how=\"outer\", indicator=True)" + "left.join(right, on=[\"key1\", \"key2\"], how=\"outer\")" ] }, { "cell_type": "markdown", - "id": "67ca901b", + "id": "b984a202", + "metadata": {}, + "source": [ + "Polars doesn't include a built-in indicator parameter to track which side contributed each row in a join. However, you can achieve the same effect by comparing the two DataFrames after the join:" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "574c4508", "metadata": {}, + "outputs": [], "source": [ - "Now we can see that the products of all key combinations are here. The `indicator=True` option has caused an extra column to be added, called '_merge', that tells us which data frame the keys on that row came from.\n", + "# Perform an outer join\n", + "joined = left.join(right, on=[\"key1\", \"key2\"], how=\"outer\")\n", "\n", + "# Add indicator column manually\n", + "# Rows with nulls in B only came from right\n", + "# Rows with nulls in D only came from left\n", + "# Rows with no nulls came from both\n", + "joined = joined.with_columns(\n", + " pl.when(pl.col(\"B\").is_null() & pl.col(\"D\").is_not_null())\n", + " .then(pl.lit(\"right_only\"))\n", + " .when(pl.col(\"B\").is_not_null() & pl.col(\"D\").is_null())\n", + " .then(pl.lit(\"left_only\"))\n", + " .otherwise(pl.lit(\"both\"))\n", + " .alias(\"_merge\")\n", + ")\n", + "joined" + ] + }, + { + "cell_type": "markdown", + "id": "67ca901b", + "metadata": {}, + "source": [ "::: {.callout-tip title=\"Exercise\"}\n", "\n", - "Merge the following two data frames using the `left_on` and `right_on` keyword arguments to specify a join on `lkey` and `rkey` respectively:\n", + "Join the following two data frames using the `left_on` and `right_on` keyword arguments to specify a join on lkey and rkey respectively:\n", "\n", "```python\n", - "df1 = pd.DataFrame({'lkey': ['foo', 'bar', 'baz', 'foo'],\n", - " 'value': [1, 2, 3, 5]})\n", - "df2 = pd.DataFrame({'rkey': ['foo', 'bar', 'baz', 'foo'],\n", - " 'value': [5, 6, 7, 8]})\n", + "df1 = pl.DataFrame({\n", + " 'lkey': ['foo', 'bar', 'baz', 'foo'],\n", + " 'value': [1, 2, 3, 5]\n", + "})\n", + "df2 = pl.DataFrame({\n", + " 'rkey': ['foo', 'bar', 'baz', 'foo'],\n", + " 'value': [5, 6, 7, 8]\n", + "})\n", "```\n", ":::\n", "\n", "::: {.callout-tip title=\"Exercise\"}\n", "\n", - "Merge the following two data frames on `\"a\"` using `how=\"left\"` as a keyword argument:\n", + "Join the following two data frames on `\"a\"` using `how=\"left\"`:\n", "\n", "```python\n", - "df1 = pd.DataFrame({'a': ['foo', 'bar'], 'b': [1, 2]})\n", - "df2 = pd.DataFrame({'a': ['foo', 'baz'], 'c': [3, 4]})\n", + "df1 = pl.DataFrame({'a': ['foo', 'bar'], 'b': [1, 2]})\n", + "df2 = pl.DataFrame({'a': ['foo', 'baz'], 'c': [3, 4]})\n", "```\n", "\n", - "What do you notice about the position `.loc[1, \"c\"]` in the merged data frame? \n", - ":::\n", + "What do you notice about the resulting row for `\"baz\"`?\n", + ":::\n" + ] + }, + { + "cell_type": "markdown", + "id": "5d09e145", + "metadata": {}, + "source": [ + "## Semi and Anti Joins\n", + "\n", + "Polars also supports semi and anti joins, which are useful for filtering:\n", + "\n", + "Semi join: Keep rows from the left where keys appear in the right:" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "33207e8a", + "metadata": {}, + "outputs": [], + "source": [ + "# Semi join - only keeps left rows that have matches in right\n", + "left.join(right, on=[\"key1\", \"key2\"], how=\"semi\")" + ] + }, + { + "cell_type": "markdown", + "id": "d55603d0", + "metadata": {}, + "source": [ + "Anti join: Keep rows from the left where keys do NOT appear in the right:" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "92218fbb", + "metadata": {}, + "outputs": [], + "source": [ + "# Anti join - only keeps left rows without matches in right\n", + "left.join(right, on=[\"key1\", \"key2\"], how=\"anti\")" + ] + }, + { + "cell_type": "markdown", + "id": "355135ce", + "metadata": {}, + "source": [ + "## Handling Duplicate Column Names\n", + "\n", + "When joining, if both DataFrames have columns with the same name (other than the join keys), polars will add suffixes to differentiate them. By default it adds `_right`:" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "a042820b", + "metadata": {}, + "outputs": [], + "source": [ + "# Create DataFrames with overlapping column names\n", + "left2 = pl.DataFrame(\n", + " {\"key\": [\"K0\", \"K1\", \"K2\"], \"value\": [1, 2, 3], \"shared\": [\"L0\", \"L1\", \"L2\"]}\n", + ")\n", + "\n", + "right2 = pl.DataFrame(\n", + " {\"key\": [\"K0\", \"K1\", \"K2\"], \"value\": [4, 5, 6], \"shared\": [\"R0\", \"R1\", \"R2\"]}\n", + ")\n", "\n", - "For more on the options for merging, see **pandas**' comprehensive [merging documentation](https://pandas.pydata.org/docs/user_guide/merging.html#database-style-dataframe-or-named-series-joining-merging)." + "# Join with default suffix\n", + "left2.join(right2, on=\"key\", suffix=\"_right\")" + ] + }, + { + "cell_type": "markdown", + "id": "aa6e6071", + "metadata": {}, + "source": [ + "For more information about the available join types and options, see **polars'** [joining documentation](https://docs.pola.rs/api/python/stable/reference/index.html)." ] } ], "metadata": { - "interpreter": { - "hash": "9d7534ecd9fbc7d385378f8400cf4d6cb9c6175408a574f1c99c5269f08771cc" - }, "jupytext": { "cell_metadata_filter": "-all", "encoding": "# -*- coding: utf-8 -*-", @@ -241,7 +407,7 @@ "main_language": "python" }, "kernelspec": { - "display_name": "Python 3 (ipykernel)", + "display_name": ".venv", "language": "python", "name": "python3" },