[Python] Polars Python: Filter list column using a boolean list column, but keeping list size

I would like to get elements from a list dtype column using another boolean list column and keeping the original size of the list (as oppose to this solution).

Starting from this dataframe:

df = pl.DataFrame({
'identity_vector': [[True, False], [False, True]],
'string_vector': [['name1', 'name2'], ['name3', 'name4']]
})

shape: (2, 2)
┌─────────────────┬────────────────────┐
│ identity_vector ┆ string_vector │
│ --- ┆ --- │
│ list[bool] ┆ list[str] │
╞═════════════════╪════════════════════╡
│ [true, false] ┆ ["name1", "name2"] │
│ [false, true] ┆ ["name3", "name4"] │
└─────────────────┴────────────────────┘


The objective is to get this output:

shape: (2, 3)
┌─────────────────┬────────────────────┬──────────────────┐
│ identity_vector ┆ string_vector ┆ filtered_strings │
│ --- ┆ --- ┆ --- │
│ list[bool] ┆ list[str] ┆ list[str] │
╞═════════════════╪════════════════════╪══════════════════╡
│ [true, false] ┆ ["name1", "name2"] ┆ ["name1", null] │
│ [false, true] ┆ ["name3", "name4"] ┆ [null, "name4"] │
└─────────────────┴────────────────────┴──────────────────┘



Which I can get using the block of code below and map_elements, but the solution is sub-optimal for performance reasons:

df.with_columns(
filtered_strings=pl.struct(["string_vector", "identity_vector"]).map_elements(
lambda row: [s if keep else None for s, keep in zip(row["string_vector"], row["identity_vector"])]
)
)



Do you have any suggestion on how to improve the performance of this process?

Continue reading...