Apply the smallest possible datatype for each column in a pandas dataframe to reduce RAM use

Question

This function is for Python Pandas users. I noticed a lot of RAM can be saved by applying smaller datatypes to columns. The function's input is a Pandas DataFrame. The function output is the compact dataframe and a list of columns that have been altered by filling out NA values with dummy values. This function iterates over every numeric column and does the following:

• Determines if the column is numeric
• Determine if the column can be represented by an integer
• Determine the range of values (min-max)
• Apply the smallest datatype that can fit that range

Important: Integer datatypes do not support NA values. Therefore, I fill NA's with the lowest value in the column, minus one. This way, I can still select an integer datatype for columns containing NA. I keep track of these columns in the NAlist. After running the function, it is important that the code-user reviews each column in the NAlist and determines the appropriate way to handle the missing values.

On my own dataset, I reduced the size from 1.3GB to 466 MB

def reduce_mem_usage(df):
    start_mem_usg = df.memory_usage().sum() / 1024**2 
    print("Memory usage of properties dataframe is :",start_mem_usg," MB")
    NAlist = [] # Keeps track of columns that have missing values filled in. 
    for col in df.columns:
        if df[col].dtype != object:  # Exclude strings

            # Print current column type
            print("******************************")
            print("Column: ",col)
            print("dtype before: ",df[col].dtype)

            # make variables for Int, max and min
            IsInt = False
            mx = df[col].max()
            mn = df[col].min()

            # Integer does not support NA, therefore, NA needs to be filled
            if not np.isfinite(df[col]).all(): 
                NAlist.append(col)
                df[col].fillna(mn-1,inplace=True)  

            # test if column can be converted to an integer
            asint = df[col].fillna(0).astype(np.int64)
            result = (df[col] - asint)
            result = result.sum()
            if result > -0.01 and result < 0.01:
                IsInt = True


            # Make Integer/unsigned Integer datatypes
            if IsInt:
                if mn >= 0:
                    if mx < 255:
                        df[col] = df[col].astype(np.uint8)
                    elif mx < 65535:
                        df[col] = df[col].astype(np.uint16)
                    elif mx < 4294967295:
                        df[col] = df[col].astype(np.uint32)
                    else:
                        df[col] = df[col].astype(np.uint64)
                else:
                    if mn > np.iinfo(np.int8).min and mx < np.iinfo(np.int8).max:
                        df[col] = df[col].astype(np.int8)
                    elif mn > np.iinfo(np.int16).min and mx < np.iinfo(np.int16).max:
                        df[col] = df[col].astype(np.int16)
                    elif mn > np.iinfo(np.int32).min and mx < np.iinfo(np.int32).max:
                        df[col] = df[col].astype(np.int32)
                    elif mn > np.iinfo(np.int64).min and mx < np.iinfo(np.int64).max:
                        df[col] = df[col].astype(np.int64)    

            # Make float datatypes
            else:
                if mn > np.finfo(np.float16).min and mx < np.finfo(np.float16).max:
                    df[col] = df[col].astype(np.float16)
                elif mn > np.finfo(np.float32).min and mx < np.finfo(np.float32).max:
                    df[col] = df[col].astype(np.float32)
                elif mn > np.finfo(np.float64).min and mx < np.finfo(np.float64).max:
                    df[col] = df[col].astype(np.float64)   

            # Print new column type
            print("dtype after: ",df[col].dtype)
            print("******************************")

    # Print final result
    print("___MEMORY USAGE AFTER COMPLETION:___")
    mem_usg = df.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 df, NAlist

Answer 1

This review is primarily about code clarity not function. But, before we get going, a couple of quick comments about function:

1. I am unsure of the veracity of encoding Nan Integers as one less than the minimum value in your data.

   You mention that people should check afterwards to insure what to do with the now encoded Nans. But this method basically requires someone to hand manage everything. I am not sure of the practicality of that.

2. The mechanism used to decided to convert to integers could have issues if all numbers were very small.

   Very small numbers will look like integers, and will all be converted to 0!

3. Floating point numbers by their nature are not really limited by the max value.

   The reason for different sized floats is to be able to encode different levels of precision. So deciding to shrink a float down in size based on the maximum value in a data set is highly unlikely to be generally useful.

But having said that I will assume the numeric functions are relevant and correct for your data, and will discuss how to make the code more Pythonic.

Stacked ifs often should be consolidated

Roughly half the code you have presented here is a large structure of stacked ifs that are looking for the range of values. Each if and the action from the if are all very similar. In many cases it is much clearer to code only the differences between the ifs, and then have the common code work with the codified differences. So for example, this:

# Make Integer/unsigned Integer datatypes
if IsInt:
    if mn >= 0:
        if mx < 255:
            df[col] = df[col].astype(np.uint8)
        elif mx < 65535:
            ....
            df[col] = df[col].astype(np.uint16)
    else:
        if mn > np.iinfo(np.int8).min and mx < np.iinfo(np.int8).max:
            df[col] = df[col].astype(np.int8)
        elif mn > np.iinfo(np.int16).min and mx < np.iinfo(np.int16).max:
            ....

# Make float datatypes
else:
    if mn > np.finfo(np.float16).min and mx < np.finfo(np.float16).max:
        df[col] = df[col].astype(np.float16)
    elif mn > np.finfo(np.float32).min and mx < np.finfo(np.float32).max:
        ....

Can be changed to:

if IsInt:
    info = np.iinfo
    # Make Integer/unsigned Integer datatypes
    if mn >= 0:
        types = (np.uint8, np.uint16, np.uint32, np.uint64)
    else:
        types = (np.int8, np.int16, np.int32, np.int64)
else:
    info = np.finfo
    types = (np.float16, np.float32, np.float64)

for t in types:
    if info(t).min <= mn and mx <= info(t).max:
        df[col] = df[col].astype(t)
        break

This code makes the differences between signed, unsigned and floats explicit. And then the action code works against a description of the differences and so the action taken is more explicit. Additionally, the resulting code is much smaller.

Chained Comparisons

I would change:

# test if column can be converted to an integer
asint = df[col].fillna(0).astype(np.int64)
result = (df[col] - asint)
result = result.sum()
if result > -0.01 and result < 0.01:
    IsInt = True

# Make Integer/unsigned Integer datatypes
if IsInt:

To:

# test if column can be converted to an integer
asint = df[col].fillna(0).astype(np.int64)
errors_from_convert_to_int = (df[col] - asint).sum()

if -0.01 < errors_from_convert_to_int < 0.01:

Changes include:

• name the result in a manner to describe what was calculated
• use chained comparisons
• remove unneeded intermediate variable IsInt

Be Careful with Ranges

In the code to encode a Nan to min-1, there is a comment about integers, but the code is applied to all columns. In addition the range needs to be expanded to include the now lower minimum value. How about:

if -0.01 < convert_to_int_errors < 0.01:
    # Integer does not support NA, therefore, NA needs to be filled
    if not np.isfinite(df[col]).all():
        na_list.append(col)
        mn -= 1
        df[col].fillna(mn, inplace=True)            

Bonus Comment:

This suggestion fits more in personal style bucket. The main loop works against each column in turn, but does want to work on object columns. That is coded as:

for col in df.columns:
    if df[col].dtype != object:   

But to know that the code always only works with non-object columns requires the reader to scroll down looking for an else. The fact that the loop only works against non-objects can be made more explicit (ie: done in one expression) with something like:

for col in (c for c in df.columns if df[c].dtype != object):