feat(util.pandas): Add optimise_df to reduce dataframe size.

aaronmussig · Apr 15, 2022 · 706baff · 706baff
1 parent 52bca3b
commit 706baff
Show file tree

Hide file tree

Showing 3 changed files with 21 additions and 0 deletions.
diff --git a/docs/source/index.rst b/docs/source/index.rst
@@ -52,6 +52,7 @@ This has been written with the intention of personal use, but feel free to use/c
    util/io
    util/accession
    util/tree
+   util/pandas
 
 
 .. toctree::

diff --git a/docs/source/util/pandas.rst b/docs/source/util/pandas.rst
@@ -0,0 +1,6 @@
+******
+Pandas
+******
+
+.. autofunction:: magna.util.pandas.optimise_df
+
diff --git a/magna/util/pandas.py b/magna/util/pandas.py
@@ -0,0 +1,14 @@
+import pandas as pd
+
+
+def optimise_df(df: pd.DataFrame):
+    """Optimise a Pandas DataFrame by using the smallest possible data type.
+
+    Args:
+        df: The Pandas DataFrame to optimise.
+    """
+
+    float_cols = df.select_dtypes('float').columns
+    int_cols = df.select_dtypes('integer').columns
+    df[float_cols] = df[float_cols].apply(pd.to_numeric, downcast='float')
+    df[int_cols] = df[int_cols].apply(pd.to_numeric, downcast='integer')