Using a Python Decorator to Print Memory Usage of a DataFrame
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.
1| import functools 2| import pandas as pd 3| 4| def print_memory_usage(f): 5| @functools.wraps(f) 6| def wrapper(*args, **kwargs): 7| df = f(*args, **kwargs) 8| print('Memory Usage:',df.memory_usage(index=False).sum(),'bytes') 9| return wrapper 10| 11| @print_memory_usage 12| def create_dataframe(*args, **kwargs): 13| df = pd.DataFrame(*args,**kwargs) 14| return df 15| 16| df = create_dataframe([[1,2,3,4],[2,3,4,6]],columns=['A','B','C','D']) 17| 18| >> Memory Usage: 64 bytes
