In Python polars convert a json string column to dict for filtering

Question

Hi have a dataframe where I have a column called tags which is a json string.

I want to filter this dataframe on the tags column so it only contains rows where a certain tag key is present or where a tag has a particular value.

I guess I could do a string contains match but think it may be more robust to have the json convert into a dict first and using has_key etc ?

What would be the recommended way to do this in python polars ?

Thanks

Answer 1

Polars does not have a generic dictionary type. Instead, dictionaries are imported/mapped as structs. Each dictionary key is mapped to a struct 'field name', and the corresponding dictionary value becomes the value of this field.

However, there are some constraints for creating a Series of type struct. Two of them are:

• all structs must have the same field names.
• the field names must be listed in the same order.

In your description, you mention has_key, which indicates that the dictionaries will not have the same keys. As such, creating a column of struct from your dictionaries will not work. (For more information, you can see this Stack Overflow response.)

json_path_match

I suggest using json_path_match, which extracts values based on some simple JSONPath syntax. Using JSONPath syntax, you should be able to query whether a key exists, and retrieve it's value. (For simple unnested dictionaries, these are the same query.)

For example, let's start with this data:

import polars as pl

json_list = [
    """{"name": "Maria",
        "position": "developer",
        "office": "Seattle"}""",
    """{"name": "Josh",
        "position": "analyst",
        "termination_date": "2020-01-01"}""",
    """{"name": "Jorge",
        "position": "architect",
        "office": "",
        "manager_st_dt": "2020-01-01"}""",
]

df = pl.DataFrame(
    {
        "tags": json_list,
    }
).with_row_count("id", 1)
df

shape: (3, 2)
┌─────┬────────────────────┐
│ id  ┆ tags               │
│ --- ┆ ---                │
│ u32 ┆ str                │
╞═════╪════════════════════╡
│ 1   ┆ {"name": "Maria",  │
│     ┆         "posit...  │
├╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┤
│ 2   ┆ {"name": "Josh",   │
│     ┆         "positi... │
├╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┤
│ 3   ┆ {"name": "Jorge",  │
│     ┆         "posit...  │
└─────┴────────────────────┘

To query for values:

df.with_columns([
    pl.col('tags').str.json_path_match(r"$.name").alias('name'),
    pl.col('tags').str.json_path_match(r"$.office").alias('location'),
    pl.col('tags').str.json_path_match(r"$.manager_st_dt").alias('manager start date'),
])

shape: (3, 5)
┌─────┬────────────────────┬───────┬──────────┬────────────────────┐
│ id  ┆ tags               ┆ name  ┆ location ┆ manager start date │
│ --- ┆ ---                ┆ ---   ┆ ---      ┆ ---                │
│ u32 ┆ str                ┆ str   ┆ str      ┆ str                │
╞═════╪════════════════════╪═══════╪══════════╪════════════════════╡
│ 1   ┆ {"name": "Maria",  ┆ Maria ┆ Seattle  ┆ null               │
│     ┆         "posit...  ┆       ┆          ┆                    │
├╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┤
│ 2   ┆ {"name": "Josh",   ┆ Josh  ┆ null     ┆ null               │
│     ┆         "positi... ┆       ┆          ┆                    │
├╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┤
│ 3   ┆ {"name": "Jorge",  ┆ Jorge ┆          ┆ 2020-01-01         │
│     ┆         "posit...  ┆       ┆          ┆                    │
└─────┴────────────────────┴───────┴──────────┴────────────────────┘

Notice the null values. This is the return value when a key is not found. We'll use this fact for the has_key functionality you mentioned.

Also, if we look at the "location" column, you'll see that json_path_match does distinguish between an empty string "office":"" and a key not found..

To filter for the presence of a key, we simply filter for null values.

df.filter(
    pl.col('tags').str.json_path_match(r"$.manager_st_dt").is_not_null()
)

shape: (1, 2)
┌─────┬───────────────────┐
│ id  ┆ tags              │
│ --- ┆ ---               │
│ u32 ┆ str               │
╞═════╪═══════════════════╡
│ 3   ┆ {"name": "Jorge", │
│     ┆         "posit... │
└─────┴───────────────────┘

The json_path_match will also work with nested structures. (See the Syntax page for details.)

One limitation, however: json_path_match will only return the first match for a query, rather than a list of matches. If your JSON strings are not lists or nested dictionaries, this won't be a problem.