Is it possible to end up with a categorical variable column after a melt operation in pandas?
If I set up the data like this:
import pandas as pd
import numpy as np
df = pd.DataFrame(
np.random.randn(3, 5),
columns=["A", "B", "C", "D", "E"]
)
df["id"] = range(1, 4)
df
| | A | B | C | D | E | id |
|----|-----------|-----------|-----------|-----------|------------|------|
| 0 | -0.406174 | -0.686917 | -0.172913 | -0.273074 | -0.0246714 | 1 |
| 1 | 0.323783 | -1.7731 | 1.57581 | -1.15671 | -1.23926 | 2 |
| 2 | -1.1426 | -0.591279 | 1.15265 | 0.326712 | -0.86374 | 3 |
and then apply
melted_df = df.melt(id_vars="id", value_vars=["A", "B", "C", "D", "E"])
melted_df
| | id | variable | value |
|----|------|------------|------------|
| 0 | 1 | A | -0.406174 |
| 1 | 2 | A | 0.323783 |
| 2 | 3 | A | -1.1426 |
| 3 | 1 | B | -0.686917 |
| 4 | 2 | B | -1.7731 |
| 5 | 3 | B | -0.591279 |
| 6 | 1 | C | -0.172913 |
| 7 | 2 | C | 1.57581 |
| 8 | 3 | C | 1.15265 |
| 9 | 1 | D | -0.273074 |
| 10 | 2 | D | -1.15671 |
| 11 | 3 | D | 0.326712 |
| 12 | 1 | E | -0.0246714 |
| 13 | 2 | E | -1.23926 |
| 14 | 3 | E | -0.86374 |
The dtype of the variable column is object
melted_df.dtypes
id int64
variable object
value float64
dtype: object
I'd like this to be category. I know, I can convert it easily by:
melted_df["variable"].astype("category")
But for large datasets, I'd like to avoid this overhead. In the documentation I didn't find such an option, but since the resulting column contains categorical data by definition, I presume there must be a possiblity.
I don't think it's possible with melt, because when it creates that column it infers the dtype and 'category' is not a dtype that pandas currently infers. (Here's a related issue where it doesn't correctly infer Int32 dtypes Why is pandas.melt messing with my dtypes?).
stack will keep the categorical dtype if you first convert the columns. stack will result in a slightly different ordering than melt, but the data will be the same. stack is also a bit clunkier with naming the resulting columns.
df = df.set_index('id')
df.columns = df.columns.astype('category')
res = (df.stack()
.rename_axis(['id', 'variable'])
.rename('value')
.reset_index())
# id variable value
#0 1 A 0.424781
#1 1 B -0.317107
#2 1 C 0.731121
#3 1 D 0.042642
#4 1 E 0.648352
#...
#13 3 D -0.889600
#14 3 E -1.822898
res.dtypes
#id int64
#variable category
#value float64
#dtype: object