I have 3 polars dataframes, one that contains 2 IDS, and the other ones contains an ID and a value. I would like to join the 3 dataframes if the ID of the main table exists on one of the other tables and bring a values from a desired column.
My current aproach its just rename each Table ID and then do a .join(how = 'left'), however, i think renaming and duplicate tables its not the correct way to approach this problem. (Due the extra code, and the wasted ram)
The first one contains 2 ID columns:
data = {
"ID1" : [1,2,3],
"ID2" : [1,4,5]
}
df = pl.DataFrame(data)
The second and third are dataframes than contains an ID and a value:
T1 = {
"ID" : [9,2,5],
"Values" : ["A","B","c"],
"Values II" : ["foo","boo","baz"]
}
T1 = pl.DataFrame(T1)
T2 = {
"ID" : [1,4,10],
"Values" : ["X","J","c"]
}
T2 = pl.DataFrame(T2)
I can check if the ID exists on the other tables like this
(
df
.with_columns(
ID1_is_on_T1 = pl.col("ID1").is_in(T1.select(pl.col("ID"))),
ID2_is_on_T1 = pl.col("ID2").is_in(T1.select(pl.col("ID"))),
ID1_is_on_T2 = pl.col("ID1").is_in(T2.select(pl.col("ID"))),
ID2_is_on_T2 = pl.col("ID2").is_in(T2.select(pl.col("ID"))),
)
)
And i’m looking to do somehting like this:
(
df
.with_columns(
pl
.when(
pl.col("ID1").is_in(T1.select(pl.col("ID")))
)
.then(
T1.select(pl.col("Values"))
)
.otherwise(0)
)
)
ValueError: can only call .item() if the dataframe is of shape (1, 1), or if explicit row/col values are provided; frame has shape (3, 1)
Current .join() approach:
T1_1 = (
T1
.rename(
{"ID": "ID1"}
)
)
T1_2 = (
T1
.rename(
{"ID": "ID2"}
)
)
Join_1 = df.join(T1_1,on = "ID1", how="left").rename({"Values" : "ID1_Values", "Values II" : "ID1_Values II"})
Join_2 = Join_1.join(T1_2, on = "ID2", how="left").rename({"Values" : "ID2_Values", "Values II" : "ID2_Values II"})
On this approach its only considering the first table, i would need to do the same for the T2 too.
4
Two joins will be most efficient here. You can avoid excessive renaming by specifing left_on and right_on separately, as well as using a suffix to clarify the duplicate names.
>>> (df.join(T1, left_on="ID1", right_on="ID", how="left")
... .join(T2, left_on="ID2", right_on="ID", how="left", suffix="_T2"))
shape: (3, 5)
┌─────┬─────┬────────┬───────────┬───────────┐
│ ID1 ┆ ID2 ┆ Values ┆ Values II ┆ Values_T2 │
│ --- ┆ --- ┆ --- ┆ --- ┆ --- │
│ i64 ┆ i64 ┆ str ┆ str ┆ str │
╞═════╪═════╪════════╪═══════════╪═══════════╡
│ 1 ┆ 1 ┆ null ┆ null ┆ X │
│ 2 ┆ 4 ┆ B ┆ boo ┆ J │
│ 3 ┆ 5 ┆ null ┆ null ┆ null │
└─────┴─────┴────────┴───────────┴───────────┘
If instead of storing T1 and T2 as their own variables, you put them in a dictionary, you could make a function that does it all for you.
Setup variables like this:
df = pl.DataFrame({
"ID1" : [1,2,3],
"ID2" : [1,4,5]
})
T={}
T["T1"]=pl.DataFrame({
"ID" : [9,2,5],
"Values" : ["A","B","c"],
"Values II" : ["foo","boo","baz"]
}
)
T["T2"]=pl.DataFrame({
"ID" : [1,4,10],
"Values" : ["X","J","c"]
})
And then the function
def big_join(df:pl.DataFrame, T:dict[str, pl.DataFrame])->pl.DataFrame:
joined=df
for id_col in df.columns:
for tbl_name, subdf in T.items():
joined=joined.join(
subdf.select("ID", pl.exclude("ID").name.prefix(f"{id_col}_{tbl_name}_")),
left_on=id_col, right_on='ID', how='left'
)
return joined
When you want to join them all up, just do
big_join(df, T)
shape: (3, 8)
┌─────┬─────┬──────────────┬──────────────┬──────────────┬─────────────┬─────────────┬─────────────┐
│ ID1 ┆ ID2 ┆ ID1_T1_Value ┆ ID1_T1_Value ┆ ID1_T2_Value ┆ ID2_T1_Valu ┆ ID2_T1_Valu ┆ ID2_T2_Valu │
│ --- ┆ --- ┆ s ┆ s II ┆ s ┆ es ┆ es II ┆ es │
│ i64 ┆ i64 ┆ --- ┆ --- ┆ --- ┆ --- ┆ --- ┆ --- │
│ ┆ ┆ str ┆ str ┆ str ┆ str ┆ str ┆ str │
╞═════╪═════╪══════════════╪══════════════╪══════════════╪═════════════╪═════════════╪═════════════╡
│ 1 ┆ 1 ┆ null ┆ null ┆ X ┆ null ┆ null ┆ X │
│ 2 ┆ 4 ┆ B ┆ boo ┆ null ┆ null ┆ null ┆ J │
│ 3 ┆ 5 ┆ null ┆ null ┆ null ┆ c ┆ baz ┆ null │
└─────┴─────┴──────────────┴──────────────┴──────────────┴─────────────┴─────────────┴─────────────┘
This is extensible to any number of T tables so long as they’re in the dictionary. Even if you wanted to maintain T1 and T2 as their own variables you could do big_join(df, {"T1":T1, "T2":T2}) instead.
Polars doesn’t copy unless it needs to so you don’t need to worry about wasted ram. For example, if you do df1=df2=df3=df it doesn’t copy df 3 times, it just has a pointer in each of the “copies” pointing back to df. The same is true for columns so if you left join df to T1 then it isn’t copying df to the new join table, it just points to the columns. I’m not a good authority on when exactly copies are made/needed but it is far fewer than other bear based DataFrame libraries.