# remove cell
import matplotlib.pyplot as plt
import matplotlib_inline.backend_inline
# Plot settings
plt.style.use("https://github.com/aeturrell/python4DS/raw/main/plot_style.txt")
matplotlib_inline.backend_inline.set_matplotlib_formats("svg")Categorical Data
Introduction
In this chapter, we’ll introduce how to work with categorical variables—that is, variables that have a fixed and known set of possible values. This chapter is enormously indebted to the polars documentation.
Prerequisites
This chapter will use the polars data analysis package.
The Category Datatype
Everything in Python has a type, even the data in polars data frame columns. While you may be more familiar with numbers and even strings, there is also a special data type for categorical data called Categorical and Enum. There are some benefits to using categorical variables (where appropriate):
- they can keep track even when elements of the category isn’t present, which can sometimes be as interesting as when they are (imagine you find no-one from a particular school goes to university)
- they can use vastly less of your computer’s memory than encoding the same information in other ways
- they can be used efficiently with modelling packages, where they will be recognised as potential ‘dummy variables’, or with plotting packages, which will treat them as discrete values
- you can order them (for example, “neutral”, “agree”, “strongly agree”) using pl.Enum
All values of categorical data for a polars column are either in the given categories or take the value null.
Polars has two categorical types:
pl.Categorical: Unordered categories (like pandas category)
pl.Enum: Ordered categories with a fixed set of values (similar to pandas ordered categorical)
Creating Categorical Data
Let’s create a categorical column of data:
import numpy as np
import polars as pl
df = pl.DataFrame({"A": ["a", "b", "c", "a"]})
df = df.with_columns(pl.col("A").cast(pl.Categorical))
df["A"]
shape: (4,)
| A |
|---|
| cat |
| "a" |
| "b" |
| "c" |
| "a" |
Notice that polars tracks the categorical type, though the display is more compact than pandas.
You can also use special expression methods, such as pl.col().cut(), to group data into discrete bins. Here’s an example where we specify the labels for each category directly.
df = pl.DataFrame({"value": np.random.randint(0, 100, 20)})
labels = [f"{i} - {i+9}" for i in range(0, 100, 10)]
df = df.with_columns(
pl.col("value")
.cut(breaks=list(range(10, 100, 10)), labels=labels, left_closed=True)
.alias("group")
)
df.head()
shape: (5, 2)
| value | group |
|---|---|
| i32 | cat |
| 27 | "20 - 29" |
| 78 | "70 - 79" |
| 51 | "50 - 59" |
| 98 | "90 - 99" |
| 36 | "30 - 39" |
In the example above, the group column is of categorical type.
Another way to create a categorical variable is directly using the pl.Categorical or pl.Enum types:
raw_cat = pl.Series(["a", "b", "c", "a", "d", "a", "c"]).cast(
pl.Enum(["b", "c", "d"]), strict=False
)
raw_cat
shape: (7,)
| enum |
| null |
| "b" |
| "c" |
| null |
| "d" |
| null |
| "c" |
pl.Enum is preferred when the specific categories are known upfront. pl.Categorical infers categories as it reads the dataframe column.
We can then enter this into a data frame:
df = pl.DataFrame({"cat_type": raw_cat})
df["cat_type"]
shape: (7,)
| cat_type |
|---|
| enum |
| null |
| "b" |
| "c" |
| null |
| "d" |
| null |
| "c" |
Note that null appears for any value that isn’t in the categories we specified—you can find more on this in Missing Values.
You can also create ordered categories:
ordered_cat = pl.Series(["a", "b", "c", "a", "d", "a", "c"]).cast(
pl.Enum(["a", "b", "c", "d"])
)
ordered_cat
shape: (7,)
| enum |
| "a" |
| "b" |
| "c" |
| "a" |
| "d" |
| "a" |
| "c" |
Working with Categories
Categorical data in polars has categories and ordering properties. For pl.Enum, the categories are fixed and ordered. For pl.Categorical, categories are not ordered and can be more flexible.
Let’s see some examples:
df["cat_type"].dtype.categories
shape: (3,)
| category |
|---|
| str |
| "b" |
| "c" |
| "d" |
isinstance(df["cat_type"].dtype, pl.Enum)TrueIf categorical data is ordered (ie pl.Enum), then the order of the categories has a meaning and certain operations are possible: you can sort values, and apply .min() and .max().
Renaming Categories
Renaming categories in polars requires re-casting with new categories:
df = df.with_columns(
pl.col("cat_type")
.cast(pl.String)
.replace({"b": "alpha", "c": "beta", "d": "gamma"})
.cast(pl.Enum(["alpha", "beta", "gamma"]))
)To add categories, you need to re-cast with an extended Enum:
df = df.with_columns(
pl.col("cat_type").cast(pl.Enum(["alpha", "beta", "gamma", "delta"]))
)
df["cat_type"]
shape: (7,)
| cat_type |
|---|
| enum |
| null |
| "alpha" |
| "beta" |
| null |
| "gamma" |
| null |
| "beta" |
Because Enum categories are fixed in Polars, you cannot remove categories directly. Instead, ensure that the data contains only the desired values, then re-cast the column to a new Enum containing just those categories.
Operations on Categories
Enum columns behave like any other Polars column, but they also retain the complete set of allowed categories. This means Polars remembers categories even if no rows currently contain them.
For example, suppose we define an Enum with five possible categories, but only use four of them:
df["cat_type"].value_counts()
shape: (4, 2)
| cat_type | count |
|---|---|
| enum | u32 |
| "gamma" | 1 |
| "alpha" | 1 |
| "beta" | 2 |
| null | 3 |
Notice that “delta” is not shown because value_counts() only reports values that actually occur in the data.
However, the Enum still knows that “delta” is a valid category.
You can inspect the complete set of categories stored in the Enum:
df["cat_type"].dtype.categories
shape: (4,)
| category |
|---|
| str |
| "alpha" |
| "beta" |
| "gamma" |
| "delta" |
If you want a frequency table that includes unused categories, combine the stored category list with the observed counts:
all_categories = pl.DataFrame(
{"cat_type": df["cat_type"].dtype.categories.cast(df["cat_type"].dtype)}
)
all_categories.join(df["cat_type"].value_counts(), on="cat_type", how="left").fill_null(
0
)
shape: (4, 2)
| cat_type | count |
|---|---|
| enum | u32 |
| "alpha" | 1 |
| "beta" | 2 |
| "gamma" | 1 |
| "delta" | 0 |
This demonstrates one advantage of Enum: the complete set of allowed categories is retained, even when some categories are not present in the data.
You can also inspect the datatype of the column:
df["cat_type"].dtypeEnum(categories=['alpha', 'beta', 'gamma', 'delta'])You can compute the most frequent category:
df["cat_type"].drop_nulls().mode()
shape: (1,)
| cat_type |
|---|
| enum |
| "beta" |
And if your categorical column happens to consist of elements that can undergo operations, those same operations will still work. For example:
import datetime
time_df = pl.DataFrame(
{
"datetime": pl.date_range(
datetime.date(2015, 5, 31), datetime.date(2015, 9, 30), "1mo", eager=True
)
.cast(pl.String)
.cast(pl.Categorical)
}
)
time_df
shape: (5, 1)
| datetime |
|---|
| cat |
| "2015-05-31" |
| "2015-06-30" |
| "2015-07-31" |
| "2015-08-31" |
| "2015-09-30" |
time_df["datetime"].cast(pl.String).str.to_date().dt.month()
shape: (5,)
| datetime |
|---|
| i8 |
| 5 |
| 6 |
| 7 |
| 8 |
| 9 |
Finally, if you ever need to translate your actual data types in your categorical column into a code, you can use .to_physical() to get unique codes for each value.
time_df["datetime"].to_physical()
shape: (5,)
| datetime |
|---|
| u32 |
| 0 |
| 1 |
| 2 |
| 3 |
| 4 |