pandas SparseSeries 可以将值存储在 float16 dtype 中吗？

Question

我想在稀疏的 pandas 容器中使用较小的数据类型的原因是为了减少内存使用。这在处理最初使用 bool（例如来自 to_dummies）或小的数字 dtypes（例如 int8）的数据时是相关的，这些数据在稀疏容器中都转换为 float64。

DataFrame 创建

提供的示例使用适度的 20k x 145 数据帧。在实践中，我正在使用 1e6 x 5e3 顺序的数据帧。

In []: bool_df.info()
<class 'pandas.core.frame.DataFrame'>
RangeIndex: 19849 entries, 0 to 19848
Columns: 145 entries, topic.party_nl.p.pvda to topic.sub_cat_Reizen
dtypes: bool(145)
memory usage: 2.7 MB

In []: bool_df.memory_usage(index=False).sum()
Out[]: 2878105

In []: bool_df.values.itemsize
Out[]: 1

这个数据帧的稀疏版本需要更少的内存，但仍然比需要的大得多，给定原始 dtype。

In []: sparse_df = bool_df.to_sparse(fill_value=False)

In []: sparse_df.info()
<class 'pandas.sparse.frame.SparseDataFrame'>
RangeIndex: 19849 entries, 0 to 19848
Columns: 145 entries, topic.party_nl.p.pvda to topic.sub_cat_Reizen
dtypes: float64(145)
memory usage: 1.1 MB

In []: sparse_df.memory_usage(index=False).sum()
Out[]: 1143456

In []: sparse_df.values.itemsize
Out[]: 8

即使这些数据相当稀疏，从 bool 到 float64 的 dtype 转换也会导致非填充值占用 8 倍以上的空间。

In []: sparse_df.memory_usage(index=False).describe()
Out[]:
count      145.000000
mean      7885.903448
std      17343.762402
min          8.000000
25%        640.000000
50%       1888.000000
75%       4440.000000
max      84688.000000

鉴于数据的稀疏性，人们希望更大幅度地减少内存大小：

In []: sparse_df.density
Out[]: 0.04966184346992205

底层存储的内存占用

SparseDataFrame 的列是SparseSeries，它们使用SparseArray 作为底层numpy.ndarray 存储的包装器。稀疏数据帧使用的字节数可以（也）直接从这些 ndarray 计算：

In []: col64_nbytes = [
.....:     sparse_df[col].values.sp_values.nbytes
.....:     for col in sparse_df
.....: ]

In []: sum(col64_nbytes)
Out[]: 1143456

ndarrays 可以转换为使用更小的浮点数，这样可以计算数据帧在使用时需要多少内存，例如浮动16s。正如人们所预料的那样，这将导致数据帧缩小 4 倍。

In []: col16_nbytes = [
.....:     sparse_df[col].values.sp_values.astype('float16').nbytes
.....:     for col in sparse_df
.....: ]

In []: sum(col16_nbytes)
Out[]: 285864

通过使用更合适的 dtype，内存使用可以减少到密集版本的 10%，而 float64 稀疏数据帧减少到 40%。对于我的数据，这可能会导致需要 20 GB 和 5 GB 的可用内存。

In []: sum(col64_nbytes) / bool_df.memory_usage(index=False).sum()
Out[]: 0.3972947477593764

In []: sum(col16_nbytes) / bool_df.memory_usage(index=False).sum()
Out[]: 0.0993236869398441

问题

很遗憾，pandas 中还没有实现稀疏容器的 dtype 转换：

In []: sparse_df.astype('float16')
---------------------------------------------------
[...]/pandas/sparse/frame.py in astype(self, dtype)
    245
    246     def astype(self, dtype):
--> 247         raise NotImplementedError
    248
    249     def copy(self, deep=True):

NotImplementedError:

如何将SparseDataFrame 中的SparseSeries 转换为使用numpy.float16 数据类型或其他每项使用少于64 个字节的数据类型，而不是默认的numpy.float64？

Answer 1

SparseArray 构造函数可用于转换其底层ndarray 的dtype。要转换数据框中的所有稀疏序列，可以遍历 df 的序列，转换它们的数组，然后用转换后的版本替换序列。

import pandas as pd
import numpy as np

def convert_sparse_series_dtype(sparse_series, dtype):
    dtype = np.dtype(dtype)
    if 'float' not in str(dtype):
        raise TypeError('Sparse containers only support float dtypes')

    sparse_array = sparse_series.values
    converted_sp_array = pd.SparseArray(sparse_array, dtype=dtype)

    converted_sp_series = pd.SparseSeries(converted_sp_array)
    return converted_sp_series


def convert_sparse_columns_dtype(sparse_dataframe, dtype):
    for col_name in sparse_dataframe:
        if isinstance(sparse_dataframe[col_name], pd.SparseSeries):
            sparse_dataframe.loc[:, col_name] = convert_sparse_series_dtype(
                 sparse_dataframe[col_name], dtype
            )

这实现了减少稀疏数据帧的内存占用的既定目的：

In []: sparse_df.info()
<class 'pandas.sparse.frame.SparseDataFrame'>
RangeIndex: 19849 entries, 0 to 19848
Columns: 145 entries, topic.party_nl.p.pvda to topic.sub_cat_Reizen
dtypes: float64(145)
memory usage: 1.1 MB

In []: convert_sparse_columns_dtype(sparse_df, 'float16')

In []: sparse_df.info()
<class 'pandas.sparse.frame.SparseDataFrame'>
RangeIndex: 19849 entries, 0 to 19848
Columns: 145 entries, topic.party_nl.p.pvda to topic.sub_cat_Reizen
dtypes: float16(145)
memory usage: 279.2 KB

In []: bool_df.equals(sparse_df.to_dense().astype('bool'))
Out[]: True

然而，这是一个有点糟糕的解决方案，因为转换后的数据帧在与其他数据帧交互时表现出不可预测的行为。例如，当转换后的稀疏数据帧与其他数据帧连接时，所有包含的序列都变为密集序列。对于未转换的稀疏数据帧，情况并非如此。它们在结果数据框中保持稀疏序列。