Repository: mars-project/mars

Issue:
Support memory_usage on DataFrame objects
Support memory_usage on md.{DataFrame,Series,Index}


Discussion:

Resolving patch:
diff --git a/mars/dataframe/base/__init__.py b/mars/dataframe/base/__init__.py
--- a/mars/dataframe/base/__init__.py
+++ b/mars/dataframe/base/__init__.py
@@ -32,6 +32,7 @@
 from .drop_duplicates import df_drop_duplicates, \
     series_drop_duplicates, index_drop_duplicates
 from .melt import melt
+from .memory_usage import df_memory_usage, series_memory_usage, index_memory_usage
 
 
 def _install():
@@ -65,6 +66,7 @@ def _install():
         setattr(t, '__delitem__', lambda df, items: df_drop(df, items, axis=1, inplace=True))
         setattr(t, 'drop_duplicates', df_drop_duplicates)
         setattr(t, 'melt', melt)
+        setattr(t, 'memory_usage', df_memory_usage)
 
     for t in SERIES_TYPE:
         setattr(t, 'to_gpu', to_gpu)
@@ -90,11 +92,13 @@ def _install():
         setattr(t, 'astype', astype)
         setattr(t, 'drop', series_drop)
         setattr(t, 'drop_duplicates', series_drop_duplicates)
+        setattr(t, 'memory_usage', series_memory_usage)
 
     for t in INDEX_TYPE:
         setattr(t, 'rechunk', rechunk)
         setattr(t, 'drop', index_drop)
         setattr(t, 'drop_duplicates', index_drop_duplicates)
+        setattr(t, 'memory_usage', index_memory_usage)
 
     for method in _string_method_to_handlers:
         if not hasattr(StringAccessor, method):
diff --git a/mars/dataframe/base/memory_usage.py b/mars/dataframe/base/memory_usage.py
new file mode 100644
--- /dev/null
+++ b/mars/dataframe/base/memory_usage.py
@@ -0,0 +1,442 @@
+# Copyright 1999-2020 Alibaba Group Holding Ltd.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import operator
+from functools import reduce
+
+import numpy as np
+import pandas as pd
+
+from ... import opcodes
+from ...config import options
+from ...operands import OperandStage
+from ...serialize import BoolField, Int64Field
+from ...utils import ceildiv, lazy_import
+from ..operands import DataFrameOperand, DataFrameOperandMixin
+from ..core import IndexValue
+from ..utils import parse_index
+
+cudf = lazy_import('cudf', globals=globals())
+
+
+class DataFrameMemoryUsage(DataFrameOperand, DataFrameOperandMixin):
+    _op_type_ = opcodes.MEMORY_USAGE
+
+    # raw arguments of memory_usage method
+    _index = BoolField('index')
+    _deep = BoolField('deep')
+
+    # size of range index, when set, the value will be prepended to the result series
+    # if the input is a dataframe, or added to the result when the input is a series
+    _range_index_size = Int64Field('range_index_size')
+
+    def __init__(self, index=None, deep=None, range_index_size=None, **kw):
+        super().__init__(_index=index, _deep=deep, _range_index_size=range_index_size, **kw)
+
+    @property
+    def index(self) -> bool:
+        return self._index
+
+    @index.setter
+    def index(self, value: bool):
+        self._index = value
+
+    @property
+    def deep(self) -> bool:
+        return self._deep
+
+    @property
+    def range_index_size(self) -> int:
+        return self._range_index_size
+
+    @range_index_size.setter
+    def range_index_size(self, value: int):
+        self._range_index_size = value
+
+    def _adapt_index(self, input_index, index=0):
+        """
+        When ``index=True`` is passed, an extra column will be prepended to the result series
+        Thus we need to update the index of initial chunk for returned dataframe chunks
+        """
+        if not self.index or index != 0:
+            return input_index
+        idx_data = input_index.to_pandas().insert(0, 'Index')
+        return parse_index(idx_data, store_data=True)
+
+    def _adapt_nsplits(self, input_nsplit):
+        """
+        When ``index=True`` is passed, the size of returned series is one element larger
+        than the number of columns, which affects ``nsplits``.
+        """
+        if not self.index:
+            return (input_nsplit[-1],)
+        nsplits_list = list(input_nsplit[-1])
+        nsplits_list[0] += 1
+        return (tuple(nsplits_list),)
+
+    def __call__(self, df_or_series):
+        """
+        Return output object of memory_usage() call
+        """
+        if df_or_series.ndim == 1:
+            # the input data is a series, a Scalar will be returned
+            return self.new_scalar([df_or_series], dtype=np.dtype(np.int_))
+        else:
+            # the input data is a DataFrame, a Scalar will be returned
+            # calculate shape of returning series given ``op.index``
+            new_shape = (df_or_series.shape[-1] + 1,) if self.index else (df_or_series.shape[-1],)
+            return self.new_series(
+                [df_or_series], index_value=self._adapt_index(df_or_series.columns_value),
+                shape=new_shape, dtype=np.dtype(np.int_))
+
+    @classmethod
+    def _tile_single(cls, op: "DataFrameMemoryUsage"):
+        """
+        Tile when input data has only one chunk on rows
+        """
+        df_or_series = op.inputs[0]
+        output = op.outputs[0]
+
+        chunks = []
+        for c in df_or_series.chunks:
+            new_op = op.copy().reset_key()
+            if c.ndim == 1:
+                # Tile for series
+                chunks.append(new_op.new_chunk([c], index=c.index, dtype=output.dtype, shape=()))
+            else:
+                # tile for dataframes
+                # only calculate with index=True on the initial chunk
+                new_op.index = op.index and c.index[-1] == 0
+
+                # calculate shape of returning chunk given ``op.index``
+                new_shape = (c.shape[-1] + 1,) if c.index[-1] == 0 and op.index else (c.shape[-1],)
+                chunks.append(new_op.new_chunk(
+                    [c], shape=new_shape, dtype=output.dtype, index=(c.index[-1],),
+                    index_value=op._adapt_index(c.columns_value, c.index[-1])))
+
+        new_op = op.copy().reset_key()
+        # return objects with chunks and nsplits (if needed)
+        if df_or_series.ndim == 1:
+            return new_op.new_scalar([df_or_series], dtype=output.dtype, chunks=chunks)
+        else:
+            return new_op.new_series([df_or_series], shape=output.shape, dtype=output.dtype,
+                                     index_value=output.index_value, chunks=chunks,
+                                     nsplits=op._adapt_nsplits(df_or_series.nsplits))
+
+    @classmethod
+    def _tile_dataframe(cls, op: "DataFrameMemoryUsage"):
+        """
+        Tile dataframes using tree reduction
+        """
+        df = op.inputs[0]
+        output = op.outputs[0]
+        is_range_index = isinstance(df.index_value.value, IndexValue.RangeIndex)
+
+        # produce map chunks
+        # allocate matrix of chunks
+        chunks_to_reduce = np.empty(shape=df.chunk_shape, dtype=np.object)
+        for c in df.chunks:
+            new_op = op.copy().reset_key()
+            new_op.stage = OperandStage.map
+
+            if op.index and is_range_index:
+                # when the index is ``pd.RangeIndex``, the size should be included
+                # after all computations are done
+                new_op.index = False
+            else:
+                # when the chunk is not the first chunk in the row, index size is not needed
+                new_op.index = op.index and c.index[-1] == 0
+
+            new_shape = (c.shape[-1] + 1,) if c.index[-1] == 0 and op.index else (c.shape[-1],)
+
+            chunks_to_reduce[c.index] = new_op.new_chunk(
+                [c], index=(c.index[-1],), dtype=output.dtype, shape=new_shape,
+                index_value=op._adapt_index(c.columns_value, c.index[-1]))
+
+        # reduce chunks using tree reduction
+        combine_size = options.combine_size
+        while chunks_to_reduce.shape[0] > 1:
+            # allocate matrix of chunks
+            new_chunks_to_reduce = np.empty((ceildiv(chunks_to_reduce.shape[0], combine_size),
+                                             chunks_to_reduce.shape[1]), dtype=np.object)
+            for idx in range(0, chunks_to_reduce.shape[0], combine_size):
+                for idx2 in range(chunks_to_reduce.shape[1]):
+                    new_op = op.copy().reset_key()
+                    new_op._stage = OperandStage.reduce
+                    chunks = list(chunks_to_reduce[idx:idx + combine_size, idx2])
+
+                    new_chunks_to_reduce[idx // combine_size, idx2] = new_op.new_chunk(
+                        chunks, index=(idx2,), dtype=output.dtype, shape=chunks[0].shape,
+                        index_value=chunks[0].index_value)
+
+            chunks_to_reduce = new_chunks_to_reduce
+
+        # handle RangeIndex at final outputs
+        if op.index and is_range_index:
+            chunks_to_reduce[0, 0].op.range_index_size = df.index_value.to_pandas().memory_usage()
+
+        # return series with chunks and nsplits
+        new_op = op.copy().reset_key()
+        return new_op.new_series([df], dtype=output.dtype, shape=output.shape,
+                                 index_value=output.index_value,
+                                 chunks=list(chunks_to_reduce[0, :]),
+                                 nsplits=op._adapt_nsplits(df.nsplits))
+
+    @classmethod
+    def _tile_series(cls, op: "DataFrameMemoryUsage"):
+        """
+        Tile series using tree reduction
+        """
+        series = op.inputs[0]
+        output = op.outputs[0]
+        is_range_index = isinstance(series.index_value.value, IndexValue.RangeIndex)
+
+        chunks_to_reduce = []
+        for c in series.chunks:
+            new_op = op.copy().reset_key()
+            new_op.stage = OperandStage.map
+
+            # when the index is ``pd.RangeIndex``, the size should be included
+            # after all computations are done
+            new_op.index = op.index and not is_range_index
+
+            chunks_to_reduce.append(
+                new_op.new_chunk([c], index=c.index, dtype=output.dtype, shape=()))
+
+        # reduce chunks using tree reduction
+        combine_size = options.combine_size
+        while len(chunks_to_reduce) > 1:
+            new_chunks_to_reduce = []
+            for idx in range(0, len(chunks_to_reduce), combine_size):
+                new_op = op.copy().reset_key()
+                new_op.stage = OperandStage.reduce
+
+                new_chunks_to_reduce.append(new_op.new_chunk(
+                    chunks_to_reduce[idx:idx + combine_size], shape=(), index=(0,),
+                    dtype=output.dtype))
+
+            chunks_to_reduce = new_chunks_to_reduce
+
+        # handle RangeIndex at final outputs
+        if op.index and is_range_index:
+            chunks_to_reduce[0].op.range_index_size = series.index_value.to_pandas().memory_usage()
+
+        # return series with chunks
+        new_op = op.copy().reset_key()
+        return new_op.new_scalar([series], dtype=output.dtype, chunks=chunks_to_reduce)
+
+    @classmethod
+    def tile(cls, op: "DataFrameMemoryUsage"):
+        df_or_series = op.inputs[0]
+        if df_or_series.chunk_shape[0] == 1:  # only one chunk in row, no aggregation needed
+            return cls._tile_single(op)
+        elif df_or_series.ndim == 1:   # series
+            return cls._tile_series(op)
+        else:  # dataframe
+            return cls._tile_dataframe(op)
+
+    @classmethod
+    def execute(cls, ctx, op: "DataFrameMemoryUsage"):
+        in_data = ctx[op.inputs[0].key]
+        # choose correct dataframe library
+        xdf = cudf if op.gpu else pd
+
+        if op.stage == OperandStage.reduce:
+            result = reduce(operator.add, (ctx[c.key] for c in op.inputs))
+            if op.range_index_size is not None:
+                if hasattr(in_data, 'ndim'):
+                    # dataframe input: prepend index size column
+                    prepend_series = xdf.Series([op.range_index_size], index=['Index'], dtype=result.dtype)
+                    result = xdf.concat([prepend_series, result])
+                else:
+                    # series input: add index size to the output
+                    result += op.range_index_size
+            ctx[op.outputs[0].key] = result
+        elif isinstance(in_data, xdf.Index):
+            ctx[op.outputs[0].key] = in_data.memory_usage(deep=op.deep)
+        else:
+            ctx[op.outputs[0].key] = in_data.memory_usage(index=op.index, deep=op.deep)
+
+
+def df_memory_usage(df, index=True, deep=False):
+    """
+    Return the memory usage of each column in bytes.
+
+    The memory usage can optionally include the contribution of
+    the index and elements of `object` dtype.
+
+    This value is displayed in `DataFrame.info` by default. This can be
+    suppressed by setting ``pandas.options.display.memory_usage`` to False.
+
+    Parameters
+    ----------
+    index : bool, default True
+        Specifies whether to include the memory usage of the DataFrame's
+        index in returned Series. If ``index=True``, the memory usage of
+        the index is the first item in the output.
+    deep : bool, default False
+        If True, introspect the data deeply by interrogating
+        `object` dtypes for system-level memory consumption, and include
+        it in the returned values.
+
+    Returns
+    -------
+    Series
+        A Series whose index is the original column names and whose values
+        is the memory usage of each column in bytes.
+
+    See Also
+    --------
+    numpy.ndarray.nbytes : Total bytes consumed by the elements of an
+        ndarray.
+    Series.memory_usage : Bytes consumed by a Series.
+    Categorical : Memory-efficient array for string values with
+        many repeated values.
+    DataFrame.info : Concise summary of a DataFrame.
+
+    Examples
+    --------
+    >>> import mars.tensor as mt
+    >>> import mars.dataframe as md
+    >>> dtypes = ['int64', 'float64', 'complex128', 'object', 'bool']
+    >>> data = dict([(t, mt.ones(shape=5000).astype(t))
+    ...              for t in dtypes])
+    >>> df = md.DataFrame(data)
+    >>> df.head().execute()
+       int64  float64            complex128  object  bool
+    0      1      1.0    1.000000+0.000000j       1  True
+    1      1      1.0    1.000000+0.000000j       1  True
+    2      1      1.0    1.000000+0.000000j       1  True
+    3      1      1.0    1.000000+0.000000j       1  True
+    4      1      1.0    1.000000+0.000000j       1  True
+
+    >>> df.memory_usage().execute()
+    Index           128
+    int64         40000
+    float64       40000
+    complex128    80000
+    object        40000
+    bool           5000
+    dtype: int64
+
+    >>> df.memory_usage(index=False).execute()
+    int64         40000
+    float64       40000
+    complex128    80000
+    object        40000
+    bool           5000
+    dtype: int64
+
+    The memory footprint of `object` dtype columns is ignored by default:
+
+    >>> df.memory_usage(deep=True).execute()
+    Index            128
+    int64          40000
+    float64        40000
+    complex128     80000
+    object        160000
+    bool            5000
+    dtype: int64
+
+    Use a Categorical for efficient storage of an object-dtype column with
+    many repeated values.
+
+    >>> df['object'].astype('category').memory_usage(deep=True).execute()
+    5216
+    """
+    op = DataFrameMemoryUsage(index=index, deep=deep)
+    return op(df)
+
+
+def series_memory_usage(series, index=True, deep=False):
+    """
+    Return the memory usage of the Series.
+
+    The memory usage can optionally include the contribution of
+    the index and of elements of `object` dtype.
+
+    Parameters
+    ----------
+    index : bool, default True
+        Specifies whether to include the memory usage of the Series index.
+    deep : bool, default False
+        If True, introspect the data deeply by interrogating
+        `object` dtypes for system-level memory consumption, and include
+        it in the returned value.
+
+    Returns
+    -------
+    int
+        Bytes of memory consumed.
+
+    See Also
+    --------
+    numpy.ndarray.nbytes : Total bytes consumed by the elements of the
+        array.
+    DataFrame.memory_usage : Bytes consumed by a DataFrame.
+
+    Examples
+    --------
+    >>> import mars.dataframe as md
+    >>> s = md.Series(range(3))
+    >>> s.memory_usage().execute()
+    152
+
+    Not including the index gives the size of the rest of the data, which
+    is necessarily smaller:
+
+    >>> s.memory_usage(index=False).execute()
+    24
+
+    The memory footprint of `object` values is ignored by default:
+
+    >>> s = md.Series(["a", "b"])
+    >>> s.values.execute()
+    array(['a', 'b'], dtype=object)
+
+    >>> s.memory_usage().execute()
+    144
+
+    >>> s.memory_usage(deep=True).execute()
+    260
+    """
+    op = DataFrameMemoryUsage(index=index, deep=deep)
+    return op(series)
+
+
+def index_memory_usage(index, deep=False):
+    """
+    Memory usage of the values.
+
+    Parameters
+    ----------
+    deep : bool
+        Introspect the data deeply, interrogate
+        `object` dtypes for system-level memory consumption.
+
+    Returns
+    -------
+    bytes used
+
+    See Also
+    --------
+    numpy.ndarray.nbytes
+
+    Notes
+    -----
+    Memory usage does not include memory consumed by elements that
+    are not components of the array if deep=False
+    """
+    op = DataFrameMemoryUsage(index=False, deep=deep)
+    return op(index)


Test patch:
diff --git a/mars/dataframe/base/tests/test_base.py b/mars/dataframe/base/tests/test_base.py
--- a/mars/dataframe/base/tests/test_base.py
+++ b/mars/dataframe/base/tests/test_base.py
@@ -839,3 +839,50 @@ def testDropDuplicates(self):
         s = df['c7']
         with self.assertRaises(ValueError):
             s.drop_duplicates(method='unknown')
+
+    def testMemoryUsage(self):
+        dtypes = ['int64', 'float64', 'complex128', 'object', 'bool']
+        data = dict([(t, np.ones(shape=500).astype(t))
+                    for t in dtypes])
+        raw = pd.DataFrame(data)
+
+        df = from_pandas_df(raw, chunk_size=(500, 2))
+        r = df.memory_usage().tiles()
+
+        self.assertIsInstance(r, SERIES_TYPE)
+        self.assertEqual(r.shape, (6,))
+        self.assertEqual(len(r.chunks), 3)
+        self.assertIsNone(r.chunks[0].op.stage)
+
+        df = from_pandas_df(raw, chunk_size=(100, 3))
+        r = df.memory_usage(index=True).tiles()
+
+        self.assertIsInstance(r, SERIES_TYPE)
+        self.assertEqual(r.shape, (6,))
+        self.assertEqual(len(r.chunks), 2)
+        self.assertEqual(r.chunks[0].op.stage, OperandStage.reduce)
+
+        r = df.memory_usage(index=False).tiles()
+
+        self.assertIsInstance(r, SERIES_TYPE)
+        self.assertEqual(r.shape, (5,))
+        self.assertEqual(len(r.chunks), 2)
+        self.assertEqual(r.chunks[0].op.stage, OperandStage.reduce)
+
+        raw = pd.Series(np.ones(shape=500).astype('object'), name='s')
+
+        series = from_pandas_series(raw)
+        r = series.memory_usage().tiles()
+
+        self.assertIsInstance(r, TENSOR_TYPE)
+        self.assertEqual(r.shape, ())
+        self.assertEqual(len(r.chunks), 1)
+        self.assertIsNone(r.chunks[0].op.stage)
+
+        series = from_pandas_series(raw, chunk_size=100)
+        r = series.memory_usage().tiles()
+
+        self.assertIsInstance(r, TENSOR_TYPE)
+        self.assertEqual(r.shape, ())
+        self.assertEqual(len(r.chunks), 1)
+        self.assertEqual(r.chunks[0].op.stage, OperandStage.reduce)
diff --git a/mars/dataframe/base/tests/test_base_execution.py b/mars/dataframe/base/tests/test_base_execution.py
--- a/mars/dataframe/base/tests/test_base_execution.py
+++ b/mars/dataframe/base/tests/test_base_execution.py
@@ -1370,3 +1370,72 @@ def testDropDuplicates(self):
             result = self.executor.execute_dataframe(series, concat=True)[0]
             expected = s.drop_duplicates()
             pd.testing.assert_series_equal(result, expected)
+
+    def testMemoryUsageExecution(self):
+        dtypes = ['int64', 'float64', 'complex128', 'object', 'bool']
+        data = dict([(t, np.ones(shape=500).astype(t))
+                     for t in dtypes])
+        raw = pd.DataFrame(data)
+
+        df = from_pandas_df(raw, chunk_size=(500, 2))
+        r = df.memory_usage(index=False)
+        pd.testing.assert_series_equal(self.executor.execute_dataframe(r, concat=True)[0],
+                                       raw.memory_usage(index=False))
+
+        df = from_pandas_df(raw, chunk_size=(500, 2))
+        r = df.memory_usage(index=True)
+        pd.testing.assert_series_equal(self.executor.execute_dataframe(r, concat=True)[0],
+                                       raw.memory_usage(index=True))
+
+        df = from_pandas_df(raw, chunk_size=(100, 3))
+        r = df.memory_usage(index=False)
+        pd.testing.assert_series_equal(self.executor.execute_dataframe(r, concat=True)[0],
+                                       raw.memory_usage(index=False))
+
+        r = df.memory_usage(index=True)
+        pd.testing.assert_series_equal(self.executor.execute_dataframe(r, concat=True)[0],
+                                       raw.memory_usage(index=True))
+
+        raw = pd.DataFrame(data, index=np.arange(500).astype('object'))
+
+        df = from_pandas_df(raw, chunk_size=(100, 3))
+        r = df.memory_usage(index=True)
+        pd.testing.assert_series_equal(self.executor.execute_dataframe(r, concat=True)[0],
+                                       raw.memory_usage(index=True))
+
+        raw = pd.Series(np.ones(shape=500).astype('object'), name='s')
+
+        series = from_pandas_series(raw)
+        r = series.memory_usage(index=True)
+        self.assertEqual(self.executor.execute_dataframe(r, concat=True)[0],
+                         raw.memory_usage(index=True))
+
+        series = from_pandas_series(raw, chunk_size=100)
+        r = series.memory_usage(index=False)
+        self.assertEqual(self.executor.execute_dataframe(r, concat=True)[0],
+                         raw.memory_usage(index=False))
+
+        series = from_pandas_series(raw, chunk_size=100)
+        r = series.memory_usage(index=True)
+        self.assertEqual(self.executor.execute_dataframe(r, concat=True)[0],
+                         raw.memory_usage(index=True))
+
+        raw = pd.Series(np.ones(shape=500).astype('object'),
+                        index=np.arange(500).astype('object'), name='s')
+
+        series = from_pandas_series(raw, chunk_size=100)
+        r = series.memory_usage(index=True)
+        self.assertEqual(self.executor.execute_dataframe(r, concat=True)[0],
+                         raw.memory_usage(index=True))
+
+        raw = pd.Index(np.arange(500), name='s')
+
+        index = from_pandas_index(raw)
+        r = index.memory_usage()
+        self.assertEqual(self.executor.execute_dataframe(r, concat=True)[0],
+                         raw.memory_usage())
+
+        index = from_pandas_index(raw, chunk_size=100)
+        r = index.memory_usage()
+        self.assertEqual(self.executor.execute_dataframe(r, concat=True)[0],
+                         raw.memory_usage())