dataframedata-cleaningpython-polars

How can I replace multiple rows in a polars dataframe with a dict mapping?


So lets say I have an excel sheet/csv of users and their salaries records. I have created an account for each user in a db and want to create salary record for each using the ids from the db.

import polars as pl

# Create the DataFrame with repeated entries for John and Jane
df = pl.DataFrame({
    "first_name": ["John", "Jane", "Alice", "Bob", "John", "Jane", "John", "Jane"],
    "middle_name": ["A.", "B.", "C.", "D.", "A.", "B.", "A.", "B."],
    "last_name": ["Doe", "Smith", "Johnson", "Brown", "Doe", "Smith", "Doe", "Smith"],
    "salary": [50000, 60000, 55000, 62000, 50000, 60000, 50000, 60000],
    "date": ["2023-01-15", "2023-02-20", "2023-03-05", "2023-04-10", "2023-05-15", "2023-06-20", "2023-07-15", "2023-08-20"]
})
print(df)
shape: (8, 5)
┌────────────┬────────────┬────────────┬────────┬────────────┐
│ first_name │ middle_name│ last_name  │ salary │ date       │
│ ---        │ ---        │ ---        │ ---    │ ---        │
│ str        │ str        │ str        │ i64    │ date       │
├────────────┼────────────┼────────────┼────────┼────────────┤
│ John       │ A.         │ Doe        │ 50000  │ 2023-01-15 │
│ Jane       │ B.         │ Smith      │ 60000  │ 2023-02-20 │
│ Alice      │ C.         │ Johnson    │ 55000  │ 2023-03-05 │
│ Bob        │ D.         │ Brown      │ 62000  │ 2023-04-10 │
│ John       │ A.         │ Doe        │ 50000  │ 2023-05-15 │
│ Jane       │ B.         │ Smith      │ 60000  │ 2023-06-20 │
│ John       │ A.         │ Doe        │ 50000  │ 2023-07-15 │
│ Jane       │ B.         │ Smith      │ 60000  │ 2023-08-20 │
└────────────┴────────────┴────────────┴────────┴────────────┘

#Get unique values
 subset_df = df.select(["first_name", "middle_name", "last_name"])
 unique_subset_df = subset_df.unique()
 for row in subset_df.select(pl.struct(pl.all()).value_counts()):
    # create acoount

The user list with their corresponding ids is as follows

users = [
    {'id': 1, 'first_name': 'John', 'middle_name': 'A.', 'last_name': 'Doe'},
    {'id': 2, 'first_name': 'Jane', 'middle_name': 'B.', 'last_name': 'Smith'},
    {'id': 3, 'first_name': 'Alice', 'middle_name': 'C.', 'last_name': 'Johnson'},
    {'id': 4, 'first_name': 'Bob', 'middle_name': 'D.', 'last_name': 'Brown'}
]
# Note:The data above can also be transformed into a list of tuples

How can I replace the values in the first_name, middle_name, and last_name columns of the dataframe with their corresponding IDs from this list of dictionaries?


Solution

  • You need to join your df with your users like this

    df.join(pl.DataFrame(users), on=['first_name', 'middle_name', 'last_name'])
    

    If you want the id column to be first then you can put a select at the end.

    (
        df
        .join(pl.DataFrame(users), on=['first_name', 'middle_name', 'last_name'])
        .select('id',pl.exclude('id'))
    )
    
    shape: (8, 6)
    ┌─────┬────────────┬─────────────┬───────────┬────────┬────────────┐
    │ id  ┆ first_name ┆ middle_name ┆ last_name ┆ salary ┆ date       │
    │ --- ┆ ---        ┆ ---         ┆ ---       ┆ ---    ┆ ---        │
    │ i64 ┆ str        ┆ str         ┆ str       ┆ i64    ┆ str        │
    ╞═════╪════════════╪═════════════╪═══════════╪════════╪════════════╡
    │ 1   ┆ John       ┆ A.          ┆ Doe       ┆ 50000  ┆ 2023-01-15 │
    │ 2   ┆ Jane       ┆ B.          ┆ Smith     ┆ 60000  ┆ 2023-02-20 │
    │ 3   ┆ Alice      ┆ C.          ┆ Johnson   ┆ 55000  ┆ 2023-03-05 │
    │ 4   ┆ Bob        ┆ D.          ┆ Brown     ┆ 62000  ┆ 2023-04-10 │
    │ 1   ┆ John       ┆ A.          ┆ Doe       ┆ 50000  ┆ 2023-05-15 │
    │ 2   ┆ Jane       ┆ B.          ┆ Smith     ┆ 60000  ┆ 2023-06-20 │
    │ 1   ┆ John       ┆ A.          ┆ Doe       ┆ 50000  ┆ 2023-07-15 │
    │ 2   ┆ Jane       ┆ B.          ┆ Smith     ┆ 60000  ┆ 2023-08-20 │
    └─────┴────────────┴─────────────┴───────────┴────────┴────────────┘