Get first n chars from a str column in Python Polars

Polars is a powerful DataFrame library designed for speed and ease of use, particularly with large datasets. If you need to extract the first n characters from a string column in a Polars DataFrame, Polars offers efficient and straightforward methods to achieve this. In this article, we will go through three good code examples demonstrating how to perform this task.

Problem Statement

When working with textual data in a DataFrame, extracting substrings from string columns is a common operation in data manipulation. Whether you're cleaning data, creating new features, or preparing data for analysis, being able to slice strings efficiently is crucial. Polars provide methods to work with string columns effectively, ensuring high performance even with large datasets.

Extracting First n chars from a String Column in Python Polars

Python Polars offers a variety of functions for string manipulation, making it easy to extract substrings from a column. Let us see a few different examples for a better understanding of the concept.

Using Apply with a Lambda Function

In this example, we will use the apply() function combined with a lambda function. This approach allows for flexible and customized operations on each element of the column. The "pl.col("text").apply(lambda x: x[:n])" applies a lambda function to each element in the "text" column, extracting the first n characters.

import polars as pl

# Create a Polars DataFrame
df = pl.DataFrame({
    "text": ["apple", "banana", "cherry", "date"]
})

# Number of characters to extract
n = 3

# Extract first n characters using apply and a lambda function
df = df.with_columns(
    pl.col("text").apply(lambda x: x[:n]).alias("first_n_chars")
)

print(df)

Output:

Using the ste.extract Function

In this example, the f"^.{{0,{n}}}" constructs a regular expression pattern to match the first n characters. Then the characters.pl.col("text").str.extract(pattern, 0) uses the str.extract method to extract the matched substring. The alias("first_n_chars") renames the resulting column to "first_n_chars".

import polars as pl

# Create a Polars DataFrame
df = pl.DataFrame({
    "text": ["apple", "banana", "cherry", "date"]
})

# Number of characters to extract
n = 3

# Extract first n characters using str.extract
pattern = f"^.{{0,{n}}}"
df = df.with_columns(
    pl.col("text").str.extract(pattern, 0).alias("first_n_chars")
)

print(df)

Output:

Using String Expression Methods

Polars string expressions offer a variety of methods to manipulate string columns. The str namespace includes a slice method, which is another way to achieve our goal. Here, pl.col("text").str.slice(0, n) is used to slice the first n characters from each element in the "text" column.

import polars as pl

# Create a Polars DataFrame
df = pl.DataFrame({
    "text": ["apple", "banana", "cherry", "date"]
})

# Number of characters to extract
n = 3

# Extract first n characters using string expression slice method
df = df.with_columns(
    pl.col("text").str.slice(0, n).alias("first_n_chars")
)

print(df)

Output:

Conclusion

Polars provides multiple efficient ways to extract the first n characters from a string column. Whether you use the apply function with a lambda, the str_slice method, or the str.slice expression, Polars ensures that the operations are performed quickly and efficiently, even on large datasets. Experiment with these methods to find the one that best fits your workflow and performance requirements.