Polars DataFrame select() Method

In Polars, the select() method is used to select specific columns from a DataFrame, which can be useful for filtering and working with only the relevant data. It allows you to work with expressions that specify which columns to include in the result. This method supports simple column selection as well as more advanced operations, including filtering, aggregation, and applying expressions.

In this article, I will explain the select() method, covering its syntax, parameters, and usage. Through detailed examples, I will demonstrate how to generate a new Polars DataFrame that includes the selected columns or transformed expressions.

Key Points –

select() is used to create a new DataFrame with selected columns or computed expressions without modifying the original DataFrame.
Columns can be selected using expressions like pl.col("column_name") or directly by column names.
Allows naming output columns explicitly by using meaningful names for computed expressions.
Does not modify the original DataFrame; instead, it returns a new DataFrame.
Apply multiple operations to the same or different columns in a single select statement.
Supports complex expressions such as arithmetic operations, conditional logic (pl.when().then().otherwise()), and function calls.
Supports single-column selection, multiple-column selection, and computed columns within the same method call.

Polars DataFrame.select() Syntax

Following is a syntax of the DataFrame.select(). This function takes exprs, and named_exprs params.


# Syntax of select()
DataFrame.select(
    *exprs: IntoExpr | Iterable[IntoExpr],
    **named_exprs: IntoExpr
) → DataFrame

Parameters of the Polars DataFrame.select()

Following are the parameters of the select() method.

*exprs - One or more expressions or column names that can either be:
- Direct column names (strings) like "A", "B", etc.
- Expressions that represent transformations on columns (e.g., pl.col("A") * 2).
- You can pass these as positional arguments (i.e., a list of expressions or column names).
**named_exprs - Named expressions, where you can assign a name to a new column created from an expression. The name of the expression will appear as the column name in the resulting DataFrame.

Return Value

This function returns a new Polars DataFrame containing the selected columns or transformed expressions.

Usage of Polars DataFrame.select() Method

The select() method in Polars is used to choose specific columns or apply expressions, transformations, or aggregations to one or more columns in a Polars DataFrame. It allows for efficient column-based operations.

To run some examples of the Polars DataFrame.select() method, let’s create a Polars DataFrame.


import polars as pl

# Creating a new Polars DataFrame
technologies = {
    'Courses': ["Spark", "Hadoop", "Python", "Pandas"],
    'Fees': [22000, 25000, 20000, 26000],
    'Duration': ['30days', '50days', '40days', '60days'],
    'Discount': [1000, 1500, 1200, 2000]
}

df = pl.DataFrame(technologies)
print("Original DataFrame:\n", df)

Yields below output.

To select a single column from a Polars DataFrame using the select() method, you can pass a single column name or an expression for that column.


# Select a single column by name
df2 = df.select("Courses")
print(df2)

Here, we select the Courses column by passing the column name as a string. The result is a DataFrame containing just the selected column.

Alternatively, to select a single column using an expression in Polars, you can utilize the pl.col() function to refer to the column, and optionally apply transformations or calculations.


# Select a single column using expression
df2 = df.select(pl.col("Courses"))
print(df2)

Here, we use pl.col() to select the Courses column. This is useful when you want to apply transformations or operations on the column in the same step. Yields the same output as above.

Select Specific Columns

To select specific columns in a Polars DataFrame using the select() method, you can directly pass a list of column names or expressions as arguments.


# Select specific columns
df2 = df.select(["Courses", "Fees"])
print(df2)

# Output:
# shape: (4, 2)
┌─────────┬───────┐
│ Courses ┆ Fees  │
│ ---     ┆ ---   │
│ str     ┆ i64   │
╞═════════╪═══════╡
│ Spark   ┆ 22000 │
│ Hadoop  ┆ 25000 │
│ Python  ┆ 20000 │
│ Pandas  ┆ 26000 │
└─────────┴───────┘

In the example above, we select just the "Courses" and "Fees" columns from the DataFrame. The result is a new DataFrame containing only the specified columns.

Using a List of Expressions

You can use a list of expressions in the select() method to define multiple transformations or selections in a single call. This approach provides flexibility for complex operations while grouping them neatly in a list.


# Select multiple columns using a list
df2 = df.select([pl.col("Courses"), pl.col("Fees"), pl.col("Discount") + 500])
print(df2)

# Output:
# shape: (4, 3)
┌─────────┬───────┬──────────┐
│ Courses ┆ Fees  ┆ Discount │
│ ---     ┆ ---   ┆ ---      │
│ str     ┆ i64   ┆ i64      │
╞═════════╪═══════╪══════════╡
│ Spark   ┆ 22000 ┆ 1500     │
│ Hadoop  ┆ 25000 ┆ 2000     │
│ Python  ┆ 20000 ┆ 1700     │
│ Pandas  ┆ 26000 ┆ 2500     │
└─────────┴───────┴──────────┘

Select Columns Using Expressions

The select() method in Polars allows you to use expressions to select and transform columns in a DataFrame. This enables powerful and flexible data manipulations like arithmetic operations, aggregations, or conditional logic applied to columns while selecting them.

You can perform transformations (like addition, multiplication, etc.) on columns while selecting them.


# Applying a transformation to a column 
# Using Expressions
df2 = df.select((pl.col("Fees") - pl.col("Discount")).alias("Net_Fees"))
print(df2)

# Output:
# shape: (4, 1)
┌──────────┐
│ Net_Fees │
│ ---      │
│ i64      │
╞══════════╡
│ 21000    │
│ 23500    │
│ 18800    │
│ 24000    │
└──────────┘

Select Multiple Columns Using Expressions

You can select multiple columns and apply transformations or operations on them simultaneously.


# Select multiple columns Using Expressions
df2 = df.select([pl.col("Courses"), (pl.col("Fees") * 1.1).alias("Fees_with_Tax")])
print(df2)
 
# Output:
# shape: (4, 2)
┌─────────┬───────────────┐
│ Courses ┆ Fees_with_Tax │
│ ---     ┆ ---           │
│ str     ┆ f64           │
╞═════════╪═══════════════╡
│ Spark   ┆ 24200.0       │
│ Hadoop  ┆ 27500.0       │
│ Python  ┆ 22000.0       │
│ Pandas  ┆ 28600.0       │
└─────────┴───────────────┘

Use Conditional Expressions

You can apply conditional logic in your expressions, such as categorizing or flagging based on a condition.


# Use conditional expressions
df2 = df.select(
    pl.when(pl.col("Fees") > 24000)
    .then(pl.lit("Expensive"))
    .otherwise(pl.lit("Affordable"))
    .alias("Fee_Category"))
print(df2)

# Output:
# shape: (4, 1)
┌──────────────┐
│ Fee_Category │
│ ---          │
│ str          │
╞══════════════╡
│ Affordable   │
│ Expensive    │
│ Affordable   │
│ Expensive    │
└──────────────┘

Using Keyword Arguments

Use keyword arguments in Polars' select() method to create a named column (e.g., threshold) based on a conditional expression.


# Select multiple columns using a list
df2 = df.select(threshold=pl.when(pl.col("Discount") > 1200).then(100).otherwise(0))
print(df2)

# Output:
# shape: (4, 1)
┌───────────┐
│ threshold │
│ ---       │
│ i32       │
╞═══════════╡
│ 0         │
│ 100       │
│ 0         │
│ 100       │
└───────────┘

Here,

pl.when(condition): Defines a conditional expression. Here, the condition checks if the values in the Discount column are greater than 1200.
.then(value): Specifies the value to return when the condition is true. If the condition (Discount > 1200) is met, it returns 100.
.otherwise(value): Specifies the value to return when the condition is false. If the condition is not met, it returns 0.
threshold=..: The resulting column from this conditional logic is named threshold using the keyword argument in select().

Conclusion

In conclusion, the select() method in Polars is a powerful and versatile tool for selecting specific columns, applying transformations, and performing complex operations using expressions. It enables efficient data manipulation by returning a new DataFrame with the desired structure and content.

Happy Learning!!

References

https://docs.pola.rs/api/python/stable/reference/dataframe/api/polars.DataFrame.select_seq.html