Python Polars select filter and with_columns expressions

When working with Polars, three methods form the backbone of most data transformations: select, filter, and with_columns. These expressions let you choose columns, filter rows, and add or modify columns in a declarative, composable way. This article focuses on python polars select filter and with_columns expressions and shows how to use them effectively in real-world data pipelines.

Understanding Expressions in Polars

In Polars, an expression is a description of a computation that operates on a column or a set of columns. Unlike pandas, where you often write operations directly on the DataFrame, Polars expressions are lazy and can be combined, reused, and optimized by the query engine. The methods select, filter, and with_columns all accept expressions as arguments, which is why they are so powerful.

For example, pl.col("age") is an expression that refers to the age column. You can apply operations to it, such as pl.col("age").mean() or pl.col("age").alias("average_age"). Expressions can be combined with operators like +, -, &, |, and >.

Selecting Columns with select

The select method returns a new DataFrame containing only the columns you specify. It accepts one or more expressions, and each expression must produce a column. The simplest use is to pass column names as strings, but you can also use pl.col for more flexibility.

import polars as pl

df = pl.DataFrame({
    "name": ["Alice", "Bob", "Charlie"],
    "age": [25, 30, 35],
    "city": ["NYC", "LA", "SF"]
})

# Select by column name
df.select("name", "age")

# Select using pl.col
df.select(pl.col("name"), pl.col("age"))

Both produce the same result. You can also create new columns in select by using expressions. For instance, to get the age in months:

df.select(pl.col("name"), (pl.col("age") * 12).alias("age_months"))

select is ideal when you want to narrow down the DataFrame to specific columns or compute derived columns without keeping the original ones.

Filtering Rows with filter

The filter method keeps rows that satisfy a boolean expression. It accepts a single expression that evaluates to a boolean mask. You can combine conditions using & (and) and | (or), but remember to wrap each condition in parentheses because of operator precedence.

# Filter rows where age > 30
df.filter(pl.col("age") > 30)

# Multiple conditions with &
df.filter((pl.col("age") > 25) & (pl.col("city") == "NYC"))

filter returns a new DataFrame with only the matching rows. It does not modify the original DataFrame, as Polars DataFrames are immutable by default.

Adding or Modifying Columns with with_columns

The with_columns method adds new columns or replaces existing ones. It accepts one or more expressions, and each expression must produce a column. If the expression uses alias, the resulting column gets that name; otherwise, the expression's name is used.

# Add a new column
df.with_columns((pl.col("age") * 2).alias("age_double"))

# Replace an existing column
df.with_columns((pl.col("age") + 1).alias("age"))

You can add multiple columns at once:

df.with_columns([
    (pl.col("age") * 2).alias("age_double"),
    pl.col("city").str.to_uppercase().alias("city_upper")
])

with_columns is the go-to method when you need to keep all existing columns and add derived ones.

Combining select, filter, and with_columns in a Pipeline

These three methods are often used together in a data transformation pipeline. Because each returns a new DataFrame, you can chain them. For example, to filter rows, add a derived column, and then select a subset of columns:

result = (df
    .filter(pl.col("age") > 25)
    .with_columns((pl.col("age") * 12).alias("age_months"))
    .select("name", "age_months"))

This is a common pattern. In eager mode, each step is executed immediately. In lazy mode, the entire pipeline is optimized before execution.

Performance Considerations: Lazy vs Eager Evaluation

Polars offers two modes: eager and lazy. In eager mode, operations like select, filter, and with_columns are executed immediately. In lazy mode, you build a query plan and execute it at the end with collect(). Lazy evaluation allows Polars to optimize the entire pipeline, such as pushing filters down and reducing the amount of data read.

# Lazy pipeline
lazy_result = (df.lazy()
    .filter(pl.col("age") > 25)
    .with_columns((pl.col("age") * 12).alias("age_months"))
    .select("name", "age_months")
    .collect())

The lazy API uses the same expressions, but the execution is deferred. This can lead to significant performance improvements on large datasets because Polars can reorder operations and eliminate unnecessary computations. For small DataFrames, the overhead of lazy evaluation may not be worth it, but for production pipelines, it is often the preferred approach.

Common Mistakes and How to Avoid Them

One frequent mistake is using select when you want to keep all existing columns. select returns only the columns you specify, so if you need to add a column without dropping others, use with_columns. Another mistake is forgetting parentheses when combining conditions in filter. Because & and | have higher precedence than comparison operators, you must wrap each condition in parentheses.

# Wrong: missing parentheses
df.filter(pl.col("age") > 25 & pl.col("city") == "NYC")  # This will raise an error

# Correct
df.filter((pl.col("age") > 25) & (pl.col("city") == "NYC"))

Also, when using with_columns to replace a column, the new expression must have the same name as the original column. If you use alias with a different name, it will add a new column instead of replacing.

Advanced Expression Patterns

Expressions can be nested and combined to perform complex transformations. For example, you can use when and otherwise to create conditional columns:

df.with_columns(
    pl.when(pl.col("age") >= 30)
      .then(pl.lit("senior"))
      .otherwise(pl.lit("junior"))
      .alias("level")
)

You can also use expressions inside select to compute aggregates, though that often requires group_by. The key is that expressions are composable, so you can build complex logic without breaking the pipeline.

When working with multiple columns, you can use pl.col("*") to select all columns, or pl.col("^name|age$") to select by regex. This is useful in select and with_columns when you need to apply a transformation to a set of columns.