Select all columns where column name starts with string

Question

Given the following dataframe, is there some way to select only columns starting with a given prefix? I know I could do e.g. pl.col(column) for column in df.columns if column.startswith("prefix_"), but I’m wondering if I can do it as part of a single expression.

df = pl.DataFrame(
    {"prefix_a": [1, 2, 3], "prefix_b": [1, 2, 3], "some_column": [3, 2, 1]}
)
df.select(pl.all().<column_name_starts_with>("prefix_"))

Would this be possible to do lazily?

Asked By: TomNorway

||

Source

Answer 1

From the documentation for polars.col, the expression can take one of the following arguments:

a single column by name

all columns by using a wildcard “*”

column by regular expression if the regex starts with ^ and ends with $

So in this case, we can use a regex expression to select for the prefix. And this does work in lazy mode.

(
    df
    .lazy()
    .select(pl.col('^prefix_.*$'))
    .collect()
)

>>> (
...     df
...     .lazy()
...     .select(pl.col('^prefix_.*$'))
...     .collect()
... 
... )
shape: (3, 2)
┌──────────┬──────────┐
│ prefix_a ┆ prefix_b │
│ ---      ┆ ---      │
│ i64      ┆ i64      │
╞══════════╪══════════╡
│ 1        ┆ 1        │
├╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌┤
│ 2        ┆ 2        │
├╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌┤
│ 3        ┆ 3        │
└──────────┴──────────┘

Note: we can also use polars.exclude with regex expressions:

(
    df
    .lazy()
    .select(pl.exclude('^prefix_.*$'))
    .collect()
)

shape: (3, 1)
┌─────────────┐
│ some_column │
│ ---         │
│ i64         │
╞═════════════╡
│ 3           │
├╌╌╌╌╌╌╌╌╌╌╌╌╌┤
│ 2           │
├╌╌╌╌╌╌╌╌╌╌╌╌╌┤
│ 1           │
└─────────────┘

Answered By: user18559875

Select all columns where column name starts with string

Question:

Answers: