使用 pandas.DataFrame.to_csv 写入 csv 文件时如何强制执行分类整数 dtype

Question

(The question has been updated to isolate the problem more strictly)

I have data in a pandas.DataFrame with categorical variables. The categories are integers. The data may have missing values.

import pandas as pd

# Define all dtypes
dtypes = {
    'var_001': pd.api.types.CategoricalDtype(
        categories=[1, 2, 3, 4],
        ordered=False,
    ),
    'var_002': pd.UInt8Dtype(),
    'var_003': pd.api.types.CategoricalDtype(
        categories=[1, 2, 3, 4, 5],
        ordered=True,
    ),
}

# Create a dataframe
df = pd.DataFrame(
    data={
        'var_001': [1, '', 3],
        'var_002': [43, 62, 99],
        'var_003': [2, 3, 3],
    },
)

# Convert to the right dtypes (btw, why this cannot be done in the construcor??)
df = df.astype(dtype=dtypes)

The dtypes seem good:

>>> print(df.dtypes)
var_001    category
var_002       UInt8
var_003    category
dtype: object

As do the data in the dataframe:

>>> print(df)
  var_001  var_002 var_003
0       1       43       2
1     NaN       62       3
2       3       99       3

However, when I write the dataframe into a csv file ( df.to_csv('data.csv', index=False) ), the values of the variable with missing values get printed as float instead of integers:

var_001, var_002, var_003
1.0,     43,      2
,        62,      3
3.0,     99,      3

Is there a way to keep the integer categories also for data with missing values when writing into a csv file?

Answer 1

Apparently, the problem lies with the non-nullable integers:

In Working with missing data, we saw that pandas primarily uses NaN to represent missing data. Because NaN is a float, this forces an array of integers with any missing values to become floating point.

https://pandas.pydata.org/pandas-docs/stable/user_guide/integer_na.html

So, to enable missing values also with categorical variables, we must define the categories as nullable integers:

import pandas as pd

# Create an array with nullable integer values
cat_0_4 = pd.array([0, 1, 2, 3, 4], dtype="Int8")

# Define an ordered categorical dtype with nullable integer values
var_dtype = pd.api.types.CategoricalDtype(
    categories=cat_0_4,
    ordered=True,
)

...