One of the most common operations in data analysis is adding new columns to a DataFrame. In Polars, this can be accomplished using the with_columns() method. This blog post will guide you through the process of adding new columns to a Polars DataFrame, demonstrating various use cases and techniques.

Basic Usage of with_columns()

The with_columns() method in Polars allows you to add one or more new columns to a DataFrame. It’s a versatile method that can handle simple constant values, calculations based on existing columns, and complex expressions.

Let’s start with a simple example:

import polars as pl
# Create a sample DataFrame
df = pl.DataFrame({
    "name": ["Alice", "Bob", "Charlie"],
    "age": [25, 30, 35]
})
# Add a new column with a constant value
df_new = df.with_columns(pl.lit("USA").alias("country"))
print(df)
print(df_new)

Output:

shape: (3, 2)
┌─────────┬─────┐
│ name    ┆ age │
│ ---     ┆ --- │
│ str     ┆ i64 │
╞═════════╪═════╡
│ Alice   ┆ 25  │
│ Bob     ┆ 30  │
│ Charlie ┆ 35  │
└─────────┴─────┘
shape: (3, 3)
┌─────────┬─────┬─────────┐
│ name    ┆ age ┆ country │
│ ---     ┆ --- ┆ ---     │
│ str     ┆ i64 ┆ str     │
╞═════════╪═════╪═════════╡
│ Alice   ┆ 25  ┆ USA     │
│ Bob     ┆ 30  ┆ USA     │
│ Charlie ┆ 35  ┆ USA     │
└─────────┴─────┴─────────┘

In this example, we used pl.lit(“USA”) to create a literal value for the new “country” column. The alias() method is used to name the new column. As you can see, the new DataFrame df_new now includes the “country” column with the value “USA” for all rows.

Adding Calculated Columns

One of the strengths of with_columns() is its ability to create new columns based on calculations involving existing columns. Here’s an example:

# Add a new column with a calculation
df_calc = df.with_columns(
    (pl.col("age") * 2).alias("double_age")
)
print(df)
print(df_calc)

Output:

shape: (3, 2)
┌─────────┬─────┐
│ name    ┆ age │
│ ---     ┆ --- │
│ str     ┆ i64 │
╞═════════╪═════╡
│ Alice   ┆ 25  │
│ Bob     ┆ 30  │
│ Charlie ┆ 35  │
└─────────┴─────┘
shape: (3, 3)
┌─────────┬─────┬────────────┐
│ name    ┆ age ┆ double_age │
│ ---     ┆ --- ┆ ---        │
│ str     ┆ i64 ┆ i64        │
╞═════════╪═════╪════════════╡
│ Alice   ┆ 25  ┆ 50         │
│ Bob     ┆ 30  ┆ 60         │
│ Charlie ┆ 35  ┆ 70         │
└─────────┴─────┴────────────┘

In this case, we’ve created a new column “double_age” by multiplying the “age” column by 2. The pl.col(“age”) syntax is used to reference existing columns in the DataFrame.

Using Expressions in with_columns()

Polars truly shines when it comes to working with expressions. Expressions in Polars are powerful constructs that allow you to perform complex operations on columns efficiently. They are lazily evaluated, which means the computation is only performed when the result is actually needed, leading to better performance, especially with large datasets.

Let’s look at a more complex example that showcases the use of expressions:

import polars as pl
# Create a sample DataFrame
df = pl.DataFrame({
    "name": ["Alice", "Bob", "Charlie"],
    "age": [25, 30, 35],
    "city": ["New York", "San Francisco", "London"]
})
# Add multiple new columns using expressions
df_expr = df.with_columns([
    pl.col("name").str.to_uppercase().alias("name_upper"),
    (pl.col("age") > 30).alias("is_over_30"),
    (pl.col("city").str.contains("New|San").cast(pl.Int8)).alias("is_us_city")
])
print(df)
print(df_expr)

Output:

shape: (3, 3)
┌─────────┬─────┬───────────────┐
│ name    ┆ age ┆ city          │
│ ---     ┆ --- ┆ ---           │
│ str     ┆ i64 ┆ str           │
╞═════════╪═════╪═══════════════╡
│ Alice   ┆ 25  ┆ New York      │
│ Bob     ┆ 30  ┆ San Francisco │
│ Charlie ┆ 35  ┆ London        │
└─────────┴─────┴───────────────┘
shape: (3, 6)
┌─────────┬─────┬───────────────┬────────────┬────────────┬────────────┐
│ name    ┆ age ┆ city          ┆ name_upper ┆ is_over_30 ┆ is_us_city │
│ ---     ┆ --- ┆ ---           ┆ ---        ┆ ---        ┆ ---        │
│ str     ┆ i64 ┆ str           ┆ str        ┆ bool       ┆ i8         │
╞═════════╪═════╪═══════════════╪════════════╪════════════╪════════════╡
│ Alice   ┆ 25  ┆ New York      ┆ ALICE      ┆ false      ┆ 1          │
│ Bob     ┆ 30  ┆ San Francisco ┆ BOB        ┆ false      ┆ 1          │
│ Charlie ┆ 35  ┆ London        ┆ CHARLIE    ┆ true       ┆ 0          │
└─────────┴─────┴───────────────┴────────────┴────────────┴────────────┘

In this example, we’re using several expressions to create multiple new columns:

1. pl.col(“name”).str.to_uppercase().alias(“name_upper”): This expression converts the “name” column to uppercase.
2. (pl.col(“age”) > 30).alias(“is_over_30”): This creates a boolean column indicating whether the age is over 30.
3. (pl.col(“city”).str.contains(“New|San”).cast(pl.Int8)).alias(“is_us_city”): This complex expression checks if the city name contains “New” or “San”, implying it’s a US city, and casts the result to an 8-bit integer.

These expressions demonstrate the power and flexibility of Polars. You can chain multiple operations, use string methods, perform boolean comparisons, and even cast data types, all within a single expression.

Conclusion

The with_columns() method in Polars is a flexible tool for adding new columns to your DataFrame. Whether you’re adding constant values, performing calculations, or using complex expressions, with_columns() provides a clean and efficient way to manipulate your data.

By leveraging the power of Polars expressions, you can perform complex operations efficiently, even on large datasets. This makes Polars a great choice for data scientists and analysts who need to work with big data and require high performance. As you continue to work with Polars, you’ll discover even more ways to use with_columns() and other methods to transform and analyze your data effectively.