Repository: mars-project/mars

Issue:
[BUG] Run DataFrame getitem leads to memory insufficiency
<!--
Thank you for your contribution!

Please review https://github.com/mars-project/mars/blob/master/CONTRIBUTING.rst before opening an issue.
-->

**Describe the bug**
When execute df.iloc[:1000000], the job hangs and warning found in server side:
```
2020-10-28 14:08:40,463 mars.worker.quota 200 WARNING  w:0:MemQuotaActor met hard memory limitation: request 16613268648, available 9270380134, hard limit 61847529062
```
this may be caused by the optimization introduced in #1328, if `n` is too large, it will lead to a huge chunk when do reductions, we should skip the optimization if `n` greater than a threshold.



Discussion:

Resolving patch:
diff --git a/mars/config.py b/mars/config.py
--- a/mars/config.py
+++ b/mars/config.py
@@ -377,6 +377,7 @@ def validate(x):
 default_options.register_option('optimize.min_stats_count', 10, validator=is_integer)
 default_options.register_option('optimize.stats_sufficient_ratio', 0.9, validator=is_float, serialize=True)
 default_options.register_option('optimize.default_disk_io_speed', 10 * 1024 ** 2, validator=is_integer)
+default_options.register_option('optimize.head_optimize_threshold', 1000, validator=is_integer)
 
 default_options.register_option('optimize_tileable_graph', True, validator=is_bool)
 
diff --git a/mars/dataframe/indexing/iloc.py b/mars/dataframe/indexing/iloc.py
--- a/mars/dataframe/indexing/iloc.py
+++ b/mars/dataframe/indexing/iloc.py
@@ -212,9 +212,10 @@ def _tile_head_tail(cls, op):
         params['chunks'] = chunks
         return new_op.new_tileables(op.inputs, kws=[params])
 
-    def is_head(self):
+    def can_be_optimized(self):
         return self._is_indexes_head_or_tail(self._indexes) and \
-               self._is_head(self._indexes[0])
+               self._is_head(self._indexes[0]) and \
+               self._indexes[0].stop <= options.optimize.head_optimize_threshold
 
     @classmethod
     def tile(cls, op):
diff --git a/mars/optimizes/runtime/dataframe.py b/mars/optimizes/runtime/dataframe.py
--- a/mars/optimizes/runtime/dataframe.py
+++ b/mars/optimizes/runtime/dataframe.py
@@ -134,7 +134,7 @@ def match(chunk, graph, keys):
         op = chunk.op
         inputs = graph.predecessors(chunk)
         if len(inputs) == 1 and isinstance(op, (DataFrameIlocGetItem, SeriesIlocGetItem)) and \
-                op.is_head() and isinstance(inputs[0].op, (DataFrameReadCSV, DataFrameReadSQL)) and \
+                op.can_be_optimized() and isinstance(inputs[0].op, (DataFrameReadCSV, DataFrameReadSQL)) and \
                 inputs[0].key not in keys:
             return True
         return False


Test patch:
diff --git a/mars/dataframe/indexing/tests/test_indexing_execution.py b/mars/dataframe/indexing/tests/test_indexing_execution.py
--- a/mars/dataframe/indexing/tests/test_indexing_execution.py
+++ b/mars/dataframe/indexing/tests/test_indexing_execution.py
@@ -888,10 +888,10 @@ def testOptimization(self):
 
             filename = os.path.join(tempdir, 'test_head.csv')
             rs = np.random.RandomState(0)
-            pd_df = pd.DataFrame({'a': rs.randint(1000, size=(100,)).astype(np.int64),
-                                  'b': rs.randint(1000, size=(100,)).astype(np.int64),
-                                  'c': ['sss' for _ in range(100)],
-                                  'd': ['eeee' for _ in range(100)]})
+            pd_df = pd.DataFrame({'a': rs.randint(1000, size=(2000,)).astype(np.int64),
+                                  'b': rs.randint(1000, size=(2000,)).astype(np.int64),
+                                  'c': ['sss' for _ in range(2000)],
+                                  'd': ['eeee' for _ in range(2000)]})
             pd_df.to_csv(filename, index=False)
 
             size = os.path.getsize(filename)
@@ -960,6 +960,19 @@ def testOptimization(self):
             pd.testing.assert_frame_equal(result.reset_index(drop=True),
                                           expected.reset_index(drop=True))
 
+            # test head number greater than limit
+            df = md.read_csv(filename, chunk_bytes=chunk_bytes)
+            r = df.head(1100)
+
+            with self.assertRaises(RuntimeError):
+                with self._inject_execute_data_source(3, DataFrameReadCSV):
+                    result = executor.execute_tileables([r])[0]
+
+            result = executor.execute_tileables([r])[0]
+            expected = pd_df.head(1100)
+            pd.testing.assert_frame_equal(result.reset_index(drop=True),
+                                          expected.reset_index(drop=True))
+
             filename = os.path.join(tempdir, 'test_sql.db')
             conn = sa.create_engine('sqlite:///' + filename)
             pd_df.to_sql('test_sql', conn)