How to Add a Column to a Polars DataFrame Using .with_columns()

The .with_columns() method in Polars allows us to add one or more columns to a DataFrame. Unlike traditional methods that modify the DataFrame in place, .with_columns() returns a new DataFrame with the added columns, preserving immutability. This method is highly versatile, allowing us to create new columns based on existing ones, use constant values, or even apply complex transformations.

In this article, we’ll explore different methods to add a new column to a Polars DataFrame using .with_columns().

Install Polars

We first make sure that we have Polars installed in our system.

pip install polars

Basic Usage of .with_columns()

To get started, let's first understand the syntax of the .with_columns() method.

In this example, we created a new column "Age_in_5_Years" by adding 5 to the "Age" column. The .with_columns() method is passed a list of expressions, each representing a column to be added.

import polars as pl

# Create a sample DataFrame
df = pl.DataFrame({
    "Name": ["Alice", "Bob", "Charlie"],
    "Age": [25, 30, 35]
})

# Add a new column
new_df = df.with_columns([
    (pl.col("Age") + 5).alias("Age_in_5_Years")
])

print(new_df)

Output

Example 1: Adding a Constant Value Column

One of the simplest use cases for .with_columns() is adding a column with a constant value. This is useful when we need to add metadata or a static identifier to each row.

import polars as pl

# Create a sample DataFrame
df = pl.DataFrame({
    "Product": ["A", "B", "C"],
    "Price": [100, 150, 200]
})

# Add a constant value column
new_df = df.with_columns([
    pl.lit("₹").alias("Currency")
])

print(new_df)

Output

Here, we added a "Currency" column with a constant value "₹".

Example 2: Creating a Column Based on Multiple Existing Columns

We can create a new column by performing operations on multiple existing columns. For instance, let's create a "Total_Cost" column by multiplying the "Price" by a "Quantity" column.

import polars as pl

# Create a sample DataFrame
df = pl.DataFrame({
    "Product": ["A", "B", "C"],
    "Price": [100, 150, 200]
})

# Add a conditional column
new_df = df.with_columns([
    pl.when(pl.col("Price") > 150)
      .then(pl.lit("Expensive"))
      .otherwise(pl.lit("Affordable"))
      .alias("Category")
])

print(new_df)

Output

This example demonstrates how to create a "Total_Cost" column by multiplying "Price" and "Quantity".

Example 3: Conditional Column Creation

The .with_columns() method allows us to create columns conditionally based on existing data. Let's create a column "Category" that categorizes products as "Expensive" if their price is above 150 and "Affordable" otherwise.

import polars as pl

# Create a sample DataFrame
df = pl.DataFrame({
    "Product": ["A", "B", "C"],
    "Price": [100, 150, 200]
})

# Add a conditional column
new_df = df.with_columns([
    pl.when(pl.col("Price") > 150)
      .then("Expensive")
      .otherwise("Affordable")
      .alias("Category")
])

print(new_df)

Output

Here, we categorized products based on their price.

Example 4: Adding Multiple Columns Simultaneously

We can add multiple columns in one go using the .with_columns() method by passing a list of expressions.

import polars as pl

# Create a sample DataFrame
df = pl.DataFrame({
    "Product": ["A", "B", "C"],
    "Price": [100, 150, 200]
})

# Add multiple columns
new_df = df.with_columns([
    (pl.col("Price") * 1.1).alias("Price_with_Tax"),
    (pl.col("Price") + 50).alias("Discounted_Price")
])

print(new_df)

Output

This example shows how to add multiple columns: "Price_with_Tax" and "Discounted_Price".

Example 5: Adding a Column with a Custom Function

Sometimes, We may need to add a column based on a custom function. We can achieve this using the .map_elements() method inside .with_columns().

import polars as pl

# Create a sample DataFrame
df = pl.DataFrame({
    "Name": ["Alice", "Bob", "Charlie"],
    "Age": [25, 30, 35]
})

# Define a custom function
def age_category(age):
    if age < 30:
        return "Young"
    else:
        return "Mature"

# Add a column using a custom function
new_df = df.with_columns([
    pl.col("Age").map_elements(age_category).alias("Category")
])

print(new_df)

Output

This example adds a "Category" column based on a custom function that categorizes people by age.

Conclusion

The .with_columns() method in Polars is a powerful and flexible way to add new columns to a DataFrame. Whether we are adding constant values, performing calculations, creating conditional columns, or applying custom functions, .with_columns() provides an intuitive and efficient interface for DataFrame manipulation. With the examples provided, we can now confidently add columns to our Polars DataFrames in a variety of scenarios.