pythonpython-polars

Polars Create Column with String Formatting


I have a polars dataframe:

df = pl.DataFrame({'schema_name': ['test_schema', 'test_schema_2'], 
                       'table_name': ['test_table', 'test_table_2'],
                       'column_name': ['test_column, test_column_2','test_column']})
schema_name table_name column_name
test_schema test_table test_column, test_column_2
test_schema_2 test_table_2 test_column

I have a string:

date_field_value_max_query = '''
    select '{0}' as schema_name, 
           '{1}' as table_name, 
           greatest({2})
    from {0}.{1}
    group by 1, 2
'''

I would like to use polars to add a column by using string formatting. The target dataframe is this:

schema_name table_name column_name query
test_schema test_table test_column, test_column_2 select test_schema, test_table, greatest(test_column, test_column_2) from test_schema.test_table group by 1, 2
test_schema_2 test_table_2 test_column select test_schema_2, test_table_2, greatest(test_column) from test_schema_2.test_table_2 group by 1, 2

In pandas, I would do something like this:

df.apply(lambda row: date_field_value_max_query.format(row['schema_name'], row['table_name'], row['column_name']), axis=1)

For polars, I tried this:

df.map_rows(lambda row: date_field_value_max_query.format(row[0], row[1], row[2]))

...but this returns only the one column, and I lose the original three columns. I know this approach is also not recommended for polars, when possible.

How can I perform string formatting across multiple dataframe columns with the output column attached to the original dataframe?


Solution

  • Another option is to use polars.format to create your string. For example:

    date_field_value_max_query = (
    '''select {} as schema_name,
           {} as table_name,
           greatest({})
        from {}.{}
        group by 1, 2
    '''
    )
    
    (
        df
        .with_columns(
            pl.format(date_field_value_max_query,
                      'schema_name',
                      'table_name',
                      'column_name',
                      'schema_name',
                      'table_name')
        )
    )
    
    shape: (2, 4)
    ┌───────────────┬──────────────┬────────────────────────────┬─────────────────────────────────────────────┐
    │ schema_name   ┆ table_name   ┆ column_name                ┆ literal                                     │
    │ ---           ┆ ---          ┆ ---                        ┆ ---                                         │
    │ str           ┆ str          ┆ str                        ┆ str                                         │
    ╞═══════════════╪══════════════╪════════════════════════════╪═════════════════════════════════════════════╡
    │ test_schema   ┆ test_table   ┆ test_column, test_column_2 ┆ select test_schema as schema_name,          │
    │               ┆              ┆                            ┆        test_table as table_name,            │
    │               ┆              ┆                            ┆        greatest(test_column, test_column_2) │
    │               ┆              ┆                            ┆     from test_schema.test_table             │
    │               ┆              ┆                            ┆     group by 1, 2                           │
    │               ┆              ┆                            ┆                                             │
    │               ┆              ┆                            ┆                                             │
    │ test_schema_2 ┆ test_table_2 ┆ test_column                ┆ select test_schema_2 as schema_name,        │
    │               ┆              ┆                            ┆        test_table_2 as table_name,          │
    │               ┆              ┆                            ┆        greatest(test_column)                │
    │               ┆              ┆                            ┆     from test_schema_2.test_table_2         │
    │               ┆              ┆                            ┆     group by 1, 2                           │
    │               ┆              ┆                            ┆                                             │
    │               ┆              ┆                            ┆                                             │
    └───────────────┴──────────────┴────────────────────────────┴─────────────────────────────────────────────┘