有效地创建带有空值的分类数据框

时间:2019-07-12 17:27:36

标签: python-3.x pandas

我想创建一个包含空值的分类数据框,并在扩展索引之前设置类别。索引非常大,我想避免内存高峰,但似乎无法做到这一点。

示例:

# memory spike
df = pd.DataFrame(index=list(range(0, 1000)), columns=['a', 'b'])
df.info(memory_usage='deep')

输出:

<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

转换为分类:

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')

输出:

<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

有没有办法避免内存高峰?

1 个答案:

答案 0 :(得分:1)

如果数据框是由DataFrame构造函数创建的,则可以将列初始化为类别类型。

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)

另一个替代解决方案如下。

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)

如果通过诸如read_csv之类的方法创建数据框,则可以使用dtype关键字参数来确保输出列具有所需的数据类型,而不是在创建数据框之后进行转换-这会导致更多的内存消耗。

df = pd.read_csv("file.csv", dtype={"a": cat_type1, "b": cat_type2})

在这里,也可以通过传入dtype={"a": "category"}从数据直接推断类别值。预先指定类别可以节省推理开销,还可以让解析器检查数据值是否与指定的类别值匹配。如果数据中未出现某些类别值,则也很有必要。