Repository: mars-project/mars

Issue:
Add memory estimation logic in `read_xxx` for DataFrame
<!--
Thank you for your contribution!

Please review https://github.com/mars-project/mars/blob/master/CONTRIBUTING.rst before opening an issue.
-->

**Is your feature request related to a problem? Please describe.**

For now, memory estimation is absent in `read_parquet`, `read_csv` and `read_sql`, we need to add the logic for them.



Discussion:

Resolving patch:
diff --git a/mars/dataframe/datasource/read_csv.py b/mars/dataframe/datasource/read_csv.py
--- a/mars/dataframe/datasource/read_csv.py
+++ b/mars/dataframe/datasource/read_csv.py
@@ -103,7 +103,8 @@ class DataFrameReadCSV(DataFrameOperand, DataFrameOperandMixin):
     def __init__(self, path=None, names=None, sep=None, header=None, index_col=None,
                  compression=None, usecols=None, offset=None, size=None, nrows=None,
                  gpu=None, keep_usecols_order=None, incremental_index=None,
-                 use_arrow_dtype=None, storage_options=None, **kw):
+                 use_arrow_dtype=None, storage_options=None,
+                 memory_scale=None, **kw):
         super().__init__(_path=path, _names=names, _sep=sep, _header=header,
                          _index_col=index_col, _compression=compression,
                          _usecols=usecols, _offset=offset, _size=size, _nrows=nrows,
@@ -111,7 +112,8 @@ def __init__(self, path=None, names=None, sep=None, header=None, index_col=None,
                          _keep_usecols_order=keep_usecols_order,
                          _use_arrow_dtype=use_arrow_dtype,
                          _storage_options=storage_options,
-                         _output_types=[OutputType.dataframe], **kw)
+                         _output_types=[OutputType.dataframe],
+                         _memory_scale=memory_scale, **kw)
 
     @property
     def path(self):
@@ -334,6 +336,10 @@ def execute(cls, ctx, op):
 
         ctx[out_df.key] = df
 
+    def estimate_size(cls, ctx, op):
+        phy_size = op.size * (op.memory_scale or 1)
+        ctx[op.outputs[0].key] = (phy_size, phy_size * 2)
+
     def __call__(self, index_value=None, columns_value=None, dtypes=None, chunk_bytes=None):
         shape = (np.nan, len(dtypes))
         return self.new_dataframe(None, shape, dtypes=dtypes, index_value=index_value,
@@ -343,7 +349,7 @@ def __call__(self, index_value=None, columns_value=None, dtypes=None, chunk_byte
 def read_csv(path, names=None, sep=',', index_col=None, compression=None, header='infer',
              dtype=None, usecols=None, nrows=None, chunk_bytes='64M', gpu=None, head_bytes='100k',
              head_lines=None, incremental_index=False, use_arrow_dtype=None,
-             storage_options=None, **kwargs):
+             storage_options=None, memory_scale=None, **kwargs):
     r"""
     Read a comma-separated values (csv) file into DataFrame.
     Also supports optionally iterating or breaking of the file
@@ -669,7 +675,7 @@ def read_csv(path, names=None, sep=',', index_col=None, compression=None, header
     op = DataFrameReadCSV(path=path, names=names, sep=sep, header=header, index_col=index_col,
                           usecols=usecols, compression=compression, gpu=gpu,
                           incremental_index=incremental_index, use_arrow_dtype=use_arrow_dtype,
-                          storage_options=storage_options,
+                          storage_options=storage_options, memory_scale=memory_scale,
                           **kwargs)
     chunk_bytes = chunk_bytes or options.chunk_store_limit
     dtypes = mini_df.dtypes
diff --git a/mars/dataframe/datasource/read_parquet.py b/mars/dataframe/datasource/read_parquet.py
--- a/mars/dataframe/datasource/read_parquet.py
+++ b/mars/dataframe/datasource/read_parquet.py
@@ -34,15 +34,20 @@
 
 from ... import opcodes as OperandDef
 from ...config import options
-from ...filesystem import get_fs, open_file, glob
+from ...filesystem import file_size, get_fs, glob, open_file
 from ...serialize import AnyField, BoolField, DictField, ListField,\
-    StringField, Int32Field, BytesField
+    StringField, Int32Field, Int64Field, BytesField
+from ...utils import is_object_dtype
 from ..arrays import ArrowStringDtype
 from ..operands import DataFrameOperandMixin, DataFrameOperand, OutputType
 from ..utils import parse_index, to_arrow_dtypes, contain_arrow_dtype, \
     standardize_range_index
 
 
+PARQUET_MEMORY_SCALE = 15
+STRING_FIELD_OVERHEAD = 50
+
+
 def check_engine(engine):
     if engine == 'auto':
         if pa is not None:
@@ -73,6 +78,9 @@ def get_engine(engine):
 
 
 class ParquetEngine:
+    def get_row_num(self, f):
+        raise NotImplementedError
+
     def read_dtypes(self, f, **kwargs):
         raise NotImplementedError
 
@@ -86,6 +94,10 @@ def read_group_to_pandas(self, f, group_index, columns=None,
 
 
 class ArrowEngine(ParquetEngine):
+    def get_row_num(self, f):
+        file = pq.ParquetFile(f)
+        return file.metadata.num_rows
+
     def read_dtypes(self, f, **kwargs):
         file = pq.ParquetFile(f)
         return file.schema_arrow.empty_table().to_pandas().dtypes
@@ -113,6 +125,10 @@ def read_group_to_pandas(self, f, group_index, columns=None,
 
 
 class FastpaquetEngine(ParquetEngine):
+    def get_row_num(self, f):
+        file = fastparquet.ParquetFile(f)
+        return file.count
+
     def read_dtypes(self, f, **kwargs):
         file = fastparquet.ParquetFile(f)
         dtypes_dict = file._dtypes()
@@ -140,15 +156,19 @@ class DataFrameReadParquet(DataFrameOperand, DataFrameOperandMixin):
     _read_kwargs = DictField('read_kwargs')
     _incremental_index = BoolField('incremental_index')
     _storage_options = DictField('storage_options')
-
     # for chunk
     _partitions = BytesField('partitions')
     _partition_keys = ListField('partition_keys')
+    # row numbers
+    _row_num = Int64Field('row_num')
+    # raw bytes of parquet files
+    _raw_bytes = Int64Field('raw_bytes')
 
     def __init__(self, path=None, engine=None, columns=None, use_arrow_dtype=None,
                  groups_as_chunks=None, group_index=None, incremental_index=None,
                  read_kwargs=None, partitions=None, partition_keys=None,
-                 storage_options=None, **kw):
+                 storage_options=None, row_num=None, raw_bytes=None,
+                 memory_scale=None, **kw):
         super().__init__(_path=path, _engine=engine, _columns=columns,
                          _use_arrow_dtype=use_arrow_dtype,
                          _groups_as_chunks=groups_as_chunks,
@@ -158,6 +178,8 @@ def __init__(self, path=None, engine=None, columns=None, use_arrow_dtype=None,
                          _partitions=partitions,
                          _partition_keys=partition_keys,
                          _storage_options=storage_options,
+                         _memory_scale=memory_scale,
+                         _row_num=row_num, _raw_bytes=raw_bytes,
                          _output_types=[OutputType.dataframe], **kw)
 
     @property
@@ -204,6 +226,14 @@ def partition_keys(self):
     def storage_options(self):
         return self._storage_options
 
+    @property
+    def row_num(self):
+        return self._row_num
+
+    @property
+    def raw_bytes(self):
+        return self._raw_bytes
+
     @classmethod
     def _to_arrow_dtypes(cls, dtypes, op):
         if op.use_arrow_dtype is None and not op.gpu and \
@@ -229,9 +259,11 @@ def _tile_partitioned(cls, op):
         out_chunks = []
         for piece in dataset.pieces:
             chunk_op = op.copy().reset_key()
-            chunk_op._path = path_prefix + piece.path
+            chunk_op._path = chunk_path = path_prefix + piece.path
             chunk_op._partitions = pickle.dumps(dataset.partitions)
             chunk_op._partition_keys = piece.partition_keys
+            chunk_op._row_num = piece.get_metadata().num_rows
+            chunk_op._raw_bytes = file_size(chunk_path, op.storage_options)
             new_chunk = chunk_op.new_chunk(
                 None, shape=shape, index=(chunk_index, 0),
                 index_value=out_df.index_value,
@@ -248,7 +280,7 @@ def _tile_partitioned(cls, op):
                                      chunks=out_chunks, nsplits=nsplits)
 
     @classmethod
-    def _tile_no_partitioned(cls, op):
+    def _tile_no_partitioned(cls, op: "DataFrameReadParquet"):
         chunk_index = 0
         out_chunks = []
         out_df = op.outputs[0]
@@ -260,11 +292,22 @@ def _tile_no_partitioned(cls, op):
             glob(op.path, storage_options=op.storage_options)
 
         for pth in paths:
+            raw_bytes = file_size(pth, op.storage_options)
+            with open_file(pth, storage_options=op.storage_options) as f:
+                row_num = get_engine(op.engine).get_row_num(f)
+
             if op.groups_as_chunks:
-                for group_idx in range(pq.ParquetFile(pth).num_row_groups):
+                num_row_groups = pq.ParquetFile(pth).num_row_groups
+                for group_idx in range(num_row_groups):
                     chunk_op = op.copy().reset_key()
                     chunk_op._path = pth
                     chunk_op._group_index = group_idx
+                    # we don't know exact row number and bytes of each group
+                    # just treat them as equal
+                    chunk_op._row_num = \
+                        np.ceil(np.divide(row_num, num_row_groups)).astype(int).item()
+                    chunk_op._raw_bytes = \
+                        np.ceil(np.divide(raw_bytes, num_row_groups)).astype(int).item()
                     new_chunk = chunk_op.new_chunk(
                         None, shape=shape, index=(chunk_index, 0),
                         index_value=out_df.index_value,
@@ -275,6 +318,8 @@ def _tile_no_partitioned(cls, op):
             else:
                 chunk_op = op.copy().reset_key()
                 chunk_op._path = pth
+                chunk_op._row_num = row_num
+                chunk_op._raw_bytes = raw_bytes
                 new_chunk = chunk_op.new_chunk(
                     None, shape=shape, index=(chunk_index, 0),
                     index_value=out_df.index_value,
@@ -331,6 +376,13 @@ def execute(cls, ctx, op):
 
             ctx[out.key] = df
 
+    @classmethod
+    def estimate_size(cls, ctx, op: "DataFrameReadParquet"):
+        phy_size = op.raw_bytes * (op.memory_scale or PARQUET_MEMORY_SCALE)
+        n_strings = len([dt for dt in op.outputs[0].dtypes if is_object_dtype(dt)])
+        pd_size = phy_size + n_strings * op.row_num * STRING_FIELD_OVERHEAD
+        ctx[op.outputs[0].key] = (pd_size, pd_size + phy_size)
+
     def __call__(self, index_value=None, columns_value=None, dtypes=None):
         shape = (np.nan, len(dtypes))
         return self.new_dataframe(None, shape, dtypes=dtypes, index_value=index_value,
@@ -340,7 +392,7 @@ def __call__(self, index_value=None, columns_value=None, dtypes=None):
 def read_parquet(path, engine: str = "auto", columns=None,
                  groups_as_chunks=False, use_arrow_dtype=None,
                  incremental_index=False, storage_options=None,
-                 **kwargs):
+                 memory_scale=None, **kwargs):
     """
     Load a parquet object from the file path, returning a DataFrame.
 
@@ -372,6 +424,8 @@ def read_parquet(path, engine: str = "auto", columns=None,
         If True, use arrow dtype to store columns.
     storage_options: dict, optional
         Options for storage connection.
+    memory_scale: int, optional
+        Scale that real memory occupation divided with raw file size.
     **kwargs
         Any additional kwargs are passed to the engine.
 
@@ -416,6 +470,7 @@ def read_parquet(path, engine: str = "auto", columns=None,
                               use_arrow_dtype=use_arrow_dtype,
                               read_kwargs=kwargs,
                               incremental_index=incremental_index,
-                              storage_options=storage_options)
+                              storage_options=storage_options,
+                              memory_scale=memory_scale)
     return op(index_value=index_value, columns_value=columns_value,
               dtypes=dtypes)
diff --git a/mars/utils.py b/mars/utils.py
--- a/mars/utils.py
+++ b/mars/utils.py
@@ -1064,17 +1064,18 @@ def __getattr__(self, item):  # pragma: no cover
         return getattr(self._file_obj, item)
 
 
+def is_object_dtype(dtype):
+    try:
+        return np.issubdtype(dtype, np.object_) \
+               or np.issubdtype(dtype, np.unicode_) \
+               or np.issubdtype(dtype, np.bytes_)
+    except TypeError:  # pragma: no cover
+        return False
+
+
 def calc_object_overhead(chunk, shape):
     from .dataframe.core import DATAFRAME_CHUNK_TYPE, SERIES_CHUNK_TYPE, INDEX_CHUNK_TYPE
 
-    def is_object_dtype(dtype):
-        try:
-            return np.issubdtype(dtype, np.object_) \
-                   or np.issubdtype(dtype, np.unicode_) \
-                   or np.issubdtype(dtype, np.bytes_)
-        except TypeError:  # pragma: no cover
-            return False
-
     if not shape or np.isnan(shape[0]) or getattr(chunk, 'dtypes', None) is None:
         return 0
 


Test patch:
diff --git a/mars/dataframe/datasource/tests/test_datasource_execution.py b/mars/dataframe/datasource/tests/test_datasource_execution.py
--- a/mars/dataframe/datasource/tests/test_datasource_execution.py
+++ b/mars/dataframe/datasource/tests/test_datasource_execution.py
@@ -319,8 +319,11 @@ def testReadCSVExecution(self):
             df.to_csv(file_path)
 
             pdf = pd.read_csv(file_path, index_col=0)
-            mdf = self.executor.execute_dataframe(md.read_csv(file_path, index_col=0), concat=True)[0]
+            r = md.read_csv(file_path, index_col=0)
+            mdf = self.executor.execute_dataframe(r, concat=True)[0]
             pd.testing.assert_frame_equal(pdf, mdf)
+            size_res = self.executor.execute_dataframe(r, mock=True)
+            self.assertEqual(sum(s[0] for s in size_res), os.stat(file_path).st_size)
 
             mdf2 = self.executor.execute_dataframe(md.read_csv(file_path, index_col=0, chunk_bytes=10),
                                                    concat=True)[0]
@@ -805,6 +808,8 @@ def testReadParquetArrow(self):
             df = md.read_parquet(file_path)
             result = self.executor.execute_dataframe(df, concat=True)[0]
             pd.testing.assert_frame_equal(result, test_df)
+            size_res = self.executor.execute_dataframe(df, mock=True)
+            self.assertGreater(sum(s[0] for s in size_res), test_df.memory_usage(deep=True).sum())
 
         with tempfile.TemporaryDirectory() as tempdir:
             file_path = os.path.join(tempdir, 'test.csv')
@@ -859,3 +864,5 @@ def testReadParquetFastParquet(self):
             df = md.read_parquet(file_path, engine='fastparquet')
             result = self.executor.execute_dataframe(df, concat=True)[0]
             pd.testing.assert_frame_equal(result, test_df)
+            size_res = self.executor.execute_dataframe(df, mock=True)
+            self.assertGreater(sum(s[0] for s in size_res), test_df.memory_usage(deep=True).sum())