Joins

Introduction

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.

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.

Prerequisites

This chapter will use the polars data analysis package. For reading Stata (.dta) files, we’ll use pandas as a bridge (pd.read_stata()) and then convert to polars with pl.from_pandas().

Concatenate

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().

For the same columns, use pl.concat() to stack rows vertically. Polars doesn’t have a row index like pandas; pl.concat(how="horizontal") joins DataFrames by position, aligning rows in order.

Here’s an example using data on two different states’ populations:

import pandas as pd
import polars as pl

base_url = (
    "https://github.com/aeturrell/coding-for-economists/raw/refs/heads/main/data/"
)
state_codes = ["ca", "il"]
end_url = "pop.dta"

# Read the Stata files and convert to polars
list_of_state_dfs = []
for state in state_codes:
    temp_df = pd.read_stata(base_url + state + end_url)
    list_of_state_dfs.append(pl.from_pandas(temp_df))

# shows example of the 2 dataframes
print(list_of_state_dfs[0])
print(list_of_state_dfs[1])

# concatenate the list of dataframes
df = pl.concat(list_of_state_dfs)
df
shape: (3, 2)
┌─────────────┬─────────┐
│ county      ┆ pop     │
│ ---         ┆ ---     │
│ str         ┆ i32     │
╞═════════════╪═════════╡
│ Los Angeles ┆ 9878554 │
│ Orange      ┆ 2997033 │
│ Ventura     ┆ 798364  │
└─────────────┴─────────┘
shape: (3, 2)
┌────────┬─────────┐
│ county ┆ pop     │
│ ---    ┆ ---     │
│ str    ┆ i32     │
╞════════╪═════════╡
│ Cook   ┆ 5285107 │
│ DeKalb ┆ 103729  │
│ Will   ┆ 673586  │
└────────┴─────────┘
shape: (6, 2)
county pop
str i32
"Los Angeles" 9878554
"Orange" 2997033
"Ventura" 798364
"Cook" 5285107
"DeKalb" 103729
"Will" 673586

Note that when concatenating, you can track the origin of rows by adding a column before concatenation:

# Add a state identifier column to each dataframe
list_with_state = []
for state, df_state in zip(state_codes, list_of_state_dfs):
    df_state = df_state.with_columns(pl.lit(state).alias("state"))
    list_with_state.append(df_state)

# Concatenate with state identifiers
df_with_state = pl.concat(list_with_state)
df_with_state
shape: (6, 3)
county pop state
str i32 str
"Los Angeles" 9878554 "ca"
"Orange" 2997033 "ca"
"Ventura" 798364 "ca"
"Cook" 5285107 "il"
"DeKalb" 103729 "il"
"Will" 673586 "il"
TipExercise

Concatenate the follow two data frames:

df1 = pl.DataFrame({
    'letter': ['a', 'b'],
    'number': [1, 2]
})

df2 = pl.DataFrame({
    'letter': ['c', 'd'],
    'number': [3, 4]

Joins

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:

The how= keyword works in the following ways:

  • how=“left” uses keys from the left data frame only to join.

  • how=“right” uses keys from the right data frame only to join.

  • how=“inner” uses keys that appear in both data frames to join.

  • how=“full” uses all keys from both data frames to join on (full outter join).

  • how=“semi” returns rows from the left where keys appear in the right.

  • how=“anti” returns rows from the left where keys do NOT appear in the right.

Let’s see examples of some of these:

left = pl.DataFrame(
    {
        "key1": ["K0", "K0", "K1", "K2"],
        "key2": ["K0", "K1", "K0", "K1"],
        "A": ["A0", "A1", "A2", "A3"],
        "B": ["B0", "B1", "B2", "B3"],
    }
)
right = pl.DataFrame(
    {
        "key1": ["K0", "K1", "K1", "K2"],
        "key2": ["K0", "K0", "K0", "K0"],
        "C": ["C0", "C1", "C2", "C3"],
        "D": ["D0", "D1", "D2", "D3"],
    }
)

Right join:

left.join(right, on=["key1", "key2"], how="right")
shape: (4, 6)
A B key1 key2 C D
str str str str str str
"A0" "B0" "K0" "K0" "C0" "D0"
"A2" "B2" "K1" "K0" "C1" "D1"
"A2" "B2" "K1" "K0" "C2" "D2"
null null "K2" "K0" "C3" "D3"

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.

What about an inner join?

left.join(right, on=["key1", "key2"], how="inner")
shape: (3, 6)
key1 key2 A B C D
str str str str str str
"K0" "K0" "A0" "B0" "C0" "D0"
"K1" "K0" "A2" "B2" "C1" "D1"
"K1" "K0" "A2" "B2" "C2" "D2"

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.

Finally, let’s take a look at an outer join:

left.join(right, on=["key1", "key2"], how="full")
shape: (6, 8)
key1 key2 A B key1_right key2_right C D
str str str str str str str str
"K0" "K0" "A0" "B0" "K0" "K0" "C0" "D0"
"K1" "K0" "A2" "B2" "K1" "K0" "C1" "D1"
"K1" "K0" "A2" "B2" "K1" "K0" "C2" "D2"
null null null null "K2" "K0" "C3" "D3"
"K0" "K1" "A1" "B1" null null null null
"K2" "K1" "A3" "B3" null null null null

This performs a full outer join, keeping all rows from both data frames.

Polars doesn’t include a built-in indicator parameter to track which side contributed each row in a join. On a full join, polars keeps the join keys from both sides with suffixes. However, you can achieve the same effect by comparing the two DataFrames after the join:

# Perform an outer join
joined = left.join(right, on=["key1", "key2"], how="full")

# Add indicator column manually
# Rows with nulls in B only came from right
# Rows with nulls in D only came from left
# Rows with no nulls came from both
joined = joined.with_columns(
    pl.when(pl.col("B").is_null() & pl.col("D").is_not_null())
    .then(pl.lit("right_only"))
    .when(pl.col("B").is_not_null() & pl.col("D").is_null())
    .then(pl.lit("left_only"))
    .otherwise(pl.lit("both"))
    .alias("_merge")
)
joined
shape: (6, 9)
key1 key2 A B key1_right key2_right C D _merge
str str str str str str str str str
"K0" "K0" "A0" "B0" "K0" "K0" "C0" "D0" "both"
"K1" "K0" "A2" "B2" "K1" "K0" "C1" "D1" "both"
"K1" "K0" "A2" "B2" "K1" "K0" "C2" "D2" "both"
null null null null "K2" "K0" "C3" "D3" "right_only"
"K2" "K1" "A3" "B3" null null null null "left_only"
"K0" "K1" "A1" "B1" null null null null "left_only"
TipExercise

Join the following two data frames using the left_on and right_on keyword arguments to specify a join on lkey and rkey respectively:

df1 = pl.DataFrame({
    'lkey': ['foo', 'bar', 'baz', 'foo'],
    'value': [1, 2, 3, 5]
})
df2 = pl.DataFrame({
    'rkey': ['foo', 'bar', 'baz', 'foo'],
    'value': [5, 6, 7, 8]
})
TipExercise

Join the following two data frames on "a" using how="left":

df1 = pl.DataFrame({'a': ['foo', 'bar'], 'b': [1, 2]})
df2 = pl.DataFrame({'a': ['foo', 'baz'], 'c': [3, 4]})

What do you notice about the resulting row for "baz"?

Semi and Anti Joins

Polars also supports semi and anti joins, which are useful for filtering:

Semi join: Keep rows from the left where keys appear in the right:

# Semi join - only keeps left rows that have matches in right
left.join(right, on=["key1", "key2"], how="semi")
shape: (2, 4)
key1 key2 A B
str str str str
"K0" "K0" "A0" "B0"
"K1" "K0" "A2" "B2"

Anti join: Keep rows from the left where keys do NOT appear in the right:

# Anti join - only keeps left rows without matches in right
left.join(right, on=["key1", "key2"], how="anti")
shape: (2, 4)
key1 key2 A B
str str str str
"K0" "K1" "A1" "B1"
"K2" "K1" "A3" "B3"

Handling Duplicate Column Names

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:

# Create DataFrames with overlapping column names
left2 = pl.DataFrame(
    {"key": ["K0", "K1", "K2"], "value": [1, 2, 3], "shared": ["L0", "L1", "L2"]}
)

right2 = pl.DataFrame(
    {"key": ["K0", "K1", "K2"], "value": [4, 5, 6], "shared": ["R0", "R1", "R2"]}
)

# Join with default suffix
left2.join(right2, on="key", suffix="_right")
shape: (3, 5)
key value shared value_right shared_right
str i64 str i64 str
"K0" 1 "L0" 4 "R0"
"K1" 2 "L1" 5 "R1"
"K2" 3 "L2" 6 "R2"

For more information about the available join types and options, see polars’ joining documentation.