How to use join in expression context?

NedDasty 0
Viewed 53

Suppose I have a mapping dataframe that I would like to join to an original dataframe:

df = pl.DataFrame({
    'A': [1, 2, 3, 2, 1],
})

mapper = pl.DataFrame({
    'key': [1, 2, 3, 4, 5],
    'value': ['a', 'b', 'c', 'd', 'e']
})

I can map A to value directly via df.join(mapper, ...), but is there a way to do this in an expression context, i.e. while building columns? As in:

df.with_columns([
    (pl.col('A')+1).join(mapper, left_on='A', right_on='key')
])

With would furnish:

shape: (5, 2)
┌─────┬───────┐
│ A   ┆ value │
│ --- ┆ ---   │
│ i64 ┆ str   │
╞═════╪═══════╡
│ 1   ┆ b     │
├╌╌╌╌╌┼╌╌╌╌╌╌╌┤
│ 1   ┆ b     │
├╌╌╌╌╌┼╌╌╌╌╌╌╌┤
│ 2   ┆ c     │
├╌╌╌╌╌┼╌╌╌╌╌╌╌┤
│ 2   ┆ c     │
├╌╌╌╌╌┼╌╌╌╌╌╌╌┤
│ 3   ┆ d     │
└─────┴───────┘
python-polars
1 Answers

Probably, yes. I just putted df.select(col('A')+1) inside.

df = df.with_columns([
    col('A'),
    df.select(col('A')+1).join(mapper, left_on='A', right_on='key')['value']
])

print(df)

df

┌─────┬───────┐
│ A   ┆ value │
│ --- ┆ ---   │
│ i64 ┆ str   │
╞═════╪═══════╡
│ 1   ┆ b     │
├╌╌╌╌╌┼╌╌╌╌╌╌╌┤
│ 2   ┆ b     │
├╌╌╌╌╌┼╌╌╌╌╌╌╌┤
│ 3   ┆ c     │
├╌╌╌╌╌┼╌╌╌╌╌╌╌┤
│ 2   ┆ c     │
├╌╌╌╌╌┼╌╌╌╌╌╌╌┤
│ 1   ┆ d     │
└─────┴───────┘
glebcom0
Related