Efficiently create categorical data frame with nulls - python-3.x

Question

I want to create a categorical data frame with nulls and set the categories before expanding the index. The index is very large and I want to avoid the memory spike and I cannot seem to do this.
Example:
# memory spike
df = pd.DataFrame(index=list(range(0, 1000)), columns=['a', 'b'])
df.info(memory_usage='deep')
Output:
<class 'pandas.core.frame.DataFrame'>
Int64Index: 1000 entries, 0 to 999
Data columns (total 2 columns):
a 0 non-null object
b 0 non-null object
dtypes: object(2)
memory usage: 70.3 KB
Convert to Categorical:
for _ in df.columns:
df[_] = df[_].astype('category')
# set categories for columns
df['a'] = df['a'].cat.add_categories(['d', 'e', 'f'])
df['b'] = df['b'].cat.add_categories(['g', 'h', 'i'])
# check memory usage
df.info(memory_usage='deep')
Output:
<class 'pandas.core.frame.DataFrame'>
Int64Index: 1000 entries, 0 to 999
Data columns (total 2 columns):
a 0 non-null category
b 0 non-null category
dtypes: category(2)
memory usage: 9.9 KB
Is there a way to do this while avoiding the memory spike?

Answer 1

If the data frame is created by the DataFrame constructor, the columns can be initialized as category types.
import numpy as np
import pandas as pd
from pandas.api.types import CategoricalDtype
cat_type1 = CategoricalDtype(["d", "e", "f"])
cat_type2 = CategoricalDtype(["g", "h", "i"])
index = pd.Index(range(1000))
df = pd.DataFrame({"a": pd.Series([np.nan] * len(index), dtype=cat_type1, index=index),
"b": pd.Series([np.nan] * len(index), dtype=cat_type2, index=index)},
index=index)
Another alternative solution is the following.
cols = ["a", "b"]
index = pd.Index(range(1000))
df = pd.DataFrame({k: [np.nan] * len(index) for k in cols}, index=index, dtype="category")
df["a"].cat.set_categories(["d", "e", "f"], inplace=True)
df["b"].cat.set_categories(["g", "h", "i"], inplace=True)
If the data frame is created via methods such as read_csv, the dtype keyword parameter can be used to make sure the output columns have desired data types rather than making conversions after the data frame is created -- which leads to more memory consumption.
df = pd.read_csv("file.csv", dtype={"a": cat_type1, "b": cat_type2})
Here, the category values can also be directly inferred from the data by passing in dtype={"a": "category"}. Specifying the categories beforehand can save the inference overhead and also let the parser check the data values match the specified category values. It is also necessary if some category values do not occur in the data.