Repository: modin-project/modin

Issue:
REFACTOR: Push more of memory_usage down to query compiler.
The current implementation creates a separate modin.pandas dataframe with the index memory usage and concatenates that to the memory usage without the index. This is convenient for partitioned dataframes as in modin on ray, because computing the extra index column in a distributed way is difficult. However, there are cases (e.g. when using a database for storage and execution) where creating a pandas dataframe and concatenating it may be expensive or impossible.


Discussion:

Resolving patch:
diff --git a/modin/core/storage_formats/pandas/query_compiler.py b/modin/core/storage_formats/pandas/query_compiler.py
--- a/modin/core/storage_formats/pandas/query_compiler.py
+++ b/modin/core/storage_formats/pandas/query_compiler.py
@@ -844,12 +844,31 @@ def is_series_like(self):
     prod = TreeReduce.register(pandas.DataFrame.prod)
     any = TreeReduce.register(pandas.DataFrame.any, pandas.DataFrame.any)
     all = TreeReduce.register(pandas.DataFrame.all, pandas.DataFrame.all)
-    memory_usage = TreeReduce.register(
+    # memory_usage adds an extra column for index usage, but we don't want to distribute
+    # the index memory usage calculation.
+    _memory_usage_without_index = TreeReduce.register(
         pandas.DataFrame.memory_usage,
         lambda x, *args, **kwargs: pandas.DataFrame.sum(x),
         axis=0,
     )
 
+    def memory_usage(self, **kwargs):
+        index = kwargs.get("index", True)
+        deep = kwargs.get("deep", False)
+        usage_without_index = self._memory_usage_without_index(index=False, deep=deep)
+        return (
+            self.from_pandas(
+                pandas.DataFrame(
+                    [self.index.memory_usage()],
+                    columns=["Index"],
+                    index=[MODIN_UNNAMED_SERIES_LABEL],
+                ),
+                data_cls=type(self._modin_frame),
+            ).concat(axis=1, other=[usage_without_index])
+            if index
+            else usage_without_index
+        )
+
     def max(self, axis, **kwargs):
         def map_func(df, **kwargs):
             return pandas.DataFrame.max(df, **kwargs)
diff --git a/modin/pandas/dataframe.py b/modin/pandas/dataframe.py
--- a/modin/pandas/dataframe.py
+++ b/modin/pandas/dataframe.py
@@ -42,7 +42,6 @@
 from typing import IO, Optional, Union, Iterator, Hashable, Sequence
 import warnings
 
-import modin
 from modin.pandas import Categorical
 from modin.error_message import ErrorMessage
 from modin.utils import (
@@ -1460,18 +1459,6 @@ def melt(
             )
         )
 
-    def memory_usage(self, index=True, deep=False):  # noqa: PR01, RT01, D200
-        """
-        Return the memory usage of each column in bytes.
-        """
-        if index:
-            result = self._reduce_dimension(
-                self._query_compiler.memory_usage(index=False, deep=deep)
-            )
-            index_value = self.index.memory_usage(deep=deep)
-            return modin.pandas.concat([Series(index_value, index=["Index"]), result])
-        return super(DataFrame, self).memory_usage(index=index, deep=deep)
-
     def merge(
         self,
         right,
diff --git a/modin/pandas/series.py b/modin/pandas/series.py
--- a/modin/pandas/series.py
+++ b/modin/pandas/series.py
@@ -1327,13 +1327,7 @@ def memory_usage(self, index=True, deep=False):  # noqa: PR01, RT01, D200
         """
         Return the memory usage of the Series.
         """
-        if index:
-            result = self._reduce_dimension(
-                self._query_compiler.memory_usage(index=False, deep=deep)
-            )
-            index_value = self.index.memory_usage(deep=deep)
-            return result + index_value
-        return super(Series, self).memory_usage(index=index, deep=deep)
+        return super(Series, self).memory_usage(index=index, deep=deep).sum()
 
     def mod(self, other, level=None, fill_value=None, axis=0):  # noqa: PR01, RT01, D200
         """


Test patch: