SNIPPET
2 Upvotes

Using a Python Decorator to Print Memory Usage of a DataFrame

Python
General Python

This decorator can be applied to any function that manipulates and then returns a dataframe. Decorators allow you to pass an existing function into a wrapper within which the function will run. This allows you to run some code before the function and then also some other code after the function.

In this example we have a function called create_dataframe that creates a dataframe from some arguments passed at the time the function is called. By applying the print_memory_usage decorator to the create_dataframe function we pass the create_dataframe to the print_memory_usage function and subsequently the wrapper function. In here the create_dataframe is called and then after we have code that prints the memory usage of the created dataframe.

Note: In this example, index=False is passed to the memory_usage function. This ensure that the memory used in the index is not summed, just the values of the dataframe.

import functools
import pandas as pd

def print_memory_usage(f):
    @functools.wraps(f)
    def wrapper(*args, **kwargs):
        df = f(*args, **kwargs)
        print('Memory Usage:',df.memory_usage(index=False).sum(),'bytes')
    return wrapper

@print_memory_usage
def create_dataframe(*args, **kwargs):
    df = pd.DataFrame(*args,**kwargs)
    return df

df = create_dataframe([[1,2,3,4],[2,3,4,6]],columns=['A','B','C','D'])

>> Memory Usage: 64 bytes

By detro - Last Updated Nov. 20, 2021, 3:50 p.m.

COMMENTS
RELATED SNIPPETS
Python For Else Loop
Python
General Python

Loop

2
1
Search Snippets by Tag: