Pandas Categorical类型的内部存储优化：内存布局、字典编码与比较操作的性能分析

Pandas Categorical类型的内部存储优化：内存布局、字典编码与比较操作的性能分析

各位朋友，大家好！今天我们来深入探讨Pandas中一个非常重要的类型：Categorical类型。它在数据分析和处理中扮演着关键角色，尤其是在处理包含重复值的字符串或数值数据时，能够显著提升性能并降低内存占用。我们将详细分析Categorical类型的内部存储机制，包括其内存布局、字典编码，以及这些机制如何影响比较操作的性能。

1. Categorical类型：背景与优势

在传统的数据分析中，我们经常会遇到一些列，其包含的值是有限且重复的。例如，一个包含城市名称的列，或者一个包含产品类别的列。如果直接使用字符串类型存储这些数据，会浪费大量的内存空间，并且在执行比较操作时效率低下。

Pandas的Categorical类型正是为了解决这个问题而设计的。它本质上是对原始数据进行了一层编码，将原始值映射到整数编码，并维护一个从整数编码到原始值的映射关系（即类别）。这样，原始数据就被压缩成了整数编码，大大减少了内存占用。此外，由于整数比较比字符串比较快得多，Categorical类型还能显著提升比较操作的性能。

2. Categorical类型的内部存储：内存布局与字典编码

Categorical类型的数据结构包含两个核心组成部分：

• codes: 一个整数类型的NumPy数组，存储了每个元素的类别编码。
• categories: 一个Index对象，存储了所有唯一的类别值。

让我们通过一个例子来更清晰地理解这一点：

import pandas as pd
import numpy as np

data = ['a', 'b', 'c', 'a', 'b', 'a']
categorical_data = pd.Categorical(data)

print("原始数据:", data)
print("Categorical类型:", categorical_data)
print("Codes:", categorical_data.codes)
print("Categories:", categorical_data.categories)

这段代码的输出如下：

原始数据: ['a', 'b', 'c', 'a', 'b', 'a']
Categorical类型: ['a', 'b', 'c', 'a', 'b', 'a']
Categories (3, object): ['a', 'b', 'c']
Codes: [0 1 2 0 1 0]

可以看到，categorical_data.codes存储的是整数编码，而categorical_data.categories存储的是原始的类别值。在这个例子中，’a’被编码为0，’b’被编码为1，’c’被编码为2。

这种存储方式带来的好处是显而易见的。原本需要存储多个字符串’a’、’b’、’c’，现在只需要存储它们的整数编码0、1、2，以及一个包含’a’、’b’、’c’的Index对象。尤其是在处理大量重复值时，这种压缩效果会非常显著。

内存布局

Categorical的内存布局是这样的：codes数组紧凑地存储了整数编码，而categories Index对象存储了唯一的类别值。codes数组通常是int8、int16、int32或int64类型，具体取决于类别数量。Pandas会自动选择能够容纳所有类别编码的最小整数类型。

3. Categorical类型的创建与使用

创建Categorical类型有多种方法：

• pd.Categorical(): 直接从一个列表或NumPy数组创建。
• pd.Series(..., dtype='category'): 创建一个Series，并指定dtype为’category’。
• df['column'].astype('category'): 将DataFrame中的一列转换为Categorical类型。

# 方法一：使用pd.Categorical()
categorical_data = pd.Categorical(['a', 'b', 'c', 'a', 'b', 'a'])

# 方法二：使用pd.Series(..., dtype='category')
series_data = pd.Series(['a', 'b', 'c', 'a', 'b', 'a'], dtype='category')

# 方法三：使用df['column'].astype('category')
df = pd.DataFrame({'col1': ['a', 'b', 'c', 'a', 'b', 'a']})
df['col1_categorical'] = df['col1'].astype('category')

print(type(categorical_data))
print(type(series_data))
print(df.dtypes)

输出结果为：

<class 'pandas.core.arrays.categorical.Categorical'>
<class 'pandas.core.series.Series'>
col1                object
col1_categorical    category
dtype: object

创建Categorical类型后，我们可以像使用其他Pandas数据结构一样使用它。例如，我们可以进行过滤、排序、分组等操作。

4. Categorical类型与内存优化

Categorical类型最显著的优势之一是内存优化。让我们通过一个实验来验证这一点：

import sys

n = 1000000
data = ['a', 'b', 'c'] * (n // 3)

# 使用字符串类型存储
string_series = pd.Series(data)
string_memory = sys.getsizeof(string_series)

# 使用Categorical类型存储
categorical_series = pd.Series(data, dtype='category')
categorical_memory = sys.getsizeof(categorical_series)

print("字符串类型内存占用:", string_memory, "bytes")
print("Categorical类型内存占用:", categorical_memory, "bytes")
print("内存节省比例:", (string_memory - categorical_memory) / string_memory * 100, "%")

运行结果显示，Categorical类型可以显著减少内存占用，通常可以节省50%以上的内存。这在处理大型数据集时非常重要。

5. Categorical类型的比较操作性能分析

除了内存优化，Categorical类型还能提升比较操作的性能。这是因为Categorical类型在内部使用整数编码进行比较，而整数比较比字符串比较快得多。

import time

n = 1000000
data = ['a', 'b', 'c'] * (n // 3)

# 使用字符串类型进行比较
string_series = pd.Series(data)
start_time = time.time()
string_series == 'a'
string_time = time.time() - start_time

# 使用Categorical类型进行比较
categorical_series = pd.Series(data, dtype='category')
start_time = time.time()
categorical_series == 'a'
categorical_time = time.time() - start_time

print("字符串类型比较时间:", string_time, "秒")
print("Categorical类型比较时间:", categorical_time, "秒")
print("Categorical类型速度提升比例:", (string_time - categorical_time) / string_time * 100, "%")

实验结果表明，Categorical类型的比较速度明显快于字符串类型，通常可以提升数倍甚至数十倍。

6. Categorical类型的排序与分组

Categorical类型在排序和分组操作中也能发挥优势。由于内部使用整数编码，排序和分组操作可以更快地完成。

排序

data = ['b', 'a', 'c', 'a', 'b']
categorical_series = pd.Series(data, dtype='category')

# 排序，默认按照类别顺序排序
sorted_series = categorical_series.sort_values()
print("排序后的Categorical Series:", sorted_series)

# 可以自定义类别顺序
categorical_series = categorical_series.cat.reorder_categories(['c', 'b', 'a'], ordered=True)
sorted_series = categorical_series.sort_values()
print("自定义类别顺序后的排序结果:", sorted_series)

分组

df = pd.DataFrame({'category': ['a', 'b', 'c', 'a', 'b'], 'value': [1, 2, 3, 4, 5]})
df['category'] = df['category'].astype('category')

# 分组求和
grouped_sum = df.groupby('category')['value'].sum()
print("分组求和结果:", grouped_sum)

7. Categorical类型的潜在问题与注意事项

尽管Categorical类型有很多优点，但也存在一些潜在的问题需要注意：

• 内存占用增加的情况: 如果类别数量接近或等于数据总数，使用Categorical类型可能不会减少内存占用，反而会因为额外的categories Index对象而增加内存占用。
• 类别顺序: Categorical类型的类别顺序可能会影响排序和比较结果。可以使用cat.reorder_categories()方法自定义类别顺序。
• 新增类别: 默认情况下，不能直接向Categorical类型添加新的类别。需要使用cat.add_categories()方法添加新的类别。
• 未知的类别: 如果数据中包含未在categories中定义的类别，这些值会被设置为NaN。可以使用categories参数指定已知的类别，或者使用ordered=True参数创建一个有序的Categorical类型。

8. CategoricalDtype与自定义数据类型

Pandas 1.0 引入了 CategoricalDtype 类，允许更精细地控制 Categorical 类型的创建和使用。 CategoricalDtype 允许你指定类别和顺序，从而创建自定义的 Categorical 类型。

from pandas.api.types import CategoricalDtype

# 定义类别和顺序
my_categories = ['low', 'medium', 'high']
my_dtype = CategoricalDtype(categories=my_categories, ordered=True)

# 创建 Categorical Series
data = ['low', 'medium', 'high', 'low', 'medium']
categorical_series = pd.Series(data, dtype=my_dtype)

print(categorical_series)
print(categorical_series.dtype)

使用 CategoricalDtype 可以确保数据具有预期的类别和顺序，避免潜在的错误。

9. 案例分析：大型数据集的Categorical优化

假设我们有一个包含100万行销售数据的DataFrame，其中包含一个"产品类别"列，类别包括"电子产品"、"服装"、"家居用品"等。

import pandas as pd
import numpy as np
import sys
import time

# 生成模拟数据
n = 1000000
categories = ['电子产品', '服装', '家居用品', '食品', '化妆品']
data = np.random.choice(categories, size=n)
df = pd.DataFrame({'产品类别': data, '销售额': np.random.rand(n)})

# 内存占用比较
string_memory = sys.getsizeof(df['产品类别'])
df['产品类别'] = df['产品类别'].astype('category')
categorical_memory = sys.getsizeof(df['产品类别'])
print("字符串类型内存占用:", string_memory, "bytes")
print("Categorical类型内存占用:", categorical_memory, "bytes")
print("内存节省比例:", (string_memory - categorical_memory) / string_memory * 100, "%")

# 分组求和性能比较
start_time = time.time()
df.groupby('产品类别')['销售额'].sum()
categorical_time = time.time() - start_time
print("Categorical类型分组求和时间:", categorical_time, "秒")

df['产品类别'] = df['产品类别'].astype(str)
start_time = time.time()
df.groupby('产品类别')['销售额'].sum()
string_time = time.time() - start_time
print("字符串类型分组求和时间:", string_time, "秒")
print("Categorical类型速度提升比例:", (string_time - categorical_time) / string_time * 100, "%")

通过将"产品类别"列转换为Categorical类型，我们可以显著减少内存占用，并提升分组求和操作的性能。

表格总结：Categorical类型的优势与劣势

特性	优势	劣势
内存占用	显著减少重复字符串或数值的内存占用	如果类别数量接近或等于数据总数，可能不会减少内存占用，反而会增加
比较操作	提升比较操作的性能，因为内部使用整数编码进行比较
排序与分组	提升排序和分组操作的性能	类别顺序可能会影响排序结果，需要注意
数据一致性	可以通过预定义的类别列表来保证数据的一致性	添加新类别需要使用专门的方法
适用场景	包含大量重复值的字符串或数值列	不适合类别数量接近或等于数据总数的列
自定义数据类型	CategoricalDtype允许更精细地控制类别和顺序，创建自定义Categorical类型

10. 总结：内存优化与性能提升的利器

通过今天的讨论，我们深入了解了Pandas Categorical类型的内部存储机制、内存优化效果、比较操作性能，以及使用注意事项。 Categorical类型是Pandas中一个强大的工具，能够显著提升数据分析和处理的效率，尤其是在处理包含重复值的字符串或数值数据时。 合理地利用Categorical类型，可以有效地减少内存占用，提升计算速度，从而更好地应对大型数据集的挑战。

更多IT精英技术系列讲座，到智猿学院