自动优化 Pandas 数据类型

Question

我正在使用 pandas 库研究一种算法。我在工作时遇到了一个有趣的问题。

当我将数据帧对象写入文件并再次读取时，数据帧会发生变化。当我排查原因时，发现是类型造成的。例如，我正在创建一个如下所示的数据框；

import pandas as pd

d = {'col1': [1, 2], 'col2': [3, 4]}
df = pd.DataFrame(d)
df.col1 = df.col1.astype('int8')

df.info()

输出看起来像这样：

<class 'pandas.core.frame.DataFrame'>
RangeIndex: 2 entries, 0 to 1
Data columns (total 2 columns):
col1    2 non-null int8
col2    2 non-null int64
dtypes: int64(1), int8(1)
memory usage: 98.0 bytes

它只有 98 字节。

我将其写入文件并再次读取。

<class 'pandas.core.frame.DataFrame'>
RangeIndex: 2 entries, 0 to 1
Data columns (total 2 columns):
col1    2 non-null int8
col2    2 non-null int64
dtypes: int64(1), int8(1)
memory usage: 98.0 bytes

输出看起来像这样：

<class 'pandas.core.frame.DataFrame'>
RangeIndex: 2 entries, 0 to 1
Data columns (total 2 columns):
col1    2 non-null int64
col2    2 non-null int64
dtypes: int64(2)
memory usage: 112.0 bytes

现在内存使用112字节。这里的问题是，当读取 csv 文件时，它读取为 int64。我在一个大型数据帧上执行此操作，我的文件大小为250 mb达到1.14 GB

我的问题是；有没有办法自动将数据帧上的列类型转换为尽可能小的大小？我尝试了函数infer_dtypes但没有得到我想要的结果。它说它应该是整数，因为它应该是这样的类型。

Answer 1

经过一些研究， to_numeric 函数工作正常。我已经实现了自己的实现，如下所示。

我从 numpy 数据类型创建了一个数据框对象。

np_types = [np.int8 ,np.int16 ,np.int32, np.int64,
           np.uint8 ,np.uint16, np.uint32, np.uint64]
np_types = [np_type.__name__ for np_type in np_types]
type_df = pd.DataFrame(data=np_types, columns=['class_type'])
type_df

结果如下：

然后我将有关类型的信息添加到数据框中

np_types = [np.int8 ,np.int16 ,np.int32, np.int64,
           np.uint8 ,np.uint16, np.uint32, np.uint64]
np_types = [np_type.__name__ for np_type in np_types]
type_df = pd.DataFrame(data=np_types, columns=['class_type'])
type_df

然后我在整数列上编写了一个函数，以找出哪种类型比最小值和最大值更合适。

def optimize_types(dataframe):
for col in dataframe.loc[:, dataframe.dtypes <= np.integer]:
    col_min = dataframe[col].min()
    col_max = dataframe[col].max()
    temp = type_df[(type_df['min_value'] <= col_min) & (type_df['max_value'] >= col_max)]
    optimized_class = temp.loc[temp['range'].idxmin(), 'class_type']
    print("Col name : {} Col min_value : {} Col max_value : {} Optimized Class : {}".format(col, col_min, col_max, optimized_class))
    dataframe[col] = dataframe[col].astype(optimized_class)
return dataframe

我有一个 2.6 GB 的数据帧。通过上述功能，它减少到了 600 mb。

当我使用 to_numeric 函数时，我得到以下结果：

Answer 2

适用于所有数字类型，有助于摆脱np.int64and np.float64：

import numbers
import pandas as pd
from typing import Optional

def auto_opt_pd_dtypes(df_: pd.DataFrame, inplace=False) -> Optional[pd.DataFrame]:
    """ Automatically downcast Number dtypes for minimal possible,
        will not touch other (datetime, str, object, etc)
        
        :param df_: dataframe
        :param inplace: if False, will return a copy of input dataset
        
        :return: `None` if `inplace=True` or dataframe if `inplace=False`
    """
    df = df_ if inplace else df_.copy()
        
    for col in df.columns:
        # integers
        if issubclass(df[col].dtypes.type, numbers.Integral):
            # unsigned integers
            if df[col].min() >= 0:
                df[col] = pd.to_numeric(df[col], downcast='unsigned')
            # signed integers
            else:
                df[col] = pd.to_numeric(df[col], downcast='integer')
        # other real numbers
        elif issubclass(df[col].dtypes.type, numbers.Real):
            df[col] = pd.to_numeric(df[col], downcast='float')
    
    if not inplace:
        return df

用法：

# return optimized copy
df_opt = auto_opt_pd_dtypes(df)
# or optimize in place
auto_opt_pd_dtypes(df, inplace=True)

Answer 3

如果列都是数字，您可以执行以下操作：

import numpy as np
df = df.astype(np.int8)

如果列不都是数字，您可以首先对它们进行切片，选择数字列，然后调用astype.