緩解pandas中DataFrame佔用內存過大的問題

本文最後修改於2018-01-21，文章有問題或者轉載請及時聯繫本人，如果對你有幫助，別忘了點下關注~

0 背景

在我們使用pandas進行數據處理的時候，有時候發現文件在本地明明不大，但是用pandas以DataFrame形式加載內存中的時候會佔用非常高的內存，本文即解決這樣的問題。

1 原因

如果是計算機相關專業的同學，你應該知道int8，int16，int64的區別。如果你忘記了，那我們可以舉一個例子。內存相當於倉庫，數字相當於貨物，數字需要裝到箱子裏才能堆到倉庫。現在有小，中，大三種箱子，我們一個個數字用小箱子就可以裝好，然後堆到倉庫去，而現在pandas的處理邏輯是，如果你不告訴用哪個箱子，我都會用最大的箱子去裝，這樣倉庫很快就滿了。OK，這就是有時候DataFrame內存佔用過高的原因。

常用的數據類型範圍如下所示：

header 1	範圍下限（含）	範圍上限（含）
unit8	0	255
unit16	0	65535
int8	-128	127
int16	-32768	32767
int32	-2147483648	2147483647
int64	–9,223,372,036,854,775,808	9,223,372,036,854,775,807

其他信息可以通過numpy中的函數來查看範圍

import numpy as np
# 查看int16的範圍
ii16 = np.iinfo(np.int16)
ii16.min
-32768

# 與iinfo相應，finfo可以查看float類型的範圍
fi16 = np.finfo(np.float16)
fin16.min
-3.4028235e+38

2 解決方法

第一種

當我們明確知道要加載數據的範圍，使用pd.read_table讀取數據時，可以用其中的dtype參數來手動指定類型。比如某一列的數據範圍肯定在0~255之中，那麼我們可以指定爲np.uint8類型，如果不手動指定的話默認爲np.int64類型，這之間的差距巨大。

第二種

如果數據列數太多，或者不清楚數據具體範圍的話這裏提供一個腳本，可以自動判斷類型，並根據類型修改數據範圍。雖然我認爲這個腳本已經可以cover大部分的情況，但是仍然強烈建議你在讀懂這個腳本的基礎上，根據你的數據修改成更適合你數據的形式。

注意：代碼最初源於這裏Reducing DataFrame memory size by ~65%，在Apache 2.0協議下，我對其中不太合理的地方做了些修改。

# @from: https://www.kaggle.com/arjanso/reducing-dataframe-memory-size-by-65/code
# @liscense: Apache 2.0
# @author: weijian
def reduce_mem_usage(props):
    # 計算當前內存
    start_mem_usg = props.memory_usage().sum() / 1024 ** 2
    print("Memory usage of the dataframe is :", start_mem_usg, "MB")
    
    # 哪些列包含空值，空值用-999填充。why：因爲np.nan當做float處理
    NAlist = []
    for col in props.columns:
        # 這裏只過濾了objectd格式，如果你的代碼中還包含其他類型，請一併過濾
        if (props[col].dtypes != object):
            
            print("**************************")
            print("columns: ", col)
            print("dtype before", props[col].dtype)
            
            # 判斷是否是int類型
            isInt = False
            mmax = props[col].max()
            mmin = props[col].min()
            
            # Integer does not support NA, therefore Na needs to be filled
            if not np.isfinite(props[col]).all():
                NAlist.append(col)
                props[col].fillna(-999, inplace=True) # 用-999填充
                
            # test if column can be converted to an integer
            asint = props[col].fillna(0).astype(np.int64)
            result = np.fabs(props[col] - asint)
            result = result.sum()
            if result < 0.01: # 絕對誤差和小於0.01認爲可以轉換的，要根據task修改
                isInt = True
            
            # make interger / unsigned Integer datatypes
            if isInt:
                if mmin >= 0: # 最小值大於0，轉換成無符號整型
                    if mmax <= 255:
                        props[col] = props[col].astype(np.uint8)
                    elif mmax <= 65535:
                        props[col] = props[col].astype(np.uint16)
                    elif mmax <= 4294967295:
                        props[col] = props[col].astype(np.uint32)
                    else:
                        props[col] = props[col].astype(np.uint64)
                else: # 轉換成有符號整型
                    if mmin > np.iinfo(np.int8).min and mmax < np.iinfo(np.int8).max:
                        props[col] = props[col].astype(np.int8)
                    elif mmin > np.iinfo(np.int16).min and mmax < np.iinfo(np.int16).max:
                        props[col] = props[col].astype(np.int16)
                    elif mmin > np.iinfo(np.int32).min and mmax < np.iinfo(np.int32).max:
                        props[col] = props[col].astype(np.int32)
                    elif mmin > np.iinfo(np.int64).min and mmax < np.iinfo(np.int64).max:
                        props[col] = props[col].astype(np.int64)  
            else: # 注意：這裏對於float都轉換成float16，需要根據你的情況自己更改
                props[col] = props[col].astype(np.float16)
            
            print("dtype after", props[col].dtype)
            print("********************************")
    print("___MEMORY USAGE AFTER COMPLETION:___")
    mem_usg = props.memory_usage().sum() / 1024**2 
    print("Memory usage is: ",mem_usg," MB")
    print("This is ",100*mem_usg/start_mem_usg,"% of the initial size")
    return props, NAlist

其中NAlist表明是含有空值的列表。經試驗，效果明顯，對於數值型特徵非常多的數據，至少可以減少50%以上的內存佔用。

Reference

1 https://www.kaggle.com/arjanso/reducing-dataframe-memory-size-by-65/code